Fix/adjust diffusion: export active-state guards, cleanup window, unload race for PR #5754

Round 41 review findings (2 P1, 5/12 reviewers consensus on the dominant one):

1. routes/export.py: load_checkpoint already refuses 409 when training
   or another export is active, but /export/{merged,base,gguf,lora} and
   /cleanup went through _export_public_window without those checks.
   A user could start training, then trigger an export (or cleanup),
   and both would double-own the GPU. Factor the training-active and
   export-active guards into _raise_if_training_active_for_export and
   _raise_if_export_active_for_export, call them inside the context
   manager so all /export/* + /cleanup share the same fail-closed
   semantics as load_checkpoint, and wrap /cleanup with the window.

2. core/inference/diffusion.py: DiffusionBackend.unload_model cleared
   _pipe / _repo_id / _family / ... under _lock BEFORE _release(old)
   and _drain_cuda_cache. Between the lock release and cache drain,
   status() reported is_loaded=False / is_loading=False, so the
   helper-busy check (which OR-s those two) could let an AI Assist
   GGUF backend start while diffusion VRAM was still being freed.
   Set _loading=True inside the lock as a busy marker before clearing
   the slot, and only clear it in a finally after release + drain
   complete.

studio/backend/core/inference/diffusion.py

@@ -1349,6 +1349,13 @@ class DiffusionBackend:
         with self._load_lock, self._generate_lock:
             with self._lock:
                 old = self._pipe
+                # Mark the slot as busy BEFORE clearing _pipe so a
+                # concurrent helper-busy check (which treats either
+                # is_loaded OR is_loading as busy) does not see a
+                # ``free`` GPU during the release + cache-drain window.
+                # is_loading is cleared in finally once the VRAM is
+                # actually freed.
+                self._loading = True
                 self._pipe = None
                 self._family = None
                 self._repo_id = None
@@ -1359,9 +1366,13 @@ class DiffusionBackend:
                 self._dtype = None
                 self._cpu_offload_enabled = False
                 self._loaded_at = None
-            _release(old)
-            old = None  # noqa: F841
-            _drain_cuda_cache()
+            try:
+                _release(old)
+                old = None  # noqa: F841
+                _drain_cuda_cache()
+            finally:
+                with self._lock:
+                    self._loading = False
         return {"is_loaded": False}
 
     # ── generation ────────────────────────────────────────────────

studio/backend/routes/export.py

@@ -53,6 +53,68 @@ logger = get_logger(__name__)
 import contextlib
 
 
+def _raise_if_training_active_for_export() -> None:
+    """409 if a training run is in flight; 503 if status check itself
+    raises. Mirrors the load_checkpoint guard so /export/* and /cleanup
+    never tear down or alter export state while training is using the
+    GPU. Missing core.training is treated as 'no tracker'."""
+    try:
+        from core.training import get_training_backend  # type: ignore
+    except Exception as e:
+        logger.debug("core.training not importable, skipping training guard: %s", e)
+        return
+    try:
+        trn = get_training_backend()
+        active = trn.is_training_active()
+    except Exception as e:
+        logger.warning("Could not verify training status before export op: %s", e)
+        raise HTTPException(
+            status_code = 503,
+            detail = (
+                "Could not verify training status before the export "
+                "operation. Try again."
+            ),
+        ) from e
+    if active:
+        raise HTTPException(
+            status_code = 409,
+            detail = (
+                "Training is currently active. Stop the training run "
+                "before starting an export operation."
+            ),
+        )
+
+
+def _raise_if_export_active_for_export() -> None:
+    """409 if another export job is already running; 503 if the status
+    check itself raises. Backends without is_export_active() are
+    treated as 'no tracker available' to stay compatible with mocked
+    backends in tests."""
+    backend = get_export_backend()
+    is_export_active_fn = getattr(backend, "is_export_active", None)
+    if is_export_active_fn is None:
+        return
+    try:
+        export_is_active = bool(is_export_active_fn())
+    except Exception as e:
+        logger.warning("Could not verify export status before export op: %s", e)
+        raise HTTPException(
+            status_code = 503,
+            detail = (
+                "Could not verify export status before starting the "
+                "export operation. Try again."
+            ),
+        ) from e
+    if export_is_active:
+        raise HTTPException(
+            status_code = 409,
+            detail = (
+                "An export job is currently active. Wait for it to "
+                "finish before starting another export operation."
+            ),
+        )
+
+
 @contextlib.asynccontextmanager
 async def _export_public_window():
     """Publish the public-load window across an /export/* operation.
@@ -64,6 +126,12 @@ async def _export_public_window():
     subprocess. Mirror the load_checkpoint guard so the pending counter
     is set for the whole export call, and the helper-busy preflight
     refuses if AI Assist is mid-handoff.
+
+    Also refuses 409 if training or another export is already active so
+    a queued /export/{merged,base,gguf,lora} or /cleanup cannot
+    double-own the GPU with a running training / export job (round 41
+    consensus: load_checkpoint already runs these checks but /export/*
+    and /cleanup were skipping them).
     """
     from routes.inference import (
         _clear_public_load_window,
@@ -72,6 +140,8 @@ async def _export_public_window():
 
     export_window_published = False
     try:
+        _raise_if_training_active_for_export()
+        _raise_if_export_active_for_export()
         _raise_if_helper_advisor_busy("export")
         export_window_published = True
         yield
@@ -258,7 +328,12 @@ async def cleanup_export_memory(
     """
     try:
         backend = get_export_backend()
-        success = await asyncio.to_thread(backend.cleanup_memory)
+        # Run the cleanup under the same public-load window /export/*
+        # uses so a queued export's handoff gap cannot race a cleanup
+        # call that tears down current_checkpoint. The window also
+        # refuses 409 if training or another export is in flight.
+        async with _export_public_window():
+            success = await asyncio.to_thread(backend.cleanup_memory)
 
         if not success:
             raise HTTPException(