N-th fix od model DL
This commit is contained in:
@@ -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 0–36) 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).
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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):
|
||||||
|
|||||||
Reference in New Issue
Block a user