Stop a background page and a stale record taking the GPU or a download with them
Five fixes from a review pass over the diffusion work. delete-finetuned rmtree'd a model the Images or Video engine was holding: every guard on that route is chat-only, and Images loads any local path, so deleting a local diffusion model under the storage root pulled the weights (and the companion VAE / text encoders sd.cpp re-reads each generation) out from under a live pipeline. The cached-model route already refuses this; the trained/exported one now does too, matching by path rather than repo id, and failing open on a chat-only install so it cannot block ordinary deletes. A staged download finishing while its page was hidden loaded the model and evicted whatever the user was actually using: both diffusion pages stay mounted behind the router and a load takes the GPU unconditionally. The pick is now held until its page is on screen again, which is also what chat does. A scoped download could report success having fetched nothing. With Hugging Face metadata unavailable no manifest is written, so verification is a no-op, and snapshot_download returns an existing snapshot folder without downloading when its own repo_info call fails. A repo already on disk from a full snapshot job (which ignores *.gguf) therefore completed with no weights and auto-loaded against them. The requested file list needs no network, so it is checked against the disk directly. The XET to HTTP retry reclaimed the job slot without the scoped file list, and that claim overwrites the stored record, so a later identical scoped start compared an empty list against the real one and 409'd instead of adopting the running download. The DiT accelerator gate probed torch.mps.is_available(), which only exists from torch 2.5 while the supported floor is 2.4. All three probes shared one try/except, so on torch 2.4 the AttributeError read as 'no block' and a CPU-only host still evicted the resident pipeline, downloaded the encoders and died in the child. Each accelerator is probed on its own now, through torch.backends.mps.
This commit is contained in:
parent
273755ab42
commit
7827679b77
11 changed files with 486 additions and 7 deletions
|
|
@ -370,12 +370,22 @@ def dit_accelerator_missing_reason(resolved_family: str) -> Optional[str]:
|
|||
try:
|
||||
import torch
|
||||
|
||||
xpu = getattr(torch, "xpu", None)
|
||||
mps = getattr(torch, "mps", None)
|
||||
def probe(owner: Any) -> bool:
|
||||
# Each accelerator is probed on its own: one missing or throwing probe must not
|
||||
# decide the other two. torch.mps.is_available() only exists from torch 2.5 and
|
||||
# the supported floor is 2.4, so a shared try/except here would swallow the
|
||||
# AttributeError and wave a CPU-only host through the gate it exists for.
|
||||
try:
|
||||
fn = getattr(owner, "is_available", None)
|
||||
return bool(fn()) if callable(fn) else False
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
|
||||
if (
|
||||
torch.cuda.is_available()
|
||||
or bool(xpu is not None and xpu.is_available())
|
||||
or bool(mps is not None and mps.is_available())
|
||||
probe(torch.cuda)
|
||||
# torch.backends.mps has carried is_available() since torch 1.12.
|
||||
or probe(getattr(torch, "xpu", None))
|
||||
or probe(getattr(torch.backends, "mps", None))
|
||||
):
|
||||
return None
|
||||
except Exception: # noqa: BLE001 -- probe failure must not block a start
|
||||
|
|
|
|||
|
|
@ -430,6 +430,11 @@ def _try_http_retry(
|
|||
cancel_marker_transport = original_metadata.transport,
|
||||
hub_cache = original_metadata.hub_cache,
|
||||
xet_cache = original_metadata.xet_cache,
|
||||
# Carry the scoped file list across the reclaim. The record it overwrites is what a
|
||||
# later start for this scope slot is compared against, so dropping it makes an
|
||||
# identical scoped start read as a different file set and 409 instead of adopting
|
||||
# the running download (the worker args below read the same list).
|
||||
scoped_files = original_metadata.scoped_files or None,
|
||||
)
|
||||
if claimed:
|
||||
break
|
||||
|
|
|
|||
|
|
@ -746,6 +746,23 @@ def _download_scoped_snapshot(
|
|||
allow_patterns = files,
|
||||
max_workers = 1,
|
||||
)
|
||||
if info is None:
|
||||
# With no metadata there is no manifest, so _verify_completed_download below is a
|
||||
# no-op -- and snapshot_download RETURNS AN EXISTING SNAPSHOT FOLDER, without
|
||||
# fetching anything, when its own repo_info call also fails. A repo already on disk
|
||||
# from a full snapshot job (which ignores *.gguf) would therefore flip this job to
|
||||
# complete having downloaded no weights, and the page would load against them. The
|
||||
# requested list needs no network, so check it against the disk directly.
|
||||
root = Path(snapshot_path)
|
||||
absent = tuple(f for f in files if not (root / f).exists())
|
||||
if absent:
|
||||
print(
|
||||
f"Could not reach Hugging Face for {repo_id} [{scope}] and the copy on disk "
|
||||
f"is incomplete ({len(absent)} file(s) missing): {_format_path_list(absent)}. "
|
||||
"Reconnect (or set a valid HF token) and resume the download.",
|
||||
file = sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
_verify_completed_download(
|
||||
"model",
|
||||
repo_id,
|
||||
|
|
|
|||
|
|
@ -2338,6 +2338,30 @@ def _loading_model_matches_deleted_path(loading_model: object, deleted_path: Pat
|
|||
return _loaded_model_matches_deleted_path(str(loading_model), deleted_path)
|
||||
|
||||
|
||||
def _active_diffusion_backend():
|
||||
"""The live Images engine, or None when this install has no diffusion stack.
|
||||
|
||||
Fails OPEN on import: a chat-only install (no diffusers) must still be able to delete
|
||||
its fine-tuned models. Reading the returned backend's state is what fails closed.
|
||||
"""
|
||||
try:
|
||||
from core.inference.diffusion_engine_router import get_active_diffusion_engine
|
||||
return get_active_diffusion_engine()
|
||||
except Exception as e:
|
||||
logger.debug(f"Images engine unavailable during delete guard: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _active_video_backend():
|
||||
"""The live Video backend, or None when this install has no video stack."""
|
||||
try:
|
||||
from core.inference.video import get_video_backend
|
||||
return get_video_backend()
|
||||
except Exception as e:
|
||||
logger.debug(f"Video backend unavailable during delete guard: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def _prune_empty_parents(start: Path, stop_at: Path) -> None:
|
||||
"""Remove empty ancestors of ``start`` up to (not including) ``stop_at``.
|
||||
|
||||
|
|
@ -2557,6 +2581,44 @@ async def delete_finetuned_model(
|
|||
detail = "Could not verify model load status before deleting",
|
||||
) from e
|
||||
|
||||
# Every guard above is chat-only, and Images / Video hold their own pipelines: a local
|
||||
# diffusion or video model under the storage root loads by path, so without this rmtree
|
||||
# would pull the weights (and the companion VAE / text encoders sd.cpp re-reads on every
|
||||
# generation) out from under a live engine. The cached-model delete route runs the same
|
||||
# check via hub.services.models.deletion; it matches by repo id, this one by path.
|
||||
for label, get_backend in (
|
||||
("Images", _active_diffusion_backend),
|
||||
("Video", _active_video_backend),
|
||||
):
|
||||
backend = get_backend()
|
||||
if backend is None:
|
||||
continue
|
||||
try:
|
||||
status = backend.status()
|
||||
held = [status.get(key) for key in ("repo_id", "base_repo")] if status.get("loaded") else []
|
||||
held += list(getattr(backend, "loaded_repo_ids", tuple)())
|
||||
if any(h and _loaded_model_matches_deleted_path(str(h), target_path) for h in held):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Unload the model before deleting",
|
||||
)
|
||||
if any(
|
||||
_loading_model_matches_deleted_path(lid, target_path)
|
||||
for lid in getattr(backend, "loading_repo_ids", tuple)()
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = "Cannot delete a model while it is loading",
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.warning("Could not check the %s model before delete: %s", label, e)
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = "Could not verify model load status before deleting",
|
||||
) from e
|
||||
|
||||
try:
|
||||
if export_type == "gguf" and gguf_variant:
|
||||
if not target_path.is_dir():
|
||||
|
|
|
|||
172
studio/backend/tests/test_delete_finetuned_diffusion_guard.py
Normal file
172
studio/backend/tests/test_delete_finetuned_diffusion_guard.py
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""DELETE /api/models/delete-finetuned must refuse a directory the Images or Video
|
||||
engine is holding.
|
||||
|
||||
Every other guard on that route is chat-only (llama.cpp + the transformers backend), so a
|
||||
local diffusion model under the storage root -- Images loads any existing local path -- used
|
||||
to be rmtree'd while a pipeline was still reading it, taking the companion VAE / text encoder
|
||||
files sd.cpp re-reads on every generation with it. The cached-model delete route already
|
||||
refuses this; these tests pin the same behaviour on the trained/exported route, which matches
|
||||
by PATH rather than by repo id.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import routes.models as models_module
|
||||
from auth.authentication import get_current_subject
|
||||
from routes.models import router as models_router
|
||||
|
||||
|
||||
class _Backend:
|
||||
"""Minimal stand-in for the Images engine / Video backend delete-guard surface."""
|
||||
|
||||
def __init__(self, *, loaded=None, base=None, loading=(), extra=()):
|
||||
self._loaded = loaded
|
||||
self._base = base
|
||||
self._loading = tuple(loading)
|
||||
self._extra = tuple(extra)
|
||||
|
||||
def status(self):
|
||||
if self._loaded is None:
|
||||
return {"loaded": False}
|
||||
return {"loaded": True, "repo_id": self._loaded, "base_repo": self._base}
|
||||
|
||||
def loaded_repo_ids(self):
|
||||
return self._extra
|
||||
|
||||
def loading_repo_ids(self):
|
||||
return self._loading
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def client(tmp_path, monkeypatch):
|
||||
outputs = tmp_path / "outputs"
|
||||
outputs.mkdir()
|
||||
monkeypatch.setattr(models_module, "outputs_root", lambda: outputs)
|
||||
# No chat model is resident, so only the diffusion / video guards can refuse.
|
||||
monkeypatch.setattr(models_module, "get_inference_backend", lambda: _NoChat())
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(models_router, prefix = "/api/models")
|
||||
app.dependency_overrides[get_current_subject] = lambda: "test-user"
|
||||
return TestClient(app), outputs
|
||||
|
||||
|
||||
class _NoChat:
|
||||
active_model_name = None
|
||||
loading_models: set = set()
|
||||
|
||||
|
||||
def _model_dir(outputs):
|
||||
d = outputs / "my-diffusion-model"
|
||||
d.mkdir()
|
||||
(d / "model_index.json").write_text("{}", encoding = "utf-8")
|
||||
return d
|
||||
|
||||
|
||||
def _delete(client, path):
|
||||
return client.request(
|
||||
"DELETE",
|
||||
"/api/models/delete-finetuned",
|
||||
json = {"model_path": str(path), "source": "training"},
|
||||
)
|
||||
|
||||
|
||||
def test_refuses_to_delete_a_model_the_images_engine_has_loaded(client, monkeypatch):
|
||||
c, outputs = client
|
||||
target = _model_dir(outputs)
|
||||
monkeypatch.setattr(
|
||||
models_module, "_active_diffusion_backend", lambda: _Backend(loaded = str(target))
|
||||
)
|
||||
monkeypatch.setattr(models_module, "_active_video_backend", lambda: None)
|
||||
|
||||
resp = _delete(c, target)
|
||||
assert resp.status_code == 400
|
||||
assert "Unload the model" in resp.json()["detail"]
|
||||
assert target.exists()
|
||||
|
||||
|
||||
def test_refuses_the_companion_base_and_the_extra_repos_the_engine_reads(client, monkeypatch):
|
||||
c, outputs = client
|
||||
target = _model_dir(outputs)
|
||||
other = outputs / "somewhere-else"
|
||||
other.mkdir()
|
||||
# The checkpoint is the loaded id; the deleted dir is only the companion base.
|
||||
monkeypatch.setattr(
|
||||
models_module,
|
||||
"_active_diffusion_backend",
|
||||
lambda: _Backend(loaded = str(other), base = str(target)),
|
||||
)
|
||||
monkeypatch.setattr(models_module, "_active_video_backend", lambda: None)
|
||||
assert _delete(c, target).status_code == 400
|
||||
|
||||
# Same for a companion the engine reports through loaded_repo_ids (sd.cpp VAE / TE).
|
||||
monkeypatch.setattr(
|
||||
models_module,
|
||||
"_active_diffusion_backend",
|
||||
lambda: _Backend(loaded = str(other), extra = (str(target),)),
|
||||
)
|
||||
assert _delete(c, target).status_code == 400
|
||||
assert target.exists()
|
||||
|
||||
|
||||
def test_refuses_while_the_video_backend_is_still_fetching_it(client, monkeypatch):
|
||||
c, outputs = client
|
||||
target = _model_dir(outputs)
|
||||
monkeypatch.setattr(models_module, "_active_diffusion_backend", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
models_module, "_active_video_backend", lambda: _Backend(loading = (str(target),))
|
||||
)
|
||||
|
||||
resp = _delete(c, target)
|
||||
assert resp.status_code == 409
|
||||
assert "loading" in resp.json()["detail"].lower()
|
||||
assert target.exists()
|
||||
|
||||
|
||||
def test_allows_the_delete_when_no_diffusion_or_video_model_holds_it(client, monkeypatch):
|
||||
c, outputs = client
|
||||
target = _model_dir(outputs)
|
||||
monkeypatch.setattr(
|
||||
models_module,
|
||||
"_active_diffusion_backend",
|
||||
lambda: _Backend(loaded = str(outputs / "another-model")),
|
||||
)
|
||||
monkeypatch.setattr(models_module, "_active_video_backend", lambda: _Backend())
|
||||
|
||||
assert _delete(c, target).status_code == 200
|
||||
assert not target.exists()
|
||||
|
||||
|
||||
def test_a_chat_only_install_can_still_delete(client, monkeypatch):
|
||||
"""No diffusion stack installed: the guard must fail OPEN, not 503 every delete."""
|
||||
c, outputs = client
|
||||
target = _model_dir(outputs)
|
||||
monkeypatch.setattr(models_module, "_active_diffusion_backend", lambda: None)
|
||||
monkeypatch.setattr(models_module, "_active_video_backend", lambda: None)
|
||||
|
||||
assert _delete(c, target).status_code == 200
|
||||
assert not target.exists()
|
||||
|
||||
|
||||
def test_an_unreadable_engine_state_fails_closed(client, monkeypatch):
|
||||
"""The engine exists but cannot report its state: refuse rather than risk the rmtree."""
|
||||
c, outputs = client
|
||||
target = _model_dir(outputs)
|
||||
|
||||
class _Broken:
|
||||
def status(self):
|
||||
raise RuntimeError("engine wedged")
|
||||
|
||||
monkeypatch.setattr(models_module, "_active_diffusion_backend", lambda: _Broken())
|
||||
monkeypatch.setattr(models_module, "_active_video_backend", lambda: None)
|
||||
|
||||
resp = _delete(c, target)
|
||||
assert resp.status_code == 503
|
||||
assert target.exists()
|
||||
|
|
@ -636,7 +636,7 @@ def test_dit_accelerator_missing_reason_and_info_hide_train_without_a_gpu(monkey
|
|||
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
||||
monkeypatch.setattr(torch.xpu, "is_available", lambda: False)
|
||||
monkeypatch.setattr(torch.mps, "is_available", lambda: False)
|
||||
monkeypatch.setattr(torch.backends.mps, "is_available", lambda: False)
|
||||
assert "GPU" in (dit_accelerator_missing_reason("flux.1") or "")
|
||||
# SDXL and unknown families keep their own paths.
|
||||
assert dit_accelerator_missing_reason("sdxl") is None
|
||||
|
|
@ -652,5 +652,31 @@ def test_dit_accelerator_missing_reason_and_info_hide_train_without_a_gpu(monkey
|
|||
assert info["precision_modes"] != [] or name not in _DIT_TRAIN_FAMILIES
|
||||
|
||||
# Any accelerator clears it (MPS here, which bitsandbytes accepts).
|
||||
monkeypatch.setattr(torch.mps, "is_available", lambda: True)
|
||||
monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True)
|
||||
assert dit_accelerator_missing_reason("flux.1") is None
|
||||
|
||||
|
||||
def test_dit_accelerator_gate_survives_a_torch_without_every_probe(monkeypatch):
|
||||
"""torch.mps.is_available() only exists from torch 2.5 and the supported floor is 2.4.
|
||||
Probing the accelerators under one shared try/except turned that AttributeError into
|
||||
"no block", so the very hosts the gate exists for (CPU-only) sailed through it."""
|
||||
import torch
|
||||
|
||||
from core.training.diffusion_train_common import dit_accelerator_missing_reason
|
||||
|
||||
class _Missing:
|
||||
"""A torch.mps that predates is_available()."""
|
||||
|
||||
class _Raising:
|
||||
@staticmethod
|
||||
def is_available():
|
||||
raise RuntimeError("driver not initialised")
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
||||
monkeypatch.setattr(torch, "xpu", _Raising)
|
||||
monkeypatch.setattr(torch.backends, "mps", _Missing)
|
||||
assert "GPU" in (dit_accelerator_missing_reason("flux.1") or "")
|
||||
|
||||
# A working probe still clears the gate even when its neighbours are broken.
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
assert dit_accelerator_missing_reason("flux.1") is None
|
||||
|
|
|
|||
|
|
@ -170,6 +170,52 @@ def test_a_different_file_set_is_not_adopted(monkeypatch):
|
|||
dl._registry.set_job(key, "complete")
|
||||
|
||||
|
||||
def test_the_http_retry_keeps_the_scoped_file_list_on_the_record(monkeypatch):
|
||||
# The retry reclaims the same slot with replace_active, and that claim OVERWRITES the
|
||||
# stored metadata. Dropping the file list there left the record claiming an empty scope,
|
||||
# so the next identical scoped start compared [] against the real list and 409'd on
|
||||
# "already fetching a different set of files" instead of adopting the running download.
|
||||
monkeypatch.setattr(dl, "_reject_if_load_in_flight", lambda repo_id: None)
|
||||
monkeypatch.setattr(dl, "resolve_cached_repo_id_case", lambda repo, **k: repo)
|
||||
monkeypatch.setattr(dl, "scoped_file_blob_hashes", lambda *a, **k: frozenset())
|
||||
monkeypatch.setattr(download_lifecycle, "launch_worker", lambda *a, **k: "running")
|
||||
|
||||
class _Proc:
|
||||
pid = 4242
|
||||
|
||||
def poll(self):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(download_lifecycle, "spawn_worker", lambda *a, **k: _Proc())
|
||||
monkeypatch.setattr(download_lifecycle, "register_worker", lambda *a, **k: True)
|
||||
|
||||
key = dl._download_job_key("black-forest-labs/FLUX.1-dev", dl._scope_variant("diffusion"))
|
||||
try:
|
||||
# The retry only exists for a job that started on XET.
|
||||
assert asyncio.run(dl.download_model_response(_request(use_xet = True)))["accepted"] is True
|
||||
|
||||
retried = download_lifecycle._try_http_retry(
|
||||
dl._registry,
|
||||
key,
|
||||
hf_token = None,
|
||||
label = "FLUX.1-dev [@diffusion]",
|
||||
log_prefix = "[test]",
|
||||
logger = download_lifecycle.logging.getLogger("test"),
|
||||
repo_type = "model",
|
||||
repo_id = "black-forest-labs/FLUX.1-dev",
|
||||
watch_name = "test",
|
||||
)
|
||||
assert retried is True
|
||||
|
||||
metadata = dl._registry.get_job_metadata(key)
|
||||
assert metadata is not None and list(metadata.scoped_files) == FILES
|
||||
# And the retried job is still adoptable by the page that asked for those files.
|
||||
again = asyncio.run(dl.download_model_response(_request()))
|
||||
assert again["accepted"] is True and again["job_key"] == key
|
||||
finally:
|
||||
dl._registry.set_job(key, "complete")
|
||||
|
||||
|
||||
def test_scope_key_stays_derivable_from_the_scope_alone():
|
||||
# The download manager builds this key client-side (it polls and cancels before any
|
||||
# server round-trip tells it a key), so the scope name alone must produce it.
|
||||
|
|
|
|||
88
studio/backend/tests/test_scoped_download_worker.py
Normal file
88
studio/backend/tests/test_scoped_download_worker.py
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""The scoped download worker must never report success without the files.
|
||||
|
||||
``snapshot_download`` returns an existing snapshot folder -- having fetched nothing -- when
|
||||
its own ``repo_info`` call fails, and with HF metadata unavailable no manifest is written,
|
||||
so the usual verification is a no-op. A repo already on disk from a full snapshot job (which
|
||||
ignores ``*.gguf``) would otherwise flip a scoped job to complete with no weights, and the
|
||||
Images page auto-loads as soon as the job completes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
from hub.workers import hf_download
|
||||
|
||||
|
||||
FILES = ["model_index.json", "transformer/diffusion_pytorch_model.safetensors"]
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def offline(monkeypatch, tmp_path):
|
||||
"""A worker whose metadata lookups all fail, pointed at a snapshot dir we control."""
|
||||
monkeypatch.setattr(
|
||||
hf_download, "_model_info_with_retry", lambda *a, **k: (_ for _ in ()).throw(OSError("no net"))
|
||||
)
|
||||
monkeypatch.setattr(hf_download, "_protected_blob_hashes", lambda: frozenset())
|
||||
monkeypatch.setattr(hf_download, "_preflight_disk_space", lambda *a, **k: None)
|
||||
|
||||
import hub.utils.download_registry as registry_mod
|
||||
|
||||
monkeypatch.setattr(registry_mod, "prepare_cache_for_transport", lambda *a, **k: 0)
|
||||
|
||||
snapshot = tmp_path / "snapshot"
|
||||
snapshot.mkdir()
|
||||
|
||||
import huggingface_hub
|
||||
|
||||
monkeypatch.setattr(huggingface_hub, "snapshot_download", lambda **k: str(snapshot))
|
||||
return snapshot
|
||||
|
||||
|
||||
def _run(scope = "diffusion"):
|
||||
hf_download._download_scoped_snapshot(
|
||||
"black-forest-labs/FLUX.1-dev", scope, list(FILES), None, "http"
|
||||
)
|
||||
|
||||
|
||||
def test_offline_scoped_download_fails_when_the_files_are_not_on_disk(offline, capsys):
|
||||
(offline / "model_index.json").write_text("{}", encoding = "utf-8") # the cheap file only
|
||||
|
||||
with pytest.raises(SystemExit) as exit_info:
|
||||
_run()
|
||||
assert exit_info.value.code == 1
|
||||
err = capsys.readouterr().err
|
||||
assert "incomplete" in err
|
||||
assert "diffusion_pytorch_model.safetensors" in err
|
||||
|
||||
|
||||
def test_offline_scoped_download_passes_when_every_file_is_present(offline):
|
||||
for rel in FILES:
|
||||
path = offline / rel
|
||||
path.parent.mkdir(parents = True, exist_ok = True)
|
||||
path.write_text("weights", encoding = "utf-8")
|
||||
|
||||
_run() # no SystemExit: everything the job asked for is on disk
|
||||
|
||||
|
||||
def test_a_dangling_symlink_does_not_count_as_present(offline, capsys):
|
||||
"""Cache entries are symlinks into blobs/; a broken one is a missing file."""
|
||||
(offline / "model_index.json").write_text("{}", encoding = "utf-8")
|
||||
target = offline / "transformer"
|
||||
target.mkdir(parents = True, exist_ok = True)
|
||||
(target / "diffusion_pytorch_model.safetensors").symlink_to(offline / "gone.bin")
|
||||
|
||||
with pytest.raises(SystemExit) as exit_info:
|
||||
_run()
|
||||
assert exit_info.value.code == 1
|
||||
assert "incomplete" in capsys.readouterr().err
|
||||
|
|
@ -1768,15 +1768,32 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
} | null>(null);
|
||||
const handleLoadRef = useRef(handleLoad);
|
||||
handleLoadRef.current = handleLoad;
|
||||
// Set when a staged download finished while this page was hidden. Both diffusion pages
|
||||
// stay mounted, and a load evicts whatever holds the GPU, so a download landing in the
|
||||
// background must not take the model out from under the page the user is actually on.
|
||||
// The pick is not dropped: it fires when this page comes back.
|
||||
const stagedLoadDeferred = useRef(false);
|
||||
const { stage } = useStagedDownload({
|
||||
scopeId: "diffusion",
|
||||
onReady: () => {
|
||||
if (!active) {
|
||||
stagedLoadDeferred.current = true;
|
||||
return;
|
||||
}
|
||||
const pending = pendingStagedLoad.current;
|
||||
pendingStagedLoad.current = null;
|
||||
if (pending) void handleLoadRef.current(pending.repoId, pending.opts);
|
||||
},
|
||||
});
|
||||
|
||||
useEffect(() => {
|
||||
if (!active || !stagedLoadDeferred.current) return;
|
||||
stagedLoadDeferred.current = false;
|
||||
const pending = pendingStagedLoad.current;
|
||||
pendingStagedLoad.current = null;
|
||||
if (pending) void handleLoadRef.current(pending.repoId, pending.opts);
|
||||
}, [active]);
|
||||
|
||||
// Stage a not-yet-downloaded hub pick, else load it directly. Returns true when the pick
|
||||
// was accepted either way, so the callers' optimistic picker state stands.
|
||||
const loadOrStage = useCallback(
|
||||
|
|
|
|||
|
|
@ -1108,15 +1108,31 @@ export function VideoPage({ active = true }: { active?: boolean }) {
|
|||
} | null>(null);
|
||||
const handleLoadRef = useRef(handleLoad);
|
||||
handleLoadRef.current = handleLoad;
|
||||
// A download finishing while this page is hidden must not evict the model the visible
|
||||
// page has loaded (both pages stay mounted and a load takes the GPU unconditionally).
|
||||
// The pick is held, not dropped, and fires when Video is on screen again.
|
||||
const stagedLoadDeferred = useRef(false);
|
||||
const { stage } = useStagedDownload({
|
||||
scopeId: "diffusion",
|
||||
onReady: () => {
|
||||
if (!active) {
|
||||
stagedLoadDeferred.current = true;
|
||||
return;
|
||||
}
|
||||
const pending = pendingStagedLoad.current;
|
||||
pendingStagedLoad.current = null;
|
||||
if (pending) void handleLoadRef.current(pending.repoId, pending.opts);
|
||||
},
|
||||
});
|
||||
|
||||
useEffect(() => {
|
||||
if (!active || !stagedLoadDeferred.current) return;
|
||||
stagedLoadDeferred.current = false;
|
||||
const pending = pendingStagedLoad.current;
|
||||
pendingStagedLoad.current = null;
|
||||
if (pending) void handleLoadRef.current(pending.repoId, pending.opts);
|
||||
}, [active]);
|
||||
|
||||
// Stage a not-yet-downloaded hub pick, else load it directly.
|
||||
const loadOrStage = useCallback(
|
||||
async (
|
||||
|
|
|
|||
|
|
@ -626,6 +626,26 @@ def test_diffusion_pages_stage_downloads_through_the_manager():
|
|||
assert "catch" in body, f"{rel}: no fallback when the plan is unavailable"
|
||||
|
||||
|
||||
def test_a_hidden_diffusion_page_does_not_load_when_its_download_lands():
|
||||
"""Images and Video stay mounted behind the router, and a load evicts whoever holds the
|
||||
GPU. A multi-GB staged download finishing while the user is on another page must not take
|
||||
the model out from under them; the pick waits for its page to be visible again."""
|
||||
for rel in ("features/images/images-page.tsx", "features/video/video-page.tsx"):
|
||||
src = _read(rel)
|
||||
ready = re.search(r"onReady: \(\) => \{.*?\n \},", src, re.S)
|
||||
assert ready, f"{rel}: staged-download onReady not found"
|
||||
assert "if (!active)" in ready.group(0), f"{rel}: a hidden page still takes the GPU"
|
||||
# Deferred, not dropped: something has to fire the held pick when the page returns.
|
||||
assert "stagedLoadDeferred" in ready.group(0), f"{rel}: the pick is discarded"
|
||||
flush = re.search(
|
||||
r"if \(!active \|\| !stagedLoadDeferred\.current\) return;.*?\n \}, \[active\]\);",
|
||||
src,
|
||||
re.S,
|
||||
)
|
||||
assert flush, f"{rel}: nothing flushes the deferred load when the page is shown"
|
||||
assert "handleLoadRef.current(" in flush.group(0), f"{rel}: deferred load never runs"
|
||||
|
||||
|
||||
def test_staged_downloads_always_scope_their_files():
|
||||
"""Every staged entry must go out as a scoped job carrying its file list, GGUF
|
||||
checkpoints included. A plain snapshot job drops *.gguf via the Hub's ignore list, so
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue