From 354b7d0aca6f02b634b23c38c37a438ab77c87aa Mon Sep 17 00:00:00 2001 From: Manan17 Date: Sun, 15 Feb 2026 00:00:22 +0000 Subject: [PATCH 1/2] feat: add cancel or save and stop training --- studio/backend/core/training/trainer.py | 42 +++++++++++++++++------- studio/backend/core/training/training.py | 14 ++++++-- studio/backend/routes/training.py | 39 +++++++++++++++++++++- 3 files changed, 79 insertions(+), 16 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 52cf60af10..013c7817b2 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -58,6 +58,7 @@ class UnslothTrainer: self.progress_callbacks = [] self.is_training = False self.should_stop = False + self.save_on_stop = True # Model state tracking self.is_vlm = False @@ -756,16 +757,32 @@ class UnslothTrainer: self.trainer.train() # ========== SAVE MODEL ========== - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nTraining completed! Model saved to {output_dir}\n") - - self._update_progress( - is_training=False, - is_completed=True, - #status_message=status_msg - status_message=f"Training completed! Model saved to {output_dir}", - ) + if self.should_stop and self.save_on_stop: + # Stopped by user — save model at current checkpoint + self.trainer.save_model() + self.tokenizer.save_pretrained(output_dir) + print(f"\nTraining stopped. Model saved to {output_dir}\n") + self._update_progress( + is_training=False, + status_message=f"Training stopped. Model saved to {output_dir}", + ) + elif self.should_stop: + # Cancelled by user — don't save + print("\nTraining cancelled.\n") + self._update_progress( + is_training=False, + status_message="Training cancelled.", + ) + else: + # Normal completion + self.trainer.save_model() + self.tokenizer.save_pretrained(output_dir) + print(f"\nTraining completed! Model saved to {output_dir}\n") + self._update_progress( + is_training=False, + is_completed=True, + status_message=f"Training completed! Model saved to {output_dir}", + ) except Exception as e: logger.error(f"Training error: {e}") @@ -774,10 +791,11 @@ class UnslothTrainer: finally: self.is_training = False - def stop_training(self): + def stop_training(self, save: bool = True): """Stop ongoing training""" - print("\nStopping training...") + 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="") diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index fa8b3b5daf..62febadc13 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -103,6 +103,7 @@ class TrainingBackend: try: # Reset stop flag and clear history self.trainer.should_stop = False + self.trainer.save_on_stop = True self.loss_history = [] self.lr_history = [] self.step_history = [] @@ -224,16 +225,19 @@ class TrainingBackend: ) return False - def stop_training(self) -> bool: + def stop_training(self, save: bool = True) -> bool: """ Stop ongoing training. + Args: + save: If True, save the model at the current checkpoint. + Returns: True if training was successfully stopped. """ try: - logger.info("Stopping training...") - self.trainer.stop_training() + logger.info(f"Stopping training (save={save})...") + self.trainer.stop_training(save=save) return True except Exception as e: logger.error(f"Error stopping training: {e}") @@ -293,6 +297,10 @@ 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: + return False + progress = self.trainer.get_training_progress() # Training is active if is_training is True # Also check if we're in loading/preparation phase (status_message indicates activity) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 20a5e1e254..1a904cd371 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -37,6 +37,11 @@ from models import ( TrainingProgress, ) from models.responses import TrainingStopResponse, TrainingMetricsResponse +from pydantic import BaseModel as PydanticBaseModel + + +class TrainingStopRequest(PydanticBaseModel): + save: bool = True router = APIRouter() logger = logging.getLogger(__name__) @@ -251,10 +256,14 @@ async def start_training( @router.post("/stop", response_model=TrainingStopResponse) async def stop_training( + body: TrainingStopRequest = TrainingStopRequest(), current_subject: str = Depends(get_current_subject), ): """ Stop the currently running training job. + + Body: + save (bool): If True (default), save the model at the current checkpoint. """ try: backend = get_training_backend() @@ -266,7 +275,7 @@ async def stop_training( ) # Call backend stop method - backend.stop_training() + backend.stop_training(save=body.save) return TrainingStopResponse( status="stopped", @@ -281,6 +290,29 @@ async def stop_training( ) +@router.post("/reset") +async def reset_training( + current_subject: str = Depends(get_current_subject), +): + """ + Reset training state so the user can return to configuration. + """ + try: + backend = get_training_backend() + backend.trainer.should_stop = False + backend.trainer.training_progress = backend.trainer.training_progress.__class__() + backend.loss_history = [] + backend.lr_history = [] + backend.step_history = [] + return {"status": "ok"} + except Exception as e: + logger.error(f"Error resetting training: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to reset training: {str(e)}", + ) + + @router.get("/status") async def get_training_status( current_subject: str = Depends(get_current_subject), @@ -313,6 +345,9 @@ async def get_training_status( ) or "Ready to train" error_message = getattr(progress, "error", None) if progress else None + # Check if training was stopped by user + trainer_stopped = getattr(backend.trainer, "should_stop", False) + # Derive high-level phase if error_message: phase = "error" @@ -326,6 +361,8 @@ async def get_training_status( phase = "configuring" else: phase = "training" + elif trainer_stopped: + phase = "stopped" elif progress and getattr(progress, "is_completed", False): phase = "completed" elif has_thread: From 55e7bd60c158362ac114a7794d856a636200249a Mon Sep 17 00:00:00 2001 From: Manan17 Date: Sun, 15 Feb 2026 00:22:28 +0000 Subject: [PATCH 2/2] feat: UI for cancel or save and stop training --- .../studio/sections/progress-section.tsx | 53 +++++++++++++++---- .../src/features/studio/studio-page.tsx | 20 +++++++ .../src/features/training/api/train-api.ts | 15 +++++- .../training/hooks/use-training-actions.ts | 16 ++++-- .../hooks/use-training-runtime-lifecycle.ts | 6 ++- .../src/features/training/lib/sync-runtime.ts | 5 ++ .../training/stores/training-runtime-store.ts | 6 ++- .../src/features/training/types/runtime.ts | 1 + 8 files changed, 104 insertions(+), 18 deletions(-) diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index b8656a7867..f1a10c2684 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -1,4 +1,14 @@ import { SectionCard } from "@/components/section-card"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; import { Button } from "@/components/ui/button"; import { Popover, @@ -104,6 +114,7 @@ export function ProgressSection(): ReactElement { ); const { stopTrainingRun } = useTrainingActions(); + const [stopDialogOpen, setStopDialogOpen] = useState(false); const localStartAtRef = useRef(null); const [, setLocalTick] = useState(0); @@ -225,15 +236,39 @@ export function ProgressSection(): ReactElement { - + + + + + Stop Training + + Choose how you want to stop the current training run. + + + + Continue Training + void stopTrainingRun(false)} + > + Cancel Training + + void stopTrainingRun(true)} + > + Stop and Save + + + + } > diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx index d5e20fc194..5811f161db 100644 --- a/studio/frontend/src/features/studio/studio-page.tsx +++ b/studio/frontend/src/features/studio/studio-page.tsx @@ -1,8 +1,12 @@ +import { Button } from "@/components/ui/button"; import { shouldShowTrainingView, + useTrainingActions, useTrainingRuntimeLifecycle, useTrainingRuntimeStore, } from "@/features/training"; +import { ArrowLeft01Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import type { ReactElement } from "react"; import { DatasetSection } from "./sections/dataset-section"; import { ModelSection } from "./sections/model-section"; @@ -14,12 +18,28 @@ export function StudioPage(): ReactElement { useTrainingRuntimeLifecycle(); const showTrainingView = useTrainingRuntimeStore(shouldShowTrainingView); const runtimeMessage = useTrainingRuntimeStore((state) => state.message); + const runtimePhase = useTrainingRuntimeStore((state) => state.phase); const isHydratingRuntime = useTrainingRuntimeStore((state) => state.isHydrating); const hasHydratedRuntime = useTrainingRuntimeStore((state) => state.hasHydrated); + const { dismissTrainingRun } = useTrainingActions(); + + const canGoBack = runtimePhase === "stopped" || runtimePhase === "error"; return (
+ {canGoBack && ( + + )} + {/* Header */}

diff --git a/studio/frontend/src/features/training/api/train-api.ts b/studio/frontend/src/features/training/api/train-api.ts index e3c53b0bc2..d2f589c298 100644 --- a/studio/frontend/src/features/training/api/train-api.ts +++ b/studio/frontend/src/features/training/api/train-api.ts @@ -41,11 +41,22 @@ export async function startTraining( return parseJson(response); } -export async function stopTraining(): Promise { - const response = await authFetch("/api/train/stop", { method: "POST" }); +export async function stopTraining(save = true): Promise { + const response = await authFetch("/api/train/stop", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ save }), + }); return parseJson(response); } +export async function resetTraining(): Promise { + const response = await authFetch("/api/train/reset", { method: "POST" }); + if (!response.ok) { + throw new Error(await readError(response)); + } +} + export async function getTrainingStatus(): Promise { const response = await authFetch("/api/train/status"); return parseJson(response); 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 510c45ed4a..4ace9dd5ed 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -1,7 +1,7 @@ import { useCallback } from "react"; import { useTrainingConfigStore } from "../stores/training-config-store"; import { useTrainingRuntimeStore } from "../stores/training-runtime-store"; -import { startTraining, stopTraining } from "../api/train-api"; +import { startTraining, stopTraining, resetTraining } from "../api/train-api"; import { buildTrainingStartPayload } from "../api/mappers"; import { syncTrainingRuntimeFromBackend } from "../lib/sync-runtime"; import { validateTrainingConfig } from "../lib/validation"; @@ -45,12 +45,12 @@ export function useTrainingActions() { } }, []); - const stopTrainingRun = useCallback(async (): Promise => { + const stopTrainingRun = useCallback(async (save = true): Promise => { const runtimeStore = useTrainingRuntimeStore.getState(); runtimeStore.setStartError(null); try { - await stopTraining(); + await stopTraining(save); await syncTrainingRuntimeFromBackend(); return true; } catch (error) { @@ -61,10 +61,20 @@ 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 + } + }, []); + return { isStarting, startError, startTrainingRun, stopTrainingRun, + dismissTrainingRun, }; } 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 3baff3a62d..8d0eb9c756 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 @@ -48,9 +48,10 @@ export function useTrainingRuntimeLifecycle(): void { }; const pollMetrics = async () => { + const gen = runtimeStore.getState().resetGeneration; try { const metrics = await getTrainingMetrics(); - if (disposed) { + if (disposed || runtimeStore.getState().resetGeneration !== gen) { return; } runtimeStore.getState().applyMetrics(metrics); @@ -62,9 +63,10 @@ export function useTrainingRuntimeLifecycle(): void { }; const pollStatus = async () => { + const gen = runtimeStore.getState().resetGeneration; try { const status = await getTrainingStatus(); - if (disposed) { + if (disposed || runtimeStore.getState().resetGeneration !== gen) { return; } diff --git a/studio/frontend/src/features/training/lib/sync-runtime.ts b/studio/frontend/src/features/training/lib/sync-runtime.ts index b5fbd0bafb..bf255f58c6 100644 --- a/studio/frontend/src/features/training/lib/sync-runtime.ts +++ b/studio/frontend/src/features/training/lib/sync-runtime.ts @@ -6,12 +6,17 @@ import { useTrainingRuntimeStore } from "../stores/training-runtime-store"; import type { TrainingStatusResponse } from "../types/runtime"; export async function syncTrainingRuntimeFromBackend(): Promise { + const gen = useTrainingRuntimeStore.getState().resetGeneration; + const [status, metrics] = await Promise.all([ getTrainingStatus(), getTrainingMetrics(), ]); const runtimeStore = useTrainingRuntimeStore.getState(); + if (runtimeStore.resetGeneration !== gen) { + return status; + } runtimeStore.applyStatus(status); runtimeStore.applyMetrics(metrics); 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 bc40019346..94db0f28f5 100644 --- a/studio/frontend/src/features/training/stores/training-runtime-store.ts +++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts @@ -34,6 +34,7 @@ const initialState: TrainingRuntimeState = { lossHistory: [], lrHistory: [], gradNormHistory: [], + resetGeneration: 0, }; function sortSeries(points: TrainingSeriesPoint[]): TrainingSeriesPoint[] { @@ -95,12 +96,13 @@ export const useTrainingRuntimeStore = create()((set) => ( setLastEventId: (value) => set({ lastEventId: value }), resetRuntime: () => - set({ + set((state) => ({ ...initialState, lossHistory: [], lrHistory: [], gradNormHistory: [], - }), + resetGeneration: state.resetGeneration + 1, + })), setStartQueued: (jobId, message) => set({ diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts index fe2afbd36d..389418680a 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -82,6 +82,7 @@ export interface TrainingRuntimeState { lossHistory: TrainingSeriesPoint[]; lrHistory: TrainingSeriesPoint[]; gradNormHistory: TrainingSeriesPoint[]; + resetGeneration: number; } export interface TrainingRuntimeActions {