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

feat: UI for cancel or save and stop training

Browse source
This commit is contained in:
Manan17 2026-02-15 00:22:28 +00:00
commit 52cbf9b699
8 changed files with 104 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

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 {