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

Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-07 14:56:08 +02:00
parent 80bd83ae0a ca49675f50
commit 5cdce1a5b0
19 changed files with 1831 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -66,6 +66,8 @@ def _run_node(cfg: dict) -> None:
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)),
torch_threads=cfg.get("torch_threads"),
torch_interop_threads=cfg.get("torch_interop_threads"),
)
except Exception as exc:
print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
@@ -151,6 +153,10 @@ def _cmd_default(args) -> int:
overrides["debug"] = True
if getattr(args, "tracker_source_disabled", False):
overrides["tracker_source_disabled"] = True
if getattr(args, "torch_threads", None) is not None:
overrides["torch_threads"] = args.torch_threads
if getattr(args, "torch_interop_threads", None) is not None:
overrides["torch_interop_threads"] = args.torch_interop_threads
if overrides:
cfg = merge_cli_overrides(cfg, **overrides)
@@ -249,6 +255,8 @@ def _cmd_start(args) -> int:
max_loaded_shards=getattr(args, "max_shards", 1),
debug=getattr(args, "debug", False),
tracker_source_disabled=getattr(args, "tracker_source_disabled", False),
torch_threads=getattr(args, "torch_threads", None),
torch_interop_threads=getattr(args, "torch_interop_threads", None),
)
except Exception as exc:
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
@@ -299,6 +307,10 @@ def main() -> None:
help="Override autodetected VRAM/RAM budget in MB used for shard assignment")
parser.add_argument("--max-shards", type=int, metavar="N", default=None,
help="Maximum shard slots this node advertises to the tracker (default 1)")
parser.add_argument("--torch-threads", type=int, metavar="N",
help="Set PyTorch intra-op CPU worker threads")
parser.add_argument("--torch-interop-threads", type=int, metavar="N",
help="Set PyTorch inter-op CPU worker threads")
parser.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
parser.add_argument("--no-tui", action="store_true", help="Plain-text output (no rich dashboard)")
parser.add_argument("--compact", action="store_true", help="Single-line status output")
@@ -334,6 +346,10 @@ def main() -> None:
help="Override autodetected VRAM/RAM budget in MB used for shard assignment")
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("--torch-threads", type=int, metavar="N",
help="Set PyTorch intra-op CPU worker threads")
start_cmd.add_argument("--torch-interop-threads", type=int, metavar="N",
help="Set PyTorch inter-op CPU worker threads")
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")

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

@@ -3,6 +3,7 @@
from __future__ import annotations
import json
import os
import socket
import sys
import threading
@@ -139,6 +140,59 @@ def _hardware_label(device: str, gpu_name: str | None = None) -> str:
return "CPU"
def _positive_int(value: int | str | None, name: str) -> int | None:
if value is None or value == "":
return None
try:
parsed = int(value)
except (TypeError, ValueError) as exc:
raise ValueError(f"{name} must be a positive integer") from exc
if parsed < 1:
raise ValueError(f"{name} must be a positive integer")
return parsed
def _configure_torch_threads(
torch_threads: int | None = None,
torch_interop_threads: int | None = None,
) -> dict[str, int]:
"""Apply PyTorch CPU thread settings before model load/benchmark."""
intra_threads = _positive_int(
torch_threads if torch_threads is not None else os.environ.get("MESHNET_TORCH_THREADS"),
"--torch-threads",
)
interop_threads = _positive_int(
torch_interop_threads
if torch_interop_threads is not None
else os.environ.get("MESHNET_TORCH_INTEROP_THREADS"),
"--torch-interop-threads",
)
if intra_threads is not None:
os.environ.setdefault("OMP_NUM_THREADS", str(intra_threads))
os.environ.setdefault("MKL_NUM_THREADS", str(intra_threads))
try:
import torch
except ModuleNotFoundError:
return {}
if intra_threads is not None:
torch.set_num_threads(intra_threads)
if interop_threads is not None:
torch.set_num_interop_threads(interop_threads)
active: dict[str, int] = {}
try:
active["torch_threads"] = int(torch.get_num_threads())
except Exception:
pass
try:
active["torch_interop_threads"] = int(torch.get_num_interop_threads())
except Exception:
pass
return active
def _max_assignable_layers(
memory_mb: int,
total_layers: int | None,
@@ -443,6 +497,8 @@ def run_startup(
max_loaded_shards: int = 1,
debug: bool = False,
tracker_source_disabled: bool = False,
torch_threads: int | None = None,
torch_interop_threads: int | None = None,
) -> StubNodeServer | TorchNodeServer:
"""Execute the full startup sequence and return a running node server.
@@ -483,6 +539,12 @@ def run_startup(
print("Detecting hardware...", flush=True)
hw = detect_hardware()
torch_thread_config = _configure_torch_threads(torch_threads, torch_interop_threads)
if torch_thread_config:
hw.update(torch_thread_config)
intra = torch_thread_config.get("torch_threads", "?")
interop = torch_thread_config.get("torch_interop_threads", "?")
print(f" PyTorch threads: intra-op={intra}, inter-op={interop}", flush=True)
device: str = hw["device"]
gpu_name: str | None = hw.get("gpu_name")
vram_mb: int = hw.get("vram_mb", 0)
@@ -835,6 +897,7 @@ def run_startup(
hf_repo: str | None = assignment.get("hf_repo")
peers: list[dict] = assignment.get("peers", [])
model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", [])
assignment_bytes_per_layer = _assignment_bytes_per_layer(assignment, quantization)
print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
# 4. Download shard
@@ -921,7 +984,7 @@ def run_startup(
f"meshnet-node ready\n"
f" Wallet: {address}\n"
f" Shard: layers {shard_start}-{shard_end} ({assigned_model})\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assignment.get('model_layers_end', shard_end) + 1, quantization)}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assignment.get('model_layers_end', shard_end) + 1, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n"
f" Endpoint: {endpoint}\n"
f" Node ID: {node_id}\n"
f" Hardware: {hw_str}\n"