From e30c5ed386b2e2df6085152402250917d114dbac Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 20:43:04 +0000 Subject: [PATCH] Fix/adjust diffusion: backend public_load_pending parity for PR #5754 Round 38 P1: R35e added the public_load_pending() check to the route-side _raise_if_helper_advisor_busy, but the backend-side _raise_if_helper_advisor_busy_for_diffusion (used by direct DiffusionBackend.load_model callers AND called transitively from the route via backend.load_model) never got the same parity check. That left a window where: * /api/training/start published the "training" pending marker via the route helper * a script/test calling DiffusionBackend.load_model() directly passed the backend's helper-busy snapshot, never checked public_load_pending(), and proceeded to destructive owner teardown + GPU allocation while training was still pending. Add the parity check with a kw-only `excluding` parameter on public_load_pending so a route-wrapped backend call can ignore the marker its own route already published (route publishes "diffusion"; backend publishes the separate "diffusion-backend" tag). load_model gains ignore_public_load_pending_workload to thread the route's tag through; the diffusion route passes "diffusion" so the backend's atomic check does not self-block on the route's own publication. Verified by smoke test: route-wrapped backend with excluding= "diffusion" allowed during route's diffusion pending; direct backend call refused with RuntimeError "Another GPU workload is mid-handoff" when training is pending. 86 backend tests pass. --- studio/backend/core/inference/diffusion.py | 20 ++++++++++++++++++++ studio/backend/routes/inference.py | 6 ++++++ studio/backend/utils/datasets/llm_assist.py | 18 +++++++++++++++--- 3 files changed, 41 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 5c4450fd15..e15c4daea2 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -763,6 +763,7 @@ class DiffusionBackend: hf_token: Optional[str] = None, family_override: Optional[str] = None, enable_model_cpu_offload: bool = True, + ignore_public_load_pending_workload: Optional[str] = None, ) -> dict[str, Any]: """Load a diffusion model. @@ -1087,6 +1088,9 @@ class DiffusionBackend: # contribute to public_load_pending(). backend_pending_published = _raise_if_helper_advisor_busy_for_diffusion( publish_pending = True, + ignore_pending_workload = ( + ignore_public_load_pending_workload + ), ) _release_other_gpu_owners_for_diffusion() _release_chat_backend_for_diffusion(check_helper_advisor = False) @@ -1562,6 +1566,7 @@ def encode_png_base64(pil_image: "Any") -> str: def _raise_if_helper_advisor_busy_for_diffusion( *, publish_pending: bool = False, + ignore_pending_workload: Optional[str] = None, ) -> bool: """Round 29 P1 #1: split the helper-busy check out of _release_chat_backend_for_diffusion so the diffusion load can @@ -1590,6 +1595,7 @@ def _raise_if_helper_advisor_busy_for_diffusion( _HELPER_ADVISOR_START_LOCK, _publish_public_load_pending, helper_advisor_busy, + public_load_pending, ) except Exception: return False @@ -1599,6 +1605,20 @@ def _raise_if_helper_advisor_busy_for_diffusion( "AI Assist (helper / advisor GGUF) is still using the GPU. " "Wait for it to finish before loading a diffusion image model." ) + # Round 38 P1: mirror the route-side _raise_if_helper_advisor_busy + # public_load_pending parity check. When publishing, refuse if + # ANOTHER public workload is already mid-handoff. Route-wrapped + # calls pass ignore_pending_workload="diffusion" so the + # route's own publish (which happened just before + # backend.load_model) does not cause the backend's atomic + # check to self-block. + if publish_pending and public_load_pending( + excluding = ignore_pending_workload + ): + raise RuntimeError( + "Another GPU workload is mid-handoff. Wait for it to " + "finish before loading a diffusion image model." + ) if publish_pending: _publish_public_load_pending("diffusion-backend") return True diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index f9e5a58ecc..547b50f567 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2462,6 +2462,12 @@ async def diffusion_load( family_override = payload.family, hf_token = payload.hf_token, enable_model_cpu_offload = payload.enable_model_cpu_offload, + # Round 38 P1: this route already published the + # "diffusion" pending marker above; tell the + # backend to ignore it so the parity check it + # now applies does not self-block on our own + # publication. + ignore_public_load_pending_workload = "diffusion", ), ) return JSONResponse(content = status) diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 2512a83e07..e3976125dd 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -133,13 +133,25 @@ def _release_public_load_pending(workload: str) -> None: _PUBLIC_LOAD_PENDING_COUNT.pop(needle, None) -def public_load_pending() -> bool: +def public_load_pending(*, excluding: str | None = None) -> bool: """True if any public GPU workload has passed its helper-busy snapshot but not yet flipped its public ownership flags. Helper / advisor starts treat this as busy so they cannot race a public - load mid-handoff.""" + load mid-handoff. + + Round 38 P1: ``excluding`` lets a route-wrapped backend call + skip the marker its own route layer already published (e.g. the + diffusion route publishes ``diffusion`` before calling into + ``backend.load_model``, which publishes ``diffusion-backend`` -- + the backend should ignore its own ``diffusion`` marker so the + parity check does not self-block) while still seeing every + OTHER in-flight public workload.""" + ignored = excluding.lower() if excluding else None with _HELPER_ADVISOR_LOCK: - return sum(_PUBLIC_LOAD_PENDING_COUNT.values()) > 0 + return any( + count > 0 and workload != ignored + for workload, count in _PUBLIC_LOAD_PENDING_COUNT.items() + ) def _strip_think_tags(text: str) -> str: