From aeba18dc6d9621b5177e222f9bb5a371e0c82936 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 19:15:14 +0000 Subject: [PATCH] Fix/adjust diffusion: public_load_pending self-check for PR #5754 Round 35 P1: _raise_if_helper_advisor_busy published a new public pending marker without first checking public_load_pending(). Two public workloads (e.g. training + diffusion) could both pass their idle helper-busy snapshot concurrently, then both run through destructive owner teardown before either flipped its own visibility flag (is_training_active, current_checkpoint, loading_model_identifier, diffusion is_loading). Add the missing self-check under _HELPER_ADVISOR_START_LOCK so the second public workload sees the first's pending marker and gets a 503 retry instead of racing for VRAM. Helper / advisor already checked public_load_pending() on its side via _gpu_workload_busy_for_helper; this closes the symmetric public -> public window. 86 backend tests pass + smoke test confirms second public load is refused with 503 while first is pending, and the next public load is permitted once the first clears. --- studio/backend/routes/inference.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 09cfc4b237..f9e5a58ecc 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -375,6 +375,7 @@ def _raise_if_helper_advisor_busy(workload: str) -> None: _HELPER_ADVISOR_START_LOCK, _publish_public_load_pending, helper_advisor_busy, + public_load_pending, ) except Exception: return @@ -402,6 +403,21 @@ def _raise_if_helper_advisor_busy(workload: str) -> None: f"Wait for it to finish before starting {workload}." ), ) + # Round 35 P1: also refuse when another public workload is + # already mid-handoff (passed its own helper-busy snapshot + # but has not yet flipped is_training_active / + # current_checkpoint / loading_model_identifier / + # diffusion is_loading). Without this two public loads can + # both pass their idle snapshots concurrently and race + # destructive owner teardown. + if public_load_pending(): + raise HTTPException( + status_code = 503, + detail = ( + f"Another GPU workload is mid-handoff. Wait for it to " + f"finish before starting {workload}." + ), + ) _publish_public_load_pending(workload)