diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index ec204f670c..72c36307c9 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -1,6 +1,8 @@ """ Datasets API routes """ +import base64 +import io import sys from pathlib import Path from fastapi import APIRouter, HTTPException @@ -30,6 +32,42 @@ if not logger.handlers: from models.datasets import CheckFormatRequest, CheckFormatResponse +def _serialize_preview_value(value): + """make it json safe for client preview ⊂(◉‿◉)つ""" + if value is None or isinstance(value, (str, int, float, bool)): + return value + + try: + from PIL.Image import Image as PILImage + if isinstance(value, PILImage): + buffer = io.BytesIO() + value.convert("RGB").save(buffer, format="JPEG", quality=85) + return { + "type": "image", + "mime": "image/jpeg", + "width": value.width, + "height": value.height, + "data": base64.b64encode(buffer.getvalue()).decode("ascii"), + } + except Exception: + pass + + if isinstance(value, dict): + return {str(key): _serialize_preview_value(item) for key, item in value.items()} + + if isinstance(value, (list, tuple)): + return [_serialize_preview_value(item) for item in value] + + return str(value) + + +def _serialize_preview_rows(rows): + return [ + {str(key): _serialize_preview_value(value) for key, value in dict(row).items()} + for row in rows + ] + + # --- Endpoints --- @router.post("/check-format", response_model=CheckFormatResponse) @@ -92,15 +130,15 @@ async def check_format(request: CheckFormatRequest): custom_format_mapping=result.get("suggested_mapping"), ) processed = format_result["dataset"] - preview_samples = [dict(row) for row in processed] + preview_samples = _serialize_preview_rows(processed) except Exception as e: logger.warning(f"Processed preview generation failed (non-fatal): {e}") # Fall back to raw samples so frontend still has something - preview_samples = [dict(row) for row in preview_slice] + preview_samples = _serialize_preview_rows(preview_slice) else: # Format detection failed — return raw samples so user can # see actual data and map columns in the frontend - preview_samples = [dict(row) for row in preview_slice] + preview_samples = _serialize_preview_rows(preview_slice) return CheckFormatResponse( requires_manual_mapping=result["requires_manual_mapping"], diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index eaf004f7b2..dcbe83d2ae 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -180,14 +180,19 @@ async def start_training( except Exception as e: logger.error(f"Error updating progress: {e}") - # Consume the generator - this actually runs the training - update_count = 0 - for _update_tuple in backend.start_training(**training_kwargs): - update_count += 1 - if update_count % 10 == 0: - logger.info(f"Training progress update #{update_count}") + # start_training returns bool (not generator) + run_result = backend.start_training(**training_kwargs) + logger.info( + "Training job %s backend.start_training returned type=%s value=%r", + job_id, + type(run_result).__name__, + run_result, + ) + if not run_result: + progress_error = backend.trainer.training_progress.error + raise RuntimeError(progress_error or "Training failed to start") - logger.info(f"Training job {job_id} completed successfully") + logger.info(f"Training job {job_id} started successfully") except Exception as e: logger.error(f"Training error in job {job_id}: {e}", exc_info=True) @@ -653,4 +658,3 @@ async def stream_training_progress( "X-Accel-Buffering": "no", } ) - 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/components/ui/terminal.tsx b/studio/frontend/src/components/ui/terminal.tsx new file mode 100644 index 0000000000..d9483eaca5 --- /dev/null +++ b/studio/frontend/src/components/ui/terminal.tsx @@ -0,0 +1,227 @@ +import { cn } from "@/lib/utils" +import { + Children, + cloneElement, + isValidElement, + useEffect, + useRef, + useState, +} from "react" +import type { ElementType, ReactElement, ReactNode } from "react" + +type TerminalProps = { + children: ReactNode + className?: string + sequence?: boolean + startOnView?: boolean +} + +type InternalLineProps = { + __isActive?: boolean + __onDone?: () => void + __sequence?: boolean +} + +function useStartOnView(enabled: boolean): { + ref: React.RefObject + started: boolean +} { + const ref = useRef(null) + const [isInView, setIsInView] = useState(false) + const started = !enabled || isInView + + useEffect(() => { + if (!enabled) { + return + } + + const node = ref.current + if (!node) { + return + } + + const observer = new IntersectionObserver( + ([entry]) => { + if (entry?.isIntersecting) { + setIsInView(true) + observer.disconnect() + } + }, + { threshold: 0.2 } + ) + + observer.observe(node) + return () => observer.disconnect() + }, [enabled]) + + return { ref, started } +} + +export function Terminal({ + children, + className, + sequence = true, + startOnView = true, +}: TerminalProps): ReactElement { + const { ref, started } = useStartOnView(startOnView) + const childElements = Children.toArray(children).filter(isValidElement) + const [activeIndex, setActiveIndex] = useState(0) + const visibleIndex = sequence + ? started + ? activeIndex + : -1 + : Number.MAX_SAFE_INTEGER + + function handleLineDone(index: number): void { + if (!sequence) { + return + } + + setActiveIndex((prev) => { + if (prev !== index) { + return prev + } + return Math.min(index + 1, childElements.length) + }) + } + + return ( +
+ {childElements.map((child, index) => + cloneElement(child, { + __sequence: sequence, + __isActive: !sequence || visibleIndex >= index, + __onDone: () => handleLineDone(index), + key: child.key ?? index, + } as InternalLineProps) + )} +
+ ) +} + +type AnimatedSpanProps = InternalLineProps & { + children: ReactNode + className?: string + delay?: number + startOnView?: boolean +} + +export function AnimatedSpan({ + children, + className, + delay = 0, + startOnView = false, + __isActive, + __sequence, + __onDone, +}: AnimatedSpanProps): ReactElement { + const { ref, started } = useStartOnView(startOnView) + const [visible, setVisible] = useState(false) + const doneRef = useRef(false) + const onDoneRef = useRef(__onDone) + const shouldStart = __sequence ? __isActive : started + + useEffect(() => { + onDoneRef.current = __onDone + }, [__onDone]) + + useEffect(() => { + if (!shouldStart || doneRef.current) { + return + } + + const timeout = window.setTimeout(() => { + setVisible(true) + doneRef.current = true + onDoneRef.current?.() + }, delay) + + return () => window.clearTimeout(timeout) + }, [delay, shouldStart]) + + return ( +
+ {children} +
+ ) +} + +type TypingAnimationProps = InternalLineProps & { + children: string + className?: string + duration?: number + delay?: number + as?: ElementType + startOnView?: boolean +} + +export function TypingAnimation({ + children, + className, + duration = 60, + delay = 0, + as: Component = "span", + startOnView = true, + __isActive, + __sequence, + __onDone, +}: TypingAnimationProps): ReactElement { + const { ref, started } = useStartOnView(startOnView) + const [typed, setTyped] = useState("") + const doneRef = useRef(false) + const onDoneRef = useRef(__onDone) + const shouldStart = __sequence ? __isActive : started + + useEffect(() => { + onDoneRef.current = __onDone + }, [__onDone]) + + useEffect(() => { + if (!shouldStart || doneRef.current) { + return + } + + let index = 0 + let intervalId: number | null = null + const startTimer = window.setTimeout(() => { + intervalId = window.setInterval(() => { + index += 1 + setTyped(children.slice(0, index)) + + if (index >= children.length) { + if (intervalId) { + window.clearInterval(intervalId) + } + doneRef.current = true + onDoneRef.current?.() + } + }, duration) + }, delay) + + return () => { + window.clearTimeout(startTimer) + if (intervalId) { + window.clearInterval(intervalId) + } + } + }, [children, delay, duration, shouldStart]) + + return ( +
+ {typed} +
+ ) +} 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 3a4d2d1d29..60c9486565 100644 --- a/studio/frontend/src/features/studio/sections/charts-content.tsx +++ b/studio/frontend/src/features/studio/sections/charts-content.tsx @@ -23,7 +23,6 @@ import { } from "@/components/ui/dropdown-menu"; import { Label } from "@/components/ui/label"; import { Slider } from "@/components/ui/slider"; -import type { TrainingMetrics } from "@/types/training"; import { ChartAverageIcon, Settings02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { type ReactElement, useMemo, useState } from "react"; @@ -44,9 +43,11 @@ const lossConfig = { const lrConfig = { lr: { label: "LR", color: "#8b5cf6" }, } satisfies ChartConfig; + const gradNormConfig = { gradNorm: { label: "Grad Norm", color: "#f97316" }, } satisfies ChartConfig; + const evalLossConfig = { loss: { label: "Eval Loss", color: "#ef4444" }, } satisfies ChartConfig; @@ -62,10 +63,81 @@ const placeholderEvalData = [ type LossHistoryItem = { step: number; loss: number }; type SmoothedLossItem = LossHistoryItem & { smoothed: number }; +interface TrainingChartSeries { + lossHistory: LossHistoryItem[]; + lrHistory: { step: number; lr: number }[]; + gradNormHistory: { step: number; gradNorm: number }[]; +} + +const CHART_SYNC_ID = "train-metrics-sync"; +const MAX_RENDER_POINTS = 800; +const DEFAULT_VISIBLE_POINTS = 160; + +function formatStepTick(value: number): string { + if (value >= 1_000_000) { + return `${(value / 1_000_000).toFixed(1)}M`; + } + if (value >= 1_000) { + return `${(value / 1_000).toFixed(1)}k`; + } + return String(Math.round(value)); +} + +function compressSeries(data: T[], maxPoints: number): T[] { + if (data.length <= maxPoints) { + return data; + } + + const stride = Math.ceil(data.length / maxPoints); + return data.filter( + (_item, index) => index % stride === 0 || index === data.length - 1, + ); +} + +function buildStepTicks(min: number, max: number, targetCount = 6): number[] { + if (!Number.isFinite(min) || !Number.isFinite(max)) { + return [0, 1]; + } + if (max <= min) { + return [min, max]; + } + + const stepSize = Math.max(1, Math.ceil((max - min) / (targetCount - 1))); + const ticks: number[] = []; + let current = min; + + while (current < max) { + ticks.push(current); + current += stepSize; + } + + ticks.push(max); + return Array.from(new Set(ticks)); +} + +function buildYDomain(values: number[]): [number, number] { + if (values.length === 0) { + return [0, 1]; + } + + const min = Math.min(...values); + const max = Math.max(...values); + + if (min === max) { + const base = Math.abs(min); + const pad = base > 0 ? base * 0.08 : 0.1; + return [min - pad, max + pad]; + } + + const pad = (max - min) * 0.12; + return [min - pad, max + pad]; +} + function ema(data: LossHistoryItem[], alpha: number): SmoothedLossItem[] { if (data.length === 0) { return []; } + let s = data[0].loss; return data.map((d) => { s = alpha * d.loss + (1 - alpha) * s; @@ -75,18 +147,113 @@ function ema(data: LossHistoryItem[], alpha: number): SmoothedLossItem[] { export function ChartsContent({ metrics, -}: { metrics: TrainingMetrics }): ReactElement { - const [smoothing, setSmoothing] = useState(0.6); +}: { metrics: TrainingChartSeries }): ReactElement { + const [smoothing, setSmoothing] = useState(0.75); const [showRaw, setShowRaw] = useState(true); const [showSmoothed, setShowSmoothed] = useState(true); const [showAvgLine, setShowAvgLine] = useState(true); const lossHistory = metrics.lossHistory; const smoothedData = useMemo( - () => (lossHistory ? ema(lossHistory, 1 - smoothing) : []), + () => (lossHistory.length > 0 ? ema(lossHistory, 1 - smoothing) : []), [lossHistory, smoothing], ); + const reducedLossData = useMemo( + () => compressSeries(smoothedData, MAX_RENDER_POINTS), + [smoothedData], + ); + + const reducedGradNormData = useMemo( + () => compressSeries(metrics.gradNormHistory, MAX_RENDER_POINTS), + [metrics.gradNormHistory], + ); + + const reducedLrData = useMemo( + () => compressSeries(metrics.lrHistory, MAX_RENDER_POINTS), + [metrics.lrHistory], + ); + + const visibleStepDomain = useMemo<[number, number]>(() => { + const allSteps = [ + ...reducedLossData.map((point) => point.step), + ...reducedGradNormData.map((point) => point.step), + ...reducedLrData.map((point) => point.step), + ].sort((a, b) => a - b); + + if (allSteps.length === 0) { + 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] ?? minStep; + if (startStep === endStep) { + return [startStep, startStep + 4]; + } + if (endStep - startStep < 6) { + return [Math.max(minStep, endStep - 6), endStep]; + } + return [startStep, endStep]; + }, [reducedGradNormData, reducedLossData, reducedLrData]); + + const xAxisTicks = useMemo( + () => buildStepTicks(visibleStepDomain[0], visibleStepDomain[1]), + [visibleStepDomain], + ); + + const visibleLossValues = useMemo( + () => + reducedLossData + .filter( + (point) => + point.step >= visibleStepDomain[0] && point.step <= visibleStepDomain[1], + ) + .map((point) => point.loss), + [reducedLossData, visibleStepDomain], + ); + + const visibleSmoothValues = useMemo( + () => + reducedLossData + .filter( + (point) => + point.step >= visibleStepDomain[0] && point.step <= visibleStepDomain[1], + ) + .map((point) => point.smoothed), + [reducedLossData, visibleStepDomain], + ); + + const visibleGradValues = useMemo( + () => + reducedGradNormData + .filter( + (point) => + point.step >= visibleStepDomain[0] && point.step <= visibleStepDomain[1], + ) + .map((point) => point.gradNorm), + [reducedGradNormData, visibleStepDomain], + ); + + const visibleLrValues = useMemo( + () => + reducedLrData + .filter( + (point) => + point.step >= visibleStepDomain[0] && point.step <= visibleStepDomain[1], + ) + .map((point) => point.lr), + [reducedLrData, visibleStepDomain], + ); + + const lossDomain = useMemo( + () => buildYDomain([...visibleLossValues, ...visibleSmoothValues]), + [visibleLossValues, visibleSmoothValues], + ); + const gradDomain = useMemo(() => buildYDomain(visibleGradValues), [visibleGradValues]); + const lrDomain = useMemo(() => buildYDomain(visibleLrValues), [visibleLrValues]); + const avg = metrics.lossHistory.length > 0 ? +( @@ -96,8 +263,7 @@ export function ChartsContent({ : 0; return ( -
- {/* Training Loss */} +
Training Loss @@ -106,7 +272,7 @@ export function ChartsContent({ @@ -155,12 +321,10 @@ export function ChartsContent({ - + @@ -168,19 +332,27 @@ export function ChartsContent({ formatStepTick(Number(value))} interval="preserveStartEnd" /> Number(value).toFixed(2)} /> )} {showSmoothed && ( )} @@ -232,7 +412,6 @@ export function ChartsContent({ - {/* Grad Norm */} Gradient Norm @@ -240,10 +419,11 @@ export function ChartsContent({ @@ -251,19 +431,27 @@ export function ChartsContent({ formatStepTick(Number(value))} interval="preserveStartEnd" /> Number(value).toFixed(2)} /> } /> @@ -288,18 +480,15 @@ export function ChartsContent({ - {/* Learning Rate */} Learning Rate - + @@ -307,20 +496,27 @@ export function ChartsContent({ formatStepTick(Number(value))} interval="preserveStartEnd" /> v.toExponential(0)} + width={52} + tickFormatter={(value) => Number(value).toExponential(0)} /> `Step ${payload?.[0]?.payload?.step ?? ""}` } - formatter={(value) => [ - Number(value).toExponential(3), - "LR", - ]} + formatter={(value) => [Number(value).toExponential(3), "LR"]} /> } /> } /> @@ -349,7 +546,6 @@ export function ChartsContent({ - {/* Eval Loss (disabled/blurred) */} @@ -360,7 +556,7 @@ export function ChartsContent({
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-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx index e9d70846be..80fbcb65ee 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx @@ -27,6 +27,14 @@ type CheckFormatResponse = { total_rows?: number | null; }; +type PreviewImagePayload = { + type: "image"; + mime?: string; + width?: number; + height?: number; + data?: string; +}; + type DatasetPreviewDialogProps = { open: boolean; onOpenChange: (open: boolean) => void; @@ -123,6 +131,36 @@ export function DatasetPreviewDialog({ ), cell: ({ getValue }: { getValue: () => unknown }) => { const value = getValue(); + const images = collectPreviewImages(value); + if (images.length > 0) { + return ( +
+ {images.slice(0, 4).map((image, index) => { + const mime = image.mime || "image/jpeg"; + const src = image.data ? `data:${mime};base64,${image.data}` : ""; + const width = image.width ?? 128; + const height = image.height ?? 128; + return ( + {`preview-${index}`} + ); + })} + {images.length > 4 && ( + + +{images.length - 4} more + + )} +
+ ); + } + const text = formatCell(value); if (!text) { return ( @@ -285,3 +323,41 @@ function formatCell(value: unknown): string { return JSON.stringify(value).slice(0, 500); return String(value); } + +function isPreviewImagePayload(value: unknown): value is PreviewImagePayload { + if (!value || typeof value !== "object") return false; + const record = value as Record; + return ( + record.type === "image" && + typeof record.data === "string" && + record.data.length > 0 + ); +} + +function collectPreviewImages(value: unknown): PreviewImagePayload[] { + const images: PreviewImagePayload[] = []; + const stack: unknown[] = [value]; + let steps = 0; + + while (stack.length > 0 && steps < 200) { + steps += 1; + const current = stack.pop(); + if (isPreviewImagePayload(current)) { + images.push(current); + continue; + } + + if (Array.isArray(current)) { + for (const item of current) stack.push(item); + continue; + } + + if (current && typeof current === "object") { + for (const nested of Object.values(current as Record)) { + stack.push(nested); + } + } + } + + return images; +} 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 0656093e6a..b8656a7867 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -5,7 +5,12 @@ import { PopoverContent, PopoverTrigger, } from "@/components/ui/popover"; -import { useWizardStore } from "@/stores/training"; +import { + useTrainingConfigStore, + useTrainingActions, + useTrainingRuntimeStore, + type TrainingPhase, +} from "@/features/training"; import { ChartAverageIcon, DashboardSpeed01Icon, @@ -16,66 +21,160 @@ 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"; -export function ProgressSection(): ReactElement | null { - const store = useWizardStore(); - const metrics = store.trainingMetrics; - if (!metrics) { - return null; +const phaseLabel: Record = { + idle: "Idle", + loading_model: "Loading model", + loading_dataset: "Loading dataset", + configuring: "Configuring", + training: "Training", + completed: "Completed", + error: "Error", + stopped: "Stopped", +}; + +const phaseColors: Record = { + idle: "bg-muted text-muted-foreground", + loading_model: "bg-amber-100 text-amber-700 dark:bg-amber-900 dark:text-amber-300", + loading_dataset: + "bg-amber-100 text-amber-700 dark:bg-amber-900 dark:text-amber-300", + configuring: "bg-blue-100 text-blue-700 dark:bg-blue-900 dark:text-blue-300", + training: + "bg-emerald-100 text-emerald-700 dark:bg-emerald-900 dark:text-emerald-300", + completed: + "bg-emerald-100 text-emerald-700 dark:bg-emerald-900 dark:text-emerald-300", + error: "bg-red-100 text-red-700 dark:bg-red-900 dark:text-red-300", + stopped: "bg-muted text-muted-foreground", +}; + +function formatDuration(seconds: number | null): string { + if (seconds == null || seconds < 0) { + return "--"; } + const total = Math.floor(seconds); + const min = Math.floor(total / 60); + const sec = total % 60; + return `${min}m ${sec}s`; +} - const pct = Math.round((metrics.currentStep / metrics.totalSteps) * 100); - const etaSec = - metrics.totalSteps > 0 - ? Math.round( - ((metrics.totalSteps - metrics.currentStep) / - Math.max(metrics.currentStep, 1)) * - metrics.elapsed, +function formatNumber(value: number | null | undefined, digits: number): string { + if (value == null || !Number.isFinite(value)) { + return "--"; + } + return value.toFixed(digits); +} + +export function ProgressSection(): ReactElement { + const runtime = useTrainingRuntimeStore( + useShallow((state) => ({ + phase: state.phase, + message: state.message, + error: state.error, + currentStep: state.currentStep, + totalSteps: state.totalSteps, + currentEpoch: state.currentEpoch, + currentLoss: state.currentLoss, + currentLearningRate: state.currentLearningRate, + currentGradNorm: state.currentGradNorm, + progressPercent: state.progressPercent, + elapsedSeconds: state.elapsedSeconds, + etaSeconds: state.etaSeconds, + currentNumTokens: state.currentNumTokens, + isTrainingRunning: state.isTrainingRunning, + })), + ); + + const config = useTrainingConfigStore( + useShallow((state) => ({ + selectedModel: state.selectedModel, + trainingMethod: state.trainingMethod, + epochs: state.epochs, + batchSize: state.batchSize, + learningRate: state.learningRate, + maxSteps: state.maxSteps, + contextLength: state.contextLength, + warmupSteps: state.warmupSteps, + loraRank: state.loraRank, + loraAlpha: state.loraAlpha, + loraDropout: state.loraDropout, + loraVariant: state.loraVariant, + })), + ); + + const { stopTrainingRun } = useTrainingActions(); + const localStartAtRef = useRef(null); + const [, setLocalTick] = useState(0); + + const pct = + runtime.totalSteps > 0 + ? Math.min( + 100, + Math.max( + 0, + Math.round((runtime.currentStep / runtime.totalSteps) * 100), + ), ) - : 0; - const fmtTime = (s: number) => { - const m = Math.floor(s / 60); - const sec = s % 60; - return `${m}m ${sec}s`; - }; + : Math.round(runtime.progressPercent); - const statusColors = { - training: - "bg-emerald-100 text-emerald-700 dark:bg-emerald-900 dark:text-emerald-300", - warmup: "bg-amber-100 text-amber-700 dark:bg-amber-900 dark:text-amber-300", - saving: - "bg-emerald-100 text-emerald-700 dark:bg-emerald-900 dark:text-emerald-300", - }; - const statusLabels = { - training: "Training", - warmup: "Warming up", - saving: "Saving checkpoint", - }; + 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]); - const modelName = store.selectedModel ?? "—"; + 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)) + : null; + const eta = runtime.etaSeconds ?? derivedEta; + + const stepsPerSecond = + elapsed != null && elapsed > 0 + ? runtime.currentStep / elapsed + : null; const configItems = [ { section: "Hyperparams", rows: [ - ["Epochs", store.epochs], - ["Batch size", store.batchSize], - ["Learning rate", store.learningRate], - ["Max steps", store.maxSteps], - ["Context length", store.contextLength], - ["Warmup steps", store.warmupSteps], + ["Epochs", config.epochs], + ["Batch size", config.batchSize], + ["Learning rate", config.learningRate], + ["Max steps", config.maxSteps], + ["Context length", config.contextLength], + ["Warmup steps", config.warmupSteps], ], }, - ...(store.trainingMethod !== "full" + ...(config.trainingMethod !== "full" ? [ { section: "LoRA", rows: [ - ["Rank", store.loraRank], - ["Alpha", store.loraAlpha], - ["Dropout", store.loraDropout], - ["Variant", store.loraVariant], + ["Rank", config.loraRank], + ["Alpha", config.loraAlpha], + ["Dropout", config.loraDropout], + ["Variant", config.loraVariant], ], }, ] @@ -86,7 +185,7 @@ export function ProgressSection(): ReactElement | null { } title="Training Progress" - description="Live training metrics" + description={runtime.message || "Live training metrics"} accent="emerald" className="shadow-border ring-1 ring-border" headerAction={ @@ -130,7 +229,8 @@ export function ProgressSection(): ReactElement | null { variant="destructive" size="sm" className="h-7 cursor-pointer px-3 text-xs" - onClick={() => store.setIsTraining(false)} + onClick={() => void stopTrainingRun()} + disabled={!runtime.isTrainingRunning} > Stop @@ -138,24 +238,22 @@ export function ProgressSection(): ReactElement | null { } >
- {/* Left: Progress */}
- {statusLabels[metrics.status]} + {phaseLabel[runtime.phase]} - Epoch {metrics.currentEpoch.toFixed(2)} / {metrics.totalEpochs} + Epoch {runtime.currentEpoch.toFixed(2)}
- {/* Progress bar */}
- Step {metrics.currentStep} / {metrics.totalSteps} + Step {runtime.currentStep} / {runtime.totalSteps || "--"} {pct}%
@@ -167,51 +265,59 @@ export function ProgressSection(): ReactElement | null {
- {/* Metrics */} -
+ {runtime.error && ( +

{runtime.error}

+ )} + +

Loss

- {metrics.currentLoss.toFixed(4)} + {runtime.currentLoss.toFixed(4)}

LR

- {metrics.currentLR.toExponential(2)} + {runtime.currentLearningRate.toExponential(2)}

Grad Norm

- {metrics.gradNorm.toFixed(3)} + {formatNumber(runtime.currentGradNorm, 3)}

Model

- {modelName} + {config.selectedModel ?? "--"}

Method

-

{store.trainingMethod}

+

+ {config.trainingMethod.toUpperCase()} +

- {/* Timings */} -
- Elapsed: {fmtTime(metrics.elapsed)} - ETA: {fmtTime(etaSec)} - {metrics.samplesPerSecond} samples/s +
+ Elapsed: {formatDuration(elapsed)} + ETA: {formatDuration(eta)} + + {stepsPerSecond == null + ? "-- steps/s" + : `${stepsPerSecond.toFixed(2)} steps/s`} + + {runtime.currentNumTokens != null && ( + Tokens: {runtime.currentNumTokens} + )}
- {/* Right: GPU */}
-

- GPU Monitor -

+

GPU Monitor

} - value={`${metrics.gpuUtil}%`} - pct={metrics.gpuUtil} + value="--" + pct={0} /> - } - value={`${metrics.gpuTemp}°C`} - pct={metrics.gpuTemp} + icon={} + value="--" + pct={0} max={100} /> } - value={`${metrics.gpuVramUsed.toFixed(1)} / ${metrics.gpuVramTotal}GB`} - pct={(metrics.gpuVramUsed / metrics.gpuVramTotal) * 100} + value="--" + pct={0} /> } - value={`${metrics.gpuPower}W`} - pct={(metrics.gpuPower / 350) * 100} + value="--" + pct={0} />
@@ -272,6 +376,7 @@ function GpuStat({ } else if (clamped < 95) { barColor = "bg-amber-500"; } + return (
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/studio/training-start-overlay.tsx b/studio/frontend/src/features/studio/training-start-overlay.tsx new file mode 100644 index 0000000000..b6c413b400 --- /dev/null +++ b/studio/frontend/src/features/studio/training-start-overlay.tsx @@ -0,0 +1,58 @@ +import { + AnimatedSpan, + Terminal, + TypingAnimation, +} from "@/components/ui/terminal" +import type { ReactElement } from "react" + +type TrainingStartOverlayProps = { + message: string + currentStep: number +} + +export function TrainingStartOverlay({ + message, + currentStep, +}: TrainingStartOverlayProps): ReactElement { + return ( +
+
+ Unsloth mascot + + + {"> unsloth training starts..."} + + + 
{`==((====))==
+\\\\   /|
+O^O/ \\_/ \\
+\\        /
+ "-____-"`}
+ + + {"> Preparing model and dataset..."} + + + {"> We are getting everything ready for your run..."} + + + {"> Did you know, Mugi is actually short for \"Mugiwara\" xd"} + + + {`> ${message || "starting training..."} | waiting for first step... (${currentStep})`} + + +
+
+ ) +} diff --git a/studio/frontend/src/features/studio/training-view.tsx b/studio/frontend/src/features/studio/training-view.tsx index f1120b0ac2..e10faa6735 100644 --- a/studio/frontend/src/features/studio/training-view.tsx +++ b/studio/frontend/src/features/studio/training-view.tsx @@ -1,164 +1,47 @@ -import { useWizardStore } from "@/stores/training"; -import type { TrainingMetrics } from "@/types/training"; -import { type ReactElement, useEffect, useRef } from "react"; +import { cn } from "@/lib/utils"; +import { useTrainingRuntimeStore } from "@/features/training"; +import type { ReactElement } from "react"; +import { useShallow } from "zustand/react/shallow"; import { ChartsSection } from "./sections/charts-section"; import { ProgressSection } from "./sections/progress-section"; - -function createInitialMetrics( - totalSteps: number, - totalEpochs: number, - lr: number, -): TrainingMetrics { - return { - currentStep: 0, - totalSteps, - currentEpoch: 0, - totalEpochs, - currentLoss: 2.5, - currentLR: lr * 0.1, - gradNorm: 0, - samplesPerSecond: 0, - lossHistory: [], - lrHistory: [], - gradNormHistory: [], - gpuUtil: 0, - gpuTemp: 45, - gpuVramUsed: 0, - gpuVramTotal: 24, - gpuPower: 50, - elapsed: 0, - status: "warmup", - }; -} +import { TrainingStartOverlay } from "./training-start-overlay"; export function TrainingView(): ReactElement { - const { maxSteps, epochs, learningRate, warmupSteps, setTrainingMetrics } = - useWizardStore(); - const metricsRef = useRef | null>(null); - const chartsRef = useRef | null>(null); + const runtime = useTrainingRuntimeStore( + useShallow((state) => ({ + phase: state.phase, + message: state.message, + currentStep: state.currentStep, + firstStepReceived: state.firstStepReceived, + isStarting: state.isStarting, + })), + ); - useEffect(() => { - const totalSteps = maxSteps || 500; - const totalEpochs = epochs || 3; - const peakLR = learningRate; - const warmup = warmupSteps || 20; - - setTrainingMetrics(createInitialMetrics(totalSteps, totalEpochs, peakLR)); - - let step = 0; - let elapsed = 0; - - const computeStep = () => { - step++; - if (step > totalSteps) { - return null; - } - elapsed++; - - let lr: number; - if (step < warmup) { - lr = peakLR * (step / warmup); - } else { - const progress = (step - warmup) / (totalSteps - warmup); - lr = peakLR * 0.5 * (1 + Math.cos(Math.PI * progress)); - } - - const baseLoss = 2.5 * Math.exp((-3 * step) / totalSteps) + 0.3; - const noise = (Math.random() - 0.5) * 0.08; - const loss = Math.max(0.1, baseLoss + noise); - const status = - step < warmup ? "warmup" : step % 100 === 0 ? "saving" : "training"; - const gradNorm = +( - 1.2 * Math.exp(-step / totalSteps) + - 0.1 + - (Math.random() - 0.5) * 0.05 - ).toFixed(3); - - return { - step, - elapsed, - lr, - loss: +loss.toFixed(4), - status: status as TrainingMetrics["status"], - gradNorm, - }; - }; - - // Top card values — update every 1s - metricsRef.current = setInterval(() => { - const s = computeStep(); - if (!s) { - if (metricsRef.current) { - clearInterval(metricsRef.current); - } - if (chartsRef.current) { - clearInterval(chartsRef.current); - } - return; - } - - const prev = useWizardStore.getState().trainingMetrics; - setTrainingMetrics({ - currentStep: s.step, - totalSteps, - currentEpoch: - Math.floor((s.step / totalSteps) * totalEpochs * 100) / 100, - totalEpochs, - currentLoss: s.loss, - currentLR: s.lr, - gradNorm: s.gradNorm, - samplesPerSecond: +(12 + (Math.random() - 0.5) * 2).toFixed(1), - lossHistory: prev?.lossHistory ?? [], - lrHistory: prev?.lrHistory ?? [], - gradNormHistory: prev?.gradNormHistory ?? [], - gpuUtil: Math.min(99, 85 + Math.round((Math.random() - 0.5) * 10)), - gpuTemp: Math.min(89, 68 + Math.round((Math.random() - 0.5) * 6)), - gpuVramUsed: +(18.2 + (Math.random() - 0.5) * 0.4).toFixed(1), - gpuVramTotal: 24, - gpuPower: Math.round(280 + (Math.random() - 0.5) * 30), - elapsed: s.elapsed, - status: s.status, - }); - }, 1000); - - // Chart history — update every 5s - chartsRef.current = setInterval(() => { - const prev = useWizardStore.getState().trainingMetrics; - if (!prev || prev.currentStep === 0) { - return; - } - - setTrainingMetrics({ - ...prev, - lossHistory: [ - ...prev.lossHistory, - { step: prev.currentStep, loss: prev.currentLoss }, - ], - lrHistory: [ - ...prev.lrHistory, - { step: prev.currentStep, lr: prev.currentLR }, - ], - gradNormHistory: [ - ...prev.gradNormHistory, - { step: prev.currentStep, gradNorm: prev.gradNorm }, - ], - }); - }, 5000); - - return () => { - if (metricsRef.current) { - clearInterval(metricsRef.current); - } - if (chartsRef.current) { - clearInterval(chartsRef.current); - } - }; - }, [epochs, learningRate, maxSteps, setTrainingMetrics, warmupSteps]); + const isPreparingPhase = + runtime.phase === "loading_model" || + runtime.phase === "loading_dataset" || + runtime.phase === "configuring"; + const isWaitingForFirstStep = + runtime.phase === "training" && !runtime.firstStepReceived; + const showOverlay = + runtime.isStarting || + isPreparingPhase || + (isWaitingForFirstStep && runtime.currentStep <= 0); return ( -
- - +
+
+ + +
+ {showOverlay ? ( + + ) : null}
); } 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..a9d6d37f42 --- /dev/null +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -0,0 +1,105 @@ +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", + partialize: (state) => { + const { modelType, ...rest } = state; + return rest; + }, + }, + ), +); 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/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index e9d3b054d0..4745b4dc7b 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -23,6 +23,28 @@ const EXCLUDED_TAGS = new Set([ "ctranslate2", ]); +function withPopularitySort( + input: Parameters[0], + init?: Parameters[1], +): ReturnType { + const rawUrl = + typeof input === "string" + ? input + : input instanceof URL + ? input.toString() + : input.url; + const url = new URL(rawUrl); + + if (!url.searchParams.has("sort")) { + url.searchParams.set("sort", "downloads"); + } + if (!url.searchParams.has("direction")) { + url.searchParams.set("direction", "-1"); + } + + return fetch(url, init); +} + function mapModel(raw: unknown): HfModelResult | null { const m = raw as { name: string; @@ -53,10 +75,10 @@ export function useHfModelSearch( listModels({ search: { ...(query.trim() ? { query } : { owner: "unsloth" }), - tags: ["transformers"], ...(task ? { task } : {}), }, additionalFields: ["safetensors", "tags"], + fetch: withPopularitySort, ...(accessToken ? { credentials: { accessToken } } : {}), }) as AsyncGenerator, [query, task, accessToken], 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/], }, }, });