Merge pull request #352 from unslothai/fix/cancel-training

Fix/cancel training

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(

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)

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:

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

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);
+                      }
                     });
                   }}
                 >

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);