5eef2f4003
Gallery clear/delete now scope to Studio-owned files: image_gallery and video_gallery skip PNGs / MP4s without a readable recipe (a hand-dropped or orphan file the listing already hides), so clear() and a guessed-id delete no longer destroy files the gallery never surfaced. Remote ControlNets now force use_safetensors: a bare owner/name reaches from_pretrained without the base trust gate, and the Hub scan fails open when unavailable, so requiring safetensors closes the pickle deserialization vector. POSIX uninstall now stops resident sd-server / sd-cli under an owned sd.cpp root before removing the tree (marker-gated), mirroring the Windows stop-before-delete scan; a live native server no longer survives unlinking its binary. Diffusion dataset containment: the training-start read path and the discovery picker route bare names through the protected resolver, so a symlinked dataset is rejected / not advertised like the caption/delete routes already do. Uploads gain the inference decode guard (oversized real images 400 before OOMing the trainer) and dataset upload/caption/delete/import are blocked with 409 while a diffusion run is active. JSONL readers (trainer + routes) tolerate non-object JSON and invalid UTF-8 instead of raising AttributeError / 500. LoRA family compatibility is enforced in the shared resolver, not only the picker, so a direct API client cannot apply a mismatched-family adapter. GPU arbiter gains release_if so the image/video unload idle-check and release are atomic against a concurrent same-owner load's registration. Native batch recipes persist the base batch_seed and restore replays from it, so a native batch_index>0 image no longer advances its seed twice. FLUX.2-klein selects its sd.cpp text encoder by variant (4B -> Qwen3-4B, 9B -> Qwen3-8B) instead of the single family default.
1608 lines
64 KiB
Python
1608 lines
64 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Tests for the diffusion LoRA training service + routes.
|
|
|
|
The service's subprocess context and target are injected with in-thread fakes, so the
|
|
full start -> event-pump -> status -> complete path is exercised without real
|
|
multiprocessing or torch. The routes are hit with the FastAPI TestClient and a mocked
|
|
service, so wiring / validation / error mapping are covered without a GPU.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import queue as _queue
|
|
import threading
|
|
import time
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from auth.authentication import authenticated_via_api_key, get_current_subject
|
|
from core.training.diffusion_training_service import DiffusionTrainingService
|
|
from routes.training import router as training_router
|
|
|
|
|
|
# ── fake spawn context (runs the "process" target on a thread) ────────────────
|
|
class _FakeQueue:
|
|
def __init__(self) -> None:
|
|
self._q: _queue.Queue = _queue.Queue()
|
|
|
|
def put(self, x):
|
|
self._q.put(x)
|
|
|
|
def get(self, timeout = None):
|
|
return self._q.get(timeout = timeout) # raises queue.Empty on timeout
|
|
|
|
def get_nowait(self):
|
|
return self._q.get_nowait()
|
|
|
|
def empty(self):
|
|
return self._q.empty()
|
|
|
|
|
|
class _FakeProc:
|
|
def __init__(self, target, kwargs, daemon):
|
|
self._target = target
|
|
self._kwargs = kwargs
|
|
self._thread: threading.Thread | None = None
|
|
self.pid = 4321
|
|
|
|
def start(self):
|
|
self._thread = threading.Thread(target = self._target, kwargs = self._kwargs, daemon = True)
|
|
self._thread.start()
|
|
|
|
def is_alive(self):
|
|
return self._thread is not None and self._thread.is_alive()
|
|
|
|
|
|
class _FakeCtx:
|
|
def Queue(self):
|
|
return _FakeQueue()
|
|
|
|
def Process(self, target, kwargs, daemon):
|
|
return _FakeProc(target, kwargs, daemon)
|
|
|
|
|
|
@pytest.fixture(autouse = True)
|
|
def _isolated_runs_dir(monkeypatch, tmp_path):
|
|
"""Terminal service events persist a run record; point the runs dir at tmp so tests
|
|
never write into a real studio home. Yields the dir for the history tests."""
|
|
import core.training.diffusion_training_service as dts
|
|
|
|
d = tmp_path / "runs" / "diffusion"
|
|
d.mkdir(parents = True, exist_ok = True)
|
|
monkeypatch.setattr(dts, "_runs_dir", lambda: d)
|
|
yield d
|
|
|
|
|
|
def _happy_target(*, event_queue, stop_queue, config):
|
|
event_queue.put({"type": "model_load_started", "num_images": 3})
|
|
event_queue.put({"type": "model_load_completed"})
|
|
event_queue.put(
|
|
{
|
|
"type": "progress",
|
|
"step": 1,
|
|
"total_steps": 2,
|
|
"loss": 0.5,
|
|
"avg_loss": 0.5,
|
|
"learning_rate": 1e-4,
|
|
}
|
|
)
|
|
event_queue.put(
|
|
{
|
|
"type": "progress",
|
|
"step": 2,
|
|
"total_steps": 2,
|
|
"loss": 0.4,
|
|
"avg_loss": 0.45,
|
|
"learning_rate": 1e-4,
|
|
}
|
|
)
|
|
event_queue.put(
|
|
{
|
|
"type": "complete",
|
|
"output_dir": config["output_dir"],
|
|
"lora_path": config["output_dir"] + "/pytorch_lora_weights.safetensors",
|
|
"stopped": False,
|
|
}
|
|
)
|
|
|
|
|
|
def _stoppable_target(*, event_queue, stop_queue, config):
|
|
event_queue.put({"type": "model_load_completed"})
|
|
stop_queue.get(timeout = 5.0) # block until stop() signals
|
|
event_queue.put(
|
|
{"type": "complete", "output_dir": config["output_dir"], "lora_path": "x", "stopped": True}
|
|
)
|
|
|
|
|
|
def _crashing_target(*, event_queue, stop_queue, config):
|
|
event_queue.put({"type": "model_load_started"})
|
|
# Exits without a terminal event -> the pump must mark it as an error.
|
|
|
|
|
|
_CFG = {"base_model": "b", "data_dir": "d", "output_dir": "/tmp/out", "train_steps": 2}
|
|
|
|
|
|
def _wait_status(
|
|
svc,
|
|
*terminal,
|
|
timeout = 3.0,
|
|
):
|
|
end = time.time() + timeout
|
|
while time.time() < end:
|
|
st = svc.status()
|
|
if st["status"] in terminal:
|
|
return st
|
|
time.sleep(0.02)
|
|
return svc.status()
|
|
|
|
|
|
def test_service_happy_path():
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
job_id = svc.start(dict(_CFG))
|
|
assert job_id
|
|
st = _wait_status(svc, "completed")
|
|
assert st["status"] == "completed"
|
|
assert st["step"] == 2 and st["total_steps"] == 2
|
|
assert st["num_images"] == 3
|
|
assert st["loss"] == 0.4 and st["avg_loss"] == 0.45
|
|
assert st["lora_path"].endswith("pytorch_lora_weights.safetensors")
|
|
assert st["active"] is False
|
|
|
|
|
|
def test_service_rejects_bad_config_before_spawn():
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
with pytest.raises(ValueError):
|
|
svc.start({**_CFG, "train_steps": 0})
|
|
# Nothing was spawned; still idle.
|
|
assert svc.status()["status"] == "idle"
|
|
|
|
|
|
def test_service_rejects_second_concurrent_job():
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _stoppable_target)
|
|
svc.start(dict(_CFG))
|
|
_wait_status(svc, "running")
|
|
with pytest.raises(RuntimeError):
|
|
svc.start(dict(_CFG))
|
|
assert svc.stop() is True
|
|
_wait_status(svc, "stopped")
|
|
|
|
|
|
def test_service_stop_marks_stopped():
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _stoppable_target)
|
|
svc.start(dict(_CFG))
|
|
_wait_status(svc, "running")
|
|
assert svc.stop() is True
|
|
st = _wait_status(svc, "stopped")
|
|
assert st["status"] == "stopped"
|
|
assert st["active"] is False
|
|
# Stopping again when idle is a no-op.
|
|
assert svc.stop() is False
|
|
|
|
|
|
def test_service_crash_without_terminal_event_is_error():
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _crashing_target)
|
|
svc.start(dict(_CFG))
|
|
st = _wait_status(svc, "error")
|
|
assert st["status"] == "error"
|
|
assert "unexpectedly" in st["message"]
|
|
|
|
|
|
def test_apply_event_transitions():
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
svc._apply_event({"type": "model_load_started", "num_images": 5})
|
|
assert svc.status()["in_model_load"] is True and svc.status()["num_images"] == 5
|
|
svc._apply_event({"type": "model_load_completed"})
|
|
assert svc.status()["in_model_load"] is False
|
|
svc._apply_event({"type": "error", "message": "boom"})
|
|
assert svc.status()["status"] == "error" and svc.status()["message"] == "boom"
|
|
|
|
|
|
def test_progress_nulls_non_finite_floats_for_strict_json():
|
|
# A divergent step (or an inf grad norm) can push loss / avg_loss / learning_rate to
|
|
# NaN or Infinity, which strict JSON forbids. The service must null those so the status
|
|
# snapshot and the metric history stay strict-JSON serializable.
|
|
import json
|
|
import math
|
|
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
svc._apply_event(
|
|
{
|
|
"type": "progress",
|
|
"step": 3,
|
|
"total_steps": 10,
|
|
"loss": float("nan"),
|
|
"avg_loss": float("inf"),
|
|
"learning_rate": float("-inf"),
|
|
"grad_norm": float("inf"),
|
|
}
|
|
)
|
|
snap = svc.status()
|
|
assert snap["loss"] is None
|
|
assert snap["avg_loss"] is None
|
|
assert snap["learning_rate"] is None
|
|
# The reviewer's exact case: an inf pre-clip grad norm must not reach the status JSON.
|
|
assert snap["grad_norm"] is None
|
|
# The non-finite point is skipped in the history, so the loss series stays clean.
|
|
assert snap["metric_loss"] == []
|
|
assert snap["metric_steps"] == []
|
|
# strict JSON (allow_nan=False) round-trips without a ValueError from NaN/Infinity.
|
|
json.dumps(snap, allow_nan = False)
|
|
|
|
# A finite point after the bad one is recorded and preserved verbatim.
|
|
svc._apply_event(
|
|
{"type": "progress", "step": 4, "total_steps": 10, "loss": 0.5, "learning_rate": 1e-4}
|
|
)
|
|
snap2 = svc.status()
|
|
assert snap2["loss"] == 0.5
|
|
assert snap2["metric_loss"] == [0.5] and snap2["metric_steps"] == [4]
|
|
assert math.isfinite(snap2["learning_rate"])
|
|
json.dumps(snap2, allow_nan = False)
|
|
|
|
|
|
def test_terminal_events_clear_model_load_flag():
|
|
# A stop or error during model load emits complete/error WITHOUT a preceding
|
|
# model_load_completed, so the terminal update must reset in_model_load or the
|
|
# client shows a stale loading indicator after the job ended.
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
svc._apply_event({"type": "model_load_started"})
|
|
assert svc.status()["in_model_load"] is True
|
|
svc._apply_event({"type": "complete", "stopped": True})
|
|
assert svc.status()["in_model_load"] is False and svc.status()["status"] == "stopped"
|
|
|
|
svc2 = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
svc2._apply_event({"type": "model_load_started"})
|
|
svc2._apply_event({"type": "error", "message": "load failed"})
|
|
assert svc2.status()["in_model_load"] is False and svc2.status()["status"] == "error"
|
|
|
|
|
|
# ── route wiring (mocked service) ─────────────────────────────────────────────
|
|
class _FakeService:
|
|
def __init__(self):
|
|
self._running = False
|
|
self._reserved = False
|
|
self.started_with = None
|
|
self.stopped_with_save = None
|
|
# Ordered log of lifecycle calls so a test can assert reserve precedes the GPU free.
|
|
self.calls: list = []
|
|
# Extra keys merged into status() so a test can inject metric history / perf fields.
|
|
self.status_extra: dict = {}
|
|
|
|
def reserve(self):
|
|
self._reserved = True
|
|
self.calls.append("reserve")
|
|
|
|
def unreserve(self):
|
|
self._reserved = False
|
|
self.calls.append("unreserve")
|
|
|
|
def is_active(self):
|
|
return self._reserved or self._running
|
|
|
|
def start(self, config):
|
|
self.started_with = config
|
|
self._running = True
|
|
self.calls.append("start")
|
|
return "job-123"
|
|
|
|
def stop(self, save = True):
|
|
self.stopped_with_save = save
|
|
was = self._running
|
|
self._running = False
|
|
return was
|
|
|
|
def status(self):
|
|
return {
|
|
"active": self._running,
|
|
"job_id": "job-123" if self._running else None,
|
|
"status": "running" if self._running else "idle",
|
|
"message": "",
|
|
"step": 1,
|
|
"total_steps": 2,
|
|
"loss": 0.5,
|
|
"avg_loss": 0.5,
|
|
"learning_rate": 1e-4,
|
|
"num_images": 3,
|
|
"in_model_load": False,
|
|
"output_dir": None,
|
|
"lora_path": None,
|
|
"started_at": None,
|
|
"updated_at": None,
|
|
**self.status_extra,
|
|
}
|
|
|
|
|
|
class _FakeLLMBackend:
|
|
def __init__(self, active = False):
|
|
self._active = active
|
|
|
|
def is_training_active(self):
|
|
return self._active
|
|
|
|
|
|
@pytest.fixture
|
|
def client(monkeypatch):
|
|
fake = _FakeService()
|
|
monkeypatch.setattr(
|
|
"core.training.diffusion_training_service.get_diffusion_training_service", lambda: fake
|
|
)
|
|
# Neutralize the LLM interlock + GPU-free for the wiring tests (their own tests below
|
|
# exercise those behaviors). The route imports get_training_backend at module scope.
|
|
import routes.training as tr
|
|
|
|
monkeypatch.setattr(tr, "get_training_backend", lambda: _FakeLLMBackend(active = False))
|
|
monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None)
|
|
# The dataset preflight runs the trainer's discovery against _BODY's fake
|
|
# data_dir; stub it here so wiring tests pass, and let the dedicated preflight
|
|
# tests below re-point it at a real tmp dataset.
|
|
monkeypatch.setattr(
|
|
"core.training.diffusion_train_common.discover_image_caption_pairs",
|
|
lambda data_dir, **kw: [("img.png", "caption")],
|
|
)
|
|
app = FastAPI()
|
|
app.include_router(training_router, prefix = "/api/train")
|
|
app.dependency_overrides[get_current_subject] = lambda: "test-user"
|
|
# Default to session (UI) auth: the API-key inference-in-flight guard is a no-op there,
|
|
# so the wiring tests below behave as before. The guard test flips this override.
|
|
app.dependency_overrides[authenticated_via_api_key] = lambda: False
|
|
c = TestClient(app)
|
|
c._fake = fake # type: ignore[attr-defined]
|
|
c._app = app # type: ignore[attr-defined]
|
|
return c
|
|
|
|
|
|
# Studio-relative paths: the route resolves/contains them before spawn.
|
|
_BODY = {
|
|
"base_model": "stabilityai/sdxl-turbo",
|
|
"data_dir": "uploads/my-images",
|
|
"output_dir": "my-lora-run",
|
|
"train_steps": 10,
|
|
}
|
|
|
|
|
|
def test_route_start_ok(client):
|
|
r = client.post("/api/train/diffusion/start", json = _BODY)
|
|
assert r.status_code == 200, r.text
|
|
assert r.json() == {"job_id": "job-123", "status": "running"}
|
|
assert client._fake.started_with["base_model"] == "stabilityai/sdxl-turbo"
|
|
# Paths were resolved to absolute Studio-contained locations before spawn.
|
|
from pathlib import Path
|
|
|
|
assert Path(client._fake.started_with["data_dir"]).is_absolute()
|
|
assert Path(client._fake.started_with["output_dir"]).is_absolute()
|
|
|
|
|
|
def test_route_start_frees_gpu_off_the_coroutine_thread(client, monkeypatch):
|
|
# The GPU cleanup can block for seconds (engine unload waits on generation locks; the
|
|
# export subprocess join can take seconds), so the async start route must offload it via
|
|
# asyncio.to_thread rather than run it inline and freeze the event loop for concurrent
|
|
# status/progress/cancel requests. Assert the cleanup runs on a DIFFERENT thread than the
|
|
# inline coroutine body (service.start), which an inline (un-offloaded) call could not.
|
|
import threading
|
|
|
|
import routes.training as tr
|
|
|
|
threads: dict = {}
|
|
|
|
def _record_cleanup():
|
|
threads["cleanup"] = threading.current_thread()
|
|
|
|
monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", _record_cleanup)
|
|
|
|
orig_start = client._fake.start
|
|
|
|
def _record_start(config):
|
|
threads["inline"] = threading.current_thread()
|
|
return orig_start(config)
|
|
|
|
monkeypatch.setattr(client._fake, "start", _record_start)
|
|
|
|
r = client.post("/api/train/diffusion/start", json = _BODY)
|
|
assert r.status_code == 200, r.text
|
|
assert threads["cleanup"] is not threads["inline"] # offloaded to a worker, not run inline
|
|
|
|
|
|
def test_route_start_reserves_before_freeing_gpu(client, monkeypatch):
|
|
# The training slot must be reserved (is_active -> true) BEFORE the route frees resident GPU
|
|
# models, so a concurrent /images/load or /video/load guard refuses during the free-then-spawn
|
|
# window instead of double-allocating the GPU. Assert the ordering: reserve is logged before
|
|
# the GPU free runs, and the service reports active while the free is in flight.
|
|
import routes.training as tr
|
|
|
|
order: list = []
|
|
|
|
def _record_free():
|
|
order.append("free")
|
|
# During the free window the service must already look active to a concurrent load guard.
|
|
order.append(f"active={client._fake.is_active()}")
|
|
|
|
monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", _record_free)
|
|
|
|
r = client.post("/api/train/diffusion/start", json = _BODY)
|
|
assert r.status_code == 200, r.text
|
|
# reserve fires before the free, the free sees an active service, then start, then unreserve.
|
|
assert client._fake.calls[0] == "reserve"
|
|
assert client._fake.calls.index("reserve") < client._fake.calls.index("start")
|
|
assert order == ["free", "active=True"]
|
|
assert "unreserve" in client._fake.calls
|
|
|
|
|
|
def test_service_reserve_marks_active_and_rolls_back():
|
|
# The real service: reserve() flips is_active true before any proc exists (so a load guard
|
|
# refuses during the free window), and unreserve() clears it without a live proc, so a failed
|
|
# start is not left permanently "active".
|
|
from core.training.diffusion_training_service import DiffusionTrainingService
|
|
|
|
svc = DiffusionTrainingService()
|
|
assert svc.is_active() is False
|
|
svc.reserve()
|
|
assert svc.is_active() is True # active with no proc, purely from the reservation
|
|
svc.unreserve()
|
|
assert svc.is_active() is False
|
|
|
|
|
|
def test_service_reserve_is_compare_and_set():
|
|
# reserve() is the concurrency gate: two /diffusion/start requests can interleave between the
|
|
# is_active() check and the reservation, so reserve() itself must reject a second reservation
|
|
# atomically. Without the compare-and-set, both callers would reserve, both would free the
|
|
# GPU's resident chat/image model, and the loser would only 409 AFTER the eviction -- exactly
|
|
# the evict-then-fail the reservation exists to prevent. A second reserve must raise; the first
|
|
# stays reserved; after unreserve the slot is claimable again.
|
|
from core.training.diffusion_training_service import DiffusionTrainingService
|
|
|
|
svc = DiffusionTrainingService()
|
|
svc.reserve()
|
|
with pytest.raises(RuntimeError, match = "already running"):
|
|
svc.reserve()
|
|
assert svc.is_active() is True # the losing reserve did not clear the winner's claim
|
|
svc.unreserve()
|
|
svc.reserve() # claimable again once released
|
|
assert svc.is_active() is True
|
|
|
|
|
|
def test_route_start_preflights_gated_base_off_the_coroutine_thread(client, monkeypatch):
|
|
# _preflight_gated_base does a blocking urlopen HEAD (up to a 5s timeout) to Hugging Face, so
|
|
# the async start route must offload it via asyncio.to_thread rather than run it inline and
|
|
# freeze the event loop for concurrent status/progress/cancel requests. Assert it runs on a
|
|
# DIFFERENT thread than the inline coroutine body (service.start), which an inline call could
|
|
# not.
|
|
import threading
|
|
|
|
import routes.training as tr
|
|
|
|
threads: dict = {}
|
|
|
|
def _record_preflight(base_model, hf_token):
|
|
threads["preflight"] = threading.current_thread()
|
|
|
|
monkeypatch.setattr(tr, "_preflight_gated_base", _record_preflight)
|
|
|
|
orig_start = client._fake.start
|
|
|
|
def _record_start(config):
|
|
threads["inline"] = threading.current_thread()
|
|
return orig_start(config)
|
|
|
|
monkeypatch.setattr(client._fake, "start", _record_start)
|
|
|
|
r = client.post("/api/train/diffusion/start", json = _BODY)
|
|
assert r.status_code == 200, r.text
|
|
assert threads["preflight"] is not threads["inline"] # offloaded to a worker, not run inline
|
|
|
|
|
|
def test_route_start_forwards_extra_training_knobs(client):
|
|
# max_grad_norm and lora_target_modules must reach the service, not be silently dropped.
|
|
body = {**_BODY, "max_grad_norm": 0.5, "lora_target_modules": ["to_q", "to_v"]}
|
|
r = client.post("/api/train/diffusion/start", json = body)
|
|
assert r.status_code == 200, r.text
|
|
assert client._fake.started_with["max_grad_norm"] == 0.5
|
|
assert client._fake.started_with["lora_target_modules"] == ["to_q", "to_v"]
|
|
|
|
|
|
def test_route_start_forwards_num_epochs(client):
|
|
# Epochs mode: the frontend omits train_steps and sends num_epochs; it must reach the
|
|
# service so the trainer can resolve it against the dataset size.
|
|
body = {k: v for k, v in _BODY.items() if k != "train_steps"}
|
|
r = client.post("/api/train/diffusion/start", json = {**body, "num_epochs": 8})
|
|
assert r.status_code == 200, r.text
|
|
assert client._fake.started_with["num_epochs"] == 8
|
|
|
|
|
|
def test_request_model_num_epochs_bounds():
|
|
# The request schema mirrors DiffusionLoraConfig's 0..1000 num_epochs range.
|
|
from pydantic import ValidationError
|
|
|
|
from models.training import DiffusionTrainingStartRequest
|
|
|
|
base = {"base_model": "b", "data_dir": "d", "output_dir": "o"}
|
|
assert DiffusionTrainingStartRequest(**base).num_epochs == 0 # default = use train_steps
|
|
assert DiffusionTrainingStartRequest(**base, num_epochs = 1000).num_epochs == 1000
|
|
for bad in (-1, 1001):
|
|
with pytest.raises(ValidationError):
|
|
DiffusionTrainingStartRequest(**base, num_epochs = bad)
|
|
|
|
|
|
def test_request_model_base_precision_accepts_mxfp8():
|
|
# The base_precision Literal now includes mxfp8 (the DiT dense speed mode); a bogus
|
|
# mode is still rejected.
|
|
from pydantic import ValidationError
|
|
|
|
from models.training import DiffusionTrainingStartRequest
|
|
|
|
base = {"base_model": "b", "data_dir": "d", "output_dir": "o"}
|
|
assert DiffusionTrainingStartRequest(**base).base_precision == "nf4" # default
|
|
assert DiffusionTrainingStartRequest(**base, base_precision = "mxfp8").base_precision == "mxfp8"
|
|
with pytest.raises(ValidationError):
|
|
DiffusionTrainingStartRequest(**base, base_precision = "bogus")
|
|
|
|
|
|
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_permutation_sampler_covers_dataset_once_per_cycle():
|
|
# Every index must appear exactly once per cycle before any repeat (epoch-style pass),
|
|
# so a short run over a small dataset never leaves images unseen the way the old
|
|
# with-replacement draw did. Consecutive cycles must be reshuffled (differ).
|
|
import random
|
|
|
|
from core.training.diffusion_train_common import PermutationBatchSampler
|
|
|
|
n = 100
|
|
sampler = PermutationBatchSampler(n, random.Random(0))
|
|
|
|
# Draw exactly one cycle in batches of 3 (n not divisible by the batch, so a batch spans
|
|
# the cycle boundary); the first n indices must be a permutation of range(n).
|
|
drawn: list[int] = []
|
|
while len(drawn) < n:
|
|
drawn.extend(sampler.next_batch(3))
|
|
first_cycle = drawn[:n]
|
|
assert sorted(first_cycle) == list(range(n)) # each index once, none missing
|
|
|
|
# The next full cycle is also a permutation, and it is reshuffled (order differs).
|
|
fresh = PermutationBatchSampler(n, random.Random(0))
|
|
cycle_a = fresh.next_batch(n)
|
|
cycle_b = fresh.next_batch(n)
|
|
assert sorted(cycle_a) == list(range(n))
|
|
assert sorted(cycle_b) == list(range(n))
|
|
assert cycle_a != cycle_b # cycles are reshuffled, not repeated in the same order
|
|
|
|
# A seed replays the exact index stream (determinism for reproducible runs).
|
|
replay = PermutationBatchSampler(n, random.Random(0))
|
|
assert replay.next_batch(n) == cycle_a
|
|
|
|
# A batch larger than the dataset refills across cycles so it never shrinks (batch shape
|
|
# preserved), even though it must then repeat indices within the batch.
|
|
big = PermutationBatchSampler(4, random.Random(1))
|
|
batch = big.next_batch(10)
|
|
assert len(batch) == 10
|
|
assert set(batch) == {0, 1, 2, 3}
|
|
|
|
|
|
def test_permutation_sampler_honors_batch_on_tiny_dataset():
|
|
# Regression for the SDXL LoRA trainer clamp: a dataset smaller than train_batch_size must
|
|
# still yield exactly train_batch_size indices (refilled across cycles), so a tiny dataset
|
|
# trains at the configured/reported effective batch instead of silently shrinking it -- the
|
|
# contract the SDXL _next_batch now relies on (matching the DiT trainer, which never clamps).
|
|
import random
|
|
|
|
from core.training.diffusion_train_common import PermutationBatchSampler
|
|
|
|
sampler = PermutationBatchSampler(2, random.Random(0)) # 2-image dataset
|
|
batch = sampler.next_batch(8) # train_batch_size = 8, not clamped to 2
|
|
assert len(batch) == 8
|
|
assert set(batch) == {0, 1}
|
|
# A whole-multiple batch draws each image equally, so the effective gradient matches the
|
|
# configured batch rather than a shrunk one.
|
|
assert batch.count(0) == 4 and batch.count(1) == 4
|
|
|
|
|
|
def test_route_start_accepts_zero_max_grad_norm(client):
|
|
# 0 is the documented "disable clipping" value (the trainer skips clip_grad_norm_);
|
|
# the request model must not reject it.
|
|
r = client.post("/api/train/diffusion/start", json = {**_BODY, "max_grad_norm": 0.0})
|
|
assert r.status_code == 200, r.text
|
|
assert client._fake.started_with["max_grad_norm"] == 0.0
|
|
|
|
|
|
def test_route_start_rejects_nonpositive_snr_gamma(client):
|
|
# gamma <= 0 zeroes/inverts the min-SNR loss weight; null is the disable value.
|
|
r = client.post("/api/train/diffusion/start", json = {**_BODY, "snr_gamma": 0})
|
|
assert r.status_code == 422
|
|
|
|
|
|
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"})
|
|
assert r.status_code == 400
|
|
|
|
|
|
def test_route_start_resolves_bare_name_under_image_dataset_root(client, monkeypatch, tmp_path):
|
|
# The upload/labeling routes manage image datasets directly under datasets_root() and
|
|
# the UI passes the bare folder name back as data_dir. The generic resolve_dataset_path
|
|
# searches the LLM uploads and recipe roots FIRST, so an unrelated upload file or recipe
|
|
# folder sharing the name would shadow the just-uploaded image dataset (preflight 400
|
|
# "not a directory", or training the wrong data). The route must prefer the image
|
|
# dataset root for a bare name that exists there.
|
|
import utils.paths as up
|
|
|
|
ds_root = tmp_path / "assets" / "datasets"
|
|
img_ds = ds_root / "my-photos"
|
|
img_ds.mkdir(parents = True)
|
|
(img_ds / "a.png").write_bytes(b"x")
|
|
# Shadowing entries the generic resolver would pick first.
|
|
(ds_root / "uploads").mkdir()
|
|
(ds_root / "uploads" / "my-photos").write_text("an LLM dataset upload, not a folder")
|
|
(ds_root / "recipes" / "my-photos").mkdir(parents = True)
|
|
monkeypatch.setattr(up, "datasets_root", lambda: ds_root)
|
|
|
|
r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "my-photos"})
|
|
assert r.status_code == 200, r.text
|
|
assert client._fake.started_with["data_dir"] == str(img_ds)
|
|
|
|
|
|
def test_route_start_blocked_by_active_llm_training(client, monkeypatch):
|
|
import routes.training as tr
|
|
|
|
monkeypatch.setattr(tr, "get_training_backend", lambda: _FakeLLMBackend(active = True))
|
|
r = client.post("/api/train/diffusion/start", json = _BODY)
|
|
assert r.status_code == 409
|
|
assert "LLM training" in r.json()["detail"]
|
|
|
|
|
|
def test_route_start_missing_required_is_422(client):
|
|
r = client.post(
|
|
"/api/train/diffusion/start", json = {"base_model": "x"}
|
|
) # no data_dir/output_dir
|
|
assert r.status_code == 422
|
|
|
|
|
|
def test_route_start_bad_config_maps_to_400(client, monkeypatch):
|
|
def _raise(_cfg):
|
|
raise ValueError("resolution must be a multiple of 8")
|
|
|
|
client._fake.start = _raise # type: ignore[assignment]
|
|
r = client.post("/api/train/diffusion/start", json = _BODY)
|
|
assert r.status_code == 400
|
|
assert "multiple of 8" in r.json()["detail"]
|
|
|
|
|
|
def test_route_start_conflict_maps_to_409(client):
|
|
def _raise(_cfg):
|
|
raise RuntimeError("A diffusion training job is already running.")
|
|
|
|
client._fake.start = _raise # type: ignore[assignment]
|
|
r = client.post("/api/train/diffusion/start", json = _BODY)
|
|
assert r.status_code == 409
|
|
|
|
|
|
def test_route_start_over_api_with_inference_in_flight_is_409(client, monkeypatch):
|
|
# An API-key client must not start diffusion training (which frees VRAM by unloading
|
|
# chat) while an inference request is streaming; it should 409 instead of killing it.
|
|
client._app.dependency_overrides[authenticated_via_api_key] = lambda: True
|
|
monkeypatch.setattr(
|
|
"core.inference.llama_keepwarm.other_inference_request_count",
|
|
lambda current_request_counted = False: 1,
|
|
)
|
|
freed = {"called": False}
|
|
import routes.training as tr
|
|
|
|
monkeypatch.setattr(
|
|
tr, "_free_gpu_for_diffusion_training", lambda: freed.__setitem__("called", True)
|
|
)
|
|
r = client.post("/api/train/diffusion/start", json = _BODY)
|
|
assert r.status_code == 409
|
|
# The guard must run BEFORE any GPU is freed, so the live inference stream survives.
|
|
assert freed["called"] is False
|
|
|
|
|
|
def test_route_start_over_api_without_inference_proceeds(client, monkeypatch):
|
|
# Same API-key path but no inference in flight: the start proceeds normally.
|
|
client._app.dependency_overrides[authenticated_via_api_key] = lambda: True
|
|
monkeypatch.setattr(
|
|
"core.inference.llama_keepwarm.other_inference_request_count",
|
|
lambda current_request_counted = False: 0,
|
|
)
|
|
r = client.post("/api/train/diffusion/start", json = _BODY)
|
|
assert r.status_code == 200
|
|
|
|
|
|
def test_route_status_and_stop(client):
|
|
client.post("/api/train/diffusion/start", json = _BODY)
|
|
s = client.get("/api/train/diffusion/status")
|
|
assert s.status_code == 200 and s.json()["status"] == "running"
|
|
st = client.post("/api/train/diffusion/stop")
|
|
assert st.status_code == 200 and st.json()["status"] == "stopping"
|
|
# After stopping, a stop with nothing running reports idle.
|
|
st2 = client.post("/api/train/diffusion/stop")
|
|
assert st2.json()["status"] == "idle"
|
|
|
|
|
|
def test_service_restart_after_completion():
|
|
# A finished job's pump is joined OUTSIDE the lock (it needs the lock for its
|
|
# final state writes), so a second start neither stalls nor deadlocks.
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
svc.start(dict(_CFG))
|
|
_wait_status(svc, "completed")
|
|
t0 = time.time()
|
|
job2 = svc.start(dict(_CFG))
|
|
assert job2
|
|
assert time.time() - t0 < 4.0 # no 5s join-under-lock stall
|
|
st = _wait_status(svc, "completed")
|
|
assert st["status"] == "completed"
|
|
|
|
|
|
def test_stale_pump_events_cannot_corrupt_new_job():
|
|
# An event carrying a superseded job's proc identity must be dropped, so a
|
|
# straggler pump can never overwrite the state of a newly started job.
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
svc.start(dict(_CFG))
|
|
_wait_status(svc, "completed")
|
|
current = svc._proc
|
|
svc._apply_event({"type": "error", "message": "stale boom"}, proc = object())
|
|
assert svc.status()["message"] != "stale boom"
|
|
# The current job's events still apply.
|
|
svc._apply_event({"type": "progress", "step": 9}, proc = current)
|
|
assert svc.status()["step"] == 9
|
|
|
|
|
|
# ── /diffusion/info + /diffusion/dataset (dataset discovery + upload) ─────────
|
|
@pytest.fixture
|
|
def dataset_roots(client, monkeypatch, tmp_path):
|
|
# The endpoints import these lazily per-request, so patching the package attr works.
|
|
import utils.paths as up
|
|
|
|
ds_root = tmp_path / "assets" / "datasets"
|
|
out_root = tmp_path / "outputs"
|
|
ds_root.mkdir(parents = True)
|
|
out_root.mkdir(parents = True)
|
|
monkeypatch.setattr(up, "datasets_root", lambda: ds_root)
|
|
monkeypatch.setattr(up, "outputs_root", lambda: out_root)
|
|
return ds_root, out_root
|
|
|
|
|
|
def test_diffusion_info_lists_image_dataset_folders(client, dataset_roots):
|
|
ds_root, out_root = dataset_roots
|
|
good = ds_root / "cat-photos"
|
|
good.mkdir()
|
|
(good / "a.png").write_bytes(b"x")
|
|
(good / "b.jpg").write_bytes(b"x")
|
|
(good / "a.txt").write_text("a cat")
|
|
(ds_root / "empty-dir").mkdir() # no images -> not a dataset
|
|
(ds_root / "stray.txt").write_text("not a folder")
|
|
|
|
r = client.get("/api/train/diffusion/info")
|
|
assert r.status_code == 200, r.text
|
|
body = r.json()
|
|
assert body["datasets_root"] == str(ds_root)
|
|
assert body["outputs_root"] == str(out_root)
|
|
assert [d["name"] for d in body["datasets"]] == ["cat-photos"]
|
|
assert body["datasets"][0]["image_count"] == 2
|
|
assert body["datasets"][0]["caption_count"] == 1
|
|
|
|
|
|
def test_diffusion_info_counts_metadata_captions(client, dataset_roots):
|
|
# A dataset captioned via metadata.jsonl must not report caption_count=0 (which the
|
|
# Train UI treats as uncaptioned); metadata rows count like sidecars, without
|
|
# double-counting an image that has both.
|
|
ds_root, _ = dataset_roots
|
|
folder = ds_root / "meta-captioned"
|
|
folder.mkdir()
|
|
(folder / "a.png").write_bytes(b"x")
|
|
(folder / "b.png").write_bytes(b"x")
|
|
(folder / "c.png").write_bytes(b"x")
|
|
# a.png + b.png via metadata; a.png also has a sidecar (must count once); c.png none.
|
|
(folder / "metadata.jsonl").write_text(
|
|
json.dumps({"file_name": "a.png", "text": "cap a"})
|
|
+ "\n"
|
|
+ json.dumps({"file_name": "b.png", "text": "cap b"})
|
|
+ "\n",
|
|
encoding = "utf-8",
|
|
)
|
|
(folder / "a.txt").write_text("edited a", encoding = "utf-8")
|
|
|
|
r = client.get("/api/train/diffusion/info")
|
|
assert r.status_code == 200, r.text
|
|
summary = next(d for d in r.json()["datasets"] if d["name"] == "meta-captioned")
|
|
assert summary["image_count"] == 3
|
|
assert summary["caption_count"] == 2
|
|
|
|
|
|
def test_diffusion_dataset_upload_accumulates(client, dataset_roots):
|
|
ds_root, _ = dataset_roots
|
|
files = [
|
|
("files", ("a.png", b"png-bytes", "image/png")),
|
|
("files", ("b.JPG", b"jpg-bytes", "image/jpeg")),
|
|
("files", ("a.txt", b"a caption", "text/plain")),
|
|
]
|
|
r = client.post("/api/train/diffusion/dataset", data = {"name": "my style"}, files = files)
|
|
assert r.status_code == 200, r.text
|
|
body = r.json()
|
|
assert body["name"] == "my style"
|
|
assert body["uploaded"] == 3
|
|
assert body["image_count"] == 2
|
|
assert body["caption_count"] == 1
|
|
assert (ds_root / "my style" / "a.png").read_bytes() == b"png-bytes"
|
|
|
|
# A second batch into the same name accumulates (large sets arrive in chunks).
|
|
r = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "my style"},
|
|
files = [("files", ("c.webp", b"w", "image/webp"))],
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
assert r.json()["uploaded"] == 1
|
|
assert r.json()["image_count"] == 3
|
|
|
|
|
|
def test_diffusion_dataset_upload_normalizes_windows_and_rejects_dotdot(client, dataset_roots):
|
|
ds_root, _ = dataset_roots
|
|
# A Windows client can send a backslash path in the multipart filename; POSIX Path.name does
|
|
# not split on backslash, so it must be folded to the true basename, or the stored name holds
|
|
# backslashes that _safe_dataset_image_path later rejects -- an image the labeling grid can
|
|
# list but never preview/caption/delete (an orphan).
|
|
r = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "winset"},
|
|
files = [("files", ("C:\\Users\\me\\pics\\cat.png", b"png-bytes", "image/png"))],
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
assert (ds_root / "winset" / "cat.png").read_bytes() == b"png-bytes"
|
|
# It is listed under the clean basename and the per-image endpoints accept it (not an orphan).
|
|
recs = client.get("/api/train/diffusion/dataset/winset/images").json()["images"]
|
|
assert any(rec["filename"] == "cat.png" for rec in recs)
|
|
assert client.get("/api/train/diffusion/dataset/winset/image/cat.png").status_code == 200
|
|
|
|
# A basename that still contains ".." (which _safe_dataset_image_path rejects) is refused at
|
|
# upload rather than persisted as an unmanageable entry.
|
|
r = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "winset"},
|
|
files = [("files", ("a..b.png", b"x", "image/png"))],
|
|
)
|
|
assert r.status_code == 400 and "Unsupported file" in r.json()["detail"]
|
|
|
|
|
|
def test_diffusion_dataset_upload_rejects_case_insensitive_stem_clash(client, dataset_roots):
|
|
# Two images whose stems differ only by case (sample.png vs Sample.jpg) map to the SAME
|
|
# <stem>.txt caption sidecar on case-insensitive filesystems (Windows / default macOS), so
|
|
# keeping both would silently share -- and corrupt -- one caption during training. The clash
|
|
# check compares stems with casefold, so it must reject the pair (the comparison is pure
|
|
# string logic, so this fires regardless of the test host filesystem).
|
|
r = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "caseset"},
|
|
files = [
|
|
("files", ("sample.png", b"p", "image/png")),
|
|
("files", ("Sample.jpg", b"j", "image/jpeg")),
|
|
],
|
|
)
|
|
assert r.status_code == 400 and "Duplicate image name" in r.json()["detail"]
|
|
|
|
# The same clash across batches (the new image collides with one already on disk).
|
|
r = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "caseset2"},
|
|
files = [("files", ("photo.png", b"p", "image/png"))],
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
r = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "caseset2"},
|
|
files = [("files", ("PHOTO.webp", b"w", "image/webp"))],
|
|
)
|
|
assert r.status_code == 400 and "Duplicate image name" in r.json()["detail"]
|
|
|
|
# A same-name case variant with the SAME extension is an overwrite (one file on a
|
|
# case-insensitive FS), not a caption clash, so it is still allowed.
|
|
r = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "caseset3"},
|
|
files = [
|
|
("files", ("pic.png", b"a", "image/png")),
|
|
("files", ("Pic.png", b"b", "image/png")),
|
|
],
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
|
|
|
|
def test_diffusion_dataset_upload_over_cap_keeps_existing_example(
|
|
client, dataset_roots, monkeypatch
|
|
):
|
|
# A re-upload that trips the size cap mid-write must not destroy an example already
|
|
# stored under the same name: the write goes to a sibling temp file and only atomically
|
|
# replaces the original on success, so the 413 leaves the prior good bytes intact.
|
|
import utils.upload_limits as ul
|
|
|
|
ds_root, _ = dataset_roots
|
|
folder = ds_root / "my style"
|
|
folder.mkdir()
|
|
(folder / "cat.png").write_bytes(b"ORIGINAL-CAT-BYTES")
|
|
|
|
monkeypatch.setattr(ul, "get_upload_limit_bytes", lambda: 8)
|
|
monkeypatch.setattr(ul, "get_upload_limit_label", lambda: "8B")
|
|
|
|
r = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "my style"},
|
|
files = [("files", ("cat.png", b"x" * 64, "image/png"))],
|
|
)
|
|
assert r.status_code == 413, r.text
|
|
# The pre-existing example survives untouched, and no temp file is left behind.
|
|
assert (folder / "cat.png").read_bytes() == b"ORIGINAL-CAT-BYTES"
|
|
assert sorted(p.name for p in folder.iterdir()) == ["cat.png"]
|
|
|
|
|
|
def test_diffusion_info_empty_sidecar_shadows_metadata_caption(client, dataset_roots):
|
|
# An empty (tombstone) .txt sidecar shadows a metadata row -- the trainer strips it and
|
|
# skips the image -- so the summary must not count it as captioned (which would report a
|
|
# dataset as captioned that the trainer would reject as having no captioned images).
|
|
ds_root, _ = dataset_roots
|
|
folder = ds_root / "tombstoned"
|
|
folder.mkdir()
|
|
(folder / "a.png").write_bytes(b"x")
|
|
(folder / "b.png").write_bytes(b"x")
|
|
(folder / "c.png").write_bytes(b"x")
|
|
(folder / "metadata.jsonl").write_text(
|
|
json.dumps({"file_name": "a.png", "text": "cap a"})
|
|
+ "\n"
|
|
+ json.dumps({"file_name": "c.png", "text": "cap c"})
|
|
+ "\n",
|
|
encoding = "utf-8",
|
|
)
|
|
# a.png: metadata caption but an empty sidecar tombstone -> uncaptioned.
|
|
(folder / "a.txt").write_text(" ", encoding = "utf-8")
|
|
# b.png: real sidecar caption. c.png: metadata only. Both captioned.
|
|
(folder / "b.txt").write_text("cap b", encoding = "utf-8")
|
|
|
|
r = client.get("/api/train/diffusion/info")
|
|
assert r.status_code == 200, r.text
|
|
summary = next(d for d in r.json()["datasets"] if d["name"] == "tombstoned")
|
|
assert summary["image_count"] == 3
|
|
assert summary["caption_count"] == 2
|
|
|
|
|
|
def _png_bytes(width: int, height: int) -> bytes:
|
|
import io
|
|
|
|
pytest.importorskip("PIL")
|
|
from PIL import Image
|
|
|
|
buf = io.BytesIO()
|
|
Image.new("RGB", (width, height), (1, 2, 3)).save(buf, format = "PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
def test_diffusion_dataset_upload_rejects_oversized_image(client, dataset_roots):
|
|
# A decompression bomb: a small compressible PNG with huge dimensions passes the byte limit but
|
|
# would OOM the trainer on decode. It must 400 at upload (dimension check reads only the header).
|
|
pytest.importorskip("PIL")
|
|
big = _png_bytes(5000, 64) # > 4096 per side
|
|
r = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "bomb"},
|
|
files = [("files", ("huge.png", big, "image/png"))],
|
|
)
|
|
assert r.status_code == 400, r.text
|
|
assert "too large" in r.json()["detail"]
|
|
# An in-bounds real image still uploads fine.
|
|
ok = _png_bytes(64, 64)
|
|
r2 = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "bomb"},
|
|
files = [("files", ("ok.png", ok, "image/png"))],
|
|
)
|
|
assert r2.status_code == 200, r2.text
|
|
|
|
|
|
def test_diffusion_info_tolerates_non_object_jsonl(client, dataset_roots):
|
|
# A metadata.jsonl line that is valid JSON but not an object ([]/null/string/number) or malformed
|
|
# must be skipped per-line, not 500 the info endpoint; a valid row in the same file still counts.
|
|
ds_root, _ = dataset_roots
|
|
folder = ds_root / "weird-meta"
|
|
folder.mkdir()
|
|
(folder / "a.png").write_bytes(b"x")
|
|
(folder / "metadata.jsonl").write_bytes(
|
|
b"[]\n"
|
|
b"null\n"
|
|
b'"just a string"\n'
|
|
b"123\n"
|
|
b"{not json\n"
|
|
+ json.dumps({"file_name": "a.png", "text": "cap a"}).encode("utf-8")
|
|
+ b"\n"
|
|
)
|
|
r = client.get("/api/train/diffusion/info")
|
|
assert r.status_code == 200, r.text
|
|
summary = next(d for d in r.json()["datasets"] if d["name"] == "weird-meta")
|
|
assert summary["caption_count"] == 1
|
|
|
|
|
|
def test_diffusion_info_tolerates_invalid_utf8_jsonl(client, dataset_roots):
|
|
# Invalid UTF-8 in a metadata file must not 500 the info endpoint; the file is skipped, not decoded.
|
|
ds_root, _ = dataset_roots
|
|
folder = ds_root / "bad-utf8-meta"
|
|
folder.mkdir()
|
|
(folder / "a.png").write_bytes(b"x")
|
|
(folder / "metadata.jsonl").write_bytes(b"\xff\xfe not valid utf-8\n")
|
|
r = client.get("/api/train/diffusion/info")
|
|
assert r.status_code == 200, r.text
|
|
summary = next(d for d in r.json()["datasets"] if d["name"] == "bad-utf8-meta")
|
|
assert summary["caption_count"] == 0
|
|
|
|
|
|
def test_diffusion_dataset_mutations_blocked_while_training_active(client, dataset_roots):
|
|
ds_root, _ = dataset_roots
|
|
folder = ds_root / "locked"
|
|
folder.mkdir()
|
|
folder.joinpath("a.png").write_bytes(b"x")
|
|
# Flip the fake diffusion service to active.
|
|
client._fake._running = True
|
|
# Upload, caption, delete, and example-import must all 409 while a run is active.
|
|
up = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "locked"},
|
|
files = [("files", ("b.png", b"x", "image/png"))],
|
|
)
|
|
assert up.status_code == 409, up.text
|
|
cap = client.put(
|
|
"/api/train/diffusion/dataset/locked/caption/a.png", json = {"caption": "hi"}
|
|
)
|
|
assert cap.status_code == 409, cap.text
|
|
dele = client.delete("/api/train/diffusion/dataset/locked/image/a.png")
|
|
assert dele.status_code == 409, dele.text
|
|
imp = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "anything"})
|
|
assert imp.status_code == 409, imp.text
|
|
|
|
|
|
def test_diffusion_info_skips_symlinked_dataset_dir(client, dataset_roots):
|
|
# A directory symlink under the datasets root must not be advertised as a dataset (the CRUD
|
|
# resolver rejects symlinked datasets, so discovery must agree).
|
|
import os
|
|
|
|
ds_root, _ = dataset_roots
|
|
outside = ds_root.parent / "outside-images"
|
|
outside.mkdir()
|
|
(outside / "a.png").write_bytes(b"x")
|
|
try:
|
|
os.symlink(outside, ds_root / "linked")
|
|
except OSError:
|
|
pytest.skip("symlinks unavailable")
|
|
r = client.get("/api/train/diffusion/info")
|
|
assert r.status_code == 200, r.text
|
|
assert "linked" not in [d["name"] for d in r.json()["datasets"]]
|
|
|
|
|
|
def test_diffusion_dataset_upload_rejects_traversal_names(client, dataset_roots):
|
|
for bad in ("../evil", "a/b", ".hidden", " "):
|
|
r = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": bad},
|
|
files = [("files", ("a.png", b"x", "image/png"))],
|
|
)
|
|
assert r.status_code == 400, f"{bad!r}: {r.status_code}"
|
|
|
|
|
|
def test_diffusion_dataset_upload_rejects_unsupported_files(client, dataset_roots):
|
|
r = client.post(
|
|
"/api/train/diffusion/dataset",
|
|
data = {"name": "ok-name"},
|
|
files = [("files", ("weights.exe", b"mz", "application/octet-stream"))],
|
|
)
|
|
assert r.status_code == 400
|
|
assert "Unsupported file" in r.json()["detail"]
|
|
|
|
|
|
def test_free_gpu_for_diffusion_training_unloads_video(monkeypatch):
|
|
# A resident Video pipeline loads under the VIDEO arbiter owner, which the Images teardown
|
|
# does not free; starting diffusion training must unload it too or the trainer OOMs against
|
|
# the still-resident video model.
|
|
import routes.training as tr
|
|
from core.inference import gpu_arbiter
|
|
|
|
class _Exp:
|
|
current_checkpoint = None
|
|
|
|
def is_export_active(self):
|
|
return False
|
|
|
|
class _Diff:
|
|
is_loaded = False
|
|
|
|
def unload(self):
|
|
pass
|
|
|
|
unloaded = {"video": False}
|
|
|
|
class _Vid:
|
|
def status(self):
|
|
return {"loaded": True}
|
|
|
|
def unload(self):
|
|
unloaded["video"] = True
|
|
|
|
released = []
|
|
monkeypatch.setattr("core.export.get_export_backend", lambda: _Exp())
|
|
monkeypatch.setattr(
|
|
"core.inference.diffusion_engine_router.get_active_diffusion_engine", lambda: _Diff()
|
|
)
|
|
monkeypatch.setattr("core.inference.video.get_video_backend", lambda: _Vid())
|
|
monkeypatch.setattr(gpu_arbiter, "release", lambda owner: released.append(owner))
|
|
|
|
tr._free_gpu_for_diffusion_training()
|
|
|
|
assert unloaded["video"] is True
|
|
assert gpu_arbiter.VIDEO in released
|
|
|
|
|
|
def test_keepwarm_tracks_image_video_generation_paths():
|
|
# The API-key training-start guard uses other_inference_request_count(), which only sees
|
|
# paths the keepwarm middleware tracks. Image/video generation must be tracked so a training
|
|
# start is refused (409) while one is in-flight rather than its unload cancelling it; the GET
|
|
# *-progress and */cancel variants must stay untracked.
|
|
from core.inference.llama_keepwarm import _is_inference_path
|
|
|
|
assert _is_inference_path("/api/inference/images/generate")
|
|
assert _is_inference_path("/v1/images/generations")
|
|
assert _is_inference_path("/api/inference/video/generate")
|
|
assert not _is_inference_path("/api/inference/images/generate-progress")
|
|
assert not _is_inference_path("/api/inference/video/generate-progress")
|
|
assert not _is_inference_path("/api/inference/video/generate/cancel")
|
|
|
|
|
|
def test_import_example_partial_failure_leaves_no_partial_dataset(
|
|
client, dataset_roots, monkeypatch
|
|
):
|
|
# A materialize that writes some images then fails must not leave a partial dataset: it stages
|
|
# into a discarded temp dir, so the target folder stays empty and a retry re-materializes
|
|
# instead of the image_count>0 idempotency check treating a truncated result as complete.
|
|
import routes.training as tr
|
|
|
|
ds_root, _ = dataset_roots
|
|
|
|
def _boom(entry, dest, cap):
|
|
(dest / "img_0000.png").write_bytes(b"x") # partial write into staging
|
|
raise RuntimeError("transient copy error")
|
|
|
|
monkeypatch.setattr(tr, "_materialize_hf_dataset", _boom)
|
|
r = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "dreambooth-dog"})
|
|
assert r.status_code == 502
|
|
folder = ds_root / "dreambooth-dog"
|
|
assert not folder.exists() or not any(folder.iterdir())
|
|
# And no leftover staging dir surfaces as a dataset.
|
|
assert not any(p.name.startswith(".dreambooth-dog.import-") for p in ds_root.iterdir())
|
|
|
|
|
|
def test_diffusion_dataset_upload_over_cap_rolls_back_whole_batch(
|
|
client, dataset_roots, monkeypatch
|
|
):
|
|
# All-or-nothing: a valid image ahead of the one that trips the size cap must NOT be
|
|
# left on disk (the 413 mid-batch rolls back every file written this request).
|
|
import utils.upload_limits as ul
|
|
|
|
monkeypatch.setattr(ul, "get_upload_limit_bytes", lambda: 100)
|
|
ds_root, _ = dataset_roots
|
|
files = [
|
|
("files", ("small.png", b"x" * 50, "image/png")),
|
|
("files", ("big.png", b"y" * 200, "image/png")),
|
|
]
|
|
r = client.post("/api/train/diffusion/dataset", data = {"name": "rollback"}, files = files)
|
|
assert r.status_code == 413, r.text
|
|
folder = ds_root / "rollback"
|
|
assert not (folder / "small.png").exists() # the earlier valid file was rolled back
|
|
assert not (folder / "big.png").exists()
|
|
|
|
|
|
def test_diffusion_dataset_upload_over_cap_preserves_existing_file(
|
|
client, dataset_roots, monkeypatch
|
|
):
|
|
# A failed batch that reuses an existing filename must NOT delete the user's
|
|
# pre-existing file (repeat uploads accumulate). Staging to a temp file keeps the
|
|
# original intact until the whole batch commits.
|
|
import utils.upload_limits as ul
|
|
|
|
monkeypatch.setattr(ul, "get_upload_limit_bytes", lambda: 100)
|
|
ds_root, _ = dataset_roots
|
|
folder = ds_root / "keep"
|
|
folder.mkdir(parents = True)
|
|
(folder / "existing.png").write_bytes(b"ORIGINAL") # from an earlier upload
|
|
files = [
|
|
("files", ("existing.png", b"NEW", "image/png")), # re-upload, small
|
|
("files", ("big.png", b"y" * 200, "image/png")), # trips the cap
|
|
]
|
|
r = client.post("/api/train/diffusion/dataset", data = {"name": "keep"}, files = files)
|
|
assert r.status_code == 413, r.text
|
|
assert (folder / "existing.png").read_bytes() == b"ORIGINAL" # untouched
|
|
assert not (folder / "big.png").exists()
|
|
assert not list(folder.glob(".*.part")) # no leftover temp files
|
|
|
|
|
|
def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypatch):
|
|
# A doomed start (non-SDXL base) must 400 BEFORE resident GPU workloads are freed,
|
|
# so a bad pick never unloads the user's working chat/Images model.
|
|
import routes.training as tr
|
|
|
|
freed = []
|
|
monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1))
|
|
r = client.post(
|
|
"/api/train/diffusion/start", json = {**_BODY, "base_model": "unsloth/FLUX.1-dev-GGUF"}
|
|
)
|
|
assert r.status_code == 400
|
|
assert "SDXL" in r.json()["detail"]
|
|
assert freed == []
|
|
assert client._fake.started_with is None
|
|
|
|
|
|
def test_route_start_refuses_non_bf16_gpu_without_freeing_gpu(client, monkeypatch):
|
|
# A DiT precision the host cannot run (no bf16 GPU, or explicit int8 without a functional
|
|
# torchao) must 400 BEFORE resident GPU workloads are freed: otherwise the host tears down the
|
|
# user's chat/Images model and the run then dies in the trainer child. The route imports
|
|
# training_precision_preflight_error locally, so patch it on its home module.
|
|
import routes.training as tr
|
|
|
|
freed = []
|
|
monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1))
|
|
monkeypatch.setattr(
|
|
"core.training.diffusion_train_common.training_precision_preflight_error",
|
|
lambda fam, prec: (
|
|
"This trainer requires a bfloat16-capable GPU (Ampere or newer)."
|
|
if fam != "sdxl"
|
|
else None
|
|
),
|
|
)
|
|
r = client.post(
|
|
"/api/train/diffusion/start",
|
|
json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"},
|
|
)
|
|
assert r.status_code == 400
|
|
assert "bfloat16" in r.json()["detail"]
|
|
assert freed == []
|
|
assert client._fake.started_with is None
|
|
|
|
# SDXL (its own mixed_precision path) is exempt: the same probe returns None, so an SDXL
|
|
# start proceeds normally past the preflight.
|
|
r2 = client.post("/api/train/diffusion/start", json = _BODY)
|
|
assert r2.status_code == 200, r2.text
|
|
|
|
|
|
# ── metric history + perf/family fields (PR A platform) ──────────────────────
|
|
def test_apply_event_records_metric_history_and_perf():
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
svc._apply_event(
|
|
{
|
|
"type": "progress",
|
|
"step": 1,
|
|
"total_steps": 10,
|
|
"loss": 0.5,
|
|
"learning_rate": 1e-4,
|
|
"samples_per_second": 3.2,
|
|
"peak_memory_gb": 7.1,
|
|
}
|
|
)
|
|
svc._apply_event(
|
|
{"type": "progress", "step": 2, "total_steps": 10, "loss": 0.4, "learning_rate": 9e-5}
|
|
)
|
|
st = svc.status()
|
|
assert st["metric_steps"] == [1, 2]
|
|
assert st["metric_loss"] == [0.5, 0.4]
|
|
assert st["metric_lr"] == [1e-4, 9e-5]
|
|
assert st["samples_per_second"] == 3.2
|
|
assert st["peak_memory_gb"] == 7.1
|
|
|
|
|
|
def test_apply_event_metric_history_skips_bad_points():
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
# step 0 (warmup / no real step), a None loss, and a NaN loss must all be skipped.
|
|
svc._apply_event({"type": "progress", "step": 0, "loss": 0.9, "learning_rate": 1e-4})
|
|
svc._apply_event({"type": "progress", "step": 1, "loss": None, "learning_rate": 1e-4})
|
|
svc._apply_event({"type": "progress", "step": 2, "loss": float("nan"), "learning_rate": 1e-4})
|
|
svc._apply_event({"type": "progress", "step": 3, "loss": 0.3, "learning_rate": None})
|
|
st = svc.status()
|
|
assert st["metric_steps"] == [3]
|
|
assert st["metric_loss"] == [0.3]
|
|
assert st["metric_lr"] == [None] # lr None is retained so the series stays index-aligned
|
|
|
|
|
|
def test_metric_history_decimates_at_cap():
|
|
from core.training import diffusion_training_service as svc_mod
|
|
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
svc._apply_event({"type": "model_load_started"}) # initialise running state
|
|
n = svc_mod._METRIC_CAP + 50
|
|
for i in range(1, n + 1):
|
|
svc._apply_event({"type": "progress", "step": i, "loss": 1.0 / i, "learning_rate": 1e-4})
|
|
st = svc.status()
|
|
# Never exceeds the cap, and stays a valid paired history with matching lengths.
|
|
assert len(st["metric_steps"]) <= svc_mod._METRIC_CAP
|
|
assert len(st["metric_steps"]) == len(st["metric_loss"]) == len(st["metric_lr"])
|
|
# Decimation keeps the curve monotonic in step (still increasing, just sparser).
|
|
assert st["metric_steps"] == sorted(st["metric_steps"])
|
|
assert st["metric_steps"][-1] == n # the latest point is always retained
|
|
|
|
|
|
def test_complete_event_records_family_and_catalog():
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
svc._apply_event(
|
|
{
|
|
"type": "complete",
|
|
"output_dir": "/o",
|
|
"lora_path": "/o/w.safetensors",
|
|
"catalog_path": "/loras/w.safetensors",
|
|
"family": "sdxl",
|
|
"base_model": "b",
|
|
}
|
|
)
|
|
st = svc.status()
|
|
assert st["status"] == "completed"
|
|
assert st["family"] == "sdxl"
|
|
assert st["base_model"] == "b"
|
|
assert st["catalog_path"] == "/loras/w.safetensors"
|
|
|
|
|
|
def test_status_route_nests_metric_history(client):
|
|
# The status route folds the service's flat arrays into a nested metric_history object.
|
|
client._fake.status_extra = {
|
|
"metric_steps": [1, 2],
|
|
"metric_loss": [0.5, 0.4],
|
|
"metric_lr": [1e-4, 9e-5],
|
|
"family": "sdxl",
|
|
"samples_per_second": 2.0,
|
|
"peak_memory_gb": 6.0,
|
|
}
|
|
r = client.get("/api/train/diffusion/status")
|
|
assert r.status_code == 200, r.text
|
|
body = r.json()
|
|
assert body["metric_history"]["steps"] == [1, 2]
|
|
assert body["metric_history"]["loss"] == [0.5, 0.4]
|
|
assert body["family"] == "sdxl"
|
|
assert body["samples_per_second"] == 2.0
|
|
|
|
|
|
# ── /diffusion/info families + gated-repo preflight (PR B) ──────────────────────
|
|
def test_info_lists_trainable_families(client):
|
|
r = client.get("/api/train/diffusion/info")
|
|
assert r.status_code == 200, r.text
|
|
families = {f["name"]: f for f in r.json()["families"]}
|
|
for fam in ("sdxl", "flux.1", "qwen-image", "z-image"):
|
|
assert fam in families
|
|
assert families["flux.1"]["default_base"] == "black-forest-labs/FLUX.1-dev"
|
|
assert families["z-image"]["defaults"]["resolution"] == 768
|
|
|
|
|
|
def test_start_gated_base_without_access_is_400_and_keeps_gpu(client, monkeypatch):
|
|
# A gated FLUX base with no valid token must 400 from the HEAD preflight BEFORE the GPU
|
|
# residents are freed, so a doomed start never evicts the user's loaded model.
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
import routes.training as tr
|
|
|
|
freed: list[int] = []
|
|
monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1))
|
|
|
|
def _fake_urlopen(req, timeout = None):
|
|
raise urllib.error.HTTPError(req.full_url, 403, "Forbidden", {}, None)
|
|
|
|
monkeypatch.setattr(urllib.request, "urlopen", _fake_urlopen)
|
|
r = client.post(
|
|
"/api/train/diffusion/start",
|
|
json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"},
|
|
)
|
|
assert r.status_code == 400
|
|
assert "gated" in r.json()["detail"].lower()
|
|
assert freed == []
|
|
assert client._fake.started_with is None
|
|
|
|
|
|
def test_start_ungated_base_preflight_is_noop(client, monkeypatch):
|
|
# A reachable base (HEAD 200) proceeds to start normally.
|
|
import urllib.request
|
|
|
|
import routes.training as tr
|
|
|
|
monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None)
|
|
monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout = None: object())
|
|
r = client.post(
|
|
"/api/train/diffusion/start",
|
|
json = {**_BODY, "base_model": "black-forest-labs/FLUX.1-dev"},
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
assert client._fake.started_with["base_model"] == "black-forest-labs/FLUX.1-dev"
|
|
|
|
|
|
# ── persisted run history ──────────────────────────────────────────────────────
|
|
def test_run_record_persisted_on_complete(_isolated_runs_dir):
|
|
# A completed run writes one JSON record: summary + scrubbed config + metric logs.
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
|
|
job_id = svc.start({**_CFG, "model_family": "z-image", "hf_token": "SECRET"})
|
|
_wait_status(svc, "completed")
|
|
# The pump persists right after the terminal event; give the thread a beat.
|
|
time.sleep(0.1)
|
|
|
|
import json
|
|
|
|
rec = json.loads((_isolated_runs_dir / f"{job_id}.json").read_text())
|
|
assert rec["job_id"] == job_id
|
|
assert rec["status"] == "completed"
|
|
assert rec["saved"] is True
|
|
assert rec["adapter"] == "out" # basename of /tmp/out
|
|
assert rec["family"] == "z-image" # falls back to the config's model_family
|
|
assert rec["step"] == 2 and rec["total_steps"] == 2
|
|
assert rec["avg_loss"] == 0.45
|
|
assert rec["metric_history"]["steps"] == [1, 2]
|
|
assert rec["metric_history"]["loss"] == [0.5, 0.4]
|
|
# Secrets never land on disk.
|
|
assert "hf_token" not in rec["config"]
|
|
assert rec["config"]["model_family"] == "z-image"
|
|
|
|
|
|
def test_run_record_no_save_stop_marks_unsaved(_isolated_runs_dir):
|
|
# A cancel (stop without save) persists too, flagged as not saved.
|
|
def _cancel_target(*, event_queue, stop_queue, config):
|
|
event_queue.put({"type": "model_load_completed"})
|
|
stop_queue.get(timeout = 5.0)
|
|
event_queue.put(
|
|
{"type": "complete", "output_dir": None, "lora_path": None, "stopped": True}
|
|
)
|
|
|
|
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _cancel_target)
|
|
job_id = svc.start(dict(_CFG))
|
|
_wait_status(svc, "running")
|
|
svc.stop(save = False)
|
|
_wait_status(svc, "stopped")
|
|
time.sleep(0.1)
|
|
|
|
import json
|
|
|
|
rec = json.loads((_isolated_runs_dir / f"{job_id}.json").read_text())
|
|
assert rec["status"] == "stopped"
|
|
assert rec["saved"] is False and rec["lora_path"] is None
|
|
|
|
|
|
def test_runs_endpoints_list_and_detail(client, _isolated_runs_dir):
|
|
# Seed two records directly (the endpoints read the persisted files, not the service).
|
|
import json
|
|
import os
|
|
|
|
a = {
|
|
"job_id": "a" * 32,
|
|
"status": "completed",
|
|
"adapter": "first",
|
|
"saved": True,
|
|
"step": 10,
|
|
"total_steps": 10,
|
|
"avg_loss": 0.4,
|
|
"config": {"train_steps": 10},
|
|
"metric_history": {"steps": [1], "loss": [0.4], "lr": [1e-4], "grad_norm": [0.2]},
|
|
}
|
|
b = {
|
|
"job_id": "b" * 32,
|
|
"status": "stopped",
|
|
"adapter": "second",
|
|
"saved": False,
|
|
"step": 3,
|
|
"total_steps": 10,
|
|
"avg_loss": 0.6,
|
|
"config": {"train_steps": 10},
|
|
"metric_history": {"steps": [1], "loss": [0.6], "lr": [1e-4], "grad_norm": [0.3]},
|
|
}
|
|
pa = _isolated_runs_dir / f"{a['job_id']}.json"
|
|
pb = _isolated_runs_dir / f"{b['job_id']}.json"
|
|
pa.write_text(json.dumps(a))
|
|
pb.write_text(json.dumps(b))
|
|
os.utime(pa, (1000, 1000))
|
|
os.utime(pb, (2000, 2000)) # b is newer -> listed first
|
|
|
|
r = client.get("/api/train/diffusion/runs")
|
|
assert r.status_code == 200, r.text
|
|
runs = r.json()["runs"]
|
|
assert [x["adapter"] for x in runs] == ["second", "first"]
|
|
# Summaries stay light: no config / metric logs.
|
|
assert "config" not in runs[0] and "metric_history" not in runs[0]
|
|
|
|
r = client.get(f"/api/train/diffusion/runs/{a['job_id']}")
|
|
assert r.status_code == 200, r.text
|
|
detail = r.json()
|
|
assert detail["adapter"] == "first"
|
|
assert detail["metric_history"]["grad_norm"] == [0.2]
|
|
assert detail["config"] == {"train_steps": 10}
|
|
|
|
# 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
|
|
|
|
|
|
def test_run_detail_route_non_object_record_is_404(client, _isolated_runs_dir):
|
|
# A valid-JSON but non-object record (a truncated / hand-edited [] file named with a real
|
|
# job id) makes DiffusionTrainingRunDetail(**rec) raise TypeError, not ValidationError; the
|
|
# detail route must shape-check like the list path and 404 instead of 500.
|
|
import json
|
|
|
|
job_id = "a" * 32
|
|
(_isolated_runs_dir / f"{job_id}.json").write_text(json.dumps([]))
|
|
|
|
r = client.get(f"/api/train/diffusion/runs/{job_id}")
|
|
assert r.status_code == 404, r.text
|