diff --git a/studio/backend/core/training/resume.py b/studio/backend/core/training/resume.py new file mode 100644 index 0000000000..165c1c2cf1 --- /dev/null +++ b/studio/backend/core/training/resume.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Helpers for validating resumable training outputs.""" + +from pathlib import Path +from typing import Optional + +from utils.paths import outputs_root, resolve_output_dir + + +def _is_under_outputs(path: Path) -> bool: + resolved = path.resolve(strict = False) + root = outputs_root().resolve(strict = False) + try: + resolved.relative_to(root) + return True + except ValueError: + return False + + +def has_resume_state(path_value: Optional[str]) -> bool: + if not path_value: + return False + return get_resume_checkpoint_path(path_value) is not None + + +def _checkpoint_step(path: Path) -> int: + try: + return int(path.name.removeprefix("checkpoint-")) + except ValueError: + return -1 + + +def get_resume_checkpoint_path(path_value: str) -> Optional[str]: + path = resolve_output_dir(path_value) + if not _is_under_outputs(path) or not path.is_dir(): + return None + if (path / "trainer_state.json").is_file(): + return str(path) + + checkpoints = [ + child + for child in path.glob("checkpoint-*") + if child.is_dir() and (child / "trainer_state.json").is_file() + ] + if not checkpoints: + return None + return str(max(checkpoints, key = _checkpoint_step)) + + +def normalize_resume_output_dir(path_value: str) -> str: + path = resolve_output_dir(path_value) + if not _is_under_outputs(path): + raise ValueError("Resume checkpoint must be inside Studio outputs.") + return str(path) + + +def can_resume_run(run: dict) -> bool: + if run.get("resumed_later"): + return False + + final_step = run.get("final_step") + total_steps = run.get("total_steps") + has_remaining_steps = ( + not isinstance(final_step, int) + or not isinstance(total_steps, int) + or total_steps <= 0 + or final_step < total_steps + ) + return ( + run.get("status") == "stopped" + and has_remaining_steps + and has_resume_state(run.get("output_dir")) + ) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index c1e2ac4a85..24d06f9f62 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -376,6 +376,7 @@ class UnslothTrainer: def _finalize_training(self, output_dir, label = ""): """Save model after training and update progress. Used by all training branches.""" if self.should_stop and self.save_on_stop: + self.trainer._save_checkpoint(self.trainer.model, trial = None) self.trainer.save_model() self.tokenizer.save_pretrained(output_dir) self._patch_adapter_config(output_dir) @@ -2828,7 +2829,9 @@ class UnslothTrainer: total_steps = total, status_message = "Starting CSM training..." ) logger.info(f"CSM training config: {config}\n") - self.trainer.train() + self.trainer.train( + resume_from_checkpoint = training_args.get("resume_from_checkpoint") + ) self._finalize_training(output_dir, "CSM") return @@ -2867,7 +2870,9 @@ class UnslothTrainer: total_steps = total, status_message = "Starting SNAC training..." ) logger.info(f"SNAC training config: {config}\n") - self.trainer.train() + self.trainer.train( + resume_from_checkpoint = training_args.get("resume_from_checkpoint") + ) self._finalize_training(output_dir, "SNAC") return @@ -2913,7 +2918,9 @@ class UnslothTrainer: total_steps = total, status_message = "Starting Whisper training..." ) logger.info(f"Whisper training config: {config}\n") - self.trainer.train() + self.trainer.train( + resume_from_checkpoint = training_args.get("resume_from_checkpoint") + ) self._finalize_training(output_dir, "Whisper") return @@ -3408,7 +3415,9 @@ class UnslothTrainer: # ========== START TRAINING ========== self._update_progress(status_message = "Starting training...") logger.info("Starting training...\n") - self.trainer.train() + self.trainer.train( + resume_from_checkpoint = training_args.get("resume_from_checkpoint") + ) # ========== SAVE MODEL ========== self._finalize_training(output_dir) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index f35c7e8ad3..5437645797 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -185,6 +185,7 @@ class TrainingBackend: "wandb_project": kwargs.get("wandb_project", "unsloth-training"), "enable_tensorboard": kwargs.get("enable_tensorboard", False), "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), + "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"), "trust_remote_code": kwargs.get("trust_remote_code", False), "gpu_ids": kwargs.get("gpu_ids"), } diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index ff046c3fc9..60b9e994ab 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -35,6 +35,15 @@ from utils.wheel_utils import ( ) +def _output_dir_from_resume_checkpoint( + resume_from_checkpoint: str | None, +) -> str | None: + if not resume_from_checkpoint: + return None + path = Path(resume_from_checkpoint) + return str(path.parent if path.name.startswith("checkpoint-") else path) + + _CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4" _CAUSAL_CONV1D_PACKAGE_VERSION = "1.6.1" _MAMBA_SSM_RELEASE_TAG = "v2.3.1" @@ -757,7 +766,10 @@ def run_training_process( return # Generate output dir - output_dir = config.get("output_dir") + resume_from_checkpoint = config.get("resume_from_checkpoint") + output_dir = config.get("output_dir") or _output_dir_from_resume_checkpoint( + resume_from_checkpoint + ) if not output_dir: output_dir = f"{model_name.replace('/', '_')}_{int(time.time())}" output_dir = str(resolve_output_dir(output_dir)) @@ -805,6 +817,7 @@ def run_training_process( max_seq_length = config.get("max_seq_length", 2048), optim = config.get("optim", "adamw_8bit"), lr_scheduler_type = config.get("lr_scheduler_type", "linear"), + resume_from_checkpoint = resume_from_checkpoint, ) _tqdm_stop.set() @@ -821,10 +834,13 @@ def run_training_process( } ) else: + saved_output_dir = ( + None if trainer.should_stop and not trainer.save_on_stop else output_dir + ) event_queue.put( { "type": "complete", - "output_dir": output_dir, + "output_dir": saved_output_dir, "status_message": progress.status_message or "Training completed", "ts": time.time(), } @@ -1109,11 +1125,15 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> ) return - output_dir = config.get("output_dir") + resume_from_checkpoint = config.get("resume_from_checkpoint") + output_dir = config.get("output_dir") or _output_dir_from_resume_checkpoint( + resume_from_checkpoint + ) if not output_dir: output_dir = str( resolve_output_dir(f"{model_name.replace('/', '_')}_{int(time.time())}") ) + output_dir = str(resolve_output_dir(output_dir)) num_epochs = config.get("num_epochs", 2) batch_size = config.get("batch_size", 256) @@ -1221,7 +1241,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> callbacks = [_EmbeddingProgressCallback()], ) - trainer.train() + trainer.train(resume_from_checkpoint = resume_from_checkpoint) except Exception as e: event_queue.put( { @@ -1247,6 +1267,8 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> _send_status(event_queue, "Saving model...") try: + if _should_stop and _save_on_stop: + trainer._save_checkpoint(trainer.model, trial = None) model.save_pretrained(output_dir) model.tokenizer.save_pretrained(output_dir) logger.info("Embedding model saved to %s", output_dir) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 07a306ca39..a9f4caa1bb 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -127,6 +127,9 @@ class TrainingStartRequest(BaseModel): wandb_project: Optional[str] = Field(None, description = "W&B project name") enable_tensorboard: bool = Field(False, description = "Enable TensorBoard logging") tensorboard_dir: Optional[str] = Field(None, description = "TensorBoard directory") + resume_from_checkpoint: Optional[str] = Field( + None, description = "Saved training output directory to resume from" + ) # GPU selection gpu_ids: Optional[List[int]] = Field( @@ -220,6 +223,8 @@ class TrainingRunSummary(BaseModel): duration_seconds: Optional[float] = None error_message: Optional[str] = None loss_sparkline: Optional[List[float]] = None + can_resume: bool = False + resumed_later: bool = False class TrainingRunListResponse(BaseModel): diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index e625408bad..e5195bb337 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -25,6 +25,12 @@ if str(backend_path) not in sys.path: # Import backend functions try: from core.training import get_training_backend + from core.training.resume import ( + can_resume_run, + get_resume_checkpoint_path, + normalize_resume_output_dir, + ) + from storage.studio_db import get_resumable_run_by_output_dir from utils.models.model_config import load_model_defaults from utils.paths import resolve_dataset_path except ImportError: @@ -33,6 +39,12 @@ except ImportError: if str(parent_backend) not in sys.path: sys.path.insert(0, str(parent_backend)) from core.training import get_training_backend + from core.training.resume import ( + can_resume_run, + get_resume_checkpoint_path, + normalize_resume_output_dir, + ) + from storage.studio_db import get_resumable_run_by_output_dir from utils.models.model_config import load_model_defaults from utils.paths import resolve_dataset_path @@ -152,6 +164,28 @@ async def start_training( request.local_eval_datasets = _validate_local_dataset_paths( request.local_eval_datasets, "Local eval dataset" ) + resume_output_dir: Optional[str] = None + if request.resume_from_checkpoint: + try: + resume_output_dir = normalize_resume_output_dir( + request.resume_from_checkpoint + ) + except ValueError as e: + raise HTTPException(status_code = 400, detail = str(e)) + + resume_run = get_resumable_run_by_output_dir(resume_output_dir) + if not resume_run or not can_resume_run(resume_run): + raise HTTPException( + status_code = 400, + detail = "Resume checkpoint must belong to a stopped run with saved trainer state.", + ) + resume_checkpoint = get_resume_checkpoint_path(resume_output_dir) + if not resume_checkpoint: + raise HTTPException( + status_code = 400, + detail = "Resume checkpoint must include saved trainer state.", + ) + request.resume_from_checkpoint = resume_checkpoint # Convert request to kwargs for backend training_kwargs = { @@ -209,6 +243,8 @@ async def start_training( "wandb_project": request.wandb_project or "", "enable_tensorboard": request.enable_tensorboard, "tensorboard_dir": request.tensorboard_dir or "", + "output_dir": resume_output_dir, + "resume_from_checkpoint": request.resume_from_checkpoint, "trust_remote_code": request.trust_remote_code, "gpu_ids": request.gpu_ids, } @@ -437,6 +473,9 @@ async def get_training_status( "loss": getattr(progress, "loss", None), "learning_rate": getattr(progress, "learning_rate", None), } + output_dir = getattr(backend, "_output_dir", None) + if output_dir: + details["output_dir"] = output_dir # Build metric history for chart recovery after SSE reconnection metric_history = None diff --git a/studio/backend/routes/training_history.py b/studio/backend/routes/training_history.py index 597c4424c0..6f34321959 100644 --- a/studio/backend/routes/training_history.py +++ b/studio/backend/routes/training_history.py @@ -11,6 +11,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query from loggers import get_logger from auth.authentication import get_current_subject +from core.training.resume import can_resume_run from models import ( TrainingRunDeleteResponse, TrainingRunDetailResponse, @@ -34,7 +35,10 @@ async def list_training_runs( """List training runs, newest first.""" result = list_runs(limit = limit, offset = offset) return TrainingRunListResponse( - runs = [TrainingRunSummary(**r) for r in result["runs"]], + runs = [ + TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)}) + for r in result["runs"] + ], total = result["total"], ) @@ -58,7 +62,12 @@ async def get_training_run_detail( metrics_data = get_run_metrics(run_id) return TrainingRunDetailResponse( - run = TrainingRunSummary(**{k: v for k, v in run.items() if k != "config_json"}), + run = TrainingRunSummary( + **{ + **{k: v for k, v in run.items() if k != "config_json"}, + "can_resume": can_resume_run(run), + } + ), config = config, metrics = TrainingRunMetrics(**metrics_data), ) diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 89f75632ef..29e787c196 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -267,10 +267,23 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: total = conn.execute("SELECT COUNT(*) FROM training_runs").fetchone()[0] rows = conn.execute( """ - SELECT id, status, model_name, dataset_name, started_at, ended_at, - total_steps, final_step, final_loss, output_dir, - duration_seconds, error_message, loss_sparkline - FROM training_runs + SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at, + r.ended_at, r.total_steps, r.final_step, r.final_loss, + r.output_dir, r.duration_seconds, r.error_message, + r.loss_sparkline, + CASE + WHEN r.status = 'stopped' + AND r.output_dir IS NOT NULL + AND EXISTS ( + SELECT 1 + FROM training_runs newer + WHERE newer.output_dir = r.output_dir + AND newer.status IN ('stopped', 'completed') + AND newer.started_at > r.started_at + ) + THEN 1 ELSE 0 + END AS resumed_later + FROM training_runs r ORDER BY started_at DESC LIMIT ? OFFSET ? """, @@ -297,7 +310,26 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict: def get_run(id: str) -> Optional[dict]: conn = get_connection() try: - row = conn.execute("SELECT * FROM training_runs WHERE id = ?", (id,)).fetchone() + row = conn.execute( + """ + SELECT r.*, + CASE + WHEN r.status = 'stopped' + AND r.output_dir IS NOT NULL + AND EXISTS ( + SELECT 1 + FROM training_runs newer + WHERE newer.output_dir = r.output_dir + AND newer.status IN ('stopped', 'completed') + AND newer.started_at > r.started_at + ) + THEN 1 ELSE 0 + END AS resumed_later + FROM training_runs r + WHERE r.id = ? + """, + (id,), + ).fetchone() if row is None: return None run = dict(row) @@ -313,6 +345,45 @@ def get_run(id: str) -> Optional[dict]: conn.close() +def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]: + conn = get_connection() + try: + row = conn.execute( + """ + SELECT r.*, + 0 AS resumed_later + FROM training_runs r + WHERE r.output_dir = ? + AND r.status = 'stopped' + AND NOT EXISTS ( + SELECT 1 + FROM training_runs newer + WHERE newer.output_dir = r.output_dir + AND newer.status IN ('stopped', 'completed') + AND newer.started_at > r.started_at + ) + ORDER BY r.started_at DESC + LIMIT 1 + """, + (output_dir,), + ).fetchone() + if row is None: + return None + run = dict(row) + sparkline = run.get("loss_sparkline") + if sparkline: + try: + run["loss_sparkline"] = json.loads(sparkline) + except (json.JSONDecodeError, TypeError): + logger.debug( + "Failed to parse loss_sparkline for output_dir %s", output_dir + ) + run["loss_sparkline"] = None + return run + finally: + conn.close() + + def get_run_metrics(id: str) -> dict: """Return metric arrays for a run, using paired step arrays per metric.""" conn = get_connection() diff --git a/studio/frontend/src/features/studio/historical-training-view.tsx b/studio/frontend/src/features/studio/historical-training-view.tsx index d3ac434bd8..e461fc5a90 100644 --- a/studio/frontend/src/features/studio/historical-training-view.tsx +++ b/studio/frontend/src/features/studio/historical-training-view.tsx @@ -60,6 +60,7 @@ function mapToViewData(detail: TrainingRunDetailResponse): TrainingViewData { currentGradNorm: metrics.grad_norm_history.at(-1) ?? null, currentEpoch: metrics.final_epoch, currentNumTokens: metrics.final_num_tokens ?? null, + outputDir: run.output_dir ?? null, progressPercent: run.total_steps && run.final_step ? (run.final_step / run.total_steps) * 100 diff --git a/studio/frontend/src/features/studio/history-card-grid.tsx b/studio/frontend/src/features/studio/history-card-grid.tsx index 9e72c9f75c..6fbdffce85 100644 --- a/studio/frontend/src/features/studio/history-card-grid.tsx +++ b/studio/frontend/src/features/studio/history-card-grid.tsx @@ -13,7 +13,12 @@ import { } from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import type { TrainingRunSummary } from "@/features/training"; -import { deleteTrainingRun, listTrainingRuns } from "@/features/training"; +import { + deleteTrainingRun, + listTrainingRuns, + useTrainingActions, + useTrainingRuntimeStore, +} from "@/features/training"; import { formatDuration } from "@/features/studio/sections/progress-section-lib"; import { cn } from "@/lib/utils"; import { Delete02Icon } from "@hugeicons/core-free-icons"; @@ -47,8 +52,28 @@ const statusBadge: Record< className: "bg-blue-100 text-blue-700 dark:bg-blue-950 dark:text-blue-400", }, + resumed_later: { + label: "Continued", + className: + "bg-emerald-100 text-emerald-700 dark:bg-emerald-950 dark:text-emerald-400", + }, }; +function wasContinuedInVisibleRuns( + run: TrainingRunSummary, + runs: TrainingRunSummary[], +): boolean { + if (run.status !== "stopped" || !run.output_dir) return false; + const startedAt = new Date(run.started_at).getTime(); + return runs.some( + (other) => + other.id !== run.id && + other.output_dir === run.output_dir && + (other.status === "stopped" || other.status === "completed") && + new Date(other.started_at).getTime() > startedAt, + ); +} + function catmullRomPath(points: { x: number; y: number }[]): string { if (points.length < 2) return ""; const d = [`M${points[0]!.x.toFixed(1)},${points[0]!.y.toFixed(1)}`]; @@ -133,10 +158,12 @@ function formatRelativeTime(isoDate: string): string { interface HistoryCardGridProps { onSelectRun: (runId: string) => void; + onResumeStarted?: () => void; } export function HistoryCardGrid({ onSelectRun, + onResumeStarted, }: HistoryCardGridProps): ReactElement { const [runs, setRuns] = useState([]); const [total, setTotal] = useState(0); @@ -144,7 +171,10 @@ export function HistoryCardGrid({ const [error, setError] = useState(null); const [deleteError, setDeleteError] = useState(null); const [deleteTarget, setDeleteTarget] = useState(null); + const [resumeTarget, setResumeTarget] = useState(null); const [manualFetchInFlight, setManualFetchInFlight] = useState(false); + const { resumeTrainingRunFromHistory } = useTrainingActions(); + const isStarting = useTrainingRuntimeStore((state) => state.isStarting); const userControllerRef = useRef(null); const pollControllerRef = useRef(null); @@ -233,6 +263,18 @@ export function HistoryCardGrid({ setDeleteTarget(null); }; + const handleResume = async (runId: string) => { + setResumeTarget(runId); + try { + const ok = await resumeTrainingRunFromHistory(runId); + if (ok) { + onResumeStarted?.(); + } + } finally { + setResumeTarget(null); + } + }; + if (!loading && error && runs.length === 0) { return (
@@ -264,18 +306,25 @@ export function HistoryCardGrid({ )}
{runs.map((run) => { - const badge = statusBadge[run.status] ?? statusBadge.error; + const wasContinued = + run.resumed_later || wasContinuedInVisibleRuns(run, runs); + const badge = wasContinued + ? statusBadge.resumed_later + : (statusBadge[run.status] ?? statusBadge.error); const isRunning = run.status === "running"; + const canResume = run.can_resume && !wasContinued; + const isResuming = resumeTarget === run.id; return (
onSelectRun(run.id)} onKeyDown={(e) => { @@ -299,6 +348,21 @@ export function HistoryCardGrid({ {formatRelativeTime(run.started_at)}
+ {canResume && ( + + )}

{run.loss_sparkline && run.loss_sparkline.length >= 2 && ( - +
+ +
)}
diff --git a/studio/frontend/src/features/studio/live-training-view.tsx b/studio/frontend/src/features/studio/live-training-view.tsx index 9f930ce77b..e355655d54 100644 --- a/studio/frontend/src/features/studio/live-training-view.tsx +++ b/studio/frontend/src/features/studio/live-training-view.tsx @@ -31,6 +31,7 @@ export function LiveTrainingView(): ReactElement { elapsedSeconds: state.elapsedSeconds, etaSeconds: state.etaSeconds, evalEnabled: state.evalEnabled, + outputDir: state.outputDir, isTrainingRunning: state.isTrainingRunning, lossHistory: state.lossHistory, lrHistory: state.lrHistory, @@ -57,6 +58,7 @@ export function LiveTrainingView(): ReactElement { currentGradNorm: runtime.currentGradNorm, currentEpoch: runtime.currentEpoch, currentNumTokens: runtime.currentNumTokens, + outputDir: runtime.outputDir, progressPercent: runtime.progressPercent, elapsedSeconds: runtime.elapsedSeconds, etaSeconds: runtime.etaSeconds, diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx index 5018dc3baa..8cfd39d5c1 100644 --- a/studio/frontend/src/features/studio/studio-page.tsx +++ b/studio/frontend/src/features/studio/studio-page.tsx @@ -209,6 +209,9 @@ export function StudioPage(): ReactElement { } else { setSelectedHistoryRunId(runId); } + }} onResumeStarted={() => { + setSelectedHistoryRunId(null); + handleTabChange("current-run"); }} /> )} diff --git a/studio/frontend/src/features/studio/training-start-overlay.tsx b/studio/frontend/src/features/studio/training-start-overlay.tsx index b6042d4ae6..27c5a13163 100644 --- a/studio/frontend/src/features/studio/training-start-overlay.tsx +++ b/studio/frontend/src/features/studio/training-start-overlay.tsx @@ -64,6 +64,22 @@ const EMPTY_DOWNLOAD_STATE: DownloadState = { cachePath: null, }; +function coerceCachedStateReady(state: DownloadState): DownloadState { + if (!state.cachePath) return state; + if (state.downloadedBytes > 0 && state.percent < 100) return state; + const totalBytes = + state.totalBytes > 0 ? state.totalBytes : state.downloadedBytes; + if (totalBytes <= 0) { + return { ...state, percent: 100 }; + } + return { + ...state, + downloadedBytes: totalBytes, + totalBytes, + percent: 100, + }; +} + type Fetcher = (repoId: string) => Promise; /** @@ -88,10 +104,12 @@ function useHfDownloadProgress( phase === "downloading_model" || phase === "downloading_dataset" || phase === "loading_model" || - phase === "loading_dataset"; + phase === "loading_dataset" || + phase === "training"; useEffect(() => { if (!repoId || !HF_REPO_REGEX.test(repoId) || !shouldPoll) { + setState(EMPTY_DOWNLOAD_STATE); return; } @@ -229,14 +247,42 @@ export function TrainingStartOverlay({ }: TrainingStartOverlayProps): ReactElement { const { stopTrainingRun, dismissTrainingRun } = useTrainingActions(); const isStarting = useTrainingRuntimeStore((s) => s.isStarting); - const selectedModel = useTrainingConfigStore((s) => s.selectedModel); + const phase = useTrainingRuntimeStore((s) => s.phase); + const startModelName = useTrainingRuntimeStore((s) => s.startModelName); + const startDatasetName = useTrainingRuntimeStore((s) => s.startDatasetName); + const startFromResume = useTrainingRuntimeStore((s) => s.startFromResume); + const configuredModel = useTrainingConfigStore((s) => s.selectedModel); const datasetSource = useTrainingConfigStore((s) => s.datasetSource); const dataset = useTrainingConfigStore((s) => s.dataset); // Only HF datasets have a download phase to track. Uploaded files are // already on disk by the time the overlay shows up. const hfDatasetName = datasetSource === "huggingface" ? dataset : null; - const modelDownload = useModelDownloadProgress(selectedModel); - const datasetDownload = useDatasetDownloadProgress(hfDatasetName); + const hasStartResources = startModelName !== null; + const useConfiguredResources = !isStarting && !hasStartResources; + const isDownloadPhase = + phase === "downloading_model" || phase === "downloading_dataset"; + const modelName = hasStartResources + ? startModelName + : useConfiguredResources + ? configuredModel + : null; + const datasetName = hasStartResources + ? startDatasetName + : useConfiguredResources + ? hfDatasetName + : null; + const displayMessage = + startFromResume && !isDownloadPhase && /^download/i.test(message) + ? "Resuming training..." + : message || "starting training..."; + const rawModelDownload = useModelDownloadProgress(modelName); + const rawDatasetDownload = useDatasetDownloadProgress(datasetName); + const modelDownload = isDownloadPhase + ? rawModelDownload + : coerceCachedStateReady(rawModelDownload); + const datasetDownload = isDownloadPhase + ? rawDatasetDownload + : coerceCachedStateReady(rawDatasetDownload); const [cancelDialogOpen, setCancelDialogOpen] = useState(false); const [cancelRequested, setCancelRequested] = useState(false); @@ -314,7 +360,7 @@ export function TrainingStartOverlay({ {"> We are getting everything ready for your run..."} - {`> ${message || "starting training..."} | waiting for first step... (${currentStep})`} + {`> ${displayMessage} | waiting for first step... (${currentStep})`} {datasetDownload.downloadedBytes > 0 || datasetDownload.cachePath ? ( diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts index b5b41e7a0d..7f275d5d82 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -3,6 +3,7 @@ import { useCallback } from "react"; import { checkDatasetFormat } from "../api/datasets-api"; +import { getTrainingRun } from "../api/history-api"; import { buildTrainingStartPayload } from "../api/mappers"; import { startTraining, stopTraining, resetTraining } from "../api/train-api"; import { syncTrainingRuntimeFromBackend } from "../lib/sync-runtime"; @@ -10,6 +11,7 @@ import { validateTrainingConfig } from "../lib/validation"; import { useDatasetPreviewDialogStore } from "../stores/dataset-preview-dialog-store"; import { useTrainingConfigStore } from "../stores/training-config-store"; import { useTrainingRuntimeStore } from "../stores/training-runtime-store"; +import type { TrainingStartRequest } from "../types/api"; import type { TrainingConfigState } from "../types/config"; import { toast } from "sonner"; @@ -48,6 +50,11 @@ export function useTrainingActions() { return false; } + runtimeStore.setStartResources( + config.selectedModel ?? null, + getHfDatasetName(config), + false, + ); runtimeStore.setStarting(true); try { @@ -114,6 +121,7 @@ export function useTrainingActions() { // Re-read config after potential store updates from dataset check const payload = buildTrainingStartPayload(useTrainingConfigStore.getState()); + runtimeStore.setStartResources(payload.model_name, payload.hf_dataset, false); const response = await startTraining(payload); if (response.status === "error") { @@ -153,6 +161,54 @@ export function useTrainingActions() { } }, []); + const resumeTrainingRunFromHistory = useCallback(async (runId: string): Promise => { + const runtimeStore = useTrainingRuntimeStore.getState(); + runtimeStore.setStartError(null); + runtimeStore.setStartResources(null, null, true); + runtimeStore.setStarting(true); + + try { + const detail = await getTrainingRun(runId); + const outputDir = detail.run.output_dir; + if (!detail.run.can_resume || !outputDir) { + throw new Error("Only stopped runs with a saved checkpoint can be resumed."); + } + + const config = useTrainingConfigStore.getState(); + const savedConfig = detail.config as Partial; + const payload = { + ...savedConfig, + hf_token: + typeof savedConfig.hf_token === "string" + ? savedConfig.hf_token + : config.hfToken.trim() || null, + wandb_token: null, + resume_from_checkpoint: outputDir, + } as TrainingStartRequest; + + runtimeStore.setStartResources(payload.model_name, payload.hf_dataset, true); + + const response = await startTraining(payload); + if (response.status === "error") { + throw new Error(response.error || response.message); + } + + runtimeStore.setStartQueued(response.job_id, response.message); + await syncTrainingRuntimeFromBackend(); + return true; + } catch (error) { + const rawMessage = + error instanceof Error ? error.message : "Failed to resume training"; + const safeMessage = normalizeTrainingStartError(rawMessage); + runtimeStore.setStartError(safeMessage); + runtimeStore.setStarting(false); + toast.error("Could not resume training", { + description: safeMessage, + }); + return false; + } + }, []); + const dismissTrainingRun = useCallback(async (): Promise => { try { await resetTraining(); @@ -173,6 +229,7 @@ export function useTrainingActions() { isStarting, startError, startTrainingRun, + resumeTrainingRunFromHistory, stopTrainingRun, dismissTrainingRun, }; @@ -184,6 +241,10 @@ function getDatasetName(config: TrainingConfigState): string | null { : config.uploadedFile; } +function getHfDatasetName(config: TrainingConfigState): string | null { + return config.datasetSource === "huggingface" ? config.dataset : null; +} + function hasManualMapping(config: TrainingConfigState, isVlm = false, isAudio = false): boolean { const mapping = config.datasetManualMapping; const roles = new Set(Object.values(mapping)); diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts index d958078330..da496be002 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -22,6 +22,9 @@ const initialState: TrainingRuntimeState = { hasHydrated: false, isStarting: false, startError: null, + startModelName: null, + startDatasetName: null, + startFromResume: false, sseConnected: false, firstStepReceived: false, lastEventId: null, @@ -35,6 +38,7 @@ const initialState: TrainingRuntimeState = { etaSeconds: null, currentGradNorm: null, currentNumTokens: null, + outputDir: null, lossHistory: [], lrHistory: [], gradNormHistory: [], @@ -120,6 +124,8 @@ export const useTrainingRuntimeStore = create()((set) => ( setHasHydrated: (value) => set({ hasHydrated: value }), setStarting: (value) => set({ isStarting: value }), setStartError: (value) => set({ startError: value }), + setStartResources: (startModelName, startDatasetName, startFromResume = false) => + set({ startModelName, startDatasetName, startFromResume }), setSseConnected: (value) => set({ sseConnected: value }), setLastEventId: (value) => set({ lastEventId: value }), @@ -156,6 +162,7 @@ export const useTrainingRuntimeStore = create()((set) => ( etaSeconds: null, currentGradNorm: null, currentNumTokens: null, + outputDir: null, lossHistory: [], lrHistory: [], gradNormHistory: [], @@ -207,6 +214,7 @@ export const useTrainingRuntimeStore = create()((set) => ( typeof detailLr === "number" ? detailLr : state.currentLearningRate, currentEpoch: typeof detailEpoch === "number" ? detailEpoch : state.currentEpoch, + outputDir: payload.details?.output_dir ?? state.outputDir, lossHistory: metricHistory.lossHistory ?? state.lossHistory, lrHistory: metricHistory.lrHistory ?? state.lrHistory, gradNormHistory: metricHistory.gradNormHistory ?? state.gradNormHistory, diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index 1d72e9d9d8..ae65d1a53c 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -54,6 +54,7 @@ export interface TrainingStartRequest { wandb_project: string | null; enable_tensorboard: boolean; tensorboard_dir: string | null; + resume_from_checkpoint?: string | null; } export interface TrainingStartResponse { diff --git a/studio/frontend/src/features/training/types/history.ts b/studio/frontend/src/features/training/types/history.ts index 8b89db539b..901dfe3266 100644 --- a/studio/frontend/src/features/training/types/history.ts +++ b/studio/frontend/src/features/training/types/history.ts @@ -12,6 +12,8 @@ export interface TrainingRunSummary { final_step: number | null; final_loss: number | null; output_dir: string | null; + can_resume: boolean; + resumed_later: boolean; duration_seconds: number | null; error_message: string | null; loss_sparkline: number[] | null; diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts index 3e4dbd3b79..a62ef9ef26 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -26,6 +26,7 @@ export interface TrainingStatusResponse { total_steps?: number; loss?: number; learning_rate?: number; + output_dir?: string; } | null; metric_history?: { steps?: number[]; @@ -80,6 +81,9 @@ export interface TrainingRuntimeState { hasHydrated: boolean; isStarting: boolean; startError: string | null; + startModelName: string | null; + startDatasetName: string | null; + startFromResume: boolean; sseConnected: boolean; firstStepReceived: boolean; lastEventId: number | null; @@ -93,6 +97,7 @@ export interface TrainingRuntimeState { etaSeconds: number | null; currentGradNorm: number | null; currentNumTokens: number | null; + outputDir: string | null; lossHistory: TrainingSeriesPoint[]; lrHistory: TrainingSeriesPoint[]; gradNormHistory: TrainingSeriesPoint[]; @@ -108,6 +113,11 @@ export interface TrainingRuntimeActions { setHasHydrated: (value: boolean) => void; setStarting: (value: boolean) => void; setStartError: (value: string | null) => void; + setStartResources: ( + modelName: string | null, + datasetName: string | null, + fromResume?: boolean, + ) => void; setSseConnected: (value: boolean) => void; setLastEventId: (value: number | null) => void; resetRuntime: () => void; @@ -131,6 +141,7 @@ export interface TrainingViewData { currentGradNorm: number | null; currentEpoch: number | null; currentNumTokens: number | null; + outputDir: string | null; progressPercent: number; elapsedSeconds: number | null; etaSeconds: number | null;