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

balancing improvements

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-07 18:30:30 +03:00
parent f220fd2210
commit 1299a6bb1c
6 changed files with 358 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -52,7 +52,7 @@ def _run_node(cfg: dict) -> None:
node = run_startup(
tracker_url=cfg["tracker_url"],
port=cfg.get("port", 7000),
model=cfg.get("model_name") or "stub-model",
model=cfg.get("model_name") or None,
model_id=cfg.get("model_hf_repo") or None,
shard_start=cfg.get("shard_start"),
shard_end=cfg.get("shard_end"),
@@ -215,8 +215,8 @@ def _cmd_start(args) -> int:
if args.tracker:
cfg["tracker_url"] = args.tracker
cfg["port"] = args.port if args.port is not None else _first_available_port(args.host)
model = args.model or cfg.get("model_hf_repo") or cfg.get("model_name") or "stub-model"
if args.model_id is None and "/" in model:
model = args.model or cfg.get("model_hf_repo") or cfg.get("model_name") or None
if args.model_id is None and model and "/" in model:
cfg["model_hf_repo"] = model
cfg["model_name"] = model.split("/")[-1]
else:

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

@@ -482,7 +482,7 @@ def _warn_virtual_network_ip(ip: str | None) -> None:
def run_startup(
tracker_url: str,
port: int = 0,
model: str = "stub-model",
model: str | None = None,
model_id: str | None = None,
shard_start: int | None = None,
shard_end: int | None = None,
@@ -755,8 +755,9 @@ def run_startup(
# 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 = {}
if model and model != "stub-model":
print(f"Model {model!r} requested explicitly — skipping network auto-join.", flush=True)
if model_id or (model and model != "stub-model"):
if 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})
@@ -767,17 +768,25 @@ def run_startup(
assigned_hf_repo: str | None = net_assignment.get("hf_repo")
_gap_found: bool = bool(net_assignment.get("gap_found", False))
if assigned_hf_repo and _gap_found:
if assigned_hf_repo:
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,
)
if _gap_found:
print(
f" Assigned gap: {assigned_hf_repo} "
f"layers {assigned_shard_start}–{assigned_shard_end} "
f"(of {assigned_num_layers})",
flush=True,
)
else:
print(
f" Assigned redundant copy: {assigned_hf_repo} "
f"layers {assigned_shard_start}–{assigned_shard_end} "
f"(of {assigned_num_layers})",
flush=True,
)
full_sources = [] if tracker_source_disabled else _full_model_sources(assigned_model_sources)
if full_sources:
print("Downloading assigned model snapshot...", flush=True)
@@ -874,10 +883,16 @@ def run_startup(
)
return node
# 3b. Shard assignment from tracker (stub-model / preset-based path)
if model != "stub-model":
raise RuntimeError(
"Tracker did not assign a model. Join a network that already serves one, "
"or start with --model <HF_REPO>."
)
# 3b. Stub preset path (tests / explicit stub-model only)
print("Querying tracker for shard assignment...", flush=True)
assign_qs = urllib.parse.urlencode({
"model": model,
"model": model or "stub-model",
"device": device,
# 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