5f65f01e2e
Run backend.start_training off the event loop with asyncio.to_thread so the synchronous diffusion/video unload calls (which wait on engine generation locks) cannot freeze concurrent requests; guard against overlapping starts with a _start_in_progress compare-and-set under the service lock. Resolve bare diffusion dataset names directly under datasets_root() before falling back to the generic resolver, so an unrelated LLM upload file or recipe folder sharing the name cannot shadow the image dataset. Reject exact duplicate filenames within one multipart upload batch: two parts staged to the same destination would let the later tmp.replace silently discard the earlier file. Case variants stay exempt per the existing stem-guard contract. Require an instance prompt in the train panel when only some images have captions, since backend discovery silently skips uncaptioned images.
1498 lines
60 KiB
Python
1498 lines
60 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 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
|