diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index e0936a1dd4..bdbedbf5a0 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -651,6 +651,17 @@ def _config_from_dict(config: dict) -> DiffusionLoraConfig: for k, v in config.items(): if k in valid: kwargs[k] = v + # Epoch-mode payloads from the generic Studio UI carry max_steps: 0 as the "use epochs" + # sentinel, which the max_steps -> train_steps alias copies as train_steps: 0. Since + # normalized() rejects train_steps < 1 before resolve_train_steps() can apply num_epochs, + # drop a falsy/0 train_steps when num_epochs > 0 so the dataclass default stands in until + # epoch resolution replaces it. + try: + _num_epochs = int(kwargs.get("num_epochs") or 0) + except (TypeError, ValueError): + _num_epochs = 0 + if _num_epochs > 0 and not kwargs.get("train_steps"): + kwargs.pop("train_steps", None) if kwargs.get("lora_target_modules"): kwargs["lora_target_modules"] = tuple(kwargs["lora_target_modules"]) if "gradient_checkpointing" in kwargs: diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 67f64f14ad..23c0cd2bbd 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -102,6 +102,14 @@ def list_diffusion_runs(limit: int = 20) -> list[dict]: rec = json.loads(p.read_text()) except Exception: # noqa: BLE001 -- a corrupt record never breaks the listing continue + # A valid-JSON file with the wrong shape (an old or hand-edited record that is not a + # dict, or is missing the required string job_id / status) would later blow up the + # route's DiffusionTrainingRunSummary(**r); skip it here so one bad record can never + # take down the whole Previous runs panel. + if not isinstance(rec, dict): + continue + if not (isinstance(rec.get("job_id"), str) and isinstance(rec.get("status"), str)): + continue rec.pop("metric_history", None) rec.pop("config", None) out.append(rec) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 1a85089439..f6f2a0509a 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -80,6 +80,7 @@ from models.training import ( ) from models.responses import TrainingStopResponse, TrainingMetricsResponse from pydantic import BaseModel as PydanticBaseModel +from pydantic import ValidationError class TrainingStopRequest(PydanticBaseModel): @@ -1285,9 +1286,17 @@ async def list_diffusion_training_runs( """Previous diffusion training runs (terminal), newest first, from the persisted per-run records. Summaries only; fetch one run for its config + metric logs.""" from core.training.diffusion_training_service import list_diffusion_runs - return DiffusionTrainingRunsResponse( - runs = [DiffusionTrainingRunSummary(**r) for r in list_diffusion_runs(limit = limit)] - ) + + summaries: list[DiffusionTrainingRunSummary] = [] + for r in list_diffusion_runs(limit = limit): + # list_diffusion_runs already skips non-dict / missing-id records, but a record with + # a wrong-typed field (e.g. a non-numeric avg_loss) would still raise here; catch it + # per record so one bad file never breaks the whole Previous runs panel. + try: + summaries.append(DiffusionTrainingRunSummary(**r)) + except ValidationError: + continue + return DiffusionTrainingRunsResponse(runs = summaries) @router.get("/diffusion/runs/{job_id}", response_model = DiffusionTrainingRunDetail) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index e4bd456f09..84446cc15c 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -381,6 +381,44 @@ def test_request_model_num_epochs_bounds(): DiffusionTrainingStartRequest(**base, num_epochs = bad) +def test_config_from_dict_epoch_mode_drops_max_steps_sentinel(): + # The generic Studio epoch-mode payload sends max_steps: 0 as the "use epochs" sentinel. + # The max_steps -> train_steps alias would copy that 0 and normalized() would reject + # train_steps < 1 before epochs are resolved; _config_from_dict must drop the falsy + # value so the default train_steps stands in until resolve_train_steps applies num_epochs. + from core.training.diffusion_train_common import DiffusionLoraConfig, _config_from_dict + + cfg = _config_from_dict( + { + "base_model": "stabilityai/stable-diffusion-xl-base-1.0", + "data_dir": "d", + "output_dir": "o", + "max_steps": 0, + "num_epochs": 2, + } + ) + # 0 was dropped: the dataclass default train_steps stands in and num_epochs carries over. + assert cfg.train_steps == DiffusionLoraConfig.train_steps + assert cfg.num_epochs == 2 + # normalized() no longer raises on the epoch-mode payload. + norm = cfg.normalized() + assert norm.num_epochs == 2 + + # An explicit non-zero max_steps in epochs mode is still honored (only the 0 sentinel is + # dropped), and a plain steps payload (no num_epochs) keeps max_steps: 0 -> train_steps 0 + # so normalized() surfaces the invalid value as before. + cfg_explicit = _config_from_dict( + { + "base_model": "stabilityai/stable-diffusion-xl-base-1.0", + "data_dir": "d", + "output_dir": "o", + "max_steps": 25, + "num_epochs": 2, + } + ) + assert cfg_explicit.train_steps == 25 + + def test_route_start_rejects_uncontained_paths(client): # An absolute path outside the Studio dataset roots is a 400, not silently accepted. r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "/etc"}) @@ -804,3 +842,49 @@ def test_runs_endpoints_list_and_detail(client, _isolated_runs_dir): # Unknown and malformed ids 404 (malformed also covers path traversal). assert client.get(f"/api/train/diffusion/runs/{'c' * 32}").status_code == 404 assert client.get("/api/train/diffusion/runs/not-a-job-id").status_code == 404 + + +def test_list_diffusion_runs_skips_wrong_shape_records(_isolated_runs_dir): + # A valid-JSON file with the wrong shape (non-dict, or missing the required string + # job_id / status) must be skipped by list_diffusion_runs so it never reaches the route's + # DiffusionTrainingRunSummary(**r) and takes down the whole Previous runs panel. + import json + + from core.training.diffusion_training_service import list_diffusion_runs + + good = {"job_id": "a" * 32, "status": "completed", "adapter": "good", "saved": True} + (_isolated_runs_dir / "good.json").write_text(json.dumps(good)) + # A JSON list (not a dict). + (_isolated_runs_dir / "not_a_dict.json").write_text(json.dumps([1, 2, 3])) + # A dict missing the required job_id / status. + (_isolated_runs_dir / "no_ids.json").write_text(json.dumps({"adapter": "orphan"})) + # A dict whose job_id / status are the wrong type. + (_isolated_runs_dir / "bad_types.json").write_text( + json.dumps({"job_id": 123, "status": None, "adapter": "typed"}) + ) + + runs = list_diffusion_runs() + adapters = [r.get("adapter") for r in runs] + assert adapters == ["good"] # only the well-shaped record survives + + +def test_runs_route_tolerates_bad_field_record(client, _isolated_runs_dir): + # A record that passes the service's shape check but has a wrong-typed field (a + # non-numeric avg_loss) would raise pydantic ValidationError in the route; the route must + # catch it per record so one bad file never breaks the panel and the good runs still list. + import json + + good = {"job_id": "a" * 32, "status": "completed", "adapter": "good", "saved": True} + bad = { + "job_id": "b" * 32, + "status": "completed", + "adapter": "bad", + "avg_loss": "not-a-number", # str where the summary expects Optional[float] + } + (_isolated_runs_dir / f"{good['job_id']}.json").write_text(json.dumps(good)) + (_isolated_runs_dir / f"{bad['job_id']}.json").write_text(json.dumps(bad)) + + r = client.get("/api/train/diffusion/runs") + assert r.status_code == 200, r.text + adapters = [x["adapter"] for x in r.json()["runs"]] + assert adapters == ["good"] # the bad-field record was skipped, the good one remained diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx index 813a6c1fd2..463ca86dc4 100644 --- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx +++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx @@ -520,13 +520,25 @@ export function DiffusionTrainPanel({ if (!active) return; if (status?.status === "running") return; let cancelled = false; - listDiffusionTrainingRuns() - .then((r) => { - if (!cancelled) setPrevRuns(r.runs); - }) - .catch(() => {}); + const refetch = () => { + listDiffusionTrainingRuns() + .then((r) => { + if (!cancelled) setPrevRuns(r.runs); + }) + .catch(() => {}); + }; + refetch(); + // The service exposes a terminal status before the pump has necessarily finished + // writing the run's JSON record, so the one-shot refetch above can win that race and + // miss the just-finished run. A short delayed second refetch after a terminal + // transition lets the record land so the newest run reliably appears. + let delayed: ReturnType | undefined; + if (status?.status === "completed" || status?.status === "stopped" || status?.status === "error") { + delayed = setTimeout(refetch, 1500); + } return () => { cancelled = true; + if (delayed !== undefined) clearTimeout(delayed); }; }, [active, status?.status]); @@ -745,7 +757,10 @@ export function DiffusionTrainPanel({ value={value} onChange={(e) => { settingsDirty.current = true; - set(Number(e.target.value) || fallback); + // Only fall back when the input parses to NaN (empty/invalid); a real 0 is a + // legal value for zero-legal fields (Seed, LR warmup steps) and must be kept. + const parsed = Number(e.target.value); + set(Number.isNaN(parsed) ? fallback : parsed); }} className="h-8 text-xs" />