popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

5-th DL fix

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-06 22:55:01 +03:00
parent 7e7682be47
commit 4bfdc814e2
6 changed files with 126 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
packages/node/meshnet_node/cli.py
View File

@@ -65,6 +65,7 @@ def _run_node(cfg: dict) -> None:
vram_mb_override=cfg.get("vram_mb_override"),
max_loaded_shards=int(cfg.get("max_loaded_shards", 1)),
debug=bool(cfg.get("debug", False)),
tracker_source_disabled=bool(cfg.get("tracker_source_disabled", False)),
)
except Exception as exc:
print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
@@ -148,6 +149,8 @@ def _cmd_default(args) -> int:
overrides["max_loaded_shards"] = args.max_shards
if args.debug:
overrides["debug"] = True
if getattr(args, "tracker_source_disabled", False):
overrides["tracker_source_disabled"] = True
if overrides:
cfg = merge_cli_overrides(cfg, **overrides)
@@ -245,6 +248,7 @@ def _cmd_start(args) -> int:
vram_mb_override=getattr(args, "memory", None),
max_loaded_shards=getattr(args, "max_shards", 1),
debug=getattr(args, "debug", False),
tracker_source_disabled=getattr(args, "tracker_source_disabled", False),
)
except Exception as exc:
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
@@ -281,6 +285,8 @@ def main() -> None:
help="Quantization level")
parser.add_argument("--download-dir", metavar="PATH", help="Model download directory")
parser.add_argument("--tracker", metavar="URL", help="Tracker URL")
parser.add_argument("--tracker-source-disabled", action="store_true",
help="Skip tracker/peer model-file sources and download from HuggingFace directly")
parser.add_argument("--wallet", metavar="PATH", help="Wallet file path")
parser.add_argument("--shard-start", type=int, metavar="N", help="Pin shard start layer")
parser.add_argument("--shard-end", type=int, metavar="N", help="Pin shard end layer")
@@ -329,6 +335,8 @@ def main() -> None:
start_cmd.add_argument("--max-shards", type=int, default=1, metavar="N",
help="Maximum shard slots this node advertises to the tracker (default 1)")
start_cmd.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
start_cmd.add_argument("--tracker-source-disabled", action="store_true",
help="Skip tracker/peer model-file sources and download from HuggingFace directly")
args = parser.parse_args()

86
packages/node/meshnet_node/downloader.py
View File

@@ -13,7 +13,6 @@ import tarfile
import tempfile
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any
@@ -147,11 +146,13 @@ def _download_model_source(
if progress:
print(f" {label}: transfer complete ({received / 1e9:.2f} GB), extracting ...", flush=True)
_safe_extract_shard(archive_path, extract_dir)
if shard_dir.exists():
shutil.rmtree(shard_dir)
shutil.move(str(extract_dir), str(shard_dir))
return shard_dir
except Exception as exc:
if progress:
print(f" {label}: download failed: {exc!r}", flush=True)
print(f" {label}: download failed ({url}): {exc!r}", flush=True)
return None
@@ -177,55 +178,6 @@ def _download_huggingface_subset(
return Path(snapshot_download(**kwargs))
def _download_from_fastest_source(
*,
model_sources: list[dict],
hf_repo: str,
cache_dir: Path,
shard_dir: Path,
progress: bool,
timeout: float,
) -> tuple[str, Path] | None:
shard_dir.parent.mkdir(parents=True, exist_ok=True)
with tempfile.TemporaryDirectory(prefix="meshnet-race-", dir=shard_dir.parent) as tmp:
tmp_root = Path(tmp)
jobs: dict[Any, tuple[str, Path]] = {}
pool = ThreadPoolExecutor(max_workers=min(4, len(model_sources) + 1))
try:
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, 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)] = (
"HuggingFace",
hf_candidate,
)
for future in as_completed(jobs):
label, candidate = jobs[future]
try:
result = future.result()
except Exception as exc:
if progress:
print(f" {label}: download failed: {exc!r}", flush=True)
continue
if result is None:
continue
if shard_dir.exists():
shutil.rmtree(shard_dir)
shutil.move(str(candidate), str(shard_dir))
if progress:
print(f" download source: {label}", flush=True)
pool.shutdown(wait=False, cancel_futures=True)
return label, shard_dir
finally:
pool.shutdown(wait=False, cancel_futures=True)
return None
def _allow_patterns_from_sources(model_sources: list[dict]) -> list[str] | None:
patterns: set[str] = set()
for source in model_sources:
@@ -336,21 +288,31 @@ def download_shard(
f" Downloading layers {shard_start}-{shard_end} from {hf_repo} ...",
flush=True,
)
if model_sources:
# Tracker (or peer) model sources are preferred outright — usually LAN-fast.
# HuggingFace is only the fallback when every advertised source fails.
for source in model_sources or []:
label = str(source.get("type") or "model-source")
if progress:
print(" Racing tracker model source against HuggingFace ...", flush=True)
raced = _download_from_fastest_source(
model_sources=model_sources,
hf_repo=hf_repo,
cache_dir=cache_dir,
shard_dir=shard_dir,
progress=progress,
print(f" Downloading from {label} model source (HuggingFace is the fallback) ...", flush=True)
fetched = _download_model_source(
source,
shard_dir,
timeout=max(peer_timeout, _MODEL_SOURCE_TIMEOUT_SECONDS),
progress=progress,
label=label,
)
if raced is not None:
return raced[1]
if fetched is not None:
if progress:
print(f" download source: {label}", flush=True)
return fetched
if model_sources and progress:
print(" All model sources failed — falling back to HuggingFace ...", flush=True)
allow_patterns = _allow_patterns_from_remote_index(hf_repo, cache_dir, shard_start, shard_end)
allow_patterns = None
if model_sources:
allow_patterns = _allow_patterns_from_sources(model_sources)
if allow_patterns is None:
allow_patterns = _allow_patterns_from_remote_index(hf_repo, cache_dir, shard_start, shard_end)
if progress:
if allow_patterns:
print(" download source: HuggingFace (layer-filtered)", flush=True)

15
packages/node/meshnet_node/startup.py
View File

@@ -332,6 +332,7 @@ def run_startup(
vram_mb_override: int | None = None,
max_loaded_shards: int = 1,
debug: bool = False,
tracker_source_disabled: bool = False,
) -> StubNodeServer | TorchNodeServer:
"""Execute the full startup sequence and return a running node server.
@@ -458,7 +459,10 @@ def run_startup(
if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
shard_start = net_asgn["shard_start"]
shard_end = net_asgn["shard_end"]
full_sources = _full_model_sources(net_asgn.get("model_sources", []))
full_sources = (
[] if tracker_source_disabled
else _full_model_sources(net_asgn.get("model_sources", []))
)
if full_sources:
cache_dir = download_shard(
model_id.split("/")[-1],
@@ -588,7 +592,7 @@ def run_startup(
f"(of {assigned_num_layers})",
flush=True,
)
full_sources = _full_model_sources(assigned_model_sources)
full_sources = [] if tracker_source_disabled else _full_model_sources(assigned_model_sources)
if full_sources:
print("Downloading assigned model snapshot...", flush=True)
cache_dir = download_shard(
@@ -676,7 +680,10 @@ def run_startup(
assign_qs = urllib.parse.urlencode({
"model": model,
"device": device,
"vram_mb": vram_mb,
# CPU-mode nodes must be sized by RAM: a detected-but-unusable GPU's
# VRAM would otherwise cap the shard (e.g. 8 GB VRAM → 3 layers on a
# 79 GB box whose Torch has no CUDA).
"vram_mb": assignment_vram_mb,
"ram_mb": ram_mb,
})
try:
@@ -690,7 +697,7 @@ def run_startup(
assigned_model: str = assignment.get("model", model)
hf_repo: str | None = assignment.get("hf_repo")
peers: list[dict] = assignment.get("peers", [])
model_sources: list[dict] = assignment.get("model_sources", [])
model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", [])
print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
# 4. Download shard