fix model DL doe 4-th time
This commit is contained in:
@@ -19,6 +19,11 @@ from typing import Any
|
||||
|
||||
_DEFAULT_CACHE = Path.home() / ".cache" / "meshnet" / "shards"
|
||||
_PEER_TIMEOUT_SECONDS = 2.0
|
||||
# Model-source tar streams are multi-GB; a short socket timeout must not kill
|
||||
# them on a transient read stall. Peer probes keep the short timeout because
|
||||
# they run sequentially before the race and may hit dead endpoints.
|
||||
_MODEL_SOURCE_TIMEOUT_SECONDS = 30.0
|
||||
_PROGRESS_INTERVAL_BYTES = 512 * 1024 * 1024
|
||||
|
||||
|
||||
def compute_shard_checksum(shard_dir: Path) -> str:
|
||||
@@ -111,6 +116,8 @@ def _download_model_source(
|
||||
source: dict,
|
||||
shard_dir: Path,
|
||||
timeout: float,
|
||||
progress: bool = False,
|
||||
label: str = "model-source",
|
||||
) -> Path | None:
|
||||
url = source.get("url")
|
||||
if not isinstance(url, str) or not url:
|
||||
@@ -125,16 +132,26 @@ def _download_model_source(
|
||||
extract_dir = tmp_root / "extract"
|
||||
extract_dir.mkdir()
|
||||
try:
|
||||
received = 0
|
||||
next_report = _PROGRESS_INTERVAL_BYTES
|
||||
with urllib.request.urlopen(url, timeout=timeout) as resp, archive_path.open("wb") as out:
|
||||
while True:
|
||||
chunk = resp.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
out.write(chunk)
|
||||
received += len(chunk)
|
||||
if progress and received >= next_report:
|
||||
print(f" {label}: {received / 1e9:.1f} GB received ...", flush=True)
|
||||
next_report += _PROGRESS_INTERVAL_BYTES
|
||||
if progress:
|
||||
print(f" {label}: transfer complete ({received / 1e9:.2f} GB), extracting ...", flush=True)
|
||||
_safe_extract_shard(archive_path, extract_dir)
|
||||
shutil.move(str(extract_dir), str(shard_dir))
|
||||
return shard_dir
|
||||
except Exception:
|
||||
except Exception as exc:
|
||||
if progress:
|
||||
print(f" {label}: download failed: {exc!r}", flush=True)
|
||||
return None
|
||||
|
||||
|
||||
@@ -178,7 +195,9 @@ def _download_from_fastest_source(
|
||||
for index, source in enumerate(model_sources):
|
||||
label = str(source.get("type") or "model-source")
|
||||
candidate = tmp_root / f"source-{index}"
|
||||
jobs[pool.submit(_download_model_source, source, candidate, timeout)] = (label, candidate)
|
||||
jobs[pool.submit(
|
||||
_download_model_source, source, candidate, timeout, progress, label,
|
||||
)] = (label, candidate)
|
||||
allow_patterns = _allow_patterns_from_sources(model_sources)
|
||||
hf_candidate = tmp_root / "huggingface"
|
||||
jobs[pool.submit(_download_huggingface_subset, hf_repo, cache_dir, hf_candidate, allow_patterns)] = (
|
||||
@@ -189,7 +208,9 @@ def _download_from_fastest_source(
|
||||
label, candidate = jobs[future]
|
||||
try:
|
||||
result = future.result()
|
||||
except Exception:
|
||||
except Exception as exc:
|
||||
if progress:
|
||||
print(f" {label}: download failed: {exc!r}", flush=True)
|
||||
continue
|
||||
if result is None:
|
||||
continue
|
||||
@@ -324,7 +345,7 @@ def download_shard(
|
||||
cache_dir=cache_dir,
|
||||
shard_dir=shard_dir,
|
||||
progress=progress,
|
||||
timeout=peer_timeout,
|
||||
timeout=max(peer_timeout, _MODEL_SOURCE_TIMEOUT_SECONDS),
|
||||
)
|
||||
if raced is not None:
|
||||
return raced[1]
|
||||
|
||||
@@ -561,13 +561,19 @@ def run_startup(
|
||||
raise ValueError("--shard-start / --shard-end require --model-id")
|
||||
|
||||
# 3a. Auto-join: query tracker for network-wide HF model assignment.
|
||||
print("Querying tracker for network assignment...", flush=True)
|
||||
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb})
|
||||
# Skipped when the user explicitly requested a model — the shard-assignment
|
||||
# query below (/v1/nodes/assign?model=…) is authoritative there, and a fresh
|
||||
# tracker would otherwise print a scary 503 for the model-less auto-join.
|
||||
net_assignment: dict = {}
|
||||
try:
|
||||
net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}")
|
||||
except Exception as exc:
|
||||
print(f" (auto-join unavailable: {exc})", flush=True)
|
||||
if model and model != "stub-model":
|
||||
print(f"Model {model!r} requested explicitly — skipping network auto-join.", flush=True)
|
||||
else:
|
||||
print("Querying tracker for network assignment...", flush=True)
|
||||
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb})
|
||||
try:
|
||||
net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}")
|
||||
except Exception as exc:
|
||||
print(f" (auto-join unavailable: {exc})", flush=True)
|
||||
assigned_hf_repo: str | None = net_assignment.get("hf_repo")
|
||||
_gap_found: bool = bool(net_assignment.get("gap_found", False))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user