Skip to content

Commit 1469019

Browse files
ashvinctrlentitycs
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 f34ed1d commit 1469019

3 files changed

Lines changed: 203 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: 23 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,10 @@
3030
which_tool,
3131
)
3232
from routes.shell_routes import TMUX_LOG_DIR
33-
from routes.cookbook_output import error_aware_output_tail
33+
from routes.cookbook_output import (
34+
error_aware_output_tail, classify_dead_download,
35+
HF_CACHE_COMPLETE_PROBE, HF_CACHE_INCOMPLETE_PROBE,
36+
)
3437

3538
logger = logging.getLogger(__name__)
3639

@@ -2636,30 +2639,20 @@ async def cookbook_tasks_status(request: Request):
26362639
def _cookbook_tasks_status_sync():
26372640
import subprocess
26382641

2639-
def _download_cache_complete(repo_id: str, remote_host: str = "", ssh_port: str = "") -> bool:
2642+
def _download_cache_complete(repo_id: str, remote_host: str = "", ssh_port: str = "", cache_root: str = "") -> bool:
26402643
"""Best-effort check for a completed HF cache entry.
26412644
26422645
tmux output can stop at a stale progress line if the pane/session
26432646
disappears before Cookbook captures the final DOWNLOAD_OK marker.
26442647
In that case, trust the cache shape: a snapshot directory with files
26452648
and no *.incomplete blobs means HuggingFace finished materializing the
2646-
model.
2649+
model. cache_root is the task's custom download dir — the runner
2650+
pointed HF_HOME there, so the cache lives under <cache_root>/hub,
2651+
not wherever this probe's environment says.
26472652
"""
26482653
if not repo_id or "/" not in repo_id:
26492654
return False
2650-
py = (
2651-
"import os,sys;"
2652-
"repo=sys.argv[1];"
2653-
"base=os.environ.get('HUGGINGFACE_HUB_CACHE') or os.path.join(os.environ.get('HF_HOME', os.path.expanduser('~/.cache/huggingface')), 'hub');"
2654-
"d=os.path.join(base,'models--'+repo.replace('/','--'));"
2655-
"snap=os.path.join(d,'snapshots');"
2656-
"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));"
2657-
"inc=False;"
2658-
"blobs=os.path.join(d,'blobs');"
2659-
"inc=os.path.isdir(blobs) and any(x.endswith('.incomplete') for x in os.listdir(blobs));"
2660-
"sys.exit(0 if ok and not inc else 1)"
2661-
)
2662-
cmd = ["python3", "-c", py, repo_id]
2655+
cmd = ["python3", "-c", HF_CACHE_COMPLETE_PROBE, repo_id, cache_root or ""]
26632656
try:
26642657
if remote_host:
26652658
ssh_base = ["ssh"]
@@ -2673,7 +2666,7 @@ def _download_cache_complete(repo_id: str, remote_host: str = "", ssh_port: str
26732666
except Exception:
26742667
return False
26752668

2676-
def _download_cache_incomplete(repo_id: str, remote_host: str = "", ssh_port: str = "") -> bool:
2669+
def _download_cache_incomplete(repo_id: str, remote_host: str = "", ssh_port: str = "", cache_root: str = "") -> bool:
26772670
"""Best-effort check for resumable HF partial blobs.
26782671
26792672
A lost SSH/tmux session can leave a real download still incomplete.
@@ -2682,16 +2675,7 @@ def _download_cache_incomplete(repo_id: str, remote_host: str = "", ssh_port: st
26822675
"""
26832676
if not repo_id or "/" not in repo_id:
26842677
return False
2685-
py = (
2686-
"import os,sys;"
2687-
"repo=sys.argv[1];"
2688-
"base=os.environ.get('HUGGINGFACE_HUB_CACHE') or os.path.join(os.environ.get('HF_HOME', os.path.expanduser('~/.cache/huggingface')), 'hub');"
2689-
"d=os.path.join(base,'models--'+repo.replace('/','--'));"
2690-
"blobs=os.path.join(d,'blobs');"
2691-
"inc=os.path.isdir(blobs) and any(x.endswith('.incomplete') for x in os.listdir(blobs));"
2692-
"sys.exit(0 if inc else 1)"
2693-
)
2694-
cmd = ["python3", "-c", py, repo_id]
2678+
cmd = ["python3", "-c", HF_CACHE_INCOMPLETE_PROBE, repo_id, cache_root or ""]
26952679
try:
26962680
if remote_host:
26972681
ssh_base = ["ssh"]
@@ -2896,7 +2880,7 @@ def _download_cache_incomplete(repo_id: str, remote_host: str = "", ssh_port: st
28962880
and (
28972881
".incomplete" in full_snapshot
28982882
or bool(re.search(r'model-\d+-of-\d+\.[A-Za-z0-9_.-]+:\s+(?:[0-9]|[1-8][0-9])%', full_snapshot))
2899-
or _download_cache_incomplete(_payload.get("repo_id") or model, remote, str(_tport or ""))
2883+
or _download_cache_incomplete(_payload.get("repo_id") or model, remote, str(_tport or ""), _payload.get("local_dir") or "")
29002884
)
29012885
)
29022886
if is_alive or (local_win_task and full_snapshot):
@@ -2937,11 +2921,19 @@ def _download_cache_incomplete(repo_id: str, remote_host: str = "", ssh_port: st
29372921
else:
29382922
status = "running"
29392923
else:
2940-
# Session is dead — check if it completed or crashed
2941-
if (
2924+
# Session is dead — check if it completed or crashed. The
2925+
# runner markers in the retained output are conclusive
2926+
# (DOWNLOAD_OK only prints after exit 0), so check them before
2927+
# the cache probe, which can't see ollama pulls at all.
2928+
marker = classify_dead_download(full_snapshot) if task_type == "download" else None
2929+
if marker is not None:
2930+
status, download_zero_files = marker
2931+
if status == "completed" and not progress_text:
2932+
progress_text = "Download complete"
2933+
elif (
29422934
task_type == "download"
29432935
and not download_has_incomplete_evidence
2944-
and _download_cache_complete(_payload.get("repo_id") or model, remote, str(_tport or ""))
2936+
and _download_cache_complete(_payload.get("repo_id") or model, remote, str(_tport or ""), _payload.get("local_dir") or "")
29452937
):
29462938
status = "completed"
29472939
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)