From a4a997eee68dad00994ae7978a9c45242c385796 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Fri, 13 Feb 2026 12:26:28 +0100 Subject: [PATCH] feat: add training feature with state management, API integration, and runtime synchronization --- studio/frontend/package.json | 2 + studio/frontend/src/features/auth/index.ts | 2 +- .../src/features/export/export-page.tsx | 18 +- .../components/steps/dataset-step.tsx | 4 +- .../components/steps/hyperparameters-step.tsx | 4 +- .../components/steps/model-selection-step.tsx | 4 +- .../components/steps/model-type-step.tsx | 4 +- .../components/steps/summary-step.tsx | 4 +- .../onboarding/components/wizard-content.tsx | 4 +- .../onboarding/components/wizard-footer.tsx | 4 +- .../onboarding/components/wizard-layout.tsx | 4 +- .../onboarding/components/wizard-sidebar.tsx | 4 +- .../components/wizard-step-item.tsx | 4 +- .../studio/sections/charts-content.tsx | 7 +- .../studio/sections/charts-section.tsx | 40 ++- .../studio/sections/dataset-section.tsx | 4 +- .../studio/sections/model-section.tsx | 4 +- .../studio/sections/params-section.tsx | 4 +- .../studio/sections/progress-section.tsx | 36 ++- .../studio/sections/training-section.tsx | 15 +- .../src/features/studio/studio-page.tsx | 22 +- .../src/features/training/api/mappers.ts | 63 +++++ .../src/features/training/api/train-api.ts | 173 +++++++++++++ .../training/hooks/use-training-actions.ts | 70 ++++++ .../hooks/use-training-runtime-lifecycle.ts | 183 ++++++++++++++ .../frontend/src/features/training/index.ts | 8 + .../src/features/training/lib/sync-runtime.ts | 19 ++ .../src/features/training/lib/validation.ts | 27 +++ .../training/stores/training-config-store.ts | 101 ++++++++ .../training/stores/training-runtime-store.ts | 229 ++++++++++++++++++ .../src/features/training/types/api.ts | 53 ++++ .../src/features/training/types/config.ts | 96 ++++++++ .../src/features/training/types/runtime.ts | 102 ++++++++ studio/frontend/src/stores/training.ts | 108 +-------- studio/frontend/src/types/training.ts | 25 -- studio/frontend/vite.config.ts | 27 ++- 36 files changed, 1286 insertions(+), 192 deletions(-) create mode 100644 studio/frontend/src/features/training/api/mappers.ts create mode 100644 studio/frontend/src/features/training/api/train-api.ts create mode 100644 studio/frontend/src/features/training/hooks/use-training-actions.ts create mode 100644 studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts create mode 100644 studio/frontend/src/features/training/index.ts create mode 100644 studio/frontend/src/features/training/lib/sync-runtime.ts create mode 100644 studio/frontend/src/features/training/lib/validation.ts create mode 100644 studio/frontend/src/features/training/stores/training-config-store.ts create mode 100644 studio/frontend/src/features/training/stores/training-runtime-store.ts create mode 100644 studio/frontend/src/features/training/types/api.ts create mode 100644 studio/frontend/src/features/training/types/config.ts create mode 100644 studio/frontend/src/features/training/types/runtime.ts diff --git a/studio/frontend/package.json b/studio/frontend/package.json index e7e10695b5..ac61c22969 100644 --- 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", diff --git a/studio/frontend/src/features/auth/index.ts b/studio/frontend/src/features/auth/index.ts index a6651b0c00..bd0b98f6be 100644 --- 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, diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 376a963169..ddea863397 100644 --- 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(null); const [exportMethod, setExportMethod] = useState(null); diff --git a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx index 930a51ec99..65d802c77e 100644 --- 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, diff --git a/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx b/studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx index e7afe5f1fb..45f9c82878 100644 --- 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, diff --git a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx index a3a2da37ec..d11ff6bac0 100644 --- 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, diff --git a/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx index a83b6486f8..8cf6e8bc07 100644 --- 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 = { 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, diff --git a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx index 8dfeffba3b..ae6edd230a 100644 --- 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, diff --git a/studio/frontend/src/features/onboarding/components/wizard-content.tsx b/studio/frontend/src/features/onboarding/components/wizard-content.tsx index 48b4371dd7..2f10be0fcd 100644 --- 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 = { }; 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]; diff --git a/studio/frontend/src/features/onboarding/components/wizard-footer.tsx b/studio/frontend/src/features/onboarding/components/wizard-footer.tsx index c167303ec6..2b3af5aa96 100644 --- 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, diff --git a/studio/frontend/src/features/onboarding/components/wizard-layout.tsx b/studio/frontend/src/features/onboarding/components/wizard-layout.tsx index 698b1c9052..3a38035fe1 100644 --- 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(null); const hasFiredRef = useRef(false); const isFinalStep = currentStep === STEPS.length; diff --git a/studio/frontend/src/features/onboarding/components/wizard-sidebar.tsx b/studio/frontend/src/features/onboarding/components/wizard-sidebar.tsx index 45390cb4e5..6cd5c36a6f 100644 --- 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 ( diff --git a/studio/frontend/src/features/onboarding/components/wizard-step-item.tsx b/studio/frontend/src/features/onboarding/components/wizard-step-item.tsx index d161c2cc5a..953b1699c7 100644 --- 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; diff --git a/studio/frontend/src/features/studio/sections/charts-content.tsx b/studio/frontend/src/features/studio/sections/charts-content.tsx index 74583bfc89..60c9486565 100644 --- 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]); diff --git a/studio/frontend/src/features/studio/sections/charts-section.tsx b/studio/frontend/src/features/studio/sections/charts-section.tsx index fcda887879..868df740d8 100644 --- 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 { } > - + ); } diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index abcc81f621..890faad2b3 100644 --- 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, diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index b6a6968d64..fcaa319e1d 100644 --- 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, diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index cc2953c005..8c753cde7e 100644 --- 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); diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index c2a70d96b3..b8656a7867 100644 --- 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 = { @@ -104,6 +104,8 @@ export function ProgressSection(): ReactElement { ); const { stopTrainingRun } = useTrainingActions(); + const localStartAtRef = useRef(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`} - - Tokens: {runtime.currentNumTokens == null ? "--" : runtime.currentNumTokens} - + {runtime.currentNumTokens != null && ( + Tokens: {runtime.currentNumTokens} + )} diff --git a/studio/frontend/src/features/studio/sections/training-section.tsx b/studio/frontend/src/features/studio/sections/training-section.tsx index a624d7fd3c..fa5c4e85f5 100644 --- 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 */} + {startError && ( +

{startError}

+ )} {/* Save / Clear */}
diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx index 7478f32c36..14cf64f3e9 100644 --- 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 (
@@ -18,13 +26,17 @@ export function StudioPage(): ReactElement { Fine-tuning Studio

- {isTraining - ? "Training in progress" + {showTrainingView + ? runtimeMessage || "Training in progress" : "Configure and start training"}

- {isTraining ? ( + {!hasHydratedRuntime && isHydratingRuntime ? ( +
+ Loading training runtime... +
+ ) : showTrainingView ? ( ) : (
diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts new file mode 100644 index 0000000000..0150c62b41 --- /dev/null +++ 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, + }; +} diff --git a/studio/frontend/src/features/training/api/train-api.ts b/studio/frontend/src/features/training/api/train-api.ts new file mode 100644 index 0000000000..e3c53b0bc2 --- /dev/null +++ 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 { + 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(response: Response): Promise { + if (!response.ok) { + throw new Error(await readError(response)); + } + return (await response.json()) as T; +} + +export async function startTraining( + payload: TrainingStartRequest, +): Promise { + const response = await authFetch("/api/train/start", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + }); + return parseJson(response); +} + +export async function stopTraining(): Promise { + const response = await authFetch("/api/train/stop", { method: "POST" }); + return parseJson(response); +} + +export async function getTrainingStatus(): Promise { + const response = await authFetch("/api/train/status"); + return parseJson(response); +} + +export async function getTrainingMetrics(): Promise { + const response = await authFetch("/api/train/metrics"); + return parseJson(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 { + 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 }; diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts new file mode 100644 index 0000000000..510c45ed4a --- /dev/null +++ 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 => { + 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 => { + 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, + }; +} diff --git a/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts b/studio/frontend/src/features/training/hooks/use-training-runtime-lifecycle.ts new file mode 100644 index 0000000000..3baff3a62d --- /dev/null +++ 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 | 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(); + }; + }, []); +} diff --git a/studio/frontend/src/features/training/index.ts b/studio/frontend/src/features/training/index.ts new file mode 100644 index 0000000000..48b9079fd3 --- /dev/null +++ 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"; diff --git a/studio/frontend/src/features/training/lib/sync-runtime.ts b/studio/frontend/src/features/training/lib/sync-runtime.ts new file mode 100644 index 0000000000..b5fbd0bafb --- /dev/null +++ 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 { + const [status, metrics] = await Promise.all([ + getTrainingStatus(), + getTrainingMetrics(), + ]); + + const runtimeStore = useTrainingRuntimeStore.getState(); + runtimeStore.applyStatus(status); + runtimeStore.applyMetrics(metrics); + + return status; +} diff --git a/studio/frontend/src/features/training/lib/validation.ts b/studio/frontend/src/features/training/lib/validation.ts new file mode 100644 index 0000000000..8e966153d3 --- /dev/null +++ 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 }; +} diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts new file mode 100644 index 0000000000..95ded4c211 --- /dev/null +++ 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()( + 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", + }, + ), +); diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts new file mode 100644 index 0000000000..bc40019346 --- /dev/null +++ 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()((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 + ); +} diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts new file mode 100644 index 0000000000..9cf789a702 --- /dev/null +++ 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; +} diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts new file mode 100644 index 0000000000..c0d93f76a7 --- /dev/null +++ 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; diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts new file mode 100644 index 0000000000..fe2afbd36d --- /dev/null +++ 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; diff --git a/studio/frontend/src/stores/training.ts b/studio/frontend/src/stores/training.ts index 469880800c..d3bd63d648 100644 --- 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( - (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 }; diff --git a/studio/frontend/src/types/training.ts b/studio/frontend/src/types/training.ts index 5e206278d0..49c1dddaa3 100644 --- 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; } diff --git a/studio/frontend/vite.config.ts b/studio/frontend/vite.config.ts index 84565f26d1..7b94b5d59f 100644 --- 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/], }, }, });