feat: add training feature with state management, API integration, and runtime synchronization

2026-02-13 12:26:28 +01:00 · 2026-02-13 12:26:28 +01:00 · bc5ce0f83b
commit bc5ce0f83b
parent 9ee81814cd
36 changed files with 1286 additions and 192 deletions
--- a/studio/frontend/package.json
+++ b/studio/frontend/package.json
@ -17,6 +17,8 @@
    "@assistant-ui/react-markdown": "^0.12.1",
    "@assistant-ui/react-streamdown": "^0.1.0",
    "@base-ui/react": "^1.1.0",
+    "@dagrejs/dagre": "^2.0.4",
+    "@dagrejs/graphlib": "^3.0.4",
    "@fontsource-variable/figtree": "^5.2.10",
    "@fontsource-variable/inter": "^5.2.8",
    "@fontsource-variable/space-grotesk": "^5.2.10",
--- a/studio/frontend/src/features/auth/index.ts
+++ b/studio/frontend/src/features/auth/index.ts
@ -1,6 +1,6 @@
 export { LoginPage } from "./login-page";
 export { SignupPage } from "./signup-page";
-export { refreshSession } from "./api";
+export { authFetch, refreshSession } from "./api";
 export {
  getPostAuthRoute,
  hasAuthToken,
--- a/studio/frontend/src/features/export/export-page.tsx
+++ b/studio/frontend/src/features/export/export-page.tsx
@ -13,7 +13,8 @@ import {
  TooltipContent,
  TooltipTrigger,
 } from "@/components/ui/tooltip";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingRuntimeStore } from "@/features/training";
+import { useTrainingConfigStore } from "@/features/training";
 import { isAdapterMethod } from "@/types/training";
 import { InformationCircleIcon, PackageIcon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
@ -36,36 +37,31 @@ export function ExportPage() {
    trainingMethod,
    selectedModel,
    saveSteps,
-    trainingMetrics,
    epochs,
    loraRank,
    hfToken,
    setHfToken,
-  } = useWizardStore(
+  } = useTrainingConfigStore(
    useShallow((s) => ({
      trainingMethod: s.trainingMethod,
      selectedModel: s.selectedModel,
      saveSteps: s.saveSteps,
-      trainingMetrics: s.trainingMetrics,
      epochs: s.epochs,
      loraRank: s.loraRank,
      hfToken: s.hfToken,
      setHfToken: s.setHfToken,
    })),
  );
+  const totalSteps = useTrainingRuntimeStore((state) => state.totalSteps);
  const isAdapter = isAdapterMethod(trainingMethod);

  const checkpoints = useMemo(() => {
    if (isAdapter) {
      const interval = saveSteps > 0 ? saveSteps : 100;
-      const total = trainingMetrics?.totalSteps ?? 500;
+      const total = totalSteps > 0 ? totalSteps : 500;
      const entries: { value: string; label: string; detail: string }[] = [];
      for (let step = interval; step <= total; step += interval) {
-        const loss = (
-          1.5 -
-          (step / total) * 0.7 +
-          Math.random() * 0.05
-        ).toFixed(2);
+        const loss = (1.5 - (step / total) * 0.7).toFixed(2);
        entries.push({
          value: `checkpoint-${step}`,
          label: `checkpoint-${step}`,
@ -81,7 +77,7 @@ export function ExportPage() {
        detail: "Full fine-tuned weights",
      },
    ];
-  }, [isAdapter, saveSteps, trainingMetrics?.totalSteps]);
+  }, [isAdapter, saveSteps, totalSteps]);

  const [checkpoint, setCheckpoint] = useState<string | null>(null);
  const [exportMethod, setExportMethod] = useState<ExportMethod | null>(null);
--- a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx
+++ b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx
@ -38,7 +38,7 @@ import {
  useInfiniteScroll,
 } from "@/hooks";
 import { cn, formatCompact } from "@/lib/utils";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import type { DatasetFormat } from "@/types/training";
 import {
  InformationCircleIcon,
@ -69,7 +69,7 @@ export function DatasetStep() {
    setDataset,
    uploadedFile,
    setUploadedFile,
-  } = useWizardStore(
+  } = useTrainingConfigStore(
    useShallow((s) => ({
      hfToken: s.hfToken,
      setHfToken: s.setHfToken,
--- a/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx
+++ b/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx
@ -20,7 +20,7 @@ import {
  TooltipTrigger,
 } from "@/components/ui/tooltip";
 import { CONTEXT_LENGTHS } from "@/config/training";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import { InformationCircleIcon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { useShallow } from "zustand/react/shallow";
@ -40,7 +40,7 @@ export function HyperparametersStep() {
    setLoraAlpha,
    loraDropout,
    setLoraDropout,
-  } = useWizardStore(
+  } = useTrainingConfigStore(
    useShallow((s) => ({
      trainingMethod: s.trainingMethod,
      epochs: s.epochs,
--- a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx
+++ b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx
@ -37,7 +37,7 @@ import {
  useInfiniteScroll,
 } from "@/hooks";
 import { formatCompact } from "@/lib/utils";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import type { TrainingMethod } from "@/types/training";
 import {
  InformationCircleIcon,
@ -57,7 +57,7 @@ export function ModelSelectionStep() {
    setTrainingMethod,
    hfToken,
    setHfToken,
-  } = useWizardStore(
+  } = useTrainingConfigStore(
    useShallow((s) => ({
      modelType: s.modelType,
      selectedModel: s.selectedModel,
--- a/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx
+++ b/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx
@ -8,7 +8,7 @@ import {
 } from "@/components/ui/tooltip";
 import { MODEL_TYPES } from "@/config/training";
 import { cn } from "@/lib/utils";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import type { ModelType } from "@/types/training";
 import {
  Database02Icon,
@ -38,7 +38,7 @@ const TYPE_TOOLTIPS: Record<ModelType, string> = {
 const COMING_SOON: ModelType[] = ["tts", "embeddings"];

 export function ModelTypeStep(): ReactElement {
-  const { modelType, setModelType } = useWizardStore(
+  const { modelType, setModelType } = useTrainingConfigStore(
    useShallow((s) => ({
      modelType: s.modelType,
      setModelType: s.setModelType,
--- a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx
+++ b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx
@ -1,7 +1,7 @@
 import { Badge } from "@/components/ui/badge";
 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
 import { Separator } from "@/components/ui/separator";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import { isAdapterMethod } from "@/types/training";
 import { GpuIcon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
@ -30,7 +30,7 @@ export function SummaryStep() {
    loraRank,
    loraAlpha,
    loraDropout,
-  } = useWizardStore(
+  } = useTrainingConfigStore(
    useShallow(
      ({
        modelType,
--- a/studio/frontend/src/features/onboarding/components/wizard-content.tsx
+++ b/studio/frontend/src/features/onboarding/components/wizard-content.tsx
@ -1,5 +1,5 @@
 import { STEPS } from "@/config/training";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import type { StepNumber } from "@/types/training";
 import { DatasetStep } from "./steps/dataset-step";
 import { HyperparametersStep } from "./steps/hyperparameters-step";
@ -24,7 +24,7 @@ const STEP_MASCOTS: Record<StepNumber, string> = {
 };

 export function WizardContent() {
-  const currentStep = useWizardStore((s) => s.currentStep);
+  const currentStep = useTrainingConfigStore((s) => s.currentStep);
  const stepConfig = STEPS[currentStep - 1];
  const StepComponent = STEP_COMPONENTS[currentStep];
  const mascotSrc = STEP_MASCOTS[currentStep];
--- a/studio/frontend/src/features/onboarding/components/wizard-footer.tsx
+++ b/studio/frontend/src/features/onboarding/components/wizard-footer.tsx
@ -1,14 +1,14 @@
 import { Button } from "@/components/ui/button";
 import { STEPS } from "@/config/training";
 import { markOnboardingDone } from "@/features/auth";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import { ArrowLeft02Icon, ArrowRight02Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { useNavigate } from "@tanstack/react-router";
 import { useShallow } from "zustand/react/shallow";

 export function WizardFooter() {
-  const { currentStep, prevStep, nextStep, canProceed } = useWizardStore(
+  const { currentStep, prevStep, nextStep, canProceed } = useTrainingConfigStore(
    useShallow((s) => ({
      currentStep: s.currentStep,
      prevStep: s.prevStep,
--- a/studio/frontend/src/features/onboarding/components/wizard-layout.tsx
+++ b/studio/frontend/src/features/onboarding/components/wizard-layout.tsx
@ -6,7 +6,7 @@ import { Suspense, lazy, useEffect, useRef, useState } from "react";
 import type { ConfettiRef } from "@/components/ui/confetti";
 import { STEPS } from "@/config/training";
 import { isOnboardingDone, markOnboardingDone } from "@/features/auth";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import { SplashScreen } from "./splash-screen";
 import { WizardContent } from "./wizard-content";
 import { WizardFooter } from "./wizard-footer";
@ -19,7 +19,7 @@ const Confetti = lazy(() =>
 export function WizardLayout() {
  const navigate = useNavigate();
  const [showSplash, setShowSplash] = useState(true);
-  const currentStep = useWizardStore((s) => s.currentStep);
+  const currentStep = useTrainingConfigStore((s) => s.currentStep);
  const confettiRef = useRef<ConfettiRef>(null);
  const hasFiredRef = useRef(false);
  const isFinalStep = currentStep === STEPS.length;
--- a/studio/frontend/src/features/onboarding/components/wizard-sidebar.tsx
+++ b/studio/frontend/src/features/onboarding/components/wizard-sidebar.tsx
@ -1,10 +1,10 @@
 import { Progress } from "@/components/ui/progress";
 import { STEPS } from "@/config/training";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import { WizardStepItem } from "./wizard-step-item";

 export function WizardSidebar() {
-  const currentStep = useWizardStore((s) => s.currentStep);
+  const currentStep = useTrainingConfigStore((s) => s.currentStep);
  const progress = ((currentStep - 1) / (STEPS.length - 1)) * 100;

  return (
--- a/studio/frontend/src/features/onboarding/components/wizard-step-item.tsx
+++ b/studio/frontend/src/features/onboarding/components/wizard-step-item.tsx
@ -1,5 +1,5 @@
 import { cn } from "@/lib/utils";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import type { StepConfig, StepNumber } from "@/types/training";
 import { useShallow } from "zustand/react/shallow";

@ -8,7 +8,7 @@ interface WizardStepItemProps {
 }

 export function WizardStepItem({ step }: WizardStepItemProps) {
-  const { currentStep, setStep } = useWizardStore(
+  const { currentStep, setStep } = useTrainingConfigStore(
    useShallow((s) => ({ currentStep: s.currentStep, setStep: s.setStep })),
  );
  const isActive = currentStep === step.number;
--- a/studio/frontend/src/features/studio/sections/charts-content.tsx
+++ b/studio/frontend/src/features/studio/sections/charts-content.tsx
@ -185,14 +185,15 @@ export function ChartsContent({
      return [0, 1];
    }

+    const minStep = allSteps[0] ?? 0;
    const endStep = allSteps[allSteps.length - 1] ?? 1;
    const startIndex = Math.max(0, allSteps.length - DEFAULT_VISIBLE_POINTS);
-    const startStep = allSteps[startIndex] ?? 0;
+    const startStep = allSteps[startIndex] ?? minStep;
    if (startStep === endStep) {
-      return [Math.max(0, startStep - 1), startStep + 4];
+      return [startStep, startStep + 4];
    }
    if (endStep - startStep < 6) {
-      return [Math.max(0, endStep - 6), endStep];
+      return [Math.max(minStep, endStep - 6), endStep];
    }
    return [startStep, endStep];
  }, [reducedGradNormData, reducedLossData, reducedLrData]);
--- a/studio/frontend/src/features/studio/sections/charts-section.tsx
+++ b/studio/frontend/src/features/studio/sections/charts-section.tsx
@ -1,5 +1,5 @@
-import { useWizardStore } from "@/stores/training";
-import { type ReactElement, Suspense, lazy } from "react";
+import { useTrainingRuntimeStore } from "@/features/training";
+import { type ReactElement, Suspense, lazy, useMemo } from "react";

 const ChartsContent = lazy(() =>
  import("./charts-content").then((module) => ({
@ -14,9 +14,39 @@ const SKELETON_KEYS = [
 ];

 export function ChartsSection(): ReactElement | null {
-  const metrics = useWizardStore((s) => s.trainingMetrics);
+  const currentStep = useTrainingRuntimeStore((state) => state.currentStep);
+  const totalSteps = useTrainingRuntimeStore((state) => state.totalSteps);
+  const lossHistoryRaw = useTrainingRuntimeStore((state) => state.lossHistory);
+  const lrHistoryRaw = useTrainingRuntimeStore((state) => state.lrHistory);
+  const gradNormHistoryRaw = useTrainingRuntimeStore(
+    (state) => state.gradNormHistory,
+  );

-  if (!metrics) {
+  const series = useMemo(
+    () => ({
+      currentStep,
+      totalSteps,
+      lossHistory: lossHistoryRaw.map((point) => ({
+        step: point.step,
+        loss: point.value,
+      })),
+      lrHistory: lrHistoryRaw.map((point) => ({
+        step: point.step,
+        lr: point.value,
+      })),
+      gradNormHistory: gradNormHistoryRaw.map((point) => ({
+        step: point.step,
+        gradNorm: point.value,
+      })),
+    }),
+    [currentStep, gradNormHistoryRaw, lossHistoryRaw, lrHistoryRaw, totalSteps],
+  );
+
+  if (
+    series.lossHistory.length === 0 &&
+    series.lrHistory.length === 0 &&
+    series.gradNormHistory.length === 0
+  ) {
    return null;
  }

@ -33,7 +63,7 @@ export function ChartsSection(): ReactElement | null {
        </div>
      }
    >
-      <ChartsContent metrics={metrics} />
+      <ChartsContent metrics={series} />
    </Suspense>
  );
 }
--- a/studio/frontend/src/features/studio/sections/dataset-section.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx
@ -28,7 +28,7 @@ import {
  useInfiniteScroll,
 } from "@/hooks";
 import { formatCompact } from "@/lib/utils";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import {
  CloudUploadIcon,
  Database02Icon,
@ -44,7 +44,7 @@ import { DatasetPreviewDialog } from "./dataset-preview-dialog";

 export function DatasetSection() {
  const { dataset, setDataset, datasetFormat, setDatasetFormat, hfToken } =
-    useWizardStore(
+    useTrainingConfigStore(
      useShallow(
        ({
          dataset,
--- a/studio/frontend/src/features/studio/sections/model-section.tsx
+++ b/studio/frontend/src/features/studio/sections/model-section.tsx
@ -32,7 +32,7 @@ import {
  useInfiniteScroll,
 } from "@/hooks";
 import { formatCompact } from "@/lib/utils";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import type { TrainingMethod } from "@/types/training";
 import {
  ChipIcon,
@ -65,7 +65,7 @@ export function ModelSection() {
    setTrainingMethod,
    hfToken,
    setHfToken,
-  } = useWizardStore(
+  } = useTrainingConfigStore(
    useShallow(
      ({
        modelType,
--- a/studio/frontend/src/features/studio/sections/params-section.tsx
+++ b/studio/frontend/src/features/studio/sections/params-section.tsx
@ -21,7 +21,7 @@ import {
  TooltipTrigger,
 } from "@/components/ui/tooltip";
 import { CONTEXT_LENGTHS, TARGET_MODULES } from "@/config/training";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingConfigStore } from "@/features/training";
 import type { GradientCheckpointing } from "@/types/training";
 import {
  ArrowDown01Icon,
@ -107,7 +107,7 @@ function SliderRow({
 }

 export function ParamsSection(): ReactElement {
-  const store = useWizardStore();
+  const store = useTrainingConfigStore();
  const isLora = store.trainingMethod !== "full";
  const isVision = store.modelType === "vision";
  const [loraOpen, setLoraOpen] = useState(false);
--- a/studio/frontend/src/features/studio/sections/progress-section.tsx
+++ b/studio/frontend/src/features/studio/sections/progress-section.tsx
@ -21,7 +21,7 @@ import {
  ZapIcon,
 } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
-import type { ReactElement, ReactNode } from "react";
+import { useEffect, useRef, useState, type ReactElement, type ReactNode } from "react";
 import { useShallow } from "zustand/react/shallow";

 const phaseLabel: Record<TrainingPhase, string> = {
@ -104,6 +104,8 @@ export function ProgressSection(): ReactElement {
  );

  const { stopTrainingRun } = useTrainingActions();
+  const localStartAtRef = useRef<number | null>(null);
+  const [, setLocalTick] = useState(0);

  const pct =
    runtime.totalSteps > 0
@ -116,7 +118,31 @@ export function ProgressSection(): ReactElement {
        )
      : Math.round(runtime.progressPercent);

-  const elapsed = runtime.elapsedSeconds;
+  useEffect(() => {
+    if (runtime.elapsedSeconds != null && runtime.elapsedSeconds >= 0) {
+      localStartAtRef.current = Date.now() - runtime.elapsedSeconds * 1000;
+      return;
+    }
+    if (runtime.currentStep > 0 && localStartAtRef.current == null) {
+      localStartAtRef.current = Date.now();
+    }
+  }, [runtime.currentStep, runtime.elapsedSeconds]);
+
+  useEffect(() => {
+    if (!runtime.isTrainingRunning) {
+      return;
+    }
+    const timer = window.setInterval(() => {
+      setLocalTick((prev) => prev + 1);
+    }, 1000);
+    return () => window.clearInterval(timer);
+  }, [runtime.isTrainingRunning]);
+
+  const elapsed =
+    runtime.elapsedSeconds ??
+    (localStartAtRef.current == null
+      ? null
+      : Math.max(0, Math.floor((Date.now() - localStartAtRef.current) / 1000)));
  const derivedEta =
    elapsed != null && pct > 0
      ? Math.round((elapsed * (100 - pct)) / Math.max(pct, 1))
@ -284,9 +310,9 @@ export function ProgressSection(): ReactElement {
                ? "-- steps/s"
                : `${stepsPerSecond.toFixed(2)} steps/s`}
            </span>
-            <span>
-              Tokens: {runtime.currentNumTokens == null ? "--" : runtime.currentNumTokens}
-            </span>
+            {runtime.currentNumTokens != null && (
+              <span>Tokens: {runtime.currentNumTokens}</span>
+            )}
          </div>
        </div>

--- a/studio/frontend/src/features/studio/sections/training-section.tsx
+++ b/studio/frontend/src/features/studio/sections/training-section.tsx
@ -9,7 +9,7 @@ import {
  CollapsibleTrigger,
 } from "@/components/ui/collapsible";
 import { Input } from "@/components/ui/input";
-import { useWizardStore } from "@/stores/training";
+import { useTrainingActions, useTrainingConfigStore } from "@/features/training";
 import {
  Archive04Icon,
  ArrowDown01Icon,
@ -35,7 +35,8 @@ const placeholderData = [
 ];

 export function TrainingSection() {
-  const store = useWizardStore();
+  const store = useTrainingConfigStore();
+  const { isStarting, startError, startTrainingRun } = useTrainingActions();
  const [logOpen, setLogOpen] = useState(false);

  return (
@ -94,11 +95,15 @@ export function TrainingSection() {
        {/* Start/Stop */}
        <Button
          className="w-full cursor-pointer bg-gradient-to-r from-emerald-500 to-teal-500 text-white hover:from-emerald-600 hover:to-teal-600"
-          onClick={() => store.setIsTraining(true)}
+          onClick={() => void startTrainingRun()}
+          disabled={isStarting}
        >
-          <HugeiconsIcon icon={Rocket01Icon} className="size-4" /> Start
-          Training
+          <HugeiconsIcon icon={Rocket01Icon} className="size-4" />
+          {isStarting ? "Starting..." : "Start Training"}
        </Button>
+        {startError && (
+          <p className="text-xs text-red-500 leading-relaxed">{startError}</p>
+        )}

        {/* Save / Clear */}
        <div className="grid grid-cols-2 gap-2">
--- a/studio/frontend/src/features/studio/studio-page.tsx
+++ b/studio/frontend/src/features/studio/studio-page.tsx
@ -1,4 +1,8 @@
-import { useWizardStore } from "@/stores/training";
+import {
+  shouldShowTrainingView,
+  useTrainingRuntimeLifecycle,
+  useTrainingRuntimeStore,
+} from "@/features/training";
 import type { ReactElement } from "react";
 import { DatasetSection } from "./sections/dataset-section";
 import { ModelSection } from "./sections/model-section";
@ -7,7 +11,11 @@ import { TrainingSection } from "./sections/training-section";
 import { TrainingView } from "./training-view";

 export function StudioPage(): ReactElement {
-  const isTraining = useWizardStore((s) => s.isTraining);
+  useTrainingRuntimeLifecycle();
+  const showTrainingView = useTrainingRuntimeStore(shouldShowTrainingView);
+  const runtimeMessage = useTrainingRuntimeStore((state) => state.message);
+  const isHydratingRuntime = useTrainingRuntimeStore((state) => state.isHydrating);
+  const hasHydratedRuntime = useTrainingRuntimeStore((state) => state.hasHydrated);

  return (
    <div className="min-h-screen bg-background">
@ -18,13 +26,17 @@ export function StudioPage(): ReactElement {
            Fine-tuning Studio
          </h1>
          <p className="text-sm text-muted-foreground">
-            {isTraining
-              ? "Training in progress"
+            {showTrainingView
+              ? runtimeMessage || "Training in progress"
              : "Configure and start training"}
          </p>
        </div>

-        {isTraining ? (
+        {!hasHydratedRuntime && isHydratingRuntime ? (
+          <div className="rounded-xl border bg-card p-8 text-sm text-muted-foreground">
+            Loading training runtime...
+          </div>
+        ) : showTrainingView ? (
          <TrainingView />
        ) : (
          <div className="grid grid-cols-12 items-start gap-6">
--- a/studio/frontend/src/features/training/api/mappers.ts
+++ b/studio/frontend/src/features/training/api/mappers.ts
@ -0,0 +1,63 @@
+import type { TrainingConfigState } from "../types/config";
+import type { TrainingStartRequest } from "../types/api";
+
+const BACKEND_LORA_TYPE = "LoRA/QLoRA";
+const BACKEND_FULL_TYPE = "Full Finetuning";
+
+export function toBackendTrainingType(trainingMethod: string): string {
+  return trainingMethod === "full" ? BACKEND_FULL_TYPE : BACKEND_LORA_TYPE;
+}
+
+export function buildTrainingStartPayload(
+  config: TrainingConfigState,
+): TrainingStartRequest {
+  const adapterMethod = config.trainingMethod !== "full";
+  const isQlorMethod = config.trainingMethod === "qlora";
+  const hfDataset = config.datasetSource === "huggingface" ? config.dataset : null;
+
+  return {
+    model_name: config.selectedModel ?? "",
+    training_type: toBackendTrainingType(config.trainingMethod),
+    hf_token: config.hfToken.trim() || null,
+    load_in_4bit: adapterMethod ? isQlorMethod : false,
+    max_seq_length: config.contextLength,
+    hf_dataset: hfDataset,
+    local_datasets: [],
+    format_type: config.datasetFormat,
+    num_epochs: config.epochs,
+    learning_rate: String(config.learningRate),
+    batch_size: config.batchSize,
+    gradient_accumulation_steps: config.gradientAccumulation,
+    warmup_steps: config.warmupSteps,
+    warmup_ratio: null,
+    max_steps: config.maxSteps,
+    save_steps: config.saveSteps,
+    weight_decay: config.weightDecay,
+    random_seed: config.randomSeed,
+    packing: config.packing,
+    optim: "adamw_8bit",
+    lr_scheduler_type: "linear",
+    use_lora: adapterMethod,
+    lora_r: config.loraRank,
+    lora_alpha: config.loraAlpha,
+    lora_dropout: config.loraDropout,
+    target_modules: adapterMethod ? config.targetModules : [],
+    gradient_checkpointing: config.gradientCheckpointing,
+    use_rslora: config.loraVariant === "rslora",
+    use_loftq: config.loraVariant === "loftq",
+    train_on_completions: config.trainOnCompletions,
+    finetune_vision_layers: config.finetuneVisionLayers,
+    finetune_language_layers: config.finetuneLanguageLayers,
+    finetune_attention_modules: config.finetuneAttentionModules,
+    finetune_mlp_modules: config.finetuneMLPModules,
+    enable_wandb: config.enableWandb,
+    wandb_token: config.enableWandb ? config.wandbToken.trim() || null : null,
+    wandb_project: config.enableWandb
+      ? config.wandbProject.trim() || null
+      : null,
+    enable_tensorboard: config.enableTensorboard,
+    tensorboard_dir: config.enableTensorboard
+      ? config.tensorboardDir.trim() || null
+      : null,
+  };
+}
--- a/studio/frontend/src/features/training/api/train-api.ts
+++ b/studio/frontend/src/features/training/api/train-api.ts
@ -0,0 +1,173 @@
+import { authFetch } from "@/features/auth";
+import type {
+  TrainingStartRequest,
+  TrainingStartResponse,
+  TrainingStopResponse,
+} from "../types/api";
+import type {
+  TrainingMetricsResponse,
+  TrainingProgressPayload,
+  TrainingStatusResponse,
+} from "../types/runtime";
+
+function isAbortError(error: unknown): boolean {
+  return error instanceof DOMException && error.name === "AbortError";
+}
+
+async function readError(response: Response): Promise<string> {
+  try {
+    const payload = (await response.json()) as { detail?: string; message?: string };
+    return payload.detail || payload.message || `Request failed (${response.status})`;
+  } catch {
+    return `Request failed (${response.status})`;
+  }
+}
+
+async function parseJson<T>(response: Response): Promise<T> {
+  if (!response.ok) {
+    throw new Error(await readError(response));
+  }
+  return (await response.json()) as T;
+}
+
+export async function startTraining(
+  payload: TrainingStartRequest,
+): Promise<TrainingStartResponse> {
+  const response = await authFetch("/api/train/start", {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify(payload),
+  });
+  return parseJson<TrainingStartResponse>(response);
+}
+
+export async function stopTraining(): Promise<TrainingStopResponse> {
+  const response = await authFetch("/api/train/stop", { method: "POST" });
+  return parseJson<TrainingStopResponse>(response);
+}
+
+export async function getTrainingStatus(): Promise<TrainingStatusResponse> {
+  const response = await authFetch("/api/train/status");
+  return parseJson<TrainingStatusResponse>(response);
+}
+
+export async function getTrainingMetrics(): Promise<TrainingMetricsResponse> {
+  const response = await authFetch("/api/train/metrics");
+  return parseJson<TrainingMetricsResponse>(response);
+}
+
+type ProgressEventName = "progress" | "heartbeat" | "complete" | "error";
+
+interface ParsedSseEvent {
+  event: ProgressEventName;
+  payload: TrainingProgressPayload;
+  id: number | null;
+}
+
+function parseSseEvent(rawEvent: string): ParsedSseEvent | null {
+  const lines = rawEvent.split(/\r?\n/);
+  let eventName: ProgressEventName = "progress";
+  let id: number | null = null;
+  const dataLines: string[] = [];
+
+  for (const line of lines) {
+    if (!line) {
+      continue;
+    }
+    if (line.startsWith("event:")) {
+      const value = line.slice(6).trim();
+      if (
+        value === "progress" ||
+        value === "heartbeat" ||
+        value === "complete" ||
+        value === "error"
+      ) {
+        eventName = value;
+      }
+      continue;
+    }
+    if (line.startsWith("id:")) {
+      const value = Number(line.slice(3).trim());
+      id = Number.isFinite(value) ? value : null;
+      continue;
+    }
+    if (line.startsWith("data:")) {
+      dataLines.push(line.slice(5).trimStart());
+    }
+  }
+
+  if (dataLines.length === 0) {
+    return null;
+  }
+
+  const parsed = JSON.parse(dataLines.join("\n")) as TrainingProgressPayload;
+  return { event: eventName, payload: parsed, id };
+}
+
+export async function streamTrainingProgress(options: {
+  signal: AbortSignal;
+  lastEventId?: number | null;
+  onOpen?: () => void;
+  onEvent: (event: ParsedSseEvent) => void;
+}): Promise<void> {
+  const headers = new Headers();
+  if (typeof options.lastEventId === "number") {
+    headers.set("Last-Event-ID", String(options.lastEventId));
+  }
+
+  const response = await authFetch("/api/train/progress", {
+    method: "GET",
+    headers,
+    signal: options.signal,
+  });
+
+  if (!response.ok) {
+    throw new Error(await readError(response));
+  }
+
+  if (!response.body) {
+    throw new Error("Progress stream unavailable");
+  }
+
+  options.onOpen?.();
+
+  const reader = response.body.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+
+  while (true) {
+    const { value, done } = await reader.read();
+    if (done) {
+      break;
+    }
+
+    buffer += decoder.decode(value, { stream: true });
+
+    let separatorIndex = buffer.search(/\r?\n\r?\n/);
+    while (separatorIndex >= 0) {
+      const rawEvent = buffer.slice(0, separatorIndex);
+      const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2;
+      buffer = buffer.slice(separatorIndex + separatorLength);
+
+      if (rawEvent.startsWith("retry:")) {
+        separatorIndex = buffer.search(/\r?\n\r?\n/);
+        continue;
+      }
+
+      try {
+        const event = parseSseEvent(rawEvent);
+        if (event) {
+          options.onEvent(event);
+        }
+      } catch (error) {
+        if (!isAbortError(error)) {
+          throw error;
+        }
+      }
+
+      separatorIndex = buffer.search(/\r?\n\r?\n/);
+    }
+  }
+}
+
+export { isAbortError };
--- a/studio/frontend/src/features/training/hooks/use-training-actions.ts
+++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts
@ -0,0 +1,70 @@
+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 { buildTrainingStartPayload } from "../api/mappers";
+import { syncTrainingRuntimeFromBackend } from "../lib/sync-runtime";
+import { validateTrainingConfig } from "../lib/validation";
+
+export function useTrainingActions() {
+  const isStarting = useTrainingRuntimeStore((state) => state.isStarting);
+  const startError = useTrainingRuntimeStore((state) => state.startError);
+
+  const startTrainingRun = useCallback(async (): Promise<boolean> => {
+    const config = useTrainingConfigStore.getState();
+    const runtimeStore = useTrainingRuntimeStore.getState();
+
+    runtimeStore.setStartError(null);
+    const validation = validateTrainingConfig(config);
+    if (!validation.ok) {
+      runtimeStore.setStartError(validation.message);
+      return false;
+    }
+
+    runtimeStore.setStarting(true);
+
+    try {
+      const payload = buildTrainingStartPayload(config);
+      const response = await startTraining(payload);
+
+      if (response.status === "error") {
+        runtimeStore.setStartError(response.error || response.message);
+        runtimeStore.setStarting(false);
+        return false;
+      }
+
+      runtimeStore.setStartQueued(response.job_id, response.message);
+      await syncTrainingRuntimeFromBackend();
+      return true;
+    } catch (error) {
+      const message =
+        error instanceof Error ? error.message : "Failed to start training";
+      runtimeStore.setStartError(message);
+      runtimeStore.setStarting(false);
+      return false;
+    }
+  }, []);
+
+  const stopTrainingRun = useCallback(async (): Promise<boolean> => {
+    const runtimeStore = useTrainingRuntimeStore.getState();
+    runtimeStore.setStartError(null);
+
+    try {
+      await stopTraining();
+      await syncTrainingRuntimeFromBackend();
+      return true;
+    } catch (error) {
+      const message =
+        error instanceof Error ? error.message : "Failed to stop training";
+      runtimeStore.setRuntimeError(message);
+      return false;
+    }
+  }, []);
+
+  return {
+    isStarting,
+    startError,
+    startTrainingRun,
+    stopTrainingRun,
+  };
+}
--- a/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts
+++ b/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts
@ -0,0 +1,183 @@
+import { useEffect } from "react";
+import {
+  getTrainingMetrics,
+  getTrainingStatus,
+  isAbortError,
+  streamTrainingProgress,
+} from "../api/train-api";
+import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
+import type { TrainingRuntimeStore } from "../types/runtime";
+
+const STATUS_POLL_INTERVAL_MS = 3000;
+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"
+  );
+}
+
+export function useTrainingRuntimeLifecycle(): void {
+  useEffect(() => {
+    let disposed = false;
+    let openingStream = false;
+    let streamController: AbortController | null = null;
+    let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
+
+    const runtimeStore = useTrainingRuntimeStore;
+
+    const clearReconnect = () => {
+      if (reconnectTimer) {
+        clearTimeout(reconnectTimer);
+        reconnectTimer = null;
+      }
+    };
+
+    const stopStream = () => {
+      clearReconnect();
+      if (streamController) {
+        streamController.abort();
+        streamController = null;
+      }
+      runtimeStore.getState().setSseConnected(false);
+    };
+
+    const pollMetrics = async () => {
+      try {
+        const metrics = await getTrainingMetrics();
+        if (disposed) {
+          return;
+        }
+        runtimeStore.getState().applyMetrics(metrics);
+      } catch (error) {
+        if (!isAbortError(error) && !disposed) {
+          runtimeStore.getState().setSseConnected(false);
+        }
+      }
+    };
+
+    const pollStatus = async () => {
+      try {
+        const status = await getTrainingStatus();
+        if (disposed) {
+          return;
+        }
+
+        runtimeStore.getState().applyStatus(status);
+
+        const nextState = runtimeStore.getState();
+        if (shouldUseLiveSync(nextState)) {
+          void ensureStream();
+        } else {
+          stopStream();
+        }
+      } catch (error) {
+        if (!isAbortError(error) && !disposed) {
+          runtimeStore.getState().setSseConnected(false);
+        }
+      }
+    };
+
+    const ensureStream = async () => {
+      const state = runtimeStore.getState();
+      if (
+        disposed ||
+        openingStream ||
+        streamController ||
+        !shouldUseLiveSync(state)
+      ) {
+        return;
+      }
+
+      clearReconnect();
+      openingStream = true;
+      const controller = new AbortController();
+      streamController = controller;
+
+      try {
+        await streamTrainingProgress({
+          signal: controller.signal,
+          lastEventId: state.lastEventId,
+          onOpen: () => {
+            runtimeStore.getState().setSseConnected(true);
+          },
+          onEvent: (event) => {
+            const liveStore = runtimeStore.getState();
+            if (typeof event.id === "number") {
+              liveStore.setLastEventId(event.id);
+            }
+
+            liveStore.applyProgress(event.payload, event.id ?? undefined);
+
+            if (event.event === "complete") {
+              void pollStatus();
+              void pollMetrics();
+              stopStream();
+            }
+
+            if (event.event === "error") {
+              liveStore.setRuntimeError("Training stream error");
+              stopStream();
+            }
+          },
+        });
+      } catch (error) {
+        if (!disposed && !controller.signal.aborted && !isAbortError(error)) {
+          runtimeStore.getState().setSseConnected(false);
+        }
+      } finally {
+        openingStream = false;
+        if (streamController === controller) {
+          streamController = null;
+        }
+        runtimeStore.getState().setSseConnected(false);
+
+        if (!disposed && !controller.signal.aborted) {
+          const liveState = runtimeStore.getState();
+          if (shouldUseLiveSync(liveState)) {
+            reconnectTimer = setTimeout(() => {
+              void ensureStream();
+            }, STREAM_RECONNECT_DELAY_MS);
+          }
+        }
+      }
+    };
+
+    const hydrate = async () => {
+      runtimeStore.getState().setHydrating(true);
+      try {
+        await Promise.all([pollStatus(), pollMetrics()]);
+      } finally {
+        if (!disposed) {
+          runtimeStore.getState().setHydrating(false);
+          runtimeStore.getState().setHasHydrated(true);
+        }
+      }
+    };
+
+    void hydrate();
+
+    const statusTimer = setInterval(() => {
+      void pollStatus();
+    }, STATUS_POLL_INTERVAL_MS);
+
+    const metricsTimer = setInterval(() => {
+      const state = runtimeStore.getState();
+      if (shouldUseLiveSync(state) || state.currentStep > 0) {
+        void pollMetrics();
+      }
+    }, METRICS_POLL_INTERVAL_MS);
+
+    return () => {
+      disposed = true;
+      clearInterval(statusTimer);
+      clearInterval(metricsTimer);
+      stopStream();
+    };
+  }, []);
+}
--- a/studio/frontend/src/features/training/index.ts
+++ b/studio/frontend/src/features/training/index.ts
@ -0,0 +1,8 @@
+export { useTrainingConfigStore } from "./stores/training-config-store";
+export {
+  shouldShowTrainingView,
+  useTrainingRuntimeStore,
+} from "./stores/training-runtime-store";
+export { useTrainingActions } from "./hooks/use-training-actions";
+export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle";
+export type { TrainingPhase } from "./types/runtime";
--- a/studio/frontend/src/features/training/lib/sync-runtime.ts
+++ b/studio/frontend/src/features/training/lib/sync-runtime.ts
@ -0,0 +1,19 @@
+import {
+  getTrainingMetrics,
+  getTrainingStatus,
+} from "../api/train-api";
+import { useTrainingRuntimeStore } from "../stores/training-runtime-store";
+import type { TrainingStatusResponse } from "../types/runtime";
+
+export async function syncTrainingRuntimeFromBackend(): Promise<TrainingStatusResponse> {
+  const [status, metrics] = await Promise.all([
+    getTrainingStatus(),
+    getTrainingMetrics(),
+  ]);
+
+  const runtimeStore = useTrainingRuntimeStore.getState();
+  runtimeStore.applyStatus(status);
+  runtimeStore.applyMetrics(metrics);
+
+  return status;
+}
--- a/studio/frontend/src/features/training/lib/validation.ts
+++ b/studio/frontend/src/features/training/lib/validation.ts
@ -0,0 +1,27 @@
+import type { TrainingConfigState } from "../types/config";
+
+export interface StartValidationResult {
+  ok: boolean;
+  message: string | null;
+}
+
+export function validateTrainingConfig(
+  config: TrainingConfigState,
+): StartValidationResult {
+  if (!config.selectedModel) {
+    return { ok: false, message: "Select a base model first." };
+  }
+
+  if (config.datasetSource !== "huggingface") {
+    return {
+      ok: false,
+      message: "Only Hugging Face dataset source is enabled right now.",
+    };
+  }
+
+  if (!config.dataset) {
+    return { ok: false, message: "Select a Hugging Face dataset first." };
+  }
+
+  return { ok: true, message: null };
+}
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@ -0,0 +1,101 @@
+import { DEFAULT_HYPERPARAMS, STEPS } from "@/config/training";
+import type { StepNumber } from "@/types/training";
+import { create } from "zustand";
+import { persist } from "zustand/middleware";
+import type { TrainingConfigState, TrainingConfigStore } from "../types/config";
+
+const MIN_STEP: StepNumber = 1;
+const MAX_STEP: StepNumber = STEPS.length as StepNumber;
+
+const initialState: TrainingConfigState = {
+  currentStep: MIN_STEP,
+  modelType: null,
+  selectedModel: null,
+  trainingMethod: "qlora",
+  hfToken: "",
+  datasetSource: "huggingface",
+  datasetFormat: "auto",
+  dataset: null,
+  uploadedFile: null,
+  ...DEFAULT_HYPERPARAMS,
+};
+
+function clampStep(step: number): StepNumber {
+  return Math.min(MAX_STEP, Math.max(MIN_STEP, step)) as StepNumber;
+}
+
+function canProceedForStep(state: TrainingConfigState): boolean {
+  switch (state.currentStep) {
+    case 1:
+      return state.modelType !== null;
+    case 2:
+      return state.selectedModel !== null;
+    case 3:
+      return state.datasetSource === "upload"
+        ? state.uploadedFile !== null
+        : state.dataset !== null;
+    case 4:
+    case 5:
+      return true;
+    default:
+      return false;
+  }
+}
+
+export const useTrainingConfigStore = create<TrainingConfigStore>()(
+  persist(
+    (set, get) => ({
+      ...initialState,
+      setStep: (step) => set({ currentStep: step }),
+      nextStep: () => set({ currentStep: clampStep(get().currentStep + 1) }),
+      prevStep: () => set({ currentStep: clampStep(get().currentStep - 1) }),
+      setModelType: (modelType) => set({ modelType, selectedModel: null }),
+      setSelectedModel: (selectedModel) => set({ selectedModel }),
+      setTrainingMethod: (trainingMethod) => set({ trainingMethod }),
+      setHfToken: (hfToken) => set({ hfToken }),
+      setDatasetSource: (datasetSource) => set({ datasetSource }),
+      setDatasetFormat: (datasetFormat) => set({ datasetFormat }),
+      setDataset: (dataset) => set({ dataset }),
+      setUploadedFile: (uploadedFile) => set({ uploadedFile }),
+      setEpochs: (epochs) => set({ epochs }),
+      setContextLength: (contextLength) => set({ contextLength }),
+      setLearningRate: (learningRate) => set({ learningRate }),
+      setLoraRank: (loraRank) => set({ loraRank }),
+      setLoraAlpha: (loraAlpha) => set({ loraAlpha }),
+      setLoraDropout: (loraDropout) => set({ loraDropout }),
+      setLoraVariant: (loraVariant) => set({ loraVariant }),
+      setBatchSize: (batchSize) => set({ batchSize }),
+      setGradientAccumulation: (gradientAccumulation) =>
+        set({ gradientAccumulation }),
+      setWeightDecay: (weightDecay) => set({ weightDecay }),
+      setWarmupSteps: (warmupSteps) => set({ warmupSteps }),
+      setMaxSteps: (maxSteps) => set({ maxSteps }),
+      setSaveSteps: (saveSteps) => set({ saveSteps }),
+      setPacking: (packing) => set({ packing }),
+      setTrainOnCompletions: (trainOnCompletions) =>
+        set({ trainOnCompletions }),
+      setGradientCheckpointing: (gradientCheckpointing) =>
+        set({ gradientCheckpointing }),
+      setRandomSeed: (randomSeed) => set({ randomSeed }),
+      setEnableWandb: (enableWandb) => set({ enableWandb }),
+      setWandbToken: (wandbToken) => set({ wandbToken }),
+      setWandbProject: (wandbProject) => set({ wandbProject }),
+      setEnableTensorboard: (enableTensorboard) => set({ enableTensorboard }),
+      setTensorboardDir: (tensorboardDir) => set({ tensorboardDir }),
+      setLogFrequency: (logFrequency) => set({ logFrequency }),
+      setFinetuneVisionLayers: (finetuneVisionLayers) =>
+        set({ finetuneVisionLayers }),
+      setFinetuneLanguageLayers: (finetuneLanguageLayers) =>
+        set({ finetuneLanguageLayers }),
+      setFinetuneAttentionModules: (finetuneAttentionModules) =>
+        set({ finetuneAttentionModules }),
+      setFinetuneMLPModules: (finetuneMLPModules) => set({ finetuneMLPModules }),
+      setTargetModules: (targetModules) => set({ targetModules }),
+      canProceed: () => canProceedForStep(get()),
+      reset: () => set(initialState),
+    }),
+    {
+      name: "unsloth_training_config_v1",
+    },
+  ),
+);
--- a/studio/frontend/src/features/training/stores/training-runtime-store.ts
+++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts
@ -0,0 +1,229 @@
+import { create } from "zustand";
+import type {
+  TrainingMetricsResponse,
+  TrainingProgressPayload,
+  TrainingRuntimeState,
+  TrainingRuntimeStore,
+  TrainingSeriesPoint,
+  TrainingStatusResponse,
+} from "../types/runtime";
+
+const initialState: TrainingRuntimeState = {
+  jobId: null,
+  phase: "idle",
+  isTrainingRunning: false,
+  message: "Ready to train",
+  error: null,
+  isHydrating: false,
+  hasHydrated: false,
+  isStarting: false,
+  startError: null,
+  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: [],
+};
+
+function sortSeries(points: TrainingSeriesPoint[]): TrainingSeriesPoint[] {
+  return [...points].sort((a, b) => a.step - b.step);
+}
+
+function toSeries(steps: number[], values: number[]): TrainingSeriesPoint[] {
+  const points: TrainingSeriesPoint[] = [];
+  for (let i = 0; i < steps.length; i += 1) {
+    const step = steps[i];
+    const value = values[i];
+    if (!Number.isFinite(step) || !Number.isFinite(value)) {
+      continue;
+    }
+    points.push({ step, value });
+  }
+  return sortSeries(points);
+}
+
+function upsertPoint(
+  points: TrainingSeriesPoint[],
+  step: number,
+  value: number,
+): TrainingSeriesPoint[] {
+  const next = points.slice();
+  const index = next.findIndex((point) => point.step === step);
+  if (index >= 0) {
+    next[index] = { step, value };
+    return next;
+  }
+  next.push({ step, value });
+  return sortSeries(next);
+}
+
+function applyMetricHistoryFromStatus(payload: TrainingStatusResponse): {
+  lossHistory: TrainingSeriesPoint[] | null;
+  lrHistory: TrainingSeriesPoint[] | null;
+} {
+  const history = payload.metric_history;
+  if (!history || !history.steps?.length) {
+    return { lossHistory: null, lrHistory: null };
+  }
+
+  const steps = history.steps;
+  const lossHistory = history.loss ? toSeries(steps, history.loss) : null;
+  const lrHistory = history.lr ? toSeries(steps, history.lr) : null;
+
+  return { lossHistory, lrHistory };
+}
+
+export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => ({
+  ...initialState,
+
+  setHydrating: (value) => set({ isHydrating: value }),
+  setHasHydrated: (value) => set({ hasHydrated: value }),
+  setStarting: (value) => set({ isStarting: value }),
+  setStartError: (value) => set({ startError: value }),
+  setSseConnected: (value) => set({ sseConnected: value }),
+  setLastEventId: (value) => set({ lastEventId: value }),
+
+  resetRuntime: () =>
+    set({
+      ...initialState,
+      lossHistory: [],
+      lrHistory: [],
+      gradNormHistory: [],
+    }),
+
+  setStartQueued: (jobId, message) =>
+    set({
+      jobId,
+      message,
+      error: null,
+      startError: null,
+      phase: "configuring",
+      isStarting: false,
+    }),
+
+  setRuntimeError: (message) =>
+    set({
+      error: message,
+      phase: "error",
+      isStarting: false,
+      startError: null,
+      sseConnected: false,
+    }),
+
+  applyStatus: (payload) =>
+    set((state) => {
+      const metricHistory = applyMetricHistoryFromStatus(payload);
+      const detailStep = payload.details?.step;
+      const detailTotal = payload.details?.total_steps;
+      const detailLoss = payload.details?.loss;
+      const detailLr = payload.details?.learning_rate;
+      const detailEpoch = payload.details?.epoch;
+
+      return {
+        ...state,
+        jobId: payload.job_id || state.jobId,
+        phase: payload.phase,
+        isTrainingRunning: payload.is_training_running,
+        message: payload.message,
+        error: payload.error,
+        startError: null,
+        currentStep:
+          typeof detailStep === "number" ? Math.max(detailStep, 0) : state.currentStep,
+        totalSteps:
+          typeof detailTotal === "number"
+            ? Math.max(detailTotal, 0)
+            : state.totalSteps,
+        currentLoss:
+          typeof detailLoss === "number" ? detailLoss : state.currentLoss,
+        currentLearningRate:
+          typeof detailLr === "number" ? detailLr : state.currentLearningRate,
+        currentEpoch:
+          typeof detailEpoch === "number" ? detailEpoch : state.currentEpoch,
+        lossHistory: metricHistory.lossHistory ?? state.lossHistory,
+        lrHistory: metricHistory.lrHistory ?? state.lrHistory,
+      };
+    }),
+
+  applyMetrics: (payload: TrainingMetricsResponse) =>
+    set((state) => {
+      const lossHistory = toSeries(payload.step_history, payload.loss_history);
+      const lrHistory = toSeries(payload.step_history, payload.lr_history);
+      const latestStep =
+        payload.current_step ??
+        (payload.step_history.length > 0
+          ? payload.step_history[payload.step_history.length - 1]
+          : null);
+
+      return {
+        ...state,
+        lossHistory: lossHistory.length > 0 ? lossHistory : state.lossHistory,
+        lrHistory: lrHistory.length > 0 ? lrHistory : state.lrHistory,
+        currentStep:
+          typeof latestStep === "number"
+            ? Math.max(latestStep, state.currentStep)
+            : state.currentStep,
+        currentLoss:
+          typeof payload.current_loss === "number"
+            ? payload.current_loss
+            : state.currentLoss,
+        currentLearningRate:
+          typeof payload.current_lr === "number"
+            ? payload.current_lr
+            : state.currentLearningRate,
+      };
+    }),
+
+  applyProgress: (payload: TrainingProgressPayload, eventId?: number) =>
+    set((state) => {
+      const step = Math.max(payload.step, 0);
+      return {
+        ...state,
+        jobId: payload.job_id || state.jobId,
+        currentStep: step,
+        totalSteps: Math.max(payload.total_steps, state.totalSteps),
+        currentLoss: payload.loss,
+        currentLearningRate: payload.learning_rate,
+        progressPercent: payload.progress_percent,
+        currentEpoch: payload.epoch ?? state.currentEpoch,
+        elapsedSeconds: payload.elapsed_seconds,
+        etaSeconds: payload.eta_seconds,
+        currentGradNorm: payload.grad_norm,
+        currentNumTokens: payload.num_tokens,
+        firstStepReceived: state.firstStepReceived || step > 0,
+        lastEventId: typeof eventId === "number" ? eventId : state.lastEventId,
+        lossHistory:
+          step > 0
+            ? upsertPoint(state.lossHistory, step, payload.loss)
+            : state.lossHistory,
+        lrHistory:
+          step > 0
+            ? upsertPoint(state.lrHistory, step, payload.learning_rate)
+            : state.lrHistory,
+        gradNormHistory:
+          step > 0 && typeof payload.grad_norm === "number"
+            ? upsertPoint(state.gradNormHistory, step, payload.grad_norm)
+            : state.gradNormHistory,
+      };
+    }),
+}));
+
+export function shouldShowTrainingView(state: TrainingRuntimeStore): boolean {
+  return (
+    state.phase !== "idle" ||
+    state.isTrainingRunning ||
+    state.isStarting ||
+    state.lossHistory.length > 0 ||
+    state.currentStep > 0
+  );
+}
--- a/studio/frontend/src/features/training/types/api.ts
+++ b/studio/frontend/src/features/training/types/api.ts
@ -0,0 +1,53 @@
+export interface TrainingStartRequest {
+  model_name: string;
+  training_type: string;
+  hf_token: string | null;
+  load_in_4bit: boolean;
+  max_seq_length: number;
+  hf_dataset: string | null;
+  local_datasets: string[];
+  format_type: string;
+  num_epochs: number;
+  learning_rate: string;
+  batch_size: number;
+  gradient_accumulation_steps: number;
+  warmup_steps: number | null;
+  warmup_ratio: number | null;
+  max_steps: number | null;
+  save_steps: number;
+  weight_decay: number;
+  random_seed: number;
+  packing: boolean;
+  optim: string;
+  lr_scheduler_type: string;
+  use_lora: boolean;
+  lora_r: number;
+  lora_alpha: number;
+  lora_dropout: number;
+  target_modules: string[];
+  gradient_checkpointing: string;
+  use_rslora: boolean;
+  use_loftq: boolean;
+  train_on_completions: boolean;
+  finetune_vision_layers: boolean;
+  finetune_language_layers: boolean;
+  finetune_attention_modules: boolean;
+  finetune_mlp_modules: boolean;
+  enable_wandb: boolean;
+  wandb_token: string | null;
+  wandb_project: string | null;
+  enable_tensorboard: boolean;
+  tensorboard_dir: string | null;
+}
+
+export interface TrainingStartResponse {
+  job_id: string;
+  status: "queued" | "error";
+  message: string;
+  error: string | null;
+}
+
+export interface TrainingStopResponse {
+  status: "stopped" | "idle";
+  message: string;
+}
--- a/studio/frontend/src/features/training/types/config.ts
+++ b/studio/frontend/src/features/training/types/config.ts
@ -0,0 +1,96 @@
+import type {
+  DatasetFormat,
+  DatasetSource,
+  GradientCheckpointing,
+  ModelType,
+  StepNumber,
+  TrainingMethod,
+} from "@/types/training";
+
+export type LoraVariant = "lora" | "rslora" | "loftq";
+
+export interface TrainingConfigState {
+  currentStep: StepNumber;
+  modelType: ModelType | null;
+  selectedModel: string | null;
+  trainingMethod: TrainingMethod;
+  hfToken: string;
+  datasetSource: DatasetSource;
+  datasetFormat: DatasetFormat;
+  dataset: string | null;
+  uploadedFile: string | null;
+  epochs: number;
+  contextLength: number;
+  learningRate: number;
+  loraRank: number;
+  loraAlpha: number;
+  loraDropout: number;
+  loraVariant: LoraVariant;
+  batchSize: number;
+  gradientAccumulation: number;
+  weightDecay: number;
+  warmupSteps: number;
+  maxSteps: number;
+  saveSteps: number;
+  packing: boolean;
+  trainOnCompletions: boolean;
+  gradientCheckpointing: GradientCheckpointing;
+  randomSeed: number;
+  enableWandb: boolean;
+  wandbToken: string;
+  wandbProject: string;
+  enableTensorboard: boolean;
+  tensorboardDir: string;
+  logFrequency: number;
+  finetuneVisionLayers: boolean;
+  finetuneLanguageLayers: boolean;
+  finetuneAttentionModules: boolean;
+  finetuneMLPModules: boolean;
+  targetModules: string[];
+}
+
+export interface TrainingConfigActions {
+  setStep: (step: StepNumber) => void;
+  nextStep: () => void;
+  prevStep: () => void;
+  setModelType: (type: ModelType) => void;
+  setSelectedModel: (model: string | null) => void;
+  setTrainingMethod: (method: TrainingMethod) => void;
+  setHfToken: (token: string) => void;
+  setDatasetSource: (source: DatasetSource) => void;
+  setDatasetFormat: (format: DatasetFormat) => void;
+  setDataset: (dataset: string | null) => void;
+  setUploadedFile: (file: string | null) => void;
+  setEpochs: (epochs: number) => void;
+  setContextLength: (length: number) => void;
+  setLearningRate: (rate: number) => void;
+  setLoraRank: (rank: number) => void;
+  setLoraAlpha: (alpha: number) => void;
+  setLoraDropout: (dropout: number) => void;
+  setLoraVariant: (variant: LoraVariant) => void;
+  setBatchSize: (value: number) => void;
+  setGradientAccumulation: (value: number) => void;
+  setWeightDecay: (value: number) => void;
+  setWarmupSteps: (value: number) => void;
+  setMaxSteps: (value: number) => void;
+  setSaveSteps: (value: number) => void;
+  setPacking: (value: boolean) => void;
+  setTrainOnCompletions: (value: boolean) => void;
+  setGradientCheckpointing: (value: GradientCheckpointing) => void;
+  setRandomSeed: (value: number) => void;
+  setEnableWandb: (value: boolean) => void;
+  setWandbToken: (value: string) => void;
+  setWandbProject: (value: string) => void;
+  setEnableTensorboard: (value: boolean) => void;
+  setTensorboardDir: (value: string) => void;
+  setLogFrequency: (value: number) => void;
+  setFinetuneVisionLayers: (value: boolean) => void;
+  setFinetuneLanguageLayers: (value: boolean) => void;
+  setFinetuneAttentionModules: (value: boolean) => void;
+  setFinetuneMLPModules: (value: boolean) => void;
+  setTargetModules: (value: string[]) => void;
+  canProceed: () => boolean;
+  reset: () => void;
+}
+
+export type TrainingConfigStore = TrainingConfigState & TrainingConfigActions;
--- a/studio/frontend/src/features/training/types/runtime.ts
+++ b/studio/frontend/src/features/training/types/runtime.ts
@ -0,0 +1,102 @@
+export type TrainingPhase =
+  | "idle"
+  | "loading_model"
+  | "loading_dataset"
+  | "configuring"
+  | "training"
+  | "completed"
+  | "error"
+  | "stopped";
+
+export interface TrainingStatusResponse {
+  job_id: string;
+  phase: TrainingPhase;
+  is_training_running: boolean;
+  message: string;
+  error: string | null;
+  details?: {
+    epoch?: number;
+    step?: number;
+    total_steps?: number;
+    loss?: number;
+    learning_rate?: number;
+  } | null;
+  metric_history?: {
+    steps?: number[];
+    loss?: number[];
+    lr?: number[];
+  } | null;
+}
+
+export interface TrainingMetricsResponse {
+  loss_history: number[];
+  lr_history: number[];
+  step_history: number[];
+  current_loss: number | null;
+  current_lr: number | null;
+  current_step: number | null;
+}
+
+export interface TrainingProgressPayload {
+  job_id: string;
+  step: number;
+  total_steps: number;
+  loss: number;
+  learning_rate: number;
+  progress_percent: number;
+  epoch: number | null;
+  elapsed_seconds: number | null;
+  eta_seconds: number | null;
+  grad_norm: number | null;
+  num_tokens: number | null;
+}
+
+export interface TrainingSeriesPoint {
+  step: number;
+  value: number;
+}
+
+export interface TrainingRuntimeState {
+  jobId: string | null;
+  phase: TrainingPhase;
+  isTrainingRunning: boolean;
+  message: string;
+  error: string | null;
+  isHydrating: boolean;
+  hasHydrated: boolean;
+  isStarting: boolean;
+  startError: string | null;
+  sseConnected: boolean;
+  firstStepReceived: boolean;
+  lastEventId: number | null;
+  currentStep: number;
+  totalSteps: number;
+  currentEpoch: number;
+  currentLoss: number;
+  currentLearningRate: number;
+  progressPercent: number;
+  elapsedSeconds: number | null;
+  etaSeconds: number | null;
+  currentGradNorm: number | null;
+  currentNumTokens: number | null;
+  lossHistory: TrainingSeriesPoint[];
+  lrHistory: TrainingSeriesPoint[];
+  gradNormHistory: TrainingSeriesPoint[];
+}
+
+export interface TrainingRuntimeActions {
+  setHydrating: (value: boolean) => void;
+  setHasHydrated: (value: boolean) => void;
+  setStarting: (value: boolean) => void;
+  setStartError: (value: string | null) => void;
+  setSseConnected: (value: boolean) => void;
+  setLastEventId: (value: number | null) => void;
+  resetRuntime: () => void;
+  applyStatus: (payload: TrainingStatusResponse) => void;
+  applyMetrics: (payload: TrainingMetricsResponse) => void;
+  applyProgress: (payload: TrainingProgressPayload, eventId?: number) => void;
+  setStartQueued: (jobId: string, message: string) => void;
+  setRuntimeError: (message: string) => void;
+}
+
+export type TrainingRuntimeStore = TrainingRuntimeState & TrainingRuntimeActions;
--- a/studio/frontend/src/stores/training.ts
+++ b/studio/frontend/src/stores/training.ts
@ -1,106 +1,4 @@
-import { DEFAULT_HYPERPARAMS } from "@/config/training";
-import type { StepNumber, WizardActions, WizardState } from "@/types/training";
-import { create } from "zustand";
+import { useTrainingConfigStore } from "@/features/training";

-const MIN_STEP: StepNumber = 1;
-const MAX_STEP: StepNumber = 5;
-
-const initialState: WizardState = {
-  isTraining: false,
-  trainingMetrics: null,
-  currentStep: MIN_STEP,
-  modelType: null,
-  selectedModel: null,
-  trainingMethod: "qlora",
-  hfToken: "",
-  datasetSource: "huggingface",
-  datasetFormat: "auto",
-  dataset: null,
-  uploadedFile: null,
-  ...DEFAULT_HYPERPARAMS,
-};
-
-function clampStep(step: number): StepNumber {
-  return Math.min(MAX_STEP, Math.max(MIN_STEP, step)) as StepNumber;
-}
-
-function canProceedForStep(state: WizardState): boolean {
-  switch (state.currentStep) {
-    case 1:
-      return state.modelType !== null;
-    case 2:
-      return state.selectedModel !== null;
-    case 3: {
-      if (state.datasetSource === "upload") {
-        return state.uploadedFile !== null;
-      }
-      return state.dataset !== null;
-    }
-    case 4:
-    case 5:
-      return true;
-    default:
-      return false;
-  }
-}
-
-export const useWizardStore = create<WizardState & WizardActions>(
-  (set, get) => ({
-    ...initialState,
-
-    setStep: (step) => set({ currentStep: step }),
-
-    nextStep: () => {
-      const { currentStep } = get();
-      set({ currentStep: clampStep(currentStep + 1) });
-    },
-
-    prevStep: () => {
-      const { currentStep } = get();
-      set({ currentStep: clampStep(currentStep - 1) });
-    },
-
-    setModelType: (type) => set({ modelType: type, selectedModel: null }),
-    setSelectedModel: (model) => set({ selectedModel: model }),
-    setTrainingMethod: (method) => set({ trainingMethod: method }),
-    setHfToken: (token) => set({ hfToken: token }),
-    setDatasetSource: (source) => set({ datasetSource: source }),
-    setDatasetFormat: (format) => set({ datasetFormat: format }),
-    setDataset: (dataset) => set({ dataset }),
-    setUploadedFile: (file) => set({ uploadedFile: file }),
-    setEpochs: (epochs) => set({ epochs }),
-    setContextLength: (length) => set({ contextLength: length }),
-    setLearningRate: (rate) => set({ learningRate: rate }),
-    setLoraRank: (rank) => set({ loraRank: rank }),
-    setLoraAlpha: (alpha) => set({ loraAlpha: alpha }),
-    setLoraDropout: (dropout) => set({ loraDropout: dropout }),
-    setLoraVariant: (v) => set({ loraVariant: v }),
-    setBatchSize: (v) => set({ batchSize: v }),
-    setGradientAccumulation: (v) => set({ gradientAccumulation: v }),
-    setWeightDecay: (v) => set({ weightDecay: v }),
-    setWarmupSteps: (v) => set({ warmupSteps: v }),
-    setMaxSteps: (v) => set({ maxSteps: v }),
-    setSaveSteps: (v) => set({ saveSteps: v }),
-    setPacking: (v) => set({ packing: v }),
-    setTrainOnCompletions: (v) => set({ trainOnCompletions: v }),
-    setGradientCheckpointing: (v) => set({ gradientCheckpointing: v }),
-    setRandomSeed: (v) => set({ randomSeed: v }),
-    setEnableWandb: (v) => set({ enableWandb: v }),
-    setWandbToken: (v) => set({ wandbToken: v }),
-    setWandbProject: (v) => set({ wandbProject: v }),
-    setEnableTensorboard: (v) => set({ enableTensorboard: v }),
-    setTensorboardDir: (v) => set({ tensorboardDir: v }),
-    setLogFrequency: (v) => set({ logFrequency: v }),
-    setFinetuneVisionLayers: (v) => set({ finetuneVisionLayers: v }),
-    setFinetuneLanguageLayers: (v) => set({ finetuneLanguageLayers: v }),
-    setFinetuneAttentionModules: (v) => set({ finetuneAttentionModules: v }),
-    setFinetuneMLPModules: (v) => set({ finetuneMLPModules: v }),
-    setTargetModules: (v) => set({ targetModules: v }),
-    setIsTraining: (v) => set({ isTraining: v }),
-    setTrainingMetrics: (v) => set({ trainingMetrics: v }),
-
-    canProceed: () => canProceedForStep(get()),
-
-    reset: () => set(initialState),
-  }),
-);
+export const useWizardStore = useTrainingConfigStore;
+export { useTrainingConfigStore };
--- a/studio/frontend/src/types/training.ts
+++ b/studio/frontend/src/types/training.ts
@ -9,30 +9,7 @@ export type DatasetSource = "huggingface" | "upload";
 export type DatasetFormat = "auto" | "alpaca" | "chatml" | "sharegpt";
 export type GradientCheckpointing = "none" | "true" | "unsloth";

-export interface TrainingMetrics {
-  currentStep: number;
-  totalSteps: number;
-  currentEpoch: number;
-  totalEpochs: number;
-  currentLoss: number;
-  currentLR: number;
-  gradNorm: number;
-  samplesPerSecond: number;
-  lossHistory: { step: number; loss: number }[];
-  lrHistory: { step: number; lr: number }[];
-  gradNormHistory: { step: number; gradNorm: number }[];
-  gpuUtil: number;
-  gpuTemp: number;
-  gpuVramUsed: number;
-  gpuVramTotal: number;
-  gpuPower: number;
-  elapsed: number;
-  status: "training" | "warmup" | "saving";
-}
-
 export interface WizardState {
-  isTraining: boolean;
-  trainingMetrics: TrainingMetrics | null;
  currentStep: StepNumber;
  modelType: ModelType | null;
  selectedModel: string | null;
@ -112,8 +89,6 @@ export interface WizardActions {
  setFinetuneAttentionModules: (v: boolean) => void;
  setFinetuneMLPModules: (v: boolean) => void;
  setTargetModules: (v: string[]) => void;
-  setIsTraining: (v: boolean) => void;
-  setTrainingMetrics: (v: TrainingMetrics | null) => void;
  canProceed: () => boolean;
  reset: () => void;
 }
--- a/studio/frontend/vite.config.ts
+++ b/studio/frontend/vite.config.ts
@ -6,18 +6,43 @@ import { defineConfig } from "vite";
 // https://vite.dev/config/
 export default defineConfig({
  plugins: [react(), tailwindcss()],
+  optimizeDeps: {
+    include: ["@dagrejs/dagre", "@dagrejs/graphlib"],
+  },
  server: {
-    allowedHosts: ["playground.wasimhub.dev"],
+    host: "0.0.0.0",
+    allowedHosts: true,
    proxy: {
      "/api": {
        target: "http://127.0.0.1:8000",
        changeOrigin: true,
      },
+      "/preview": {
+        target: "http://127.0.0.1:8004",
+        changeOrigin: true,
+      },
+      "/validate": {
+        target: "http://127.0.0.1:8004",
+        changeOrigin: true,
+      },
+      "/tools": {
+        target: "http://127.0.0.1:8004",
+        changeOrigin: true,
+      },
    },
  },
  resolve: {
    alias: {
      "@": path.resolve(__dirname, "./src"),
+      "@dagrejs/dagre": path.resolve(
+        __dirname,
+        "./node_modules/@dagrejs/dagre/dist/dagre.cjs.js",
+      ),
+    },
+  },
+  build: {
+    commonjsOptions: {
+      include: [/node_modules/, /@dagrejs\/dagre/, /@dagrejs\/graphlib/],
    },
  },
 });