ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Wire diffusion LoRA training into the Studio API

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-01 14:47:06 +00:00
commit fbeb6dfc6f
5 changed files with 608 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

224
studio/backend/core/training/diffusion_training_service.py Normal file
View file

@ -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

6
studio/backend/models/__init__.py
View file

@ -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",

63
studio/backend/models/training.py
View file

@ -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

51
studio/backend/routes/training.py
View file

@ -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())

264
studio/backend/tests/test_diffusion_training.py Normal file
View file

@ -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"