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.

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.
 

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

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