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.
This commit is contained in:
Daniel Han-Chen 2026-05-25 19:15:14 +00:00
commit aeba18dc6d

View file

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