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:
parent
39ee329b26
commit
123b3ea4d6
2 changed files with 28 additions and 1 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue