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

Merge remote-tracking branch 'origin/diffusion-lora-training-api' into diffusion-lora-training-ui

Browse source
This commit is contained in:
oobabooga 2026-07-01 14:33:17 -03:00
commit 92963164ca
4 changed files with 108 additions and 71 deletions
Download patch file Download diff file Expand all files Collapse all files

65
studio/backend/core/training/diffusion_training_service.py
View file

@ -34,8 +34,7 @@ _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)
run_diffusion_training_process(event_queue = event_queue, stop_queue = stop_queue, config = config)
def _idle_state() -> dict[str, Any]:
@ -94,20 +93,23 @@ class DiffusionTrainingService:
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)
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,
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
@ -115,11 +117,15 @@ class DiffusionTrainingService:
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,
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
target = self._pump_loop, args = (event_queue, self._proc), daemon = True
)
self._pump.start()
return job_id
@ -149,7 +155,7 @@ class DiffusionTrainingService:
def _pump_loop(self, event_queue: Any, proc: Any) -> None:
while True:
try:
ev = event_queue.get(timeout=1.0)
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.
@ -163,9 +169,10 @@ class DiffusionTrainingService:
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(),
active = False,
status = "error",
message = "Training process exited unexpectedly.",
updated_at = time.time(),
)
_ = drained
return
@ -182,33 +189,33 @@ class DiffusionTrainingService:
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...")
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...")
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...",
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)."
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")))
s.update(active = False, status = "error", message = str(ev.get("message", "error")))
_service: Optional[DiffusionTrainingService] = None

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

@ -675,28 +675,30 @@ class DiffusionTrainingStartRequest(BaseModel):
rest carry the trainer's defaults.
"""
model_config = ConfigDict(protected_namespaces=())
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")
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"
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)
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")
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)
lr_warmup_steps: int = Field(0, ge = 0)
center_crop: bool = Field(False)
random_flip: bool = Field(True)
caption_column: str = Field("text")

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

@ -1031,8 +1031,7 @@ async def stream_training_progress(
@router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse)
async def start_diffusion_training(
body: DiffusionTrainingStartRequest,
current_subject: str = Depends(get_current_subject),
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
@ -1047,8 +1046,11 @@ async def start_diffusion_training(
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,
e,
500,
"Failed to start diffusion training",
event = "diffusion_training.start_failed",
log = logger,
)
return DiffusionTrainingStartResponse(job_id = job_id, status = "running")
@ -1066,5 +1068,4 @@ async def stop_diffusion_training(current_subject: str = Depends(get_current_sub
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())

69
studio/backend/tests/test_diffusion_training.py
View file

@ -32,8 +32,8 @@ class _FakeQueue:
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(self, timeout = None):
return self._q.get(timeout = timeout) # raises queue.Empty on timeout
def get_nowait(self):
return self._q.get_nowait()
@ -50,7 +50,7 @@ class _FakeProc:
self.pid = 4321
def start(self):
self._thread = threading.Thread(target=self._target, kwargs=self._kwargs, daemon=True)
self._thread = threading.Thread(target = self._target, kwargs = self._kwargs, daemon = True)
self._thread.start()
def is_alive(self):
@ -69,10 +69,24 @@ 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}
{
"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}
{
"type": "progress",
"step": 2,
"total_steps": 2,
"loss": 0.4,
"avg_loss": 0.45,
"learning_rate": 1e-4,
}
)
event_queue.put(
{
@ -86,8 +100,10 @@ def _happy_target(*, event_queue, stop_queue, config):
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})
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):
@ -98,7 +114,11 @@ def _crashing_target(*, event_queue, stop_queue, config):
_CFG = {"base_model": "b", "data_dir": "d", "output_dir": "/tmp/out", "train_steps": 2}
def _wait_status(svc, *terminal, timeout=3.0):
def _wait_status(
svc,
*terminal,
timeout = 3.0,
):
end = time.time() + timeout
while time.time() < end:
st = svc.status()
@ -109,7 +129,7 @@ def _wait_status(svc, *terminal, timeout=3.0):
def test_service_happy_path():
svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_happy_target)
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
job_id = svc.start(dict(_CFG))
assert job_id
st = _wait_status(svc, "completed")
@ -122,7 +142,7 @@ def test_service_happy_path():
def test_service_rejects_bad_config_before_spawn():
svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_happy_target)
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target)
with pytest.raises(ValueError):
svc.start({**_CFG, "train_steps": 0})
# Nothing was spawned; still idle.
@ -130,7 +150,7 @@ def test_service_rejects_bad_config_before_spawn():
def test_service_rejects_second_concurrent_job():
svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_stoppable_target)
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _stoppable_target)
svc.start(dict(_CFG))
_wait_status(svc, "running")
with pytest.raises(RuntimeError):
@ -140,7 +160,7 @@ def test_service_rejects_second_concurrent_job():
def test_service_stop_marks_stopped():
svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_stoppable_target)
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _stoppable_target)
svc.start(dict(_CFG))
_wait_status(svc, "running")
assert svc.stop() is True
@ -152,7 +172,7 @@ def test_service_stop_marks_stopped():
def test_service_crash_without_terminal_event_is_error():
svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_crashing_target)
svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _crashing_target)
svc.start(dict(_CFG))
st = _wait_status(svc, "error")
assert st["status"] == "error"
@ -160,7 +180,7 @@ def test_service_crash_without_terminal_event_is_error():
def test_apply_event_transitions():
svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_happy_target)
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"})
@ -212,25 +232,32 @@ def client(monkeypatch):
"core.training.diffusion_training_service.get_diffusion_training_service", lambda: fake
)
app = FastAPI()
app.include_router(training_router, prefix="/api/train")
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}
_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)
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
r = client.post(
"/api/train/diffusion/start", json = {"base_model": "x"}
) # no data_dir/output_dir
assert r.status_code == 422
@ -239,7 +266,7 @@ def test_route_start_bad_config_maps_to_400(client, monkeypatch):
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)
r = client.post("/api/train/diffusion/start", json = _BODY)
assert r.status_code == 400
assert "multiple of 8" in r.json()["detail"]
@ -249,12 +276,12 @@ def test_route_start_conflict_maps_to_409(client):
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)
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)
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")