Studio: Add checkpoint resume for stopped training runs (#5255)
* feat: add checkpoint resume for stopped training runs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix:add resume checkpoint helpers * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: use checkpoint parent as resume output dir * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: save optimizer and scheduler state on stop-and-save Use Trainer._save_checkpoint instead of save_state so resume restores optimizer momentum and LR-schedule position via the checkpoint-NNN/ subdir written by HF's official path. * fix: clean up resume training history and startup progress * fix: preserve resume output dirs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: tighten resume run lookup * fix: remove stale output-dir lookup * fix: preserve startup download progress --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
This commit is contained in:
Lee Jackson
GitHub
parent
8cbd16786b
commit
2de17c0a96
18 changed files with 456 additions and 24 deletions
75
studio/backend/core/training/resume.py
Normal file
75
studio/backend/core/training/resume.py
Normal file
|
|
@ -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"))
|
||||
)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<TrainingRunSummary[]>([]);
|
||||
const [total, setTotal] = useState(0);
|
||||
|
|
@ -144,7 +171,10 @@ export function HistoryCardGrid({
|
|||
const [error, setError] = useState<string | null>(null);
|
||||
const [deleteError, setDeleteError] = useState<string | null>(null);
|
||||
const [deleteTarget, setDeleteTarget] = useState<string | null>(null);
|
||||
const [resumeTarget, setResumeTarget] = useState<string | null>(null);
|
||||
const [manualFetchInFlight, setManualFetchInFlight] = useState(false);
|
||||
const { resumeTrainingRunFromHistory } = useTrainingActions();
|
||||
const isStarting = useTrainingRuntimeStore((state) => state.isStarting);
|
||||
|
||||
const userControllerRef = useRef<AbortController | null>(null);
|
||||
const pollControllerRef = useRef<AbortController | null>(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 (
|
||||
<div className="flex flex-col items-center gap-2 py-16 text-center">
|
||||
|
|
@ -264,18 +306,25 @@ export function HistoryCardGrid({
|
|||
)}
|
||||
<div className="grid grid-cols-1 gap-4 sm:grid-cols-2 xl:grid-cols-3">
|
||||
{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 (
|
||||
<div
|
||||
role="button"
|
||||
tabIndex={0}
|
||||
key={run.id}
|
||||
className={cn(
|
||||
"group relative flex cursor-pointer flex-col gap-3 rounded-xl border bg-card p-4 text-left transition-colors hover:border-border hover:bg-accent/30",
|
||||
"group relative flex h-[11.5rem] cursor-pointer flex-col gap-3 rounded-xl border bg-card p-4 text-left transition-colors hover:border-border hover:bg-accent/30",
|
||||
isRunning
|
||||
? "border-blue-400/50 dark:border-blue-500/30"
|
||||
: "border-border/60",
|
||||
canResume && "gap-2",
|
||||
)}
|
||||
onClick={() => onSelectRun(run.id)}
|
||||
onKeyDown={(e) => {
|
||||
|
|
@ -299,6 +348,21 @@ export function HistoryCardGrid({
|
|||
{formatRelativeTime(run.started_at)}
|
||||
</span>
|
||||
</div>
|
||||
{canResume && (
|
||||
<Button
|
||||
type="button"
|
||||
size="xs"
|
||||
variant="outline"
|
||||
className="absolute bottom-3 left-4 h-6 rounded-full px-2.5 text-[11px] leading-none shadow-sm"
|
||||
disabled={isStarting || isResuming}
|
||||
onClick={(e) => {
|
||||
e.stopPropagation();
|
||||
void handleResume(run.id);
|
||||
}}
|
||||
>
|
||||
{isResuming ? "Resuming..." : "Resume training"}
|
||||
</Button>
|
||||
)}
|
||||
<div className="min-w-0">
|
||||
<p
|
||||
className="truncate text-sm font-medium"
|
||||
|
|
@ -311,7 +375,9 @@ export function HistoryCardGrid({
|
|||
</p>
|
||||
</div>
|
||||
{run.loss_sparkline && run.loss_sparkline.length >= 2 && (
|
||||
<Sparkline values={run.loss_sparkline} id={run.id} />
|
||||
<div className={cn(canResume && "h-7 overflow-hidden")}>
|
||||
<Sparkline values={run.loss_sparkline} id={run.id} />
|
||||
</div>
|
||||
)}
|
||||
<div className="flex flex-wrap gap-x-4 gap-y-1 text-[11px] text-muted-foreground">
|
||||
<span>
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -209,6 +209,9 @@ export function StudioPage(): ReactElement {
|
|||
} else {
|
||||
setSelectedHistoryRunId(runId);
|
||||
}
|
||||
}} onResumeStarted={() => {
|
||||
setSelectedHistoryRunId(null);
|
||||
handleTabChange("current-run");
|
||||
}} />
|
||||
)}
|
||||
</TabsContent>
|
||||
|
|
|
|||
|
|
@ -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<DownloadProgressResponse>;
|
||||
|
||||
/**
|
||||
|
|
@ -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..."}
|
||||
</TypingAnimation>
|
||||
<AnimatedSpan className="mt-2 text-muted-foreground">
|
||||
{`> ${message || "starting training..."} | waiting for first step... (${currentStep})`}
|
||||
{`> ${displayMessage} | waiting for first step... (${currentStep})`}
|
||||
</AnimatedSpan>
|
||||
{datasetDownload.downloadedBytes > 0 || datasetDownload.cachePath ? (
|
||||
<AnimatedSpan className="mt-3">
|
||||
|
|
|
|||
|
|
@ -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<boolean> => {
|
||||
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<TrainingStartRequest>;
|
||||
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<void> => {
|
||||
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));
|
||||
|
|
|
|||
|
|
@ -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<TrainingRuntimeStore>()((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<TrainingRuntimeStore>()((set) => (
|
|||
etaSeconds: null,
|
||||
currentGradNorm: null,
|
||||
currentNumTokens: null,
|
||||
outputDir: null,
|
||||
lossHistory: [],
|
||||
lrHistory: [],
|
||||
gradNormHistory: [],
|
||||
|
|
@ -207,6 +214,7 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((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,
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue