diff --git a/ainode/models/api_routes.py b/ainode/models/api_routes.py index e1be56e..a25a872 100644 --- a/ainode/models/api_routes.py +++ b/ainode/models/api_routes.py @@ -946,23 +946,67 @@ async def _poll_progress(): try: def _do_download(): - from huggingface_hub import snapshot_download + # File-by-file (NOT snapshot_download) so a cancel is real: snapshot_download + # has no cancel hook, so the old _cancel flag was a no-op and a cancelled + # pull ran to completion. We resolve the repo ONCE, then fetch each file via + # hf_hub_download into the SAME local_dir — snapshot_download internally + # does exactly this per file, so the on-disk layout is identical (load- + # bearing for on-disk-serve) — checking the cancel flag between files. + # + # Two properties we must preserve from snapshot_download (it gave them for + # free; a naive loop drops both): + # 1. Commit pinning. snapshot_download resolves repo_info.sha ONCE and + # passes revision= to every per-file fetch, so the whole snapshot + # comes from one commit even if `main` moves mid-pull (these are often + # actively-maintained community/quant repos, and a big pull is a wide + # window). We do the same: repo_info once → pin its .sha → thread it + # through every hf_hub_download. Without this, a push mid-download could + # 404 on a renamed file or silently mix artifacts from two commits. + # 2. Intra-repo parallelism. snapshot_download fetches with a worker pool + # (max_workers). Serial file-by-file would multiply wall-clock on the + # multi-shard, multi-hundred-GB repos this endpoint targets. We keep a + # bounded ThreadPoolExecutor (_download_max_workers), still checking the + # cancel flag before each file is submitted. + import os + from concurrent.futures import ThreadPoolExecutor, as_completed + from huggingface_hub import HfApi, hf_hub_download - def _progress_callback(info): - # Check cancel flag on every chunk callback + from ainode.models.registry import _download_max_workers + + token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or None + + def _check_cancel(): if jobs.get(job_id, {}).get("_cancel"): raise _DownloadCancelled("Download cancelled by user") - from ainode.models.registry import _download_max_workers - try: - snapshot_download( + _check_cancel() + # Resolve files + pinned commit in one call so every file is from one commit. + info = HfApi(token=token).repo_info(repo_id=hf_repo) + revision = info.sha + files = [s.rfilename for s in (info.siblings or [])] + + def _download_one(rfilename: str) -> None: + _check_cancel() # stop between files — no mid-file hook exists + hf_hub_download( repo_id=hf_repo, + filename=rfilename, + revision=revision, # pin to the commit resolved above local_dir=str(target), - tqdm_class=None, - max_workers=_download_max_workers(), # don't monopolise the uplink + token=token, ) - except _DownloadCancelled: - raise + + _check_cancel() + max_workers = max(1, min(_download_max_workers(), len(files) or 1)) + with ThreadPoolExecutor(max_workers=max_workers) as ex: + futures = [ex.submit(_download_one, f) for f in files] + try: + for fut in as_completed(futures): + fut.result() # propagate first error (incl. _DownloadCancelled) + except BaseException: + for fut in futures: + fut.cancel() # drop not-yet-started files; in-flight finish + raise + _check_cancel() # a cancel arriving after the last file still counts return str(target) # Serialize downloads (one fat pull at a time) so two concurrent model diff --git a/ainode/web/static/js/app.js b/ainode/web/static/js/app.js index 2f01e7d..f0eafcb 100644 --- a/ainode/web/static/js/app.js +++ b/ainode/web/static/js/app.js @@ -888,6 +888,9 @@ const AINode = { if (dot.classList.contains('active') && nodeSelector.querySelectorAll('.node-dot.active').length <= 1) return; dot.classList.toggle('active'); + // The user hand-picked nodes for this pending launch — auto-recommend + // (fired on model-select onchange) must NOT clobber that choice. + self._launchNodesUserPicked = true; updateLaunchHint(); }); }); @@ -1006,8 +1009,11 @@ const AINode = { }); // Set node selection; the fit-aware updater renders the hint (and survives polling). + // But if the user already hand-picked nodes for this launch, DON'T clobber their + // dots — a "pick node, then model" flow was landing the load on the head. We still + // recompute the hint text below so the recommendation is reflected. var ids = rec ? [head.node_id].concat(rec.peers.map(function (n) { return n.node_id; })) : [head.node_id]; - if (this._selectNodeIds) this._selectNodeIds(ids); + if (!this._launchNodesUserPicked && this._selectNodeIds) this._selectNodeIds(ids); if (this._launchHintUpdater) this._launchHintUpdater(); }, @@ -1139,6 +1145,9 @@ const AINode = { this.toast(data.error, 'error'); } else { this.toast('Launched: ' + model, 'success'); + // Launch submitted — the hand-picked-nodes intent is consumed, so the next + // model pick auto-recommends again. + this._launchNodesUserPicked = false; this.refresh(); } } catch (err) { diff --git a/tests/test_models.py b/tests/test_models.py index b4576e7..9fed2dd 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -7,6 +7,8 @@ from __future__ import annotations +import os +import threading from pathlib import Path from unittest.mock import patch, MagicMock @@ -530,3 +532,158 @@ def test_reload_stacked_under_cap_excludes_own_reservation(self): res = append_solo_instance(app, "B", 0.4, persist=False) assert res["ok"] is True and res["stacked"] is True b_backend.stop.assert_called_once() # old instance freed AFTER admission + + +# --------------------------------------------------------------------------- +# M1: cancellable file-by-file repo download (_run_download_repo) +# --------------------------------------------------------------------------- + +import asyncio as _asyncio # noqa: E402 +import sys as _sys # noqa: E402 +from types import SimpleNamespace as _SNS # noqa: E402 + +from ainode.models import api_routes as _api_routes # noqa: E402 + + +class _FakeHfHub: + """Stand-in for the huggingface_hub module: file-by-file download that + writes real bytes into local_dir (same layout snapshot_download produces) + and can invoke a hook after each file (to simulate a mid-download cancel). + + Mirrors the API _do_download actually uses: HfApi().repo_info() to resolve + the file list + a pinned commit sha ONCE, then per-file hf_hub_download with + a revision= kwarg. Records the revision passed to each fetch so tests can + assert every file came from the single pinned commit.""" + + def __init__(self, files: dict, after_file=None, sha="deadbeefcafe"): + # files: {relative_path: bytes} + self._files = files + self._after_file = after_file + self.sha = sha + self.downloaded = [] + self.revisions = [] + self._lock = threading.Lock() + hub = self + + class _FakeHfApi: + def __init__(self, token=None): + pass + + def repo_info(self, repo_id): + return _SNS( + sha=hub.sha, + siblings=[_SNS(rfilename=f) for f in hub._files], + ) + + self.HfApi = _FakeHfApi + + def hf_hub_download(self, repo_id, filename, local_dir, revision=None, token=None): + dest = Path(local_dir) / filename + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(self._files[filename]) + with self._lock: # workers run in parallel now — guard the records + self.downloaded.append(filename) + self.revisions.append(revision) + if self._after_file is not None: + self._after_file(filename) + return str(dest) + + +class TestCancellableRepoDownload: + def _run(self, tmp_path, fake, hf_repo, total_bytes): + manager = _SNS(models_dir=tmp_path) + jobs = {} + job_id = "job-1" + jobs[job_id] = {"model_id": hf_repo, "status": "downloading", + "error": None, "finished_at": None} + with patch.dict(_sys.modules, {"huggingface_hub": fake}), \ + patch.object(_api_routes, "_get_repo_total_bytes", return_value=total_bytes): + _asyncio.run(_api_routes._run_download_repo(manager, hf_repo, job_id, jobs)) + return jobs[job_id], job_id + + def test_completion_layout_identical_and_bytes_tracked(self, tmp_path): + files = { + "config.json": b'{"ok": true}', + "model.safetensors": b"\x00" * 2048, + "nested/tokenizer.json": b"abc", # nested path must survive verbatim + } + total = sum(len(v) for v in files.values()) + fake = _FakeHfHub(files) + hf_repo = "org/model-x" + job, _jid = self._run(tmp_path, fake, hf_repo, total) + + assert job["status"] == "completed" + assert job["progress"] == 100.0 + assert job["downloaded_bytes"] == total + assert job["total_bytes"] == total + + # Layout: files land at target/, IDENTICAL to what + # snapshot_download(local_dir=...) produces (load-bearing for on-disk-serve). + target = tmp_path / hf_repo.replace("/", "--") + for rel, data in files.items(): + f = target / rel + assert f.is_file(), f"missing {rel}" + assert f.read_bytes() == data + # Every listed file was actually fetched. + assert set(fake.downloaded) == set(files.keys()) + # Commit pinning: every file was fetched at the SINGLE sha resolved once + # from repo_info — not an unpinned per-file resolve of a moving `main`. + assert set(fake.revisions) == {fake.sha} + + def test_cancel_between_files_sets_cancelled_and_removes_dir(self, tmp_path): + files = {"a.bin": b"1" * 512, "b.bin": b"2" * 512, "c.bin": b"3" * 512} + hf_repo = "org/model-y" + target = tmp_path / hf_repo.replace("/", "--") + + # Simulate the user hitting Cancel right after the first file lands: the + # loop's between-files check must catch it before fetching the second. + state = {} + + def _after(filename): + if filename == "a.bin": + state["jobs"][state["job_id"]]["_cancel"] = True + + fake = _FakeHfHub(files, after_file=_after) + + manager = _SNS(models_dir=tmp_path) + jobs = {} + job_id = "job-cancel" + jobs[job_id] = {"model_id": hf_repo, "status": "downloading", + "error": None, "finished_at": None} + state["jobs"] = jobs + state["job_id"] = job_id + + # Force serial (1 worker) so the between-files cancel check is deterministic: + # with the default parallel pool, files b/c could already be in flight past + # the cancel gate when the a.bin hook fires. Cancellation is best-effort per + # in-flight file; this pins the pool to 1 to assert the stop-before-next path. + with patch.dict(_sys.modules, {"huggingface_hub": fake}), \ + patch.dict(os.environ, {"AINODE_DOWNLOAD_MAX_WORKERS": "1"}), \ + patch.object(_api_routes, "_get_repo_total_bytes", return_value=1536): + _asyncio.run(_api_routes._run_download_repo(manager, hf_repo, job_id, jobs)) + + assert jobs[job_id]["status"] == "cancelled" + assert jobs[job_id]["finished_at"] is not None + # Only the first file was fetched; the loop stopped before the rest. + assert fake.downloaded == ["a.bin"] + # Partial dir is removed on cancel. + assert not target.exists() + + def test_cancel_before_first_file(self, tmp_path): + # Cancel flag already set when the job starts → nothing downloads. + files = {"a.bin": b"x" * 256} + hf_repo = "org/model-z" + target = tmp_path / hf_repo.replace("/", "--") + fake = _FakeHfHub(files) + + manager = _SNS(models_dir=tmp_path) + job_id = "job-pre" + jobs = {job_id: {"model_id": hf_repo, "status": "downloading", + "error": None, "finished_at": None, "_cancel": True}} + with patch.dict(_sys.modules, {"huggingface_hub": fake}), \ + patch.object(_api_routes, "_get_repo_total_bytes", return_value=256): + _asyncio.run(_api_routes._run_download_repo(manager, hf_repo, job_id, jobs)) + + assert jobs[job_id]["status"] == "cancelled" + assert fake.downloaded == [] + assert not target.exists()