diff --git a/cli/commands/export.py b/cli/commands/export.py index ba487aeff3..71aa95b7f0 100644 --- a/cli/commands/export.py +++ b/cli/commands/export.py @@ -25,8 +25,11 @@ def list_checkpoints( typer.echo("No checkpoints found.") raise typer.Exit() - for display, path in checkpoints: - typer.echo(f"{display}: {path}") + for model_name, ckpt_list, metadata in checkpoints: + typer.echo(f"\n{model_name}:") + for display, path, loss in ckpt_list: + loss_str = f" (loss: {loss:.4f})" if loss is not None else "" + typer.echo(f" {display}{loss_str}: {path}") def export( diff --git a/cli/commands/train.py b/cli/commands/train.py index 2e78db7375..be124407b3 100644 --- a/cli/commands/train.py +++ b/cli/commands/train.py @@ -81,7 +81,7 @@ def train( ) raise typer.Exit(code=2) - from studio.backend.core.training import UnslothTrainer + from studio.backend.core.training.trainer import UnslothTrainer trainer = UnslothTrainer() @@ -101,18 +101,20 @@ def train( typer.echo("Model preparation failed", err=True) raise typer.Exit(code=1) - ds = trainer.load_and_format_dataset( + result = trainer.load_and_format_dataset( dataset_source=cfg.data.dataset or "", format_type=cfg.data.format_type, local_datasets=cfg.data.local_dataset, ) - if ds is None: + if result is None: typer.echo("Dataset load failed", err=True) raise typer.Exit(code=1) + ds, eval_ds = result + training_kwargs = cfg.training_kwargs() training_kwargs["wandb_token"] = wandb_token # CLI/env takes precedence - started = trainer.start_training(dataset=ds, **training_kwargs) + started = trainer.start_training(dataset=ds, eval_dataset=eval_ds, **training_kwargs) if not started: typer.echo("Training failed to start", err=True) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 101ba13621..14ff57bbd7 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -71,6 +71,7 @@ class TrainingBackend: # Progress state (updated by pump thread from subprocess events) self._progress = TrainingProgress() self._should_stop = False + self._cancel_requested = False # True only for stop(save=False) # Training Metrics (consumed by routes for SSE and /metrics) self.loss_history: list = [] @@ -114,6 +115,7 @@ class TrainingBackend: # Reset state self._should_stop = False + self._cancel_requested = False self._progress = TrainingProgress(is_training=True, status_message="Initializing training...") self.loss_history.clear() self.lr_history.clear() @@ -213,6 +215,8 @@ class TrainingBackend: def stop_training(self, save: bool = True) -> bool: """Send stop signal to the training subprocess.""" self._should_stop = True + if not save: + self._cancel_requested = True with self._lock: if self._stop_queue is not None: try: @@ -226,6 +230,20 @@ class TrainingBackend: ) return True + def force_terminate(self) -> None: + """Force-kill the training subprocess so state can be reset immediately.""" + with self._lock: + if self._proc is not None and self._proc.is_alive(): + logger.info("Force-terminating training subprocess (pid=%s)", self._proc.pid) + self._proc.terminate() + proc = self._proc + + if proc is not None: + proc.join(timeout=5.0) + if proc.is_alive(): + proc.kill() + proc.join(timeout=2.0) + def is_training_active(self) -> bool: """Check if training is currently active.""" with self._lock: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 3deabb5253..c9dd63a6a4 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -323,11 +323,16 @@ async def reset_training( is_active = backend.is_training_active() if is_active: - logger.warning("Rejected reset while training active: is_active=%s", is_active) - raise HTTPException( - status_code=409, - detail="Training is still running. Stop training and wait for it to finish before resetting.", - ) + if backend._cancel_requested: + # Cancel (save=False) was requested — force-terminate so we can reset immediately + logger.info("Force-terminating subprocess for immediate reset (cancel path)") + backend.force_terminate() + else: + logger.warning("Rejected reset while training active: is_active=%s", is_active) + raise HTTPException( + status_code=409, + detail="Training is still running. Stop training and wait for it to finish before resetting.", + ) logger.info("Reset training state: clearing runtime + metric history") backend._should_stop = False # Clear stop flag so status returns to idle diff --git a/studio/frontend/src/features/studio/training-start-overlay.tsx b/studio/frontend/src/features/studio/training-start-overlay.tsx index e2d660343a..898e47297a 100644 --- a/studio/frontend/src/features/studio/training-start-overlay.tsx +++ b/studio/frontend/src/features/studio/training-start-overlay.tsx @@ -31,7 +31,7 @@ export function TrainingStartOverlay({ message, currentStep, }: TrainingStartOverlayProps): ReactElement { - const { stopTrainingRun } = useTrainingActions(); + const { stopTrainingRun, dismissTrainingRun } = useTrainingActions(); const isStarting = useTrainingRuntimeStore((s) => s.isStarting); const [cancelDialogOpen, setCancelDialogOpen] = useState(false); const [cancelRequested, setCancelRequested] = useState(false); @@ -77,7 +77,11 @@ export function TrainingStartOverlay({ setCancelDialogOpen(false); useTrainingRuntimeStore.getState().setStopRequested(true); void stopTrainingRun(false).then((ok) => { - if (!ok) setCancelRequested(false); + if (ok) { + void dismissTrainingRun(); + } else { + setCancelRequested(false); + } }); }} > 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 b748daa907..773c8f267b 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -105,6 +105,12 @@ export function useTrainingActions() { } } + // Abort if cancel was requested during dataset check + if (useTrainingRuntimeStore.getState().stopRequested) { + runtimeStore.setStarting(false); + return false; + } + // Re-read config after potential store updates from dataset check const payload = buildTrainingStartPayload(useTrainingConfigStore.getState()); const response = await startTraining(payload);