ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #81 from unslothai/feature/early-stop-or-cancel-training-ui

Browse source 
feat: early stop or cancel training UI
This commit is contained in:
Roland Tannous 2026-02-15 08:28:39 +04:00 • committed by GitHub
commit 1a3da626e2
11 changed files with 183 additions and 34 deletions
Download patch file Download diff file Expand all files Collapse all files

42
studio/backend/core/training/trainer.py
View file

@ -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="")

14
studio/backend/core/training/training.py
View file

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

39
studio/backend/routes/training.py
View file

@ -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:

53
studio/frontend/src/features/studio/sections/progress-section.tsx
View file

@ -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<number | null>(null);
const [, setLocalTick] = useState(0);
@ -225,15 +236,39 @@ export function ProgressSection(): ReactElement {
</div>
</PopoverContent>
</Popover>
<Button
variant="destructive"
size="sm"
className="h-7 cursor-pointer px-3 text-xs"
onClick={() => void stopTrainingRun()}
disabled={!runtime.isTrainingRunning}
>
<HugeiconsIcon icon={StopIcon} className="size-3" /> Stop
</Button>
<AlertDialog open={stopDialogOpen} onOpenChange={setStopDialogOpen}>
<Button
variant="destructive"
size="sm"
className="h-7 cursor-pointer px-3 text-xs"
onClick={() => setStopDialogOpen(true)}
disabled={!runtime.isTrainingRunning}
>
<HugeiconsIcon icon={StopIcon} className="size-3" /> Stop
</Button>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>Stop Training</AlertDialogTitle>
<AlertDialogDescription>
Choose how you want to stop the current training run.
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel>Continue Training</AlertDialogCancel>
<AlertDialogAction
variant="destructive"
onClick={() => void stopTrainingRun(false)}
>
Cancel Training
</AlertDialogAction>
<AlertDialogAction
onClick={() => void stopTrainingRun(true)}
>
Stop and Save
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</div>
}
>

20
studio/frontend/src/features/studio/studio-page.tsx
View file

@ -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 (
<div className="min-h-screen bg-background">
<main className="mx-auto max-w-7xl px-6 py-4">
{canGoBack && (
<Button
variant="ghost"
size="sm"
className="mb-2 cursor-pointer gap-1.5 text-muted-foreground"
onClick={() => void dismissTrainingRun()}
>
<HugeiconsIcon icon={ArrowLeft01Icon} className="size-4" />
Back to configuration
</Button>
)}
{/* Header */}
<div className="mb-8 flex flex-col gap-0.5">
<h1 className="text-2xl font-semibold tracking-tight">

15
studio/frontend/src/features/training/api/train-api.ts
View file

@ -41,11 +41,22 @@ export async function startTraining(
return parseJson<TrainingStartResponse>(response);
}
export async function stopTraining(): Promise<TrainingStopResponse> {
const response = await authFetch("/api/train/stop", { method: "POST" });
export async function stopTraining(save = true): Promise<TrainingStopResponse> {
const response = await authFetch("/api/train/stop", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ save }),
});
return parseJson<TrainingStopResponse>(response);
}
export async function resetTraining(): Promise<void> {
const response = await authFetch("/api/train/reset", { method: "POST" });
if (!response.ok) {
throw new Error(await readError(response));
}
}
export async function getTrainingStatus(): Promise<TrainingStatusResponse> {
const response = await authFetch("/api/train/status");
return parseJson<TrainingStatusResponse>(response);

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

@ -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<boolean> => {
const stopTrainingRun = useCallback(async (save = true): Promise<boolean> => {
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<void> => {
useTrainingRuntimeStore.getState().resetRuntime();
try {
await resetTraining();
} catch {
// Frontend already reset; backend will catch up on next poll
}
}, []);
return {
isStarting,
startError,
startTrainingRun,
stopTrainingRun,
dismissTrainingRun,
};
}

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

@ -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;
}

5
studio/frontend/src/features/training/lib/sync-runtime.ts
View file

@ -6,12 +6,17 @@ import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
import type { TrainingStatusResponse } from "../types/runtime";
export async function syncTrainingRuntimeFromBackend(): Promise<TrainingStatusResponse> {
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);

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

@ -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<TrainingRuntimeStore>()((set) => (
setLastEventId: (value) => set({ lastEventId: value }),
resetRuntime: () =>
set({
set((state) => ({
...initialState,
lossHistory: [],
lrHistory: [],
gradNormHistory: [],
}),
resetGeneration: state.resetGeneration + 1,
})),
setStartQueued: (jobId, message) =>
set({

1
studio/frontend/src/features/training/types/runtime.ts
View file

@ -82,6 +82,7 @@ export interface TrainingRuntimeState {
lossHistory: TrainingSeriesPoint[];
lrHistory: TrainingSeriesPoint[];
gradNormHistory: TrainingSeriesPoint[];
resetGeneration: number;
}
export interface TrainingRuntimeActions {