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

dual billing; tracker to node model sharing

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-06 17:31:11 +03:00
parent ccb69c41e3
commit 2e696be80f
14 changed files with 1092 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -34,6 +34,21 @@ def _memory_budget(device: str, vram_mb: int, ram_mb: int, shared_vram_mb: int =
return max(0, ram_mb), "RAM"
def _full_model_sources(model_sources: list[dict]) -> list[dict]:
"""Use tracker full-snapshot URLs for real HF model loading."""
full_sources: list[dict] = []
for source in model_sources:
full_url = source.get("full_url")
if isinstance(full_url, str) and full_url:
full_sources.append({
**source,
"url": full_url,
"files": [],
"type": f"{source.get('type') or 'model-source'}-full",
})
return full_sources
def _hardware_label(device: str, gpu_name: str | None = None) -> str:
if device == "cuda":
return "CUDA"
@@ -443,6 +458,16 @@ 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", []))
if full_sources:
cache_dir = download_shard(
model_id.split("/")[-1],
shard_start,
shard_end,
cache_dir=cache_dir or Path.home() / ".cache" / "meshnet" / "shards",
hf_repo=model_id,
model_sources=full_sources,
)
print(
f" Tracker found uncovered shard: "
f"layers {shard_start}–{shard_end} (of {detected})",
@@ -550,12 +575,24 @@ def run_startup(
assigned_shard_start: int = net_assignment["shard_start"]
assigned_shard_end: int = net_assignment["shard_end"]
assigned_num_layers: int = net_assignment["num_layers"]
assigned_model_sources: list[dict] = net_assignment.get("model_sources", [])
print(
f" Assigned: {assigned_hf_repo} "
f"layers {assigned_shard_start}–{assigned_shard_end} "
f"(of {assigned_num_layers})",
flush=True,
)
full_sources = _full_model_sources(assigned_model_sources)
if full_sources:
print("Downloading assigned model snapshot...", flush=True)
cache_dir = download_shard(
assigned_hf_repo.split("/")[-1],
assigned_shard_start,
assigned_shard_end,
cache_dir=cache_dir or Path.home() / ".cache" / "meshnet" / "shards",
hf_repo=assigned_hf_repo,
model_sources=full_sources,
)
print("Loading real PyTorch model shard...", flush=True)
node = TorchNodeServer(
host=host,
@@ -647,6 +684,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", [])
print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
# 4. Download shard
@@ -658,6 +696,8 @@ def run_startup(
dl_kwargs["hf_repo"] = hf_repo
if peers:
dl_kwargs["peers"] = peers
if model_sources:
dl_kwargs["model_sources"] = model_sources
shard_path = download_shard(assigned_model, shard_start, shard_end, **dl_kwargs)
shard_checksum = compute_shard_checksum(shard_path)
print(f" Cached at: {shard_path}", flush=True)