feat: enhance training stop and reset flow with detailed checks

studio/backend/core/training/trainer.py

@@ -586,8 +586,14 @@ class UnslothTrainer:
         )
 
         self.should_stop = False
-        self.training_thread.start()
-        return True
+        self.is_training = True
+        try:
+            self.training_thread.start()
+            return True
+        except Exception as e:
+            self.is_training = False
+            logger.error(f"Failed to start training thread: {e}")
+            return False
 
     def _train_worker(self, dataset: Dataset, **training_args):
         """Worker function for training (runs in separate thread)"""
@@ -990,9 +996,12 @@ class UnslothTrainer:
         print(f"\nStopping training (save={save})...")
         self.should_stop = True
         self.save_on_stop = save
-        self.is_training = False
-        # Clear the status message so timer doesn't show stale status
-        self._update_progress(is_training=False, status_message="")
+        stop_msg = (
+            "Stopping training and saving checkpoint..."
+            if save
+            else "Cancelling training..."
+        )
+        self._update_progress(status_message=stop_msg)
 
         # If trainer exists, try to stop it gracefully
         if self.trainer:

studio/backend/core/training/training.py

@@ -369,8 +369,13 @@ class TrainingBackend:
             True if training is in progress, False otherwise
         """
         try:
-            # If user requested stop, training is no longer considered active
-            if self.trainer.should_stop:
+            training_thread = getattr(self.trainer, "training_thread", None)
+            if training_thread and training_thread.is_alive():
+                return True
+
+            # Stop requested and worker already exited => inactive.
+            # This allows UI to show stopped state + "Back to configuration".
+            if getattr(self.trainer, "should_stop", False):
                 return False
 
             progress = self.trainer.get_training_progress()
@@ -381,7 +386,23 @@ class TrainingBackend:
             # but haven't completed or errored yet
             if not is_active and not progress.is_completed and not progress.error:
                 status = progress.status_message or ""
-                if any(keyword in status.lower() for keyword in ["loading", "preparing", "training"]):
+                status_lower = status.lower()
+                if any(
+                    keyword in status_lower
+                    for keyword in ["cancelled", "canceled", "stopped", "completed", "ready to train"]
+                ):
+                    return False
+                if any(
+                    keyword in status_lower
+                    for keyword in [
+                        "loading",
+                        "preparing",
+                        "training",
+                        "configuring",
+                        "tokenizing",
+                        "starting",
+                    ]
+                ):
                     is_active = True
             return is_active
         except Exception as e:

studio/backend/routes/training.py

@@ -288,21 +288,31 @@ async def stop_training(
     """
     try:
         backend = get_training_backend()
-        
-        if not backend.is_training_active():
+        trainer_thread = getattr(getattr(backend, "trainer", None), "training_thread", None)
+        thread_alive = bool(trainer_thread and trainer_thread.is_alive())
+        is_active = backend.is_training_active()
+        logger.info(
+            "Stop requested: save=%s is_active=%s thread_alive=%s should_stop=%s",
+            body.save,
+            is_active,
+            thread_alive,
+            getattr(getattr(backend, "trainer", None), "should_stop", None),
+        )
+
+        if not is_active and not thread_alive:
             return TrainingStopResponse(
                 status="idle",
                 message="No training job is currently running"
             )
-        
+
         # Call backend stop method
         backend.stop_training(save=body.save)
-        
+
         return TrainingStopResponse(
             status="stopped",
-            message="Training job stopped successfully"
+            message="Stop requested. Training will stop at the next safe step."
         )
-        
+
     except Exception as e:
         logger.error(f"Error stopping training: {e}", exc_info=True)
         raise HTTPException(
@@ -320,6 +330,23 @@ async def reset_training(
     """
     try:
         backend = get_training_backend()
+        trainer_thread = getattr(getattr(backend, "trainer", None), "training_thread", None)
+        thread_alive = bool(trainer_thread and trainer_thread.is_alive())
+        is_active = backend.is_training_active()
+
+        if is_active or thread_alive:
+            logger.warning(
+                "Rejected reset while training active: is_active=%s thread_alive=%s should_stop=%s",
+                is_active,
+                thread_alive,
+                getattr(getattr(backend, "trainer", None), "should_stop", None),
+            )
+            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.trainer.should_stop = False
         backend.trainer.training_progress = backend.trainer.training_progress.__class__()
         backend.loss_history = []
@@ -328,6 +355,8 @@ async def reset_training(
         backend.grad_norm_history = []
         backend.grad_norm_step_history = []
         return {"status": "ok"}
+    except HTTPException:
+        raise
     except Exception as e:
         logger.error(f"Error resetting training: {e}", exc_info=True)
         raise HTTPException(

studio/frontend/src/features/studio/studio-page.tsx

@@ -25,6 +25,7 @@ const STUDIO_TOUR_KEY = "tour:studio:v1";
 export function StudioPage(): ReactElement {
   useTrainingRuntimeLifecycle();
   const showTrainingView = useTrainingRuntimeStore(shouldShowTrainingView);
+  const isTrainingRunning = useTrainingRuntimeStore((state) => state.isTrainingRunning);
   const runtimeMessage = useTrainingRuntimeStore((state) => state.message);
   const runtimePhase = useTrainingRuntimeStore((state) => state.phase);
   const isHydratingRuntime = useTrainingRuntimeStore((state) => state.isHydrating);
@@ -41,7 +42,11 @@ export function StudioPage(): ReactElement {
   const dialogInitial = useDatasetPreviewDialogStore((s) => s.initialData);
   const closeDialog = useDatasetPreviewDialogStore((s) => s.close);
 
-  const canGoBack = runtimePhase === "stopped" || runtimePhase === "error" || runtimePhase === "completed";
+  const canGoBack =
+    showTrainingView &&
+    !isTrainingRunning &&
+    !isHydratingRuntime &&
+    (runtimePhase === "stopped" || runtimePhase === "error" || runtimePhase === "completed" || runtimePhase === "idle");
   const tourEnabled = hasHydratedRuntime && !isHydratingRuntime;
   const isConfigTour = !showTrainingView;
   const tourSteps = showTrainingView ? studioTrainingTourSteps : studioTourSteps;

studio/frontend/src/features/training/hooks/use-training-actions.ts

@@ -8,6 +8,7 @@ import { useDatasetPreviewDialogStore } from "../stores/dataset-preview-dialog-s
 import { useTrainingConfigStore } from "../stores/training-config-store";
 import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
 import type { TrainingConfigState } from "../types/config";
+import { toast } from "sonner";
 
 export function useTrainingActions() {
   const isStarting = useTrainingRuntimeStore((state) => state.isStarting);
@@ -99,11 +100,18 @@ export function useTrainingActions() {
   }, []);
 
   const dismissTrainingRun = useCallback(async (): Promise<void> => {
-    useTrainingRuntimeStore.getState().resetRuntime();
     try {
       await resetTraining();
-    } catch {
-      // Frontend already reset; backend will catch up on next poll
+      useTrainingRuntimeStore.getState().resetRuntime();
+    } catch (error) {
+      const message =
+        error instanceof Error
+          ? error.message
+          : "Stop training first, then return to configuration.";
+      toast.error("Training still active", {
+        description: message,
+      });
+      await syncTrainingRuntimeFromBackend();
     }
   }, []);
 

studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts

@@ -14,13 +14,7 @@ const METRICS_POLL_INTERVAL_MS = 5000;
 const STREAM_RECONNECT_DELAY_MS = 1500;
 
 function shouldUseLiveSync(state: TrainingRuntimeStore): boolean {
-  return (
-    state.isTrainingRunning ||
-    state.phase === "loading_model" ||
-    state.phase === "loading_dataset" ||
-    state.phase === "configuring" ||
-    state.phase === "training"
-  );
+  return state.isTrainingRunning || state.phase === "training";
 }
 
 export function useTrainingRuntimeLifecycle(): void {

studio/frontend/src/features/training/stores/training-runtime-store.ts

@@ -128,14 +128,33 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
     })),
 
   setStartQueued: (jobId, message) =>
-    set({
+    set((state) => ({
+      ...state,
       jobId,
       message,
       error: null,
       startError: null,
       phase: "configuring",
       isStarting: false,
-    }),
+      sseConnected: false,
+      firstStepReceived: false,
+      lastEventId: null,
+      currentStep: 0,
+      totalSteps: 0,
+      currentEpoch: 0,
+      currentLoss: 0,
+      currentLearningRate: 0,
+      progressPercent: 0,
+      elapsedSeconds: null,
+      etaSeconds: null,
+      currentGradNorm: null,
+      currentNumTokens: null,
+      lossHistory: [],
+      lrHistory: [],
+      gradNormHistory: [],
+      evalLossHistory: [],
+      resetGeneration: state.resetGeneration + 1,
+    })),
 
   setRuntimeError: (message) =>
     set({