diff --git a/routes/cookbook_helpers.py b/routes/cookbook_helpers.py index 856c8bdb5..8c8b5caea 100644 --- a/routes/cookbook_helpers.py +++ b/routes/cookbook_helpers.py @@ -14,6 +14,7 @@ from fastapi import HTTPException from pydantic import BaseModel from routes._validators import validate_remote_host, validate_ssh_port +from routes.cookbook_output import HF_PARTIAL_BLOB_HELPER from core.platform_compat import _ssh_exec_argv logger = logging.getLogger(__name__) @@ -378,6 +379,7 @@ def _cached_model_scan_script(model_dirs: list[str] | None = None, add_hf_cache: """ lines = [ "import json, os, re, shutil, subprocess, urllib.request", + *HF_PARTIAL_BLOB_HELPER.splitlines(), "models = []", "seen = set()", "BLOCKED_ROOTS = ('/sys', '/proc', '/dev', '/run', '/var/run')", @@ -436,8 +438,10 @@ def _cached_model_scan_script(model_dirs: list[str] | None = None, add_hf_cache: " sz, nf, ic = 0, 0, False", " if os.path.isdir(blobs):", " for f in os.scandir(blobs):", + " if f.name.endswith('.incomplete'):", + " if not is_partial_blob(f.path): continue", + " ic = True", " if f.is_file(): nf += 1; sz += f.stat().st_size", - " if f.name.endswith('.incomplete'): ic = True", " snap = os.path.join(cache, d, 'snapshots')", " def snapshot_size():", " total, count, incomplete = 0, 0, False", diff --git a/routes/cookbook_output.py b/routes/cookbook_output.py index b30b18536..a8e8ed528 100644 --- a/routes/cookbook_output.py +++ b/routes/cookbook_output.py @@ -8,6 +8,20 @@ import re _FETCHING_ZERO_FILES_RE = re.compile(r"Fetching\s+0\s+files", re.IGNORECASE) +# Shared source embedded in local/remote inventory and status probes. A blank +# retry marker beside its completed, nonempty content-addressed blob is stale; +# a blank marker alone may be a download that has not received its first byte. +HF_PARTIAL_BLOB_HELPER = """def is_partial_blob(path): + if not path.endswith('.incomplete'): + return False + try: + completed = path[:-len('.incomplete')] + return (os.path.getsize(path) != 0 or not os.path.isfile(completed) + or os.path.getsize(completed) == 0) + except OSError: + return True +""" + # Probe scripts for the dead-session download check, run as # `python3 -c ` (locally or over SSH). # cache_root is the task's custom download dir, '' for the default HF cache. @@ -15,7 +29,7 @@ _FETCHING_ZERO_FILES_RE = re.compile(r"Fetching\s+0\s+files", re.IGNORECASE) # HF_HOME=, so that task's cache lives under /hub, and # the probe process's own environment knows nothing about it. HF_CACHE_COMPLETE_PROBE = ( - "import os,sys;" + "import os,sys\n" + HF_PARTIAL_BLOB_HELPER + "repo=sys.argv[1];" "root=os.path.expanduser(sys.argv[2]) if len(sys.argv)>2 and sys.argv[2] else '';" "base=os.path.join(root,'hub') if root else (os.environ.get('HUGGINGFACE_HUB_CACHE') or os.path.join(os.environ.get('HF_HOME', os.path.expanduser('~/.cache/huggingface')), 'hub'));" @@ -24,18 +38,18 @@ HF_CACHE_COMPLETE_PROBE = ( "ok=os.path.isdir(snap) and any(os.path.isdir(os.path.join(snap,x)) and os.listdir(os.path.join(snap,x)) for x in os.listdir(snap));" "inc=False;" "blobs=os.path.join(d,'blobs');" - "inc=os.path.isdir(blobs) and any(x.endswith('.incomplete') for x in os.listdir(blobs));" + "inc=os.path.isdir(blobs) and any(is_partial_blob(os.path.join(blobs,x)) for x in os.listdir(blobs));" "sys.exit(0 if ok and not inc else 1)" ) HF_CACHE_INCOMPLETE_PROBE = ( - "import os,sys;" + "import os,sys\n" + HF_PARTIAL_BLOB_HELPER + "repo=sys.argv[1];" "root=os.path.expanduser(sys.argv[2]) if len(sys.argv)>2 and sys.argv[2] else '';" "base=os.path.join(root,'hub') if root else (os.environ.get('HUGGINGFACE_HUB_CACHE') or os.path.join(os.environ.get('HF_HOME', os.path.expanduser('~/.cache/huggingface')), 'hub'));" "d=os.path.join(base,'models--'+repo.replace('/','--'));" "blobs=os.path.join(d,'blobs');" - "inc=os.path.isdir(blobs) and any(x.endswith('.incomplete') for x in os.listdir(blobs));" + "inc=os.path.isdir(blobs) and any(is_partial_blob(os.path.join(blobs,x)) for x in os.listdir(blobs));" "sys.exit(0 if inc else 1)" ) diff --git a/routes/cookbook_routes.py b/routes/cookbook_routes.py index 02b419894..a8e61e080 100644 --- a/routes/cookbook_routes.py +++ b/routes/cookbook_routes.py @@ -4343,7 +4343,7 @@ def setup_cookbook_routes() -> APIRouter: tmux output can stop at a stale progress line if the pane/session disappears before Cookbook captures the final DOWNLOAD_OK marker. In that case, trust the cache shape: a snapshot directory with files - and no *.incomplete blobs means HuggingFace finished materializing the + and no unfinished blobs means HuggingFace finished materializing the model. cache_root is the task's custom download dir — the runner pointed HF_HOME there, so the cache lives under /hub, not wherever this probe's environment says. @@ -4368,7 +4368,8 @@ def setup_cookbook_routes() -> APIRouter: """Best-effort check for resumable HF partial blobs. A lost SSH/tmux session can leave a real download still incomplete. - Treat any *.incomplete blob as stronger evidence than stale + Treat unfinished blobs (excluding empty markers paired with a + completed blob) as stronger evidence than stale "100%" lines in the captured pane output. """ if not repo_id or "/" not in repo_id: diff --git a/tests/test_cookbook_stale_cache_markers.py b/tests/test_cookbook_stale_cache_markers.py new file mode 100644 index 000000000..6439ce963 --- /dev/null +++ b/tests/test_cookbook_stale_cache_markers.py @@ -0,0 +1,83 @@ +"""A completed blob's empty retry marker must not keep its model stalled.""" +import json +import os +import subprocess +import sys + +import pytest + +from routes.cookbook_helpers import _cached_model_scan_script +from routes.cookbook_output import HF_CACHE_COMPLETE_PROBE, HF_CACHE_INCOMPLETE_PROBE + +_REPO = "fixture/stale-marker" +_STATES = ["complete", "stale", "empty-unpaired", "partial", "partial-paired", "empty-paired-empty"] + + +@pytest.fixture(params=["custom", "hf-home", "hub-env"]) +def cache(request, tmp_path): + root = tmp_path / "cache with spaces" + model = root / "hub" / ("models--" + _REPO.replace("/", "--")) + blobs = model / "blobs" + blobs.mkdir(parents=True) + snapshot = model / "snapshots" / "revision" + snapshot.mkdir(parents=True) + (snapshot / "config.json").write_text("{}") + env = {key: value for key, value in os.environ.items() + if key not in {"HF_HOME", "HUGGINGFACE_HUB_CACHE", "HF_HUB_CACHE"}} + if request.param == "custom": + argument = str(root) + scan = _cached_model_scan_script(add_hf_cache=str(root / "hub")) + else: + argument = "" + env["HF_HOME" if request.param == "hf-home" else "HUGGINGFACE_HUB_CACHE"] = str( + root if request.param == "hf-home" else root / "hub" + ) + scan = _cached_model_scan_script() + return blobs, argument, env, scan + + +def _seed(blobs, state): + if state in {"complete", "stale", "partial-paired"}: + (blobs / "sha").write_bytes(b"completed model data") + elif state == "empty-paired-empty": + (blobs / "sha").touch() + if state != "complete": + (blobs / "sha.incomplete").write_bytes(b"partial" if state.startswith("partial") else b"") + + +@pytest.mark.parametrize("state", _STATES) +@pytest.mark.parametrize("probe", [HF_CACHE_COMPLETE_PROBE, HF_CACHE_INCOMPLETE_PROBE], ids=["complete", "incomplete"]) +def test_cache_status_probes_distinguish_stale_and_partial(cache, state, probe): + blobs, argument, env, _ = cache + _seed(blobs, state) + before = {path.name: path.read_bytes() for path in blobs.iterdir()} + result = subprocess.run([sys.executable, "-c", probe, _REPO, argument], + env=env, capture_output=True, text=True, timeout=15) + complete = state in {"complete", "stale"} + expected = complete if probe == HF_CACHE_COMPLETE_PROBE else not complete + assert result.returncode == (0 if expected else 1), result.stderr + assert {path.name: path.read_bytes() for path in blobs.iterdir()} == before + + +@pytest.mark.parametrize("state", _STATES) +def test_generated_inventory_agrees_without_mutating_cache(cache, state): + blobs, _, env, scan = cache + _seed(blobs, state) + before = {path.name: path.read_bytes() for path in blobs.iterdir()} + # Inventory's unrelated Ollama discovery stays off the network. Execute + # the actual generated scanner, including its cache-root resolution. + harness = ( + "import urllib.request,shutil\n" + "def offline(*args,**kwargs): raise OSError('Ollama outside this fixture')\n" + "urllib.request.urlopen=offline\nshutil.which=lambda name: None\n" + + scan + ) + result = subprocess.run([sys.executable, "-c", harness], env=env, + capture_output=True, text=True, timeout=15) + assert result.returncode == 0, result.stderr + row = next(item for item in json.loads(result.stdout) if item["repo_id"] == _REPO) + assert row["has_incomplete"] == (state not in {"complete", "stale"}) + if state in {"complete", "stale"}: + assert row["nb_files"] == 1 + assert row["size_bytes"] == len(b"completed model data") + assert {path.name: path.read_bytes() for path in blobs.iterdir()} == before