Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 109 additions & 0 deletions studio/backend/tests/test_training_config_popover_source.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Source-level regression guards for the Training Config popover data source
(#6853).

The live Training Progress popover used to read the editable form store
(useTrainingConfigStore) while a run was active, so it showed stale/static
values whenever the user touched the form after starting the run; only the
History view read the run's saved config snapshot. These guards pin the fixed
wiring: both views feed ProgressSection a config override mapped from
GET /api/train/runs/{id}, and ProgressSection prefers that override whenever
one is present -- not only for historical views.
"""

from __future__ import annotations

from pathlib import Path

_STUDIO_FRONTEND = Path(__file__).resolve().parents[2] / "frontend" / "src" / "features" / "studio"


def _read(rel: str) -> str:
return (_STUDIO_FRONTEND / rel).read_text(encoding = "utf-8")


def test_progress_section_prefers_override_over_form_store():
src = _read("sections/progress-section.tsx")
# Fields key on the override's presence, not isHistorical: a live view passing
# an override wins over the store; without one, live keeps the store while
# History shows blanks rather than unrelated live form values.
assert "const cfg = configOverride ?? (isHistorical ? undefined : config)" in src
assert "const cfgEpochs = cfg?.epochs" in src
assert "isHistorical ? configOverride?.epochs" not in src


def test_live_view_fetches_the_active_run_config():
src = _read("live-training-view.tsx")
# Live view resolves the run's saved config snapshot by job id...
assert "getTrainingRun(" in src
assert "mapRunConfigToOverride(" in src
# ...and hands it to the popover.
assert "configOverride={runConfigOverride}" in src


def test_live_view_fetches_as_soon_as_the_job_id_exists():
# start_training() inserts the run row BEFORE the pump consumes any event, so
# the saved config is available during configuring/loading/downloading. The
# job id is therefore the whole readiness condition: gating on a first step
# or a terminal phase would show the wrong config for the entire pre-step
# window of a long load, or for a run adopted from another client.
src = _read("live-training-view.tsx")
assert "if (!runtime.jobId) {" in src
assert "[runtime.jobId, fetchedRunConfig, fetchAttempt]" in src
# No step/phase readiness gate may creep back in.
assert "runRowReady" not in src


def test_live_view_retries_the_transient_row_miss():
# start_training() creates the row before the pump, but a lookup racing that
# commit can still 404. Nothing else in the effect deps changes on failure, so
# the retry must be explicit and bounded, else a genuinely absent row would
# poll forever instead of falling back to the form store.
src = _read("live-training-view.tsx")
assert "RUN_CONFIG_FETCH_RETRIES" in src
assert "RUN_CONFIG_FETCH_RETRY_MS" in src
assert "setFetchAttempt(" in src
assert "attempts >= RUN_CONFIG_FETCH_RETRIES" in src
# The budget is keyed by job so a new run always starts fresh.
assert "fetchAttempt?.jobId === jobId ? fetchAttempt.count : 0" in src
# The pending retry must be cancelled with the effect.
assert "clearTimeout(retryTimer)" in src


def test_live_view_prefers_saved_training_method():
# The method label / LoRA-row visibility must come from the run snapshot,
# not the editable form (which may have changed since the run started).
src = _read("live-training-view.tsx")
assert "runConfigOverride?.trainingMethod ?? config.trainingMethod" in src


def test_history_view_uses_the_shared_mapper():
src = _read("historical-training-view.tsx")
# Shared mapper, not a re-inlined field-by-field copy that could drift.
assert "mapRunConfigToOverride(detail.config)" in src
assert "num_epochs" not in src


def test_shared_mapper_matches_backend_config_keys():
src = _read("sections/run-config-override.ts")
# The mapper reads the run config JSON the backend snapshots at job start;
# keep the key set pinned so a silent rename breaks loudly here.
for key in (
"training_type",
"load_in_4bit",
"num_epochs",
"batch_size",
"learning_rate",
"max_steps",
"max_seq_length",
"warmup_steps",
"optim",
"lora_r",
"lora_alpha",
"lora_dropout",
"use_rslora",
"use_loftq",
):
assert key in src, f"run-config mapper lost backend key {key}"
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import { parseBackendTrainingMethod } from "@/features/training/lib/training-met
import { type ReactElement, useEffect, useState } from "react";
import { ChartsSection } from "./sections/charts-section";
import { ProgressSection } from "./sections/progress-section";
import { mapRunConfigToOverride } from "./sections/run-config-override";
import { translate, useT } from "@/i18n";

type StudioT = ReturnType<typeof useT>;
Expand Down Expand Up @@ -147,25 +148,7 @@ export function HistoricalTrainingView({
}

const viewData = mapToViewData(detail, t);
const configOverride = detail.config
? {
epochs: detail.config.num_epochs as number | undefined,
batchSize: detail.config.batch_size as number | undefined,
learningRate: detail.config.learning_rate as string | undefined,
maxSteps: detail.config.max_steps as number | undefined,
contextLength: detail.config.max_seq_length as number | undefined,
warmupSteps: detail.config.warmup_steps as number | undefined,
optimizerType: detail.config.optim as string | undefined,
loraRank: detail.config.lora_r as number | undefined,
loraAlpha: detail.config.lora_alpha as number | undefined,
loraDropout: detail.config.lora_dropout as number | undefined,
loraVariant: detail.config.use_rslora
? "rslora"
: detail.config.use_loftq
? "loftq"
: "lora",
}
: undefined;
const configOverride = mapRunConfigToOverride(detail.config);

return (
<div className="flex flex-col gap-6">
Expand Down
91 changes: 88 additions & 3 deletions studio/frontend/src/features/studio/live-training-view.tsx
Original file line number Diff line number Diff line change
@@ -1,18 +1,42 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

import { cn } from "@/lib/utils";
import {
getTrainingRun,
useTrainingConfigStore,
useTrainingRuntimeStore,
} from "@/features/training";
import type { TrainingViewData } from "@/features/training";
import { cn } from "@/lib/utils";
import type { ReactElement } from "react";
import { useEffect, useState } from "react";
import { useShallow } from "zustand/react/shallow";
import { ChartsSection } from "./sections/charts-section";
import { ProgressSection } from "./sections/progress-section";
import {
type RunConfigOverride,
mapRunConfigToOverride,
} from "./sections/run-config-override";
import { TrainingStartOverlay } from "./training-start-overlay";

/** Retry budget for the run-config lookup. The row is inserted at
* start_training(), but a lookup issued in the same instant can still miss it;
* a few short retries cover that without polling a genuinely absent row. */
const RUN_CONFIG_FETCH_RETRIES = 5;
const RUN_CONFIG_FETCH_RETRY_MS = 1000;

/** The fetched run config only applies while it belongs to the active job;
* a stale record from a previous run falls back to the form store. */
function activeRunOverride(
fetched: { jobId: string; override: RunConfigOverride | undefined } | null,
jobId: string | null,
): RunConfigOverride | undefined {
if (fetched === null || fetched.jobId !== jobId) {
return undefined;
}
return fetched.override;
}

export function LiveTrainingView(): ReactElement {
const runtime = useTrainingRuntimeStore(
useShallow((state) => ({
Expand Down Expand Up @@ -52,6 +76,59 @@ export function LiveTrainingView(): ReactElement {
})),
);

// Show the ACTIVE run's saved config, not the editable form store the user may
// have changed since starting (#6853). start_training() commits the run row
// before the pump, so the job id alone gates the fetch; the bounded retry below
// covers the narrow uncommitted window, and until it loads ProgressSection falls
// back to the form store. The result is keyed by job id and filtered at render.
const [fetchedRunConfig, setFetchedRunConfig] = useState<{
jobId: string;
override: RunConfigOverride | undefined;
} | null>(null);
// Retry budget for the transient 404 below, keyed by job so a new run always
// starts with a fresh budget.
const [fetchAttempt, setFetchAttempt] = useState<{
jobId: string;
count: number;
} | null>(null);
useEffect(() => {
if (!runtime.jobId) {
return;
}
const jobId = runtime.jobId;
if (fetchedRunConfig !== null && fetchedRunConfig.jobId === jobId) {
return; // already resolved for this job
}
const attempts = fetchAttempt?.jobId === jobId ? fetchAttempt.count : 0;
const controller = new AbortController();
let retryTimer: ReturnType<typeof setTimeout> | undefined;
getTrainingRun(jobId, controller.signal)
.then((detail) => {
setFetchedRunConfig({
jobId,
override: mapRunConfigToOverride(detail.config),
Comment thread
hakanbaysal marked this conversation as resolved.
});
})
.catch(() => {
// A lookup racing the row commit can miss transiently; nothing else in
// the deps changes on failure, so retry explicitly. Bounded so a genuinely
// absent row falls back to the form store instead of polling forever.
if (controller.signal.aborted || attempts >= RUN_CONFIG_FETCH_RETRIES) {
return;
}
retryTimer = setTimeout(() => {
setFetchAttempt({ jobId, count: attempts + 1 });
}, RUN_CONFIG_FETCH_RETRY_MS);
});
Comment thread
hakanbaysal marked this conversation as resolved.
return () => {
controller.abort();
if (retryTimer !== undefined) {
clearTimeout(retryTimer);
}
};
}, [runtime.jobId, fetchedRunConfig, fetchAttempt]);
const runConfigOverride = activeRunOverride(fetchedRunConfig, runtime.jobId);

const activeProjectName =
runtime.startProjectName !== null
? runtime.startProjectName.trim() || null
Expand All @@ -76,7 +153,11 @@ export function LiveTrainingView(): ReactElement {
isTrainingRunning: runtime.isTrainingRunning,
modelName: runtime.startModelName ?? config.selectedModel ?? "",
projectName: activeProjectName,
trainingMethod: config.trainingMethod ?? "",
// Prefer the saved run's method: the form may have been edited (e.g. LoRA
// -> Full) after the run started, which would relabel the run and hide its
// saved LoRA rows in the popover.
trainingMethod:
runConfigOverride?.trainingMethod ?? config.trainingMethod ?? "",
lossHistory: runtime.lossHistory,
lrHistory: runtime.lrHistory,
gradNormHistory: runtime.gradNormHistory,
Expand Down Expand Up @@ -105,7 +186,11 @@ export function LiveTrainingView(): ReactElement {
)}
>
<div data-tour="studio-training-progress">
<ProgressSection key={runtime.jobId ?? "no-job"} data={viewData} />
<ProgressSection
key={runtime.jobId ?? "no-job"}
data={viewData}
configOverride={runConfigOverride}
/>
</div>
<ChartsSection
currentStep={viewData.currentStep}
Expand Down
40 changes: 16 additions & 24 deletions studio/frontend/src/features/studio/sections/progress-section.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import {
} from "@/features/training";
import { getTrainingMethodLabel } from "@/features/training/lib/training-methods";
import type { TrainingViewData } from "@/features/training";
import type { RunConfigOverride } from "./run-config-override";
import { useGpuUtilization } from "@/hooks";
import type { GpuUtilization } from "@/hooks/use-gpu-utilization";
import { cn } from "@/lib/utils";
Expand Down Expand Up @@ -81,19 +82,7 @@ function configRow(
interface ProgressSectionProps {
data: TrainingViewData;
isHistorical?: boolean;
configOverride?: {
epochs?: number;
batchSize?: number;
learningRate?: string;
maxSteps?: number;
contextLength?: number;
warmupSteps?: number;
optimizerType?: string;
loraRank?: number;
loraAlpha?: number;
loraDropout?: number;
loraVariant?: string;
};
configOverride?: RunConfigOverride;
}

export function ProgressSection({
Expand Down Expand Up @@ -183,17 +172,20 @@ export function ProgressSection({
? data.currentGradNorm
: (lastValue(data.gradNormHistory) ?? data.currentGradNorm);

const cfgEpochs = isHistorical ? configOverride?.epochs : config.epochs;
const cfgBatchSize = isHistorical ? configOverride?.batchSize : config.batchSize;
const cfgLearningRate = isHistorical ? configOverride?.learningRate : config.learningRate;
const cfgMaxSteps = isHistorical ? configOverride?.maxSteps : config.maxSteps;
const cfgContextLength = isHistorical ? configOverride?.contextLength : config.contextLength;
const cfgWarmupSteps = isHistorical ? configOverride?.warmupSteps : config.warmupSteps;
const cfgOptimizerType = isHistorical ? configOverride?.optimizerType : config.optimizerType;
const cfgLoraRank = isHistorical ? configOverride?.loraRank : config.loraRank;
const cfgLoraAlpha = isHistorical ? configOverride?.loraAlpha : config.loraAlpha;
const cfgLoraDropout = isHistorical ? configOverride?.loraDropout : config.loraDropout;
const cfgLoraVariant = isHistorical ? configOverride?.loraVariant : config.loraVariant;
// Prefer the run's saved snapshot when present (#6853). Live falls back to the
// editable form store until it loads; History shows blanks, never live form values.
const cfg = configOverride ?? (isHistorical ? undefined : config);
const cfgEpochs = cfg?.epochs;
const cfgBatchSize = cfg?.batchSize;
const cfgLearningRate = cfg?.learningRate;
const cfgMaxSteps = cfg?.maxSteps;
const cfgContextLength = cfg?.contextLength;
const cfgWarmupSteps = cfg?.warmupSteps;
const cfgOptimizerType = cfg?.optimizerType;
const cfgLoraRank = cfg?.loraRank;
const cfgLoraAlpha = cfg?.loraAlpha;
const cfgLoraDropout = cfg?.loraDropout;
const cfgLoraVariant = cfg?.loraVariant;

const optimizerLabel =
OPTIMIZER_OPTIONS.find((o) => o.value === cfgOptimizerType)?.label ??
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

import { parseBackendTrainingMethod } from "@/features/training";

/** Shape of the Training Config popover's data when it is driven by a saved
* run snapshot instead of the editable form store. */
export interface RunConfigOverride {
trainingMethod?: string;
epochs?: number;
batchSize?: number;
learningRate?: string;
maxSteps?: number;
contextLength?: number;
warmupSteps?: number;
optimizerType?: string;
loraRank?: number;
loraAlpha?: number;
loraDropout?: number;
loraVariant?: string;
}

/** Map a saved run's config (GET /api/train/runs/{id} `detail.config`) into the
* Training Config popover's override shape. Shared by the History view and the
* live Current Run view so both read the same authoritative run snapshot
* instead of the editable form store (#6853). */
export function mapRunConfigToOverride(
config: Record<string, unknown> | null | undefined,
): RunConfigOverride | undefined {
if (!config) {
return undefined;
}
return {
trainingMethod: parseBackendTrainingMethod(
config.training_type,
config.load_in_4bit,
),
epochs: config.num_epochs as number | undefined,
batchSize: config.batch_size as number | undefined,
learningRate: config.learning_rate as string | undefined,
maxSteps: config.max_steps as number | undefined,
contextLength: config.max_seq_length as number | undefined,
warmupSteps: config.warmup_steps as number | undefined,
optimizerType: config.optim as string | undefined,
loraRank: config.lora_r as number | undefined,
loraAlpha: config.lora_alpha as number | undefined,
loraDropout: config.lora_dropout as number | undefined,
loraVariant: config.use_rslora
? "rslora"
: config.use_loftq
? "loftq"
: "lora",
};
}
1 change: 1 addition & 0 deletions studio/frontend/src/features/training/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ export {
getTrainingRunDisplayTitle,
getTrainingRunModelSubtitle,
} from "./lib/run-display";
export { parseBackendTrainingMethod } from "./lib/training-methods";
export { useTrainingHistorySidebarItems } from "./hooks/use-training-history-sidebar";
export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle";
export { useTrainingCompletionWatch } from "./hooks/use-training-completion-watch";
Expand Down
Loading