From c2b25feaeefe80236a3ff7ed1ffe3d4c52965531 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Jul 2026 05:47:50 +0000 Subject: [PATCH] Guard inference loads and worker lifetime against diffusion training Teach the chat and image load guards about an active diffusion (SDXL) LoRA job: a chat load is refused (its footprint cannot be fit-checked against the trainer) and an image load is refused outright, mirroring the existing LLM training guards, so a load can no longer allocate GPU memory alongside the trainer and undo the pre-start cleanup. Bind the diffusion trainer subprocess to the parent's lifetime and scrub the native path lease secret from it by running the child through run_without_native_path_secret, matching the inference/export/LLM workers, so a Studio crash or kill no longer leaves the trainer holding the GPU. Reset in_model_load on the complete and error terminal events: a stop or failure during model loading otherwise leaves the status reporting a stale loading indicator after the job has ended. --- .../training/diffusion_training_service.py | 29 +++++++++++++-- studio/backend/routes/inference.py | 36 ++++++++++++++++--- .../backend/tests/test_diffusion_training.py | 16 +++++++++ 3 files changed, 75 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 92621931f7..bc11b38a37 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -31,12 +31,26 @@ _CTX = mp.get_context("spawn") _TERMINAL = ("complete", "error") -def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None: +def _run_diffusion_child(*, 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 _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None: + # First thing in the spawned child (before torch is imported): bind to the parent's + # death on Linux and scrub the native path lease secret, exactly like the inference / + # export / LLM-training workers. multiprocessing children cannot be given a + # parent-set preexec_fn, so the child must self-bind; otherwise a Studio crash or + # kill leaves this trainer holding the GPU. Tests inject their own target, so this + # binding only runs for the real production spawn. + from utils.native_path_leases import run_without_native_path_secret + + run_without_native_path_secret( + _run_diffusion_child, event_queue = event_queue, stop_queue = stop_queue, config = config + ) + + def _idle_state() -> dict[str, Any]: return { "active": False, @@ -225,8 +239,12 @@ class DiffusionTrainingService: message = "Training...", ) elif etype == "complete": + # Reset in_model_load: a stop during model load emits complete without a + # preceding model_load_completed, which would otherwise leave a stale + # loading indicator after the job ended. s.update( active = False, + in_model_load = False, status = "stopped" if ev.get("stopped") else "completed", output_dir = ev.get("output_dir"), lora_path = ev.get("lora_path"), @@ -235,7 +253,14 @@ class DiffusionTrainingService: else "Training complete.", ) elif etype == "error": - s.update(active = False, status = "error", message = str(ev.get("message", "error"))) + # Reset in_model_load too: an error raised during model loading has no + # model_load_completed, so the terminal state must clear it explicitly. + s.update( + active = False, + in_model_load = False, + status = "error", + message = str(ev.get("message", "error")), + ) _service: Optional[DiffusionTrainingService] = None diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5f1a5103a4..648f07ea2f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2823,12 +2823,26 @@ def _guard_chat_load_against_training( from routes.training_vram import can_load_chat_during_training try: - if not get_training_backend().is_training_active(): - return + llm_active = get_training_backend().is_training_active() except Exception as e: logger.warning("Could not check training state for chat-load guard: %s", e) return + if not llm_active: + # An SDXL LoRA trainer runs in its own subprocess and its VRAM can't be cheaply + # fit-checked here, so refuse the chat load outright while one is active rather + # than risk OOMing the run. Symmetric with the image-load guard. + if _diffusion_training_active(): + raise HTTPException( + status_code = 409, + detail = ( + "Can't load this model while diffusion (Images) training is running: " + "its GPU memory use can't be verified against the trainer, so the load " + "was refused to protect the run. Try again after training finishes." + ), + ) + return + is_gguf = bool(getattr(config, "is_gguf", False)) required_override_gb = ( _estimate_gguf_required_gb( @@ -11021,6 +11035,16 @@ async def _openai_passthrough_non_streaming( # ────────────────────────────────────────────────────────────────────────── +def _diffusion_training_active() -> bool: + """Whether a diffusion (SDXL) LoRA job is running. Best-effort so a load is never + blocked just because the training service could not be imported/read.""" + try: + from core.training.diffusion_training_service import get_diffusion_training_service + return get_diffusion_training_service().is_active() + except Exception: # noqa: BLE001 + return False + + def _guard_diffusion_load_against_training() -> None: """Refuse loading an image model while a training run is active. Unlike chat, a diffusion pipeline's VRAM can't be cheaply estimated before the load, so the @@ -11029,11 +11053,15 @@ def _guard_diffusion_load_against_training() -> None: from core.training import get_training_backend try: - if not get_training_backend().is_training_active(): - return + llm_active = get_training_backend().is_training_active() except Exception as e: logger.warning("Could not check training state for image-load guard: %s", e) return + # An SDXL LoRA trainer runs in its own subprocess on the same GPU, so an image + # load must be refused while one is active too -- otherwise the resident pipeline + # competes with the trainer for VRAM. Symmetric with the diffusion-start interlock. + if not llm_active and not _diffusion_training_active(): + return raise HTTPException( status_code = 409, detail = ( diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index c96dabb76f..35ff816046 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -189,6 +189,22 @@ def test_apply_event_transitions(): assert svc.status()["status"] == "error" and svc.status()["message"] == "boom" +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):