diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index cb3dafd440..96f94d0457 100644 --- a/studio/backend/routes/inference.py +++ b/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)) diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 9ccb3237e0..a053c9fac7 100644 --- a/studio/backend/routes/video.py +++ b/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) diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 530be7b583..f913d6e9cb 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/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"} diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 99e52d75f6..f8fa490453 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/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