Keep GPU ownership across an in-flight reload on unload

The images and video unload routes dropped their arbiter claim whenever the
backend was not committed-loaded, but a concurrent /load re-acquires the owner
and starts a background load that is not is_loaded for its whole download and
finalize window. Releasing during that window cleared the newer load's claim, so
a later chat/image load saw no owner, skipped eviction, and could allocate a
second heavy model on the GPU. Gate the release on loading_repo_ids() too (both
diffusion engines and the video backend expose it), not just the committed
state, so an overlapping load keeps ownership. Add regression tests for the
in-flight case on both routes.

studio/backend/routes/inference.py

@@ -12713,12 +12713,15 @@ async def unload_diffusion_model(current_subject: str = Depends(get_current_subj
     from core.inference.gpu_arbiter import release, DIFFUSION
 
     status_dict = await asyncio.to_thread(get_active_diffusion_engine().unload)
-    # Drop DIFFUSION ownership only if nothing is resident again: a concurrent /images/load
-    # that re-acquired DIFFUSION while this (slow) unload ran must keep ownership, or a later
-    # chat load would see no owner, skip eviction, and OOM against the newly resident pipeline.
-    # release() is owner-guarded and identity-less, so an unconditional release here would clear
-    # the newer load's claim.
-    if not get_active_diffusion_engine().is_loaded:
+    # Drop DIFFUSION ownership only if nothing is resident AND no new load is in flight: a
+    # concurrent /images/load that re-acquired DIFFUSION while this (slow) unload ran must keep
+    # ownership, or a later chat load would see no owner, skip eviction, and OOM against the newly
+    # resident pipeline. An in-flight load has is_loaded False for its whole download/finalize
+    # window, so gate on loading_repo_ids() too (both engines expose it), not just the committed
+    # state. release() is owner-guarded and identity-less, so an unconditional release here would
+    # clear the newer load's claim.
+    engine = get_active_diffusion_engine()
+    if not engine.loading_repo_ids() and not engine.is_loaded:
         release(DIFFUSION)
     return DiffusionStatusResponse(**annotate_status(status_dict))
 

studio/backend/routes/video.py

@@ -238,8 +238,16 @@ async def unload_video_model(current_subject: str = Depends(get_current_subject)
     from core.inference.gpu_arbiter import VIDEO, release
     from core.inference.video import get_video_backend
 
-    status_dict = await asyncio.to_thread(get_video_backend().unload)
-    release(VIDEO)
+    backend = get_video_backend()
+    status_dict = await asyncio.to_thread(backend.unload)
+    # Drop VIDEO ownership only if nothing is resident AND no new load is in flight: a concurrent
+    # /video/load that re-acquired VIDEO while this (slow) unload ran must keep ownership, or a
+    # later chat/image load would see no owner, skip eviction, and OOM against the newly resident
+    # (or still in-flight) video pipeline. release() is owner-guarded and identity-less, so an
+    # unconditional release here would clear the newer load's claim. Mirrors the images-route
+    # guard (inference.py), plus the in-flight check the committed-loaded state cannot cover.
+    if not backend.loading_repo_ids() and not backend.status()["loaded"]:
+        release(VIDEO)
     return VideoStatusResponse(**status_dict)
 
 

studio/backend/tests/test_diffusion_routes.py

@@ -24,11 +24,17 @@ from routes.inference import studio_router
 class _FakeBackend:
     def __init__(self) -> None:
         self.loaded = False
+        # Repo ids of in-flight (not yet committed) loads; empty tuple = none. The unload
+        # route reads this to keep DIFFUSION ownership while a concurrent load is still loading.
+        self.loading: tuple = ()
 
     @property
     def is_loaded(self) -> bool:
         return self.loaded
 
+    def loading_repo_ids(self) -> tuple:
+        return tuple(self.loading)
+
     def validate_load_request(
         self,
         model_path,
@@ -267,6 +273,25 @@ def test_unload_keeps_ownership_when_a_model_is_still_resident(client, monkeypat
     assert gpu_arbiter.current_owner() is None
 
 
+def test_unload_keeps_ownership_when_a_load_is_in_flight(client, monkeypatch):
+    # A concurrent /images/load re-acquires DIFFUSION and starts a background load, so the
+    # engine is NOT is_loaded yet (the pipeline commits later) but a load IS in flight. The
+    # unload route must keep ownership on the in-flight state alone, or a later chat load would
+    # see no owner, skip eviction, and OOM against the newly resident pipeline. is_loaded stays
+    # False the whole download/finalize window, so the loaded-only check is insufficient here.
+    backend = diffusion_module.get_diffusion_backend()
+    gpu_arbiter._owner = gpu_arbiter.DIFFUSION
+
+    backend.loaded = False
+    backend.loading = ("unsloth/z-image-turbo",)
+    monkeypatch.setattr(backend, "unload", lambda: {**_unloaded_status(), "loaded": False})
+    r = client.post("/api/inference/images/unload")
+    assert r.status_code == 200
+    assert gpu_arbiter.current_owner() == gpu_arbiter.DIFFUSION  # ownership retained for the load
+
+    backend.loading = ()
+
+
 def test_generate_batch_size_persists_each_image(client):
     client.post(
         "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}

studio/backend/tests/test_video_routes.py

@@ -64,6 +64,12 @@ class _FakeBackend:
     def __init__(self) -> None:
         self.loaded = False
         self.last_load_kwargs: dict = {}
+        # Repo ids of in-flight (not yet committed) loads; empty tuple = none. The unload
+        # route reads this to keep VIDEO ownership while a concurrent load is still loading.
+        self.loading: tuple = ()
+
+    def loading_repo_ids(self) -> tuple:
+        return tuple(self.loading)
 
     def validate_load_request(
         self,
@@ -493,6 +499,24 @@ def test_unload_releases_arbiter(client, monkeypatch):
     assert gpu_arbiter.current_owner() is None
 
 
+def test_unload_keeps_ownership_when_a_load_is_in_flight(client, monkeypatch):
+    # A concurrent /video/load re-acquires VIDEO and starts a background load, so the backend
+    # is NOT loaded yet (the pipeline commits later) but a load IS in flight. The unload route
+    # must keep ownership on the in-flight state alone, or a later chat/image load would see no
+    # owner, skip eviction, and OOM against the newly resident pipeline. The committed-loaded
+    # state stays False the whole load window, so the loaded-only check is insufficient here.
+    backend = video_module.get_video_backend()
+    monkeypatch.setattr(gpu_arbiter, "_owner", gpu_arbiter.VIDEO)
+
+    backend.loaded = False
+    backend.loading = ("unsloth/ltx-video-2b",)
+    resp = client.post("/api/inference/video/unload")
+    assert resp.status_code == 200 and resp.json()["loaded"] is False
+    assert gpu_arbiter.current_owner() == gpu_arbiter.VIDEO  # ownership retained for the load
+
+    backend.loading = ()
+
+
 def test_load_refused_during_training(client, monkeypatch):
     # A video load while training is active is refused (409) before the GPU is taken.
     import core.training as core_training