From 123b3ea4d60fc3a195b43bbd3eeef83c9dac580c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 11:57:37 +0000 Subject: [PATCH] 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. --- .../training/diffusion_training_service.py | 10 +++++++++- .../backend/tests/test_diffusion_training.py | 19 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index fe54bc1793..c02445c4f1 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -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: diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 015478769d..c0cd6b129b 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -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