N-th fix od model DL

This commit is contained in:
Dobromir Popov
2026-07-06 23:41:06 +03:00
parent 4bfdc814e2
commit d83224a62f
5 changed files with 224 additions and 18 deletions

View File

@@ -82,3 +82,34 @@ sites), falling back to `ram_mb`.
`test_mining_cli` case, present on clean master). `test_mining_cli` case, present on clean master).
- [ ] Live two-machine retest: Windows node downloads only from tracker at - [ ] Live two-machine retest: Windows node downloads only from tracker at
LAN speed and is assigned a RAM-sized shard. LAN speed and is assigned a RAM-sized shard.
## Round 3 (2026-07-06, after live retest showed mid-stream RST)
Live retest: RAM sizing worked (layers 036) and the failure finally printed —
`ConnectionResetError(10054)` ~70 s into the tar stream. Local reproduction
cleared the tracker: it streams the full 72 GB tar at ~900 MB/s, survives a
3-minute slow reader, and logs aborts in one line. The RST comes from the
network path (Windows laptop, likely WiFi + firewall/AV) — and a 72 GB
single-TCP-stream tar is inherently fragile there.
Fix: per-file downloads (design principle: nodes must be able to fetch any
missing shard or the complete model from the tracker alone — no hard HF
dependency):
- Tracker: `/v1/model-files/download?...&file=<rel>` streams one file with
`Content-Length` (rel must be in the requested shard/full set; traversal
rejected). `model_sources` now advertises `full_files` and a `file_sizes`
manifest.
- Node: `_download_source_files` fetches per file into
`<shard>.partial/`, retries each file 3×, verifies against
`Content-Length`, and reuses already-complete files (hardlink from the
existing shard) via the size manifest — so restarts and drops cost at most
one file. Tar stream remains the fallback for old trackers
(detected via Content-Type) and sources without a file list.
- `_full_model_sources` passes `full_files` through, so full-snapshot
downloads for the torch path get the same robustness.
Verified live against a local tracker: 14.7 GB shard in 7.6 s per-file;
re-run over a complete shard instant; corrupt + deleted file recovered in
1.5 s re-fetching only those two. 114 tests pass (node_startup +
tracker_routing).

View File

@@ -8,9 +8,11 @@ test suite never touches the network.
import hashlib import hashlib
import json import json
import os
import shutil import shutil
import tarfile import tarfile
import tempfile import tempfile
import time
import urllib.parse import urllib.parse
import urllib.request import urllib.request
from pathlib import Path from pathlib import Path
@@ -23,6 +25,7 @@ _PEER_TIMEOUT_SECONDS = 2.0
# they run sequentially before the race and may hit dead endpoints. # they run sequentially before the race and may hit dead endpoints.
_MODEL_SOURCE_TIMEOUT_SECONDS = 30.0 _MODEL_SOURCE_TIMEOUT_SECONDS = 30.0
_PROGRESS_INTERVAL_BYTES = 512 * 1024 * 1024 _PROGRESS_INTERVAL_BYTES = 512 * 1024 * 1024
_FILE_RETRY_ATTEMPTS = 3
def compute_shard_checksum(shard_dir: Path) -> str: def compute_shard_checksum(shard_dir: Path) -> str:
@@ -111,6 +114,121 @@ def _download_shard_from_peer(
return False return False
class _TarOnlySource(Exception):
"""The server ignored ?file= and streamed a tar — no single-file support."""
def _fetch_source_file(
file_url: str,
dest: Path,
timeout: float,
) -> int:
"""Download one file; skip if *dest* already matches the remote size.
Returns the file's final size. Raises on any failure, including a
short read against the server's Content-Length.
"""
with urllib.request.urlopen(file_url, timeout=timeout) as resp:
ctype = resp.getheader("Content-Type") if hasattr(resp, "getheader") else None
if isinstance(ctype, str) and "x-tar" in ctype:
raise _TarOnlySource(file_url)
length = resp.getheader("Content-Length") if hasattr(resp, "getheader") else None
expected = int(length) if isinstance(length, str) and length.isdigit() else None
if expected is not None and dest.exists() and dest.stat().st_size == expected:
return expected # complete from an earlier attempt/run
received = 0
with dest.open("wb") as out:
while True:
chunk = resp.read(1024 * 1024)
if not chunk:
break
out.write(chunk)
received += len(chunk)
if expected is not None and received != expected:
raise OSError(f"short read: got {received} of {expected} bytes")
return received
def _reuse_local_file(expected: int, dest: Path, final: Path) -> bool:
"""Reuse an already-complete copy of a file instead of re-downloading."""
if dest.exists() and dest.stat().st_size == expected:
return True
if final.exists() and final.stat().st_size == expected:
try:
os.link(final, dest)
except OSError:
shutil.copy2(final, dest)
return True
return False
def _download_source_files(
source: dict,
shard_dir: Path,
timeout: float,
progress: bool,
label: str,
) -> Path | None:
"""Per-file download from a model source — retries and resumes per file.
Far more robust than one multi-GB tar stream on flaky links: a dropped
connection costs at most one file, and completed files are skipped when
the download is retried or the node is restarted.
"""
url = source.get("url")
rel_files = [
rel for rel in (source.get("files") or [])
if isinstance(rel, str) and rel and not rel.startswith("/") and ".." not in Path(rel).parts
]
if not isinstance(url, str) or not url or not rel_files:
return None
sizes = source.get("file_sizes")
if not isinstance(sizes, dict):
sizes = {}
partial_dir = shard_dir.parent / f"{shard_dir.name}.partial"
partial_dir.mkdir(parents=True, exist_ok=True)
sep = "&" if "?" in url else "?"
for index, rel in enumerate(rel_files, start=1):
dest = partial_dir / rel
dest.parent.mkdir(parents=True, exist_ok=True)
known_size = sizes.get(rel)
if isinstance(known_size, int) and _reuse_local_file(known_size, dest, shard_dir / rel):
if progress:
print(
f" {label}: [{index}/{len(rel_files)}] {rel} "
f"(already complete, {known_size / 1e9:.2f} GB)",
flush=True,
)
continue
file_url = f"{url}{sep}{urllib.parse.urlencode({'file': rel})}"
for attempt in range(1, _FILE_RETRY_ATTEMPTS + 1):
try:
size = _fetch_source_file(file_url, dest, timeout)
if progress:
print(
f" {label}: [{index}/{len(rel_files)}] {rel} ({size / 1e9:.2f} GB)",
flush=True,
)
break
except _TarOnlySource:
if progress:
print(f" {label}: no single-file support — using tar stream", flush=True)
return None
except Exception as exc:
if progress:
print(
f" {label}: {rel} attempt {attempt}/{_FILE_RETRY_ATTEMPTS} failed: {exc!r}",
flush=True,
)
if attempt == _FILE_RETRY_ATTEMPTS:
return None
time.sleep(1.0 * attempt)
if shard_dir.exists():
shutil.rmtree(shard_dir)
shutil.move(str(partial_dir), str(shard_dir))
return shard_dir
def _download_model_source( def _download_model_source(
source: dict, source: dict,
shard_dir: Path, shard_dir: Path,
@@ -118,6 +236,13 @@ def _download_model_source(
progress: bool = False, progress: bool = False,
label: str = "model-source", label: str = "model-source",
) -> Path | None: ) -> Path | None:
# Prefer per-file transfers whenever the source advertises its file list;
# fall through to the tar stream if the server lacks single-file support
# or per-file transfers keep failing.
if source.get("files"):
fetched = _download_source_files(source, shard_dir, timeout, progress, label)
if fetched is not None:
return fetched
url = source.get("url") url = source.get("url")
if not isinstance(url, str) or not url: if not isinstance(url, str) or not url:
endpoint = source.get("endpoint") endpoint = source.get("endpoint")

View File

@@ -43,7 +43,9 @@ def _full_model_sources(model_sources: list[dict]) -> list[dict]:
full_sources.append({ full_sources.append({
**source, **source,
"url": full_url, "url": full_url,
"files": [], # full_files (when advertised) enables robust per-file download
# of the whole snapshot; empty list falls back to the tar stream.
"files": source.get("full_files") or [],
"type": f"{source.get('type') or 'model-source'}-full", "type": f"{source.get('type') or 'model-source'}-full",
}) })
return full_sources return full_sources

View File

@@ -3836,12 +3836,21 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
"shard_end": shard_end, "shard_end": shard_end,
}) })
full_query = urllib.parse.urlencode({"model": model, "full": 1}) full_query = urllib.parse.urlencode({"model": model, "full": 1})
full_files = _snapshot_regular_files(snapshot_dir)
file_sizes: dict[str, int] = {}
for rel in set(files) | set(full_files):
try:
file_sizes[rel] = (snapshot_dir / rel).resolve().stat().st_size
except OSError:
continue
return [{ return [{
"type": "tracker", "type": "tracker",
"endpoint": base_url, "endpoint": base_url,
"url": f"{base_url}/v1/model-files/download?{query}", "url": f"{base_url}/v1/model-files/download?{query}",
"full_url": f"{base_url}/v1/model-files/download?{full_query}", "full_url": f"{base_url}/v1/model-files/download?{full_query}",
"files": files, "files": files,
"full_files": full_files,
"file_sizes": file_sizes,
}] }]
def _handle_model_files_download(self, parsed: urllib.parse.ParseResult) -> None: def _handle_model_files_download(self, parsed: urllib.parse.ParseResult) -> None:
@@ -3877,6 +3886,42 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
if not rel_files: if not rel_files:
self._send_json(404, {"error": "no local files matched the assigned shard"}) self._send_json(404, {"error": "no local files matched the assigned shard"})
return return
# Single-file mode: ?file=<rel> streams one file with Content-Length so
# clients can verify completeness and retry per file instead of
# restarting one giant tar stream after a network hiccup.
single = params.get("file", [""])[0]
if single:
if single not in set(rel_files):
self._send_json(404, {"error": f"file not in requested shard: {single!r}"})
return
file_path = snapshot_dir / single
try:
# resolve() dereferences HF snapshot symlinks into blobs/.
size = file_path.resolve().stat().st_size
except OSError:
self._send_json(404, {"error": f"file unreadable: {single!r}"})
return
self.send_response(200)
self.send_header("Content-Type", "application/octet-stream")
self.send_header("Content-Length", str(size))
self.send_header("X-Meshnet-Model-Source", "tracker")
self.end_headers()
try:
with file_path.open("rb") as f:
while True:
chunk = f.read(1024 * 1024)
if not chunk:
break
self.wfile.write(chunk)
except (BrokenPipeError, ConnectionResetError):
print(
f"model-file download aborted by client "
f"({self.client_address[0]}, model={resolved_name}, file={single})",
flush=True,
)
return
self.send_response(200) self.send_response(200)
self.send_header("Content-Type", "application/x-tar") self.send_header("Content-Type", "application/x-tar")
self.send_header("X-Meshnet-Model-Source", "tracker") self.send_header("X-Meshnet-Model-Source", "tracker")

View File

@@ -412,21 +412,15 @@ def test_download_shard_prefers_tracker_model_source_over_huggingface(
monkeypatch, monkeypatch,
): ):
"""A working tracker model source is used exclusively — HF is never contacted.""" """A working tracker model source is used exclusively — HF is never contacted."""
source_dir = tmp_path / "source" contents = {
source_dir.mkdir() "config.json": b"{}",
(source_dir / "config.json").write_text("{}") "model-00002-of-00004.safetensors": b"tracker",
(source_dir / "model-00002-of-00004.safetensors").write_text("tracker") }
archive = io.BytesIO()
with tarfile.open(fileobj=archive, mode="w") as tf:
tf.add(source_dir / "config.json", arcname="config.json")
tf.add(
source_dir / "model-00002-of-00004.safetensors",
arcname="model-00002-of-00004.safetensors",
)
class FakeTrackerResponse: class FakeFileResponse:
def __init__(self, payload: bytes): def __init__(self, payload: bytes):
self._payload = io.BytesIO(payload) self._payload = io.BytesIO(payload)
self._length = len(payload)
def __enter__(self): def __enter__(self):
return self return self
@@ -434,14 +428,23 @@ def test_download_shard_prefers_tracker_model_source_over_huggingface(
def __exit__(self, exc_type, exc, tb): def __exit__(self, exc_type, exc, tb):
return False return False
def getheader(self, name: str):
if name == "Content-Length":
return str(self._length)
if name == "Content-Type":
return "application/octet-stream"
return None
def read(self, size: int = -1) -> bytes: def read(self, size: int = -1) -> bytes:
return self._payload.read(size) return self._payload.read(size)
monkeypatch.setattr( def fake_urlopen(url, *args, **kwargs):
urllib.request, query = urllib.parse.parse_qs(urllib.parse.urlparse(url).query)
"urlopen", rel = query.get("file", [None])[0]
lambda *args, **kwargs: FakeTrackerResponse(archive.getvalue()), assert rel in contents, f"unexpected per-file request: {url}"
) return FakeFileResponse(contents[rel])
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
hf_calls = [] hf_calls = []
def fake_snapshot_download(**kwargs): def fake_snapshot_download(**kwargs):