From f134a19b8756f86d9a5ffca6a5a715924d7d64f7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 22:55:36 +0000 Subject: [PATCH] Validate a local video checkpoint's own suffix; unload the old engine before publishing the new one Video: a local FILE picked for a gguf/single_file load is handed straight to the loader (the resolver returns the file itself, ignoring gguf_filename), so its own suffix must match the kind. A .gguf picked as single_file (or a .safetensors picked as gguf) slipped past the gguf_filename checks, evicting the resident GPU owner before failing in from_single_file. Reject the mismatch in validate, before the handoff. Engine router: publish the newly selected engine only after the old one finishes unloading. The arbiter's diffusion evictor unloads the active engine, so flipping the active name to the new (empty) engine first let a concurrent chat/video acquire evict that empty engine and take the GPU while the old model was still freeing VRAM. --- .../core/inference/diffusion_engine_router.py | 16 ++++++- studio/backend/core/inference/video.py | 20 +++++++- .../tests/test_diffusion_engine_router.py | 44 ++++++++++++++++++ studio/backend/tests/test_video_backend.py | 46 +++++++++++++++++++ 4 files changed, 123 insertions(+), 3 deletions(-) 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,