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.
This commit is contained in:
parent
3050196f6b
commit
f134a19b87
4 changed files with 123 additions and 3 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue