From 39ee329b2639a5b561ef4a1d6505f2e804cb8e55 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 10:58:13 +0000 Subject: [PATCH] Reserve the diffusion-training slot before freeing GPU residents start_diffusion_training freed resident GPU models and only then called service.start(config), which is where is_active() first flips true. During that free-then-spawn window a concurrent /images/load or /video/load saw training as inactive, passed its training guard, acquired the GPU, and began a background load, so the trainer and that pipeline both allocated VRAM. Add reserve()/unreserve() to the training service (is_active() also reports the reservation) and reserve BEFORE the free, in a try/finally so a failed start rolls the reservation back. An overlapping load's guard now refuses during the window. Regression tests: the route reserves before the free (and the free sees an active service), and the service reservation marks active then rolls back. --- .../training/diffusion_training_service.py | 22 ++++++++ studio/backend/routes/training.py | 25 ++++++--- .../backend/tests/test_diffusion_training.py | 54 +++++++++++++++++++ 3 files changed, 94 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 43946acefe..fe54bc1793 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -222,6 +222,11 @@ class DiffusionTrainingService: self._ctx = ctx if ctx is not None else _CTX self._target = target if target is not None else _default_target self._lock = threading.Lock() + # Set True by reserve() while a start is in flight, BEFORE the route frees resident GPU + # models, so the image/video load guards (which read is_active) refuse a concurrent load + # during the free-then-spawn window rather than double-allocate the GPU. Cleared by + # unreserve() once the proc is live (or the start failed). + self._reserved = False self._proc: Any = None self._stop_queue: Any = None self._pump: Optional[threading.Thread] = None @@ -232,8 +237,25 @@ class DiffusionTrainingService: # ── lifecycle ──────────────────────────────────────────────────────────── def is_active(self) -> bool: with self._lock: + if self._reserved: + return True return self._proc is not None and self._proc.is_alive() + def reserve(self) -> None: + """Mark a diffusion-training start as in flight so the image/video load guards (which + 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'.""" + with self._lock: + self._reserved = True + + def unreserve(self) -> None: + """Clear the reservation set by reserve(). Only touches the reservation flag, never + _proc, so a live job stays active on success and a failed start is fully rolled back.""" + with self._lock: + self._reserved = False + def start(self, config: dict) -> str: """Validate ``config``, spawn the trainer, and start pumping its events. diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 57932e7cd8..81a53f8c3b 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1327,21 +1327,28 @@ async def start_diffusion_training( except (FileNotFoundError, ValueError) as e: raise HTTPException(status_code = 400, detail = str(e)) - # Free resident GPU workloads (export / Images pipeline / chat) before the trainer - # loads its own pipeline. Offload the blocking teardown (engine unload waits on the - # generation locks; the export subprocess join can take seconds) to a worker thread so - # the event loop stays free for concurrent status/progress/cancel requests, as the - # inference routes do for their blocking load/unload calls. - await asyncio.to_thread(_free_gpu_for_diffusion_training) - service = get_diffusion_training_service() + # Reserve the training slot BEFORE freeing residents: is_active() otherwise flips true only + # at service.start(), after the free below, so a concurrent /images/load or /video/load would + # pass its training guard during the free-then-spawn window, acquire the GPU, and double- + # allocate VRAM against the trainer. Released in the finally so a failed start never leaves + # training stuck "active". + service.reserve() try: + # Free resident GPU workloads (export / Images pipeline / chat) before the trainer + # loads its own pipeline. Offload the blocking teardown (engine unload waits on the + # generation locks; the export subprocess join can take seconds) to a worker thread so + # the event loop stays free for concurrent status/progress/cancel requests, as the + # inference routes do for their blocking load/unload calls. + await asyncio.to_thread(_free_gpu_for_diffusion_training) job_id = service.start(config) except ValueError as e: raise HTTPException(status_code = 400, detail = str(e)) except RuntimeError as e: # A job is already running. raise HTTPException(status_code = 409, detail = str(e)) + except HTTPException: + raise except Exception as e: raise log_and_http_error( e, @@ -1350,6 +1357,10 @@ async def start_diffusion_training( event = "diffusion_training.start_failed", log = logger, ) + finally: + # On success the now-live proc keeps is_active() true; on any failure this clears the + # reservation so training is not left permanently "active". + service.unreserve() return DiffusionTrainingStartResponse(job_id = job_id, status = "running") diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 8e898a41f4..015478769d 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -264,14 +264,29 @@ def test_terminal_events_clear_model_load_flag(): class _FakeService: def __init__(self): self._running = False + self._reserved = False self.started_with = None self.stopped_with_save = None + # Ordered log of lifecycle calls so a test can assert reserve precedes the GPU free. + self.calls: list = [] # Extra keys merged into status() so a test can inject metric history / perf fields. self.status_extra: dict = {} + def reserve(self): + self._reserved = True + self.calls.append("reserve") + + def unreserve(self): + self._reserved = False + self.calls.append("unreserve") + + def is_active(self): + return self._reserved or self._running + def start(self, config): self.started_with = config self._running = True + self.calls.append("start") return "job-123" def stop(self, save = True): @@ -391,6 +406,45 @@ def test_route_start_frees_gpu_off_the_coroutine_thread(client, monkeypatch): assert threads["cleanup"] is not threads["inline"] # offloaded to a worker, not run inline +def test_route_start_reserves_before_freeing_gpu(client, monkeypatch): + # The training slot must be reserved (is_active -> true) BEFORE the route frees resident GPU + # models, so a concurrent /images/load or /video/load guard refuses during the free-then-spawn + # window instead of double-allocating the GPU. Assert the ordering: reserve is logged before + # the GPU free runs, and the service reports active while the free is in flight. + import routes.training as tr + + order: list = [] + + def _record_free(): + order.append("free") + # During the free window the service must already look active to a concurrent load guard. + order.append(f"active={client._fake.is_active()}") + + monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", _record_free) + + r = client.post("/api/train/diffusion/start", json = _BODY) + assert r.status_code == 200, r.text + # reserve fires before the free, the free sees an active service, then start, then unreserve. + assert client._fake.calls[0] == "reserve" + assert client._fake.calls.index("reserve") < client._fake.calls.index("start") + assert order == ["free", "active=True"] + assert "unreserve" in client._fake.calls + + +def test_service_reserve_marks_active_and_rolls_back(): + # The real service: reserve() flips is_active true before any proc exists (so a load guard + # refuses during the free window), and unreserve() clears it without a live proc, so a failed + # start is not left permanently "active". + from core.training.diffusion_training_service import DiffusionTrainingService + + svc = DiffusionTrainingService() + assert svc.is_active() is False + svc.reserve() + assert svc.is_active() is True # active with no proc, purely from the reservation + svc.unreserve() + assert svc.is_active() is False + + 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