From ca68fd5d131c758c5cd45df5ddd58856b57f4b83 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 22:18:02 +0000 Subject: [PATCH] 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 | 17 ++++- studio/backend/routes/export.py | 77 +++++++++++++++++++++- 2 files changed, 90 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index e3bb476dbb..9921d0543d 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/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 ──────────────────────────────────────────────── diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 07548d4508..1a69de749a 100644 --- a/studio/backend/routes/export.py +++ b/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(