N-th fix od model DL
This commit is contained in:
@@ -8,9 +8,11 @@ test suite never touches the network.
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tarfile
|
||||
import tempfile
|
||||
import time
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
@@ -23,6 +25,7 @@ _PEER_TIMEOUT_SECONDS = 2.0
|
||||
# they run sequentially before the race and may hit dead endpoints.
|
||||
_MODEL_SOURCE_TIMEOUT_SECONDS = 30.0
|
||||
_PROGRESS_INTERVAL_BYTES = 512 * 1024 * 1024
|
||||
_FILE_RETRY_ATTEMPTS = 3
|
||||
|
||||
|
||||
def compute_shard_checksum(shard_dir: Path) -> str:
|
||||
@@ -111,6 +114,121 @@ def _download_shard_from_peer(
|
||||
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(
|
||||
source: dict,
|
||||
shard_dir: Path,
|
||||
@@ -118,6 +236,13 @@ def _download_model_source(
|
||||
progress: bool = False,
|
||||
label: str = "model-source",
|
||||
) -> 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")
|
||||
if not isinstance(url, str) or not url:
|
||||
endpoint = source.get("endpoint")
|
||||
|
||||
@@ -43,7 +43,9 @@ def _full_model_sources(model_sources: list[dict]) -> list[dict]:
|
||||
full_sources.append({
|
||||
**source,
|
||||
"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",
|
||||
})
|
||||
return full_sources
|
||||
|
||||
Reference in New Issue
Block a user