Skip to content

Commit 2fa51ab

Browse files
ashvinctrlbuluma
authored andcommitted
fix(cookbook): report dead finished downloads as completed instead of stopped (odysseus-dev#4025)
When a download's tmux pane is gone, the status endpoint trusted only the HF-cache probe to tell completed from stopped. The probe derives its cache root from its own environment, but the download runner exports HF_HOME=<local_dir> (the odysseus-dev#2722 fix), so custom-dir downloads land in <local_dir>/hub where the probe never looks - and ollama pulls don't touch the HF cache at all. Finished downloads were reported as stopped forever, and tasks already persisted as completed were demoted back to stopped on the next poll. This is the backend half of odysseus-dev#3897, deliberately left out of the frontend fix in odysseus-dev#4000. - honor the conclusive runner markers first: DOWNLOAD_OK -> completed (keeping the "Fetching 0 files" error guard), DOWNLOAD_FAILED -> error - pass the task's local_dir through to the cache probes so they check the cache the download actually wrote to, keeping the env-var fallback for default-cache downloads - move the probe scripts and marker classification into routes/cookbook_output.py (dependency-free) with behavioral tests Fixes odysseus-dev#4017
1 parent d4a0b4b commit 2fa51ab

3 files changed

Lines changed: 199 additions & 31 deletions

File tree

routes/cookbook_output.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,62 @@
44
unit-tested without standing up the whole app.
55
"""
66

7+
import re
8+
9+
_FETCHING_ZERO_FILES_RE = re.compile(r"Fetching\s+0\s+files", re.IGNORECASE)
10+
11+
# Probe scripts for the dead-session download check, run as
12+
# `python3 -c <PROBE> <repo_id> <cache_root>` (locally or over SSH).
13+
# cache_root is the task's custom download dir, '' for the default HF cache.
14+
# It has to be passed explicitly: the download runner exports
15+
# HF_HOME=<local_dir>, so that task's cache lives under <local_dir>/hub, and
16+
# the probe process's own environment knows nothing about it.
17+
HF_CACHE_COMPLETE_PROBE = (
18+
"import os,sys;"
19+
"repo=sys.argv[1];"
20+
"root=os.path.expanduser(sys.argv[2]) if len(sys.argv)>2 and sys.argv[2] else '';"
21+
"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'));"
22+
"d=os.path.join(base,'models--'+repo.replace('/','--'));"
23+
"snap=os.path.join(d,'snapshots');"
24+
"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));"
25+
"inc=False;"
26+
"blobs=os.path.join(d,'blobs');"
27+
"inc=os.path.isdir(blobs) and any(x.endswith('.incomplete') for x in os.listdir(blobs));"
28+
"sys.exit(0 if ok and not inc else 1)"
29+
)
30+
31+
HF_CACHE_INCOMPLETE_PROBE = (
32+
"import os,sys;"
33+
"repo=sys.argv[1];"
34+
"root=os.path.expanduser(sys.argv[2]) if len(sys.argv)>2 and sys.argv[2] else '';"
35+
"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'));"
36+
"d=os.path.join(base,'models--'+repo.replace('/','--'));"
37+
"blobs=os.path.join(d,'blobs');"
38+
"inc=os.path.isdir(blobs) and any(x.endswith('.incomplete') for x in os.listdir(blobs));"
39+
"sys.exit(0 if inc else 1)"
40+
)
41+
42+
43+
def classify_dead_download(full_snapshot: str):
44+
"""Resolve a dead download session's status from its runner markers.
45+
46+
The runner prints DOWNLOAD_OK only after exiting 0 (and DOWNLOAD_FAILED
47+
otherwise), so the markers stay trustworthy after the tmux pane is gone.
48+
Returns (status, zero_files), or None when the snapshot carries no marker
49+
and the caller has to fall back to the cache probe. Same precedence as
50+
the live-session branch: DOWNLOAD_OK wins, except a "Fetching 0 files"
51+
run is an error (nothing matched the include/quant pattern).
52+
"""
53+
if not full_snapshot:
54+
return None
55+
if "DOWNLOAD_OK" in full_snapshot:
56+
if _FETCHING_ZERO_FILES_RE.search(full_snapshot):
57+
return ("error", True)
58+
return ("completed", False)
59+
if "DOWNLOAD_FAILED" in full_snapshot:
60+
return ("error", False)
61+
return None
62+
763

864
def error_aware_output_tail(full_snapshot: str, status: str) -> str:
965
"""Return the trailing slice of a task log for the status response.

routes/cookbook_routes.py

Lines changed: 19 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@
4444
HF_CACHE_INCOMPLETE_PROBE,
4545
)
4646

47-
4847
logger = logging.getLogger(__name__)
4948

5049
from routes.cookbook_helpers import (
@@ -2786,30 +2785,20 @@ async def cookbook_tasks_status(request: Request):
27862785
def _cookbook_tasks_status_sync():
27872786
import subprocess
27882787

2789-
def _download_cache_complete(repo_id: str, remote_host: str = "", ssh_port: str = "") -> bool:
2788+
def _download_cache_complete(repo_id: str, remote_host: str = "", ssh_port: str = "", cache_root: str = "") -> bool:
27902789
"""Best-effort check for a completed HF cache entry.
27912790
27922791
tmux output can stop at a stale progress line if the pane/session
27932792
disappears before Cookbook captures the final DOWNLOAD_OK marker.
27942793
In that case, trust the cache shape: a snapshot directory with files
27952794
and no *.incomplete blobs means HuggingFace finished materializing the
2796-
model.
2795+
model. cache_root is the task's custom download dir — the runner
2796+
pointed HF_HOME there, so the cache lives under <cache_root>/hub,
2797+
not wherever this probe's environment says.
27972798
"""
27982799
if not repo_id or "/" not in repo_id:
27992800
return False
2800-
py = (
2801-
"import os,sys;"
2802-
"repo=sys.argv[1];"
2803-
"base=os.environ.get('HUGGINGFACE_HUB_CACHE') or os.path.join(os.environ.get('HF_HOME', os.path.expanduser('~/.cache/huggingface')), 'hub');"
2804-
"d=os.path.join(base,'models--'+repo.replace('/','--'));"
2805-
"snap=os.path.join(d,'snapshots');"
2806-
"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));"
2807-
"inc=False;"
2808-
"blobs=os.path.join(d,'blobs');"
2809-
"inc=os.path.isdir(blobs) and any(x.endswith('.incomplete') for x in os.listdir(blobs));"
2810-
"sys.exit(0 if ok and not inc else 1)"
2811-
)
2812-
cmd = ["python3", "-c", py, repo_id]
2801+
cmd = ["python3", "-c", HF_CACHE_COMPLETE_PROBE, repo_id, cache_root or ""]
28132802
try:
28142803
if remote_host:
28152804
ssh_base = ["ssh"]
@@ -2823,7 +2812,7 @@ def _download_cache_complete(repo_id: str, remote_host: str = "", ssh_port: str
28232812
except Exception:
28242813
return False
28252814

2826-
def _download_cache_incomplete(repo_id: str, remote_host: str = "", ssh_port: str = "") -> bool:
2815+
def _download_cache_incomplete(repo_id: str, remote_host: str = "", ssh_port: str = "", cache_root: str = "") -> bool:
28272816
"""Best-effort check for resumable HF partial blobs.
28282817
28292818
A lost SSH/tmux session can leave a real download still incomplete.
@@ -2832,16 +2821,7 @@ def _download_cache_incomplete(repo_id: str, remote_host: str = "", ssh_port: st
28322821
"""
28332822
if not repo_id or "/" not in repo_id:
28342823
return False
2835-
py = (
2836-
"import os,sys;"
2837-
"repo=sys.argv[1];"
2838-
"base=os.environ.get('HUGGINGFACE_HUB_CACHE') or os.path.join(os.environ.get('HF_HOME', os.path.expanduser('~/.cache/huggingface')), 'hub');"
2839-
"d=os.path.join(base,'models--'+repo.replace('/','--'));"
2840-
"blobs=os.path.join(d,'blobs');"
2841-
"inc=os.path.isdir(blobs) and any(x.endswith('.incomplete') for x in os.listdir(blobs));"
2842-
"sys.exit(0 if inc else 1)"
2843-
)
2844-
cmd = ["python3", "-c", py, repo_id]
2824+
cmd = ["python3", "-c", HF_CACHE_INCOMPLETE_PROBE, repo_id, cache_root or ""]
28452825
try:
28462826
if remote_host:
28472827
ssh_base = ["ssh"]
@@ -3046,7 +3026,7 @@ def _download_cache_incomplete(repo_id: str, remote_host: str = "", ssh_port: st
30463026
and (
30473027
".incomplete" in full_snapshot
30483028
or bool(re.search(r'model-\d+-of-\d+\.[A-Za-z0-9_.-]+:\s+(?:[0-9]|[1-8][0-9])%', full_snapshot))
3049-
or _download_cache_incomplete(_payload.get("repo_id") or model, remote, str(_tport or ""))
3029+
or _download_cache_incomplete(_payload.get("repo_id") or model, remote, str(_tport or ""), _payload.get("local_dir") or "")
30503030
)
30513031
)
30523032
if is_alive or (local_win_task and full_snapshot):
@@ -3087,11 +3067,19 @@ def _download_cache_incomplete(repo_id: str, remote_host: str = "", ssh_port: st
30873067
else:
30883068
status = "running"
30893069
else:
3090-
# Session is dead — check if it completed or crashed
3091-
if (
3070+
# Session is dead — check if it completed or crashed. The
3071+
# runner markers in the retained output are conclusive
3072+
# (DOWNLOAD_OK only prints after exit 0), so check them before
3073+
# the cache probe, which can't see ollama pulls at all.
3074+
marker = classify_dead_download(full_snapshot) if task_type == "download" else None
3075+
if marker is not None:
3076+
status, download_zero_files = marker
3077+
if status == "completed" and not progress_text:
3078+
progress_text = "Download complete"
3079+
elif (
30923080
task_type == "download"
30933081
and not download_has_incomplete_evidence
3094-
and _download_cache_complete(_payload.get("repo_id") or model, remote, str(_tport or ""))
3082+
and _download_cache_complete(_payload.get("repo_id") or model, remote, str(_tport or ""), _payload.get("local_dir") or "")
30953083
):
30963084
status = "completed"
30973085
if not progress_text:
Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
"""Behavioral guards for dead-session download classification (issue #4017).
2+
3+
A download whose tmux pane is gone must not be reported as stopped when its
4+
retained output carries DOWNLOAD_OK, or when the files landed in a custom
5+
download dir. The runner exports HF_HOME=<local_dir>, so the cache lives
6+
under <local_dir>/hub — the probe only finds it if the task's dir is passed
7+
in explicitly rather than read from the probe process's environment.
8+
"""
9+
import os
10+
import subprocess
11+
import sys
12+
13+
from routes.cookbook_output import (
14+
classify_dead_download,
15+
HF_CACHE_COMPLETE_PROBE,
16+
HF_CACHE_INCOMPLETE_PROBE,
17+
)
18+
19+
REPO = "org/some-model-GGUF"
20+
21+
22+
# ── Marker classification ──
23+
24+
25+
def test_download_ok_resolves_completed():
26+
snap = "Fetching 4 files: 100%|####| 4/4\nDownload complete\n\nDOWNLOAD_OK\n$"
27+
assert classify_dead_download(snap) == ("completed", False)
28+
29+
30+
def test_download_failed_resolves_error():
31+
snap = "some progress\n\nDOWNLOAD_FAILED (exit 1 after 3 attempts)"
32+
assert classify_dead_download(snap) == ("error", False)
33+
34+
35+
def test_download_ok_with_zero_files_resolves_error():
36+
# A DOWNLOAD_OK from a run that matched no files (bad include/quant
37+
# pattern) is still a failure — same guard as the live-session branch.
38+
snap = "Fetching 0 files: 0it [00:00, ?it/s]\n\nDOWNLOAD_OK"
39+
assert classify_dead_download(snap) == ("error", True)
40+
41+
42+
def test_no_marker_returns_none():
43+
# Mid-download tail with no terminal marker — caller must fall back to
44+
# the cache probe.
45+
assert classify_dead_download("Downloading model.gguf: 42%") is None
46+
assert classify_dead_download("") is None
47+
48+
49+
def test_ollama_pull_output_resolves_completed():
50+
snap = "pulling manifest\npulling 8f39d1c3...: 100%\nsuccess\n\nDOWNLOAD_OK"
51+
assert classify_dead_download(snap) == ("completed", False)
52+
53+
54+
# ── Cache probe scripts ──
55+
56+
57+
def _make_cache(root, repo=REPO, incomplete=False, empty_snapshot=False):
58+
d = os.path.join(root, "hub", "models--" + repo.replace("/", "--"))
59+
snap = os.path.join(d, "snapshots", "abc123")
60+
os.makedirs(snap)
61+
if not empty_snapshot:
62+
with open(os.path.join(snap, "model.gguf"), "w") as f:
63+
f.write("x")
64+
if incomplete:
65+
blobs = os.path.join(d, "blobs")
66+
os.makedirs(blobs)
67+
with open(os.path.join(blobs, "deadbeef.incomplete"), "w") as f:
68+
f.write("x")
69+
70+
71+
def _run_probe(probe, repo, cache_root, env=None):
72+
# Strip the HF cache vars so the probe can't accidentally find a real
73+
# cache on the machine running the tests.
74+
full_env = {k: v for k, v in os.environ.items()
75+
if k not in ("HF_HOME", "HUGGINGFACE_HUB_CACHE", "HF_HUB_CACHE")}
76+
full_env.update(env or {})
77+
return subprocess.run(
78+
[sys.executable, "-c", probe, repo, cache_root],
79+
env=full_env, capture_output=True, timeout=30,
80+
).returncode
81+
82+
83+
def test_complete_probe_finds_custom_dir_cache(tmp_path):
84+
# Model materialized under <local_dir>/hub — found only via the explicit
85+
# cache_root argument (issue #4017).
86+
root = str(tmp_path)
87+
_make_cache(root)
88+
assert _run_probe(HF_CACHE_COMPLETE_PROBE, REPO, root) == 0
89+
90+
91+
def test_complete_probe_misses_without_cache_root(tmp_path):
92+
# Same on-disk layout, but without the cache_root argument the probe
93+
# falls back to the default cache and misses it.
94+
_make_cache(str(tmp_path))
95+
assert _run_probe(HF_CACHE_COMPLETE_PROBE, REPO, "") == 1
96+
97+
98+
def test_complete_probe_rejects_incomplete_blobs(tmp_path):
99+
root = str(tmp_path)
100+
_make_cache(root, incomplete=True)
101+
assert _run_probe(HF_CACHE_COMPLETE_PROBE, REPO, root) == 1
102+
103+
104+
def test_complete_probe_rejects_empty_snapshot(tmp_path):
105+
root = str(tmp_path)
106+
_make_cache(root, empty_snapshot=True)
107+
assert _run_probe(HF_CACHE_COMPLETE_PROBE, REPO, root) == 1
108+
109+
110+
def test_complete_probe_env_fallback_still_works(tmp_path):
111+
# No custom dir on the task — the probe must keep honoring the standard
112+
# HF env vars so default-cache downloads classify as before.
113+
root = str(tmp_path)
114+
_make_cache(root)
115+
hub = os.path.join(root, "hub")
116+
assert _run_probe(HF_CACHE_COMPLETE_PROBE, REPO, "", env={"HUGGINGFACE_HUB_CACHE": hub}) == 0
117+
118+
119+
def test_incomplete_probe_sees_custom_dir_partials(tmp_path):
120+
root = str(tmp_path)
121+
_make_cache(root, incomplete=True)
122+
assert _run_probe(HF_CACHE_INCOMPLETE_PROBE, REPO, root) == 0
123+
# Clean cache → no resumable partials.
124+
assert _run_probe(HF_CACHE_INCOMPLETE_PROBE, "org/other-model", root) == 1

0 commit comments

Comments
 (0)