From dc0cec772dc5d903adc8dfacc46498a6169ff9d0 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Tue, 17 Feb 2026 23:32:22 +0100 Subject: [PATCH] feat: enhance training stop and reset flow with detailed checks --- studio/backend/core/training/trainer.py | 19 ++++++--- studio/backend/core/training/training.py | 27 ++++++++++-- studio/backend/routes/training.py | 41 ++++++++++++++++--- .../src/features/studio/studio-page.tsx | 7 +++- .../training/hooks/use-training-actions.ts | 14 +++++-- .../hooks/use-training-runtime-lifecycle.ts | 8 +--- .../training/stores/training-runtime-store.ts | 23 ++++++++++- 7 files changed, 112 insertions(+), 27 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 8ada3d30f4..9468e282ef 100644 --- a/studio/backend/core/training/trainer.py +++ b/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: diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 9b589d215d..6fe08c2b9e 100644 --- a/studio/backend/core/training/training.py +++ b/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: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 1770999284..f8de2f639f 100644 --- a/studio/backend/routes/training.py +++ b/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( diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx index 973a7b66d0..e9f01901b7 100644 --- a/studio/frontend/src/features/studio/studio-page.tsx +++ b/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; 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 1a2a1aef8c..534240155f 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/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 => { - 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(); } }, []); diff --git a/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts b/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts index 4bf329a33e..ad7f61cc5e 100644 --- a/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts +++ b/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 { 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 6a3756a8ca..a80425188d 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -128,14 +128,33 @@ export const useTrainingRuntimeStore = create()((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({