From fbeb6dfc6f73a6c2f6d278dbd782b099732970c0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 14:47:06 +0000 Subject: [PATCH] Wire diffusion LoRA training into the Studio API Make the SDXL LoRA trainer reachable from the app with a small, self-contained job service and JSON routes, deliberately separate from the LLM TrainingBackend (whose lifecycle -- LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat- inference -- is text-training specific and would mis-handle a diffusion run). core/training/diffusion_training_service.py: DiffusionTrainingService runs one job at a time -- validate the config cheaply (before any spawn), spawn the trainer subprocess (spawn context, parent-lifetime bound), pump its events (model_load_* / progress / complete / error) into an in-memory status snapshot, and support a clean stop. The subprocess context and target are injectable so the full start -> pump -> status -> complete path is unit-tested without real multiprocessing or torch. routes/training.py: POST /api/train/diffusion/start (400 on a bad config, 409 when a job is already running), POST /api/train/diffusion/stop, GET /api/train/diffusion/status (JSON poll). models/training.py: DiffusionTrainingStartRequest + response schemas mirroring DiffusionLoraConfig, so model_dump() passes straight through. Tests: test_diffusion_training.py -- service happy path, bad-config-before-spawn, concurrent-job rejection, clean stop, crash-without-terminal-event, event transitions; plus route wiring via the FastAPI TestClient (start / 422 / 400 / 409 / status / stop) with a mocked service. The diffusion trainer's progress events already use the field names this path expects. --- .../training/diffusion_training_service.py | 224 +++++++++++++++ studio/backend/models/__init__.py | 6 + studio/backend/models/training.py | 63 +++++ studio/backend/routes/training.py | 51 ++++ .../backend/tests/test_diffusion_training.py | 264 ++++++++++++++++++ 5 files changed, 608 insertions(+) create mode 100644 studio/backend/core/training/diffusion_training_service.py create mode 100644 studio/backend/tests/test_diffusion_training.py diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py new file mode 100644 index 0000000000..166d2a9193 --- /dev/null +++ b/studio/backend/core/training/diffusion_training_service.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""A small job service for diffusion LoRA training. + +This is deliberately separate from the LLM ``TrainingBackend``: that backend's +lifecycle (LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat- +inference) is specific to text training and would mis-handle a diffusion run. This +service does only what a diffusion job needs -- spawn the trainer subprocess, pump its +events (``model_load_*`` / ``progress`` / ``complete`` / ``error``) into an in-memory +status snapshot, and support stop -- and is polled over JSON by the route layer. + +The subprocess context, target, and queues are injectable so the service can be unit +tested without real multiprocessing or torch: tests pass a fake context whose Process +runs a scripted target on a thread. +""" + +from __future__ import annotations + +import multiprocessing as mp +import threading +import time +import uuid +from typing import Any, Callable, Optional + +# Spawn (not fork): a fresh interpreter, matching the LLM training worker, so CUDA/torch +# state from the parent never leaks into the trainer. +_CTX = mp.get_context("spawn") + +# Terminal event types after which the pump stops. +_TERMINAL = ("complete", "error") + + +def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None: + # Imported lazily so this module (and the route layer) stays torch-free at import. + from .diffusion_lora_trainer import run_diffusion_training_process + + run_diffusion_training_process(event_queue=event_queue, stop_queue=stop_queue, config=config) + + +def _idle_state() -> dict[str, Any]: + return { + "active": False, + "job_id": None, + "status": "idle", + "message": "", + "step": 0, + "total_steps": 0, + "loss": None, + "avg_loss": None, + "learning_rate": None, + "num_images": None, + "in_model_load": False, + "output_dir": None, + "lora_path": None, + "started_at": None, + "updated_at": None, + } + + +class DiffusionTrainingService: + """One diffusion LoRA training job at a time, spawned as a subprocess.""" + + def __init__( + self, + *, + ctx: Any = None, + target: Optional[Callable[..., None]] = None, + ) -> None: + self._ctx = ctx if ctx is not None else _CTX + self._target = target if target is not None else _default_target + self._lock = threading.Lock() + self._proc: Any = None + self._stop_queue: Any = None + self._pump: Optional[threading.Thread] = None + self._state: dict[str, Any] = _idle_state() + + # ── lifecycle ──────────────────────────────────────────────────────────── + def is_active(self) -> bool: + with self._lock: + return self._proc is not None and self._proc.is_alive() + + def start(self, config: dict) -> str: + """Validate ``config``, spawn the trainer, and start pumping its events. + + Raises ValueError for an unusable config (before any spawn) and RuntimeError if a + job is already running. Returns the new job id.""" + # Validate cheaply BEFORE spawning so a bad request fails fast with a clear error. + from .diffusion_lora_trainer import _config_from_dict + + _config_from_dict(config).normalized() + + with self._lock: + if self._proc is not None and self._proc.is_alive(): + raise RuntimeError("A diffusion training job is already running.") + if self._pump is not None and self._pump.is_alive(): + self._pump.join(timeout=5.0) + + job_id = uuid.uuid4().hex + event_queue = self._ctx.Queue() + self._stop_queue = self._ctx.Queue() + self._proc = self._ctx.Process( + target=self._target, + kwargs={"event_queue": event_queue, "stop_queue": self._stop_queue, "config": config}, + daemon=True, + ) + self._proc.start() + try: + from utils.process_lifetime import adopt_pid + + adopt_pid(self._proc.pid) # bind to parent lifetime (no zombie on exit) + except Exception: # noqa: BLE001 -- lifetime binding is best-effort + pass + + now = time.time() + self._state = _idle_state() + self._state.update( + active=True, job_id=job_id, status="running", + message="Starting diffusion LoRA training...", started_at=now, updated_at=now, + ) + self._pump = threading.Thread( + target=self._pump_loop, args=(event_queue, self._proc), daemon=True + ) + self._pump.start() + return job_id + + def stop(self) -> bool: + """Request a clean stop (the trainer finishes the current step and saves a partial + adapter). Returns True if a stop was signalled, False if nothing was running.""" + with self._lock: + if self._proc is None or not self._proc.is_alive() or self._stop_queue is None: + return False + try: + self._stop_queue.put(True) + except Exception: # noqa: BLE001 + return False + self._state["message"] = "Stop requested; finishing the current step..." + self._state["updated_at"] = time.time() + return True + + def status(self) -> dict[str, Any]: + with self._lock: + snap = dict(self._state) + # Keep ``active`` honest even if the process died between events. + snap["active"] = self._proc is not None and self._proc.is_alive() + return snap + + # ── event pump ─────────────────────────────────────────────────────────── + def _pump_loop(self, event_queue: Any, proc: Any) -> None: + while True: + try: + ev = event_queue.get(timeout=1.0) + except Exception: # noqa: BLE001 -- Empty (timeout) or a closed queue + if not proc.is_alive(): + # Drain anything buffered, then decide if it exited cleanly. + drained = False + while True: + try: + self._apply_event(event_queue.get_nowait()) + drained = True + except Exception: # noqa: BLE001 + break + with self._lock: + if self._state.get("status") not in ("completed", "stopped", "error"): + self._state.update( + active=False, status="error", + message="Training process exited unexpectedly.", + updated_at=time.time(), + ) + _ = drained + return + continue + self._apply_event(ev) + if ev.get("type") in _TERMINAL: + return + + def _apply_event(self, ev: dict[str, Any]) -> None: + """Fold one trainer event into the status snapshot. Pure state update -- unit + tested by feeding events directly.""" + etype = ev.get("type") + with self._lock: + s = self._state + s["updated_at"] = time.time() + if etype == "model_load_started": + s.update(in_model_load=True, status="running", message="Loading base model...") + if ev.get("num_images") is not None: + s["num_images"] = ev.get("num_images") + elif etype == "model_load_completed": + s.update(in_model_load=False, message="Training...") + elif etype == "progress": + s.update( + status="running", + step=ev.get("step", s["step"]), + total_steps=ev.get("total_steps", s["total_steps"]), + loss=ev.get("loss", s["loss"]), + avg_loss=ev.get("avg_loss", s["avg_loss"]), + learning_rate=ev.get("learning_rate", s["learning_rate"]), + message="Training...", + ) + elif etype == "complete": + s.update( + active=False, + status="stopped" if ev.get("stopped") else "completed", + output_dir=ev.get("output_dir"), + lora_path=ev.get("lora_path"), + message="Stopped (partial adapter saved)." + if ev.get("stopped") + else "Training complete.", + ) + elif etype == "error": + s.update(active=False, status="error", message=str(ev.get("message", "error"))) + + +_service: Optional[DiffusionTrainingService] = None +_service_lock = threading.Lock() + + +def get_diffusion_training_service() -> DiffusionTrainingService: + """Process-wide singleton used by the route layer.""" + global _service + with _service_lock: + if _service is None: + _service = DiffusionTrainingService() + return _service diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py index dbedeeca00..8ece0a569e 100644 --- a/studio/backend/models/__init__.py +++ b/studio/backend/models/__init__.py @@ -14,6 +14,9 @@ from .training import ( TrainingRunDetailResponse, TrainingRunDeleteResponse, TrainingRunUpdateRequest, + DiffusionTrainingStartRequest, + DiffusionTrainingStartResponse, + DiffusionTrainingStatusResponse, ) from .models import ( CheckpointInfo, @@ -72,6 +75,9 @@ from .data_recipe import ( __all__ = [ # Training schemas "TrainingStartRequest", + "DiffusionTrainingStartRequest", + "DiffusionTrainingStartResponse", + "DiffusionTrainingStatusResponse", "TrainingJobResponse", "TrainingStatus", "TrainingProgress", diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index ff815a2fa9..85f0949810 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -665,3 +665,66 @@ class TrainingRunDeleteResponse(BaseModel): status: str message: str + + +class DiffusionTrainingStartRequest(BaseModel): + """Request to start a diffusion (SDXL) LoRA training job. + + Field names mirror ``core.training.diffusion_lora_trainer.DiffusionLoraConfig`` so the + service can pass ``model_dump()`` straight through. Only the paths are required; the + rest carry the trainer's defaults. + """ + + model_config = ConfigDict(protected_namespaces=()) + + base_model: str = Field(..., description="HF repo id or local path to an SDXL pipeline") + data_dir: str = Field(..., description="Folder of training images (+ captions)") + output_dir: str = Field(..., description="Directory to write the LoRA .safetensors into") + instance_prompt: Optional[str] = Field( + None, description="Dreambooth caption applied to images without their own caption" + ) + resolution: int = Field(1024, ge=64, le=2048, description="Square training resolution (multiple of 8)") + train_steps: int = Field(500, ge=1, le=100000) + learning_rate: float = Field(1e-4, gt=0) + train_batch_size: int = Field(1, ge=1, le=64) + gradient_accumulation_steps: int = Field(1, ge=1, le=256) + lora_rank: int = Field(16, ge=1, le=320) + lora_alpha: Optional[int] = Field(None, ge=1, le=640, description="Defaults to lora_rank") + lora_dropout: float = Field(0.0, ge=0.0, le=1.0) + seed: int = Field(42) + mixed_precision: Literal["bf16", "fp16", "no"] = Field("bf16") + snr_gamma: Optional[float] = Field(5.0, description="Min-SNR loss weighting; null disables") + gradient_checkpointing: bool = Field(True) + lr_scheduler: str = Field("constant") + lr_warmup_steps: int = Field(0, ge=0) + center_crop: bool = Field(False) + random_flip: bool = Field(True) + caption_column: str = Field("text") + hf_token: Optional[str] = Field(None) + + +class DiffusionTrainingStartResponse(BaseModel): + """Response for starting a diffusion training job.""" + + job_id: str + status: str + + +class DiffusionTrainingStatusResponse(BaseModel): + """A snapshot of the current diffusion training job (or idle).""" + + active: bool + job_id: Optional[str] = None + status: str + message: str = "" + step: int = 0 + total_steps: int = 0 + loss: Optional[float] = None + avg_loss: Optional[float] = None + learning_rate: Optional[float] = None + num_images: Optional[int] = None + in_model_load: bool = False + output_dir: Optional[str] = None + lora_path: Optional[str] = None + started_at: Optional[float] = None + updated_at: Optional[float] = None diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index d5ab4ae29a..b2ea6dfaa4 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -56,6 +56,9 @@ from models import ( TrainingJobResponse, TrainingStatus, TrainingProgress, + DiffusionTrainingStartRequest, + DiffusionTrainingStartResponse, + DiffusionTrainingStatusResponse, ) from models.responses import TrainingStopResponse, TrainingMetricsResponse from pydantic import BaseModel as PydanticBaseModel @@ -1017,3 +1020,51 @@ async def stream_training_progress( "X-Accel-Buffering": "no", }, ) + + +# ── Diffusion (SDXL) LoRA training ──────────────────────────────────────────── +# A separate, lightweight job path from the LLM training endpoints above: diffusion +# runs are driven by DiffusionTrainingService (its own subprocess + event pump), not +# the LLM TrainingBackend, so the two never contend and diffusion never triggers LLM +# lifecycle (DB run rows, plots, transfer-to-chat-inference). + + +@router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse) +async def start_diffusion_training( + body: DiffusionTrainingStartRequest, + current_subject: str = Depends(get_current_subject), +): + """Start an SDXL LoRA training job from an image + caption dataset.""" + from core.training.diffusion_training_service import get_diffusion_training_service + + service = get_diffusion_training_service() + try: + job_id = service.start(body.model_dump()) + except ValueError as e: + raise HTTPException(status_code = 400, detail = str(e)) + except RuntimeError as e: + # A job is already running. + raise HTTPException(status_code = 409, detail = str(e)) + except Exception as e: + raise log_and_http_error( + e, 500, "Failed to start diffusion training", + event = "diffusion_training.start_failed", log = logger, + ) + return DiffusionTrainingStartResponse(job_id = job_id, status = "running") + + +@router.post("/diffusion/stop") +async def stop_diffusion_training(current_subject: str = Depends(get_current_subject)): + """Request a clean stop of the running diffusion training job (partial adapter saved).""" + from core.training.diffusion_training_service import get_diffusion_training_service + + stopped = get_diffusion_training_service().stop() + return {"status": "stopping" if stopped else "idle"} + + +@router.get("/diffusion/status", response_model = DiffusionTrainingStatusResponse) +async def diffusion_training_status(current_subject: str = Depends(get_current_subject)): + """Poll the current diffusion training job's status/progress (JSON).""" + from core.training.diffusion_training_service import get_diffusion_training_service + + return DiffusionTrainingStatusResponse(**get_diffusion_training_service().status()) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py new file mode 100644 index 0000000000..ae95e6ce7e --- /dev/null +++ b/studio/backend/tests/test_diffusion_training.py @@ -0,0 +1,264 @@ +# 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 queue as _queue +import threading +import time + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from auth.authentication import 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) + + +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" + + +# ── route wiring (mocked service) ───────────────────────────────────────────── +class _FakeService: + def __init__(self): + self._running = False + self.started_with = None + + def start(self, config): + self.started_with = config + self._running = True + return "job-123" + + def stop(self): + 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, + } + + +@pytest.fixture +def client(monkeypatch): + fake = _FakeService() + monkeypatch.setattr( + "core.training.diffusion_training_service.get_diffusion_training_service", lambda: fake + ) + app = FastAPI() + app.include_router(training_router, prefix="/api/train") + app.dependency_overrides[get_current_subject] = lambda: "test-user" + c = TestClient(app) + c._fake = fake # type: ignore[attr-defined] + return c + + +_BODY = {"base_model": "stabilityai/sdxl-turbo", "data_dir": "/data", "output_dir": "/out", "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" + + +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_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"