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: