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

@@ -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

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)

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: