diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index 957df10ef7..ecdcb47ecb 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -103,13 +103,25 @@ def _activate(name: str, reason: Optional[str]) -> Any: if name != _active_engine_name: engine_to_unload = get_active_diffusion_engine() old_name = _active_engine_name - _active_engine_name = name - _fallback_reason = reason if name == ENGINE_DIFFUSERS else None + else: + # No engine change: publish the (possibly refreshed) fallback reason now. + _fallback_reason = reason if name == ENGINE_DIFFUSERS else None if engine_to_unload is not None: + # Publish the new engine only AFTER the old one has finished unloading. The + # arbiter's diffusion evictor unloads get_active_diffusion_engine(), so flipping + # _active_engine_name to the new (empty) engine first would let a concurrent + # chat/video acquire_for evict that empty engine and take the GPU while the old + # model is still freeing VRAM in this thread -- two large models briefly resident. + # Keeping the OLD engine published as the evict target means a concurrent evict + # serializes on its unload() (idempotent under the backend's _lock/_generate_lock) + # and the GPU is granted only once its VRAM is freed. try: engine_to_unload.unload() except Exception as exc: # noqa: BLE001 -- best-effort; never block the switch logger.warning("failed to unload previous engine %s: %s", old_name, exc) + with _lock: + _active_engine_name = name + _fallback_reason = reason if name == ENGINE_DIFFUSERS else None if name == ENGINE_SD_CPP: logger.info("diffusion engine: sd_cpp") else: diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 8839a8ed36..dc536bc4d6 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -475,7 +475,25 @@ class VideoBackend: resolve_local_gguf_child(root, gguf_filename or "") except Exception as exc: # noqa: BLE001 -- surface as client input error raise ValueError(str(exc)) from exc - elif path_shaped and not root.is_file(): + elif root.is_file(): + # The loader hands a local FILE straight to the gguf/single_file loader + # (_resolve_checkpoint_path returns the file itself, ignoring gguf_filename), + # so the file's OWN suffix must match the kind. Otherwise a .gguf picked as + # single_file (or a .safetensors picked as gguf) slips past the gguf_filename + # checks above, evicts the resident model in the route, and only then fails + # in from_single_file / the GGUF reader. Reject it here, before the handoff. + suffix = root.suffix.lower() + if kind == "gguf" and suffix != ".gguf": + raise ValueError( + f"Local checkpoint '{repo_id}' is not a .gguf file; a 'gguf' load " + f"needs a .gguf checkpoint." + ) + if kind == "single_file" and suffix != ".safetensors": + raise ValueError( + f"Local checkpoint '{repo_id}' is not a .safetensors file; a " + f"'single_file' load needs a .safetensors checkpoint." + ) + elif path_shaped: raise ValueError(f"Local model path '{repo_id}' does not exist.") # A local pipeline pick must be a real diffusers directory (model_index.json), or it # would only fail deep in from_pretrained AFTER the route evicted the resident model. diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index f4cebc893d..6667d53044 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -194,3 +194,47 @@ def test_active_status_injects_engine_and_reason(monkeypatch): st = r.active_status() assert st["engine"] == ENGINE_DIFFUSERS assert st["fallback_reason"] and "binary unavailable" in st["fallback_reason"] + + +# ── engine-switch eviction ordering ─────────────────────────────────────────── + + +def test_switch_unloads_old_engine_before_publishing_new(monkeypatch): + # The arbiter's diffusion evictor unloads get_active_diffusion_engine(); if the router + # published the new (empty) engine BEFORE the old one finished unloading, a concurrent + # chat/video acquire_for could evict the empty engine and take the GPU while the old + # model was still resident (two large models briefly co-resident -> OOM). Assert the + # OLD engine stays the published active target until its unload() completes, then the + # new engine is published. + seen = {} + + def _fake_engine(): + return SimpleNamespace( + unload = lambda: seen.__setitem__("active_during_unload", r.active_engine_name()), + status = lambda: {"loaded": False, "repo_id": None}, + ) + + monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: _fake_engine()) + r._active_engine_name = ENGINE_SD_CPP + r._activate(ENGINE_DIFFUSERS, "switch test") + assert seen["active_during_unload"] == ENGINE_SD_CPP + assert r.active_engine_name() == ENGINE_DIFFUSERS + + +def test_no_switch_keeps_engine_and_refreshes_reason(monkeypatch): + # When the engine does not change, _activate must not spuriously unload anything and must + # still refresh the recorded fallback reason (the diffusers-only steady state). + calls = {"unload": 0} + + def _fake_engine(): + return SimpleNamespace( + unload = lambda: calls.__setitem__("unload", calls["unload"] + 1), + status = lambda: {"loaded": False, "repo_id": None}, + ) + + monkeypatch.setattr(r, "get_active_diffusion_engine", lambda: _fake_engine()) + r._active_engine_name = ENGINE_DIFFUSERS + r._activate(ENGINE_DIFFUSERS, "still diffusers") + assert calls["unload"] == 0 + assert r.active_engine_name() == ENGINE_DIFFUSERS + assert r.active_status()["fallback_reason"] == "still diffusers" diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 1a25813c5d..c6928ba936 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -490,6 +490,52 @@ def test_validate_rejects_kind_extension_mismatch(tmp_path): ) +def test_validate_rejects_local_file_suffix_kind_mismatch(tmp_path): + backend = VideoBackend() + # A local FILE is handed straight to the gguf/single_file loader: _resolve_checkpoint_path + # returns the file itself, IGNORING gguf_filename, so the file's OWN suffix must match the + # kind. A .gguf file picked as single_file (or a .safetensors file picked as gguf) slips + # past the gguf_filename suffix checks, so it must be rejected HERE, before the route evicts + # the resident GPU owner and then fails deep in from_single_file / the GGUF reader. + gguf_file = tmp_path / "ltx.gguf" + gguf_file.write_bytes(b"weights") + safetensors_file = tmp_path / "ltx.safetensors" + safetensors_file.write_bytes(b"weights") + with pytest.raises(ValueError, match = "not a .safetensors file"): + backend.validate_load_request( + str(gguf_file), + gguf_filename = "ltx.safetensors", + model_kind = "single_file", + family_override = "ltx-2", + ) + with pytest.raises(ValueError, match = "not a .gguf file"): + backend.validate_load_request( + str(safetensors_file), + gguf_filename = "ltx.gguf", + model_kind = "gguf", + family_override = "ltx-2", + ) + # Matching pairs still validate: the local file's suffix agrees with the resolved kind. + assert ( + backend.validate_load_request( + str(gguf_file), + gguf_filename = "ltx.gguf", + model_kind = "gguf", + family_override = "ltx-2", + ).name + == "ltx-2" + ) + assert ( + backend.validate_load_request( + str(safetensors_file), + gguf_filename = "ltx.safetensors", + model_kind = "single_file", + family_override = "ltx-2", + ).name + == "ltx-2" + ) + + def test_validate_rejects_windows_shaped_missing_checkpoint(tmp_path): backend = VideoBackend() # A missing Windows-shaped local pick (backslash path, or a C:/ drive path) must fail HERE,