Merge pull request #218 from unslothai/fix/stop-startup-modal

Added cancel training button on the overlay

studio/frontend/src/features/studio/sections/progress-section.tsx

@@ -237,6 +237,7 @@ export function ProgressSection(): ReactElement {
                   onClick={() => {
                     setStopRequested(true);
                     setStopDialogOpen(false);
+                    useTrainingRuntimeStore.getState().setStopRequested(true);
                     void stopTrainingRun(false).then((ok) => {
                       if (!ok) setStopRequested(false);
                     });
@@ -248,6 +249,7 @@ export function ProgressSection(): ReactElement {
                   onClick={() => {
                     setStopRequested(true);
                     setStopDialogOpen(false);
+                    useTrainingRuntimeStore.getState().setStopRequested(true);
                     void stopTrainingRun(true).then((ok) => {
                       if (!ok) setStopRequested(false);
                     });

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

@@ -45,11 +45,16 @@ export function StudioPage(): ReactElement {
   const dialogInitial = useDatasetPreviewDialogStore((s) => s.initialData);
   const closeDialog = useDatasetPreviewDialogStore((s) => s.close);
 
+  const stopRequested = useTrainingRuntimeStore((state) => state.stopRequested);
   const canGoBack =
     showTrainingView &&
-    !isTrainingRunning &&
     !isHydratingRuntime &&
-    (runtimePhase === "stopped" || runtimePhase === "error" || runtimePhase === "completed" || runtimePhase === "idle");
+    (stopRequested ||
+      (!isTrainingRunning &&
+        (runtimePhase === "stopped" ||
+          runtimePhase === "error" ||
+          runtimePhase === "completed" ||
+          runtimePhase === "idle")));
   const tourEnabled = hasHydratedRuntime && !isHydratingRuntime;
   const isConfigTour = !showTrainingView;
   const tourSteps = showTrainingView ? studioTrainingTourSteps : studioTourSteps;

studio/frontend/src/features/studio/training-start-overlay.tsx

@@ -1,9 +1,23 @@
+import {
+  AlertDialog,
+  AlertDialogAction,
+  AlertDialogCancel,
+  AlertDialogContent,
+  AlertDialogDescription,
+  AlertDialogFooter,
+  AlertDialogHeader,
+  AlertDialogTitle,
+} from "@/components/ui/alert-dialog";
+import { Button } from "@/components/ui/button";
 import {
   AnimatedSpan,
   Terminal,
   TypingAnimation,
-} from "@/components/ui/terminal"
-import type { ReactElement } from "react"
+} from "@/components/ui/terminal";
+import { useTrainingActions, useTrainingRuntimeStore } from "@/features/training";
+import { Cancel01Icon } from "@hugeicons/core-free-icons";
+import { HugeiconsIcon } from "@hugeicons/react";
+import { useEffect, useState, type ReactElement } from "react";
 
 type TrainingStartOverlayProps = {
   message: string
@@ -14,18 +28,65 @@ export function TrainingStartOverlay({
   message,
   currentStep,
 }: TrainingStartOverlayProps): ReactElement {
+  const { stopTrainingRun } = useTrainingActions();
+  const isStarting = useTrainingRuntimeStore((s) => s.isStarting);
+  const [cancelDialogOpen, setCancelDialogOpen] = useState(false);
+  const [cancelRequested, setCancelRequested] = useState(false);
+
+  useEffect(() => {
+    if (!isStarting) {
+      setCancelRequested(false);
+    }
+  }, [isStarting]);
+
   return (
     <div className="pointer-events-none absolute inset-0 z-30 flex items-center justify-center rounded-2xl bg-background/45 backdrop-blur-[1px]">
-      <div className="flex w-[860px] max-w-[calc(100%-2rem)] flex-col items-center gap-4">
+      <div className="pointer-events-auto relative flex w-[860px] max-w-[calc(100%-2rem)] flex-col items-center gap-4">
         <img
           src="/Sloth emojis/large sloth wave.png"
           alt="Unsloth mascot"
           className="size-24 animate-bounce object-contain"
         />
-        <Terminal
-          className="w-full min-h-[390px] rounded-2xl px-7 py-6 text-left"
-          startOnView={false}
-        >
+        <div className="relative w-full">
+          <AlertDialog open={cancelDialogOpen} onOpenChange={setCancelDialogOpen}>
+            <Button
+              variant="ghost"
+              size="icon"
+              className="absolute right-3 top-3 z-10 size-7 cursor-pointer rounded-full text-muted-foreground/60 hover:bg-destructive/10 hover:text-destructive"
+              onClick={() => setCancelDialogOpen(true)}
+              disabled={cancelRequested}
+            >
+              <HugeiconsIcon icon={Cancel01Icon} className="size-3.5" />
+            </Button>
+            <AlertDialogContent overlayClassName="bg-background/40 supports-backdrop-filter:backdrop-blur-[1px]">
+              <AlertDialogHeader>
+                <AlertDialogTitle>Cancel Training</AlertDialogTitle>
+                <AlertDialogDescription>
+                  Do you want to cancel the current training run?
+                </AlertDialogDescription>
+              </AlertDialogHeader>
+              <AlertDialogFooter>
+                <AlertDialogCancel>Continue Training</AlertDialogCancel>
+                <AlertDialogAction
+                  variant="destructive"
+                  onClick={() => {
+                    setCancelRequested(true);
+                    setCancelDialogOpen(false);
+                    useTrainingRuntimeStore.getState().setStopRequested(true);
+                    void stopTrainingRun(false).then((ok) => {
+                      if (!ok) setCancelRequested(false);
+                    });
+                  }}
+                >
+                  Cancel Training
+                </AlertDialogAction>
+              </AlertDialogFooter>
+            </AlertDialogContent>
+          </AlertDialog>
+          <Terminal
+            className="w-full min-h-[390px] rounded-2xl px-7 py-6 text-left"
+            startOnView={false}
+          >
           <TypingAnimation
             duration={36}
             className="bg-gradient-to-r from-emerald-300 via-lime-300 to-teal-300 bg-clip-text font-semibold text-transparent"
@@ -51,7 +112,8 @@ O^O/ \\_/ \\
           <AnimatedSpan className="mt-2 text-muted-foreground">
             {`> ${message || "starting training..."} | waiting for first step... (${currentStep})`}
           </AnimatedSpan>
-        </Terminal>
+          </Terminal>
+        </div>
       </div>
     </div>
   )

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

@@ -37,6 +37,7 @@ const initialState: TrainingRuntimeState = {
   gradNormHistory: [],
   evalLossHistory: [],
   resetGeneration: 0,
+  stopRequested: false,
 };
 
 function sortSeries(points: TrainingSeriesPoint[]): TrainingSeriesPoint[] {
@@ -110,6 +111,7 @@ function applyMetricHistoryFromStatus(payload: TrainingStatusResponse): {
 export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => ({
   ...initialState,
 
+  setStopRequested: (value) => set({ stopRequested: value }),
   setHydrating: (value) => set({ isHydrating: value }),
   setHasHydrated: (value) => set({ hasHydrated: value }),
   setStarting: (value) => set({ isStarting: value }),
@@ -173,12 +175,15 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
       const detailLoss = payload.details?.loss;
       const detailLr = payload.details?.learning_rate;
       const detailEpoch = payload.details?.epoch;
+      const stopRequested =
+        payload.is_training_running ? state.stopRequested : false;
 
       return {
         ...state,
         jobId: payload.job_id || state.jobId,
         phase: payload.phase,
         isTrainingRunning: payload.is_training_running,
+        stopRequested,
         evalEnabled: payload.eval_enabled ?? state.evalEnabled,
         message: payload.message,
         error: payload.error,

studio/frontend/src/features/training/types/runtime.ts

@@ -95,9 +95,11 @@ export interface TrainingRuntimeState {
   gradNormHistory: TrainingSeriesPoint[];
   evalLossHistory: TrainingSeriesPoint[];
   resetGeneration: number;
+  stopRequested: boolean;
 }
 
 export interface TrainingRuntimeActions {
+  setStopRequested: (value: boolean) => void;
   setHydrating: (value: boolean) => void;
   setHasHydrated: (value: boolean) => void;
   setStarting: (value: boolean) => void;