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.

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

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 = (

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):