Make the diffusion-training reservation a compare-and-set

Two overlapping /diffusion/start requests can interleave between the is_active()
check and the reservation, so reserve() itself must reject a second reservation
atomically. Otherwise both callers reserve, both free the GPU's resident chat or
image model, and the loser only 409s after the eviction -- the evict-then-fail the
reservation exists to prevent. reserve() now raises under the lock if a start is
already reserved or a job is already running.
This commit is contained in:
Daniel Han 2026-07-07 11:57:37 +00:00
commit 123b3ea4d6
2 changed files with 28 additions and 1 deletions

View file

@ -246,8 +246,16 @@ class DiffusionTrainingService:
read is_active) refuse a concurrent load BEFORE the route frees resident GPU models.
Without this the training becomes active only at start(), after the free, so an
overlapping load passes its guard, acquires the GPU, and both workloads allocate VRAM.
Paired with unreserve() in a finally, so a failed start never leaves training 'active'."""
Compare-and-set: raise if a start is already reserved or a job is already running, so a
second overlapping /diffusion/start is rejected (409) BEFORE it frees GPU residents,
instead of both requests tearing down residents and racing to start() (whichever finishes
first wins, so a double-click or a retry with different parameters could start the wrong
config). Paired with unreserve() in a finally by the reserving caller, so a failed start
never leaves training 'active'."""
with self._lock:
if self._reserved or (self._proc is not None and self._proc.is_alive()):
raise RuntimeError("A diffusion training job is already running.")
self._reserved = True
def unreserve(self) -> None:

View file

@ -445,6 +445,25 @@ def test_service_reserve_marks_active_and_rolls_back():
assert svc.is_active() is False
def test_service_reserve_is_compare_and_set():
# reserve() is the concurrency gate: two /diffusion/start requests can interleave between the
# is_active() check and the reservation, so reserve() itself must reject a second reservation
# atomically. Without the compare-and-set, both callers would reserve, both would free the
# GPU's resident chat/image model, and the loser would only 409 AFTER the eviction -- exactly
# the evict-then-fail the reservation exists to prevent. A second reserve must raise; the first
# stays reserved; after unreserve the slot is claimable again.
from core.training.diffusion_training_service import DiffusionTrainingService
svc = DiffusionTrainingService()
svc.reserve()
with pytest.raises(RuntimeError, match = "already running"):
svc.reserve()
assert svc.is_active() is True # the losing reserve did not clear the winner's claim
svc.unreserve()
svc.reserve() # claimable again once released
assert svc.is_active() is True
def test_route_start_preflights_gated_base_off_the_coroutine_thread(client, monkeypatch):
# _preflight_gated_base does a blocking urlopen HEAD (up to a 5s timeout) to Hugging Face, so
# the async start route must offload it via asyncio.to_thread rather than run it inline and