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

feat: default quantization int8, GB display, shard heal cycle test

Browse Source 
- cli.py: change default --quantization from bfloat16 to int8; saves
  ~50% VRAM/RAM for new nodes that don't specify a quantization
- startup.py: display memory budget and GPU info in GB (e.g. 124.9 GB RAM)
  instead of MB; show remaining headroom after full model load
- test_tracker_routing.py: add test_shard_heal_cycle_surviving_node_covers_dead_peers_gap
  — end-to-end proof that:
    1. tracker purges expired node A and queues LOAD_SHARD for node B
    2. node B receives directive on next heartbeat
    3. TorchNodeServer.apply_tracker_directives hot-swaps the backend
    4. node B re-registers covering the full model; coverage gap closed
  Test runs in <1s with monkeypatched _load_backend (no GPU needed)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-06-30 23:08:34 +03:00
parent ff4115f611
commit b6272db93d
3 changed files with 196 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -23,7 +23,7 @@ def _run_node(cfg: dict) -> None:
model_id=cfg.get("model_hf_repo") or None,
shard_start=cfg.get("shard_start"),
shard_end=cfg.get("shard_end"),
quantization=cfg.get("quantization", "bfloat16").replace("bf16", "bfloat16"),
quantization=cfg.get("quantization", "int8").replace("bf16", "bfloat16"),
wallet_path=Path(cfg["wallet_path"]) if cfg.get("wallet_path") else None,
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
host=cfg.get("host", "0.0.0.0"),
@@ -278,7 +278,7 @@ def main() -> None:
start_cmd.add_argument("--model-id", help="HuggingFace repo ID")
start_cmd.add_argument("--shard-start", type=int)
start_cmd.add_argument("--shard-end", type=int)
start_cmd.add_argument("--quantization", choices=["bfloat16", "int8", "nf4", "bf16"], default="bfloat16")
start_cmd.add_argument("--quantization", choices=["bfloat16", "int8", "nf4", "bf16"], default="int8")
start_cmd.add_argument("--host", default="0.0.0.0")
start_cmd.add_argument("--advertise-host")
start_cmd.add_argument("--tracker-mode", action="store_true")

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

@@ -21,6 +21,40 @@ from .torch_server import TorchNodeServer
from .wallet import load_or_create_wallet
_DEFAULT_BYTES_PER_LAYER = 30 * 1024 * 1024
def _memory_budget(vram_mb: int, ram_mb: int) -> tuple[int, str]:
"""Return the capacity budget in MB and whether it came from VRAM or RAM."""
if vram_mb > 0:
return vram_mb, "VRAM"
return max(0, ram_mb), "RAM"
def _max_assignable_layers(memory_mb: int, total_layers: int | None) -> int:
if total_layers is None or total_layers <= 0 or memory_mb <= 0:
return 0
budget_bytes = memory_mb * 1024 * 1024
return min(total_layers, int((budget_bytes * 0.8) // _DEFAULT_BYTES_PER_LAYER))
def _shard_budget_line(memory_mb: int, memory_source: str, total_layers: int | None, quantization: str) -> str:
memory_gb = memory_mb / 1024
gb_str = f"{memory_gb:.1f} GB"
if total_layers is None or total_layers <= 0:
return f"Memory budget: {gb_str} {memory_source}; shard budget: unknown model layer count"
max_layers = _max_assignable_layers(memory_mb, total_layers)
# Remaining capacity after one full model load (rough estimate)
shard_bytes = max_layers * _DEFAULT_BYTES_PER_LAYER
remaining_gb = (memory_mb * 1024 * 1024 - shard_bytes) / (1024 ** 3)
remaining_str = f"; {remaining_gb:.1f} GB remaining after full load" if remaining_gb > 1 else ""
return (
f"Memory budget: {gb_str} {memory_source}; "
f"Shard budget: up to {max_layers}/{total_layers} layers at {quantization}"
f"{remaining_str}"
)
def _post_json(url: str, payload: dict, timeout: float = 10.0) -> dict:
data = json.dumps(payload).encode()
req = urllib.request.Request(
@@ -129,7 +163,11 @@ def _start_heartbeat(
uptime = time.monotonic() - _start_time
stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"}
if node_ref is not None:
stats["total_requests"] = getattr(node_ref, "total_requests", 0)
stats["total_requests"] = getattr(
node_ref,
"total_requests",
getattr(node_ref, "chat_completion_count", 0),
)
stats["failed_requests"] = getattr(node_ref, "failed_requests", 0)
stats["queue_depth"] = getattr(node_ref, "queue_depth", 0)
return stats
@@ -310,20 +348,24 @@ def run_startup(
device: str = hw["device"]
gpu_name: str | None = hw.get("gpu_name")
vram_mb: int = hw.get("vram_mb", 0)
ram_mb: int = hw.get("ram_mb", 16 * 1024)
if vram_mb_override is not None:
vram_mb = vram_mb_override
print(f" Memory budget overridden to {vram_mb} MB via --memory", flush=True)
print(f" Memory budget overridden to {vram_mb / 1024:.1f} GB via --memory", flush=True)
elif device == "cpu":
print(" WARNING: No CUDA GPU detected — running in CPU mode", flush=True)
print(f" WARNING: No CUDA GPU detected — running in CPU mode ({ram_mb / 1024:.1f} GB RAM)", flush=True)
else:
print(f" GPU: {gpu_name} ({vram_mb} MB VRAM)", flush=True)
print(f" GPU: {gpu_name} ({vram_mb / 1024:.1f} GB VRAM, {ram_mb / 1024:.1f} GB RAM)", flush=True)
memory_budget_mb, memory_budget_source = _memory_budget(vram_mb, ram_mb)
print(f" Memory budget: {memory_budget_mb} MB {memory_budget_source}", flush=True)
registration_capabilities = {
"vram_bytes": max(0, int(vram_mb)) * 1024 * 1024,
"ram_bytes": max(0, int(ram_mb)) * 1024 * 1024,
"max_loaded_shards": max_loaded_shards,
}
if vram_mb_override is not None or vram_mb > 0:
registration_capabilities["vram_bytes"] = max(0, int(vram_mb)) * 1024 * 1024
# 2. Wallet
print("Loading wallet...", flush=True)
wallet_kwargs: dict = {}
@@ -349,7 +391,7 @@ def run_startup(
if shard_start is None and shard_end is None:
try:
qs = urllib.parse.urlencode({
"device": device, "vram_mb": vram_mb, "hf_repo": model_id,
"device": device, "vram_mb": vram_mb, "ram_mb": ram_mb, "hf_repo": model_id,
})
net_asgn = _get_json(f"{tracker_url}/v1/network/assign?{qs}", timeout=5.0)
if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
@@ -432,6 +474,7 @@ def run_startup(
f" Wallet: {address}\n"
f" Model ID: {model_id}\n"
f" Shard: {shard_label}\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization)}\n"
f" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n"
f" Node ID: {tracker_node_id or 'unregistered'}\n"
@@ -445,7 +488,7 @@ def run_startup(
# 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": vram_mb})
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": vram_mb, "ram_mb": ram_mb})
net_assignment: dict = {}
try:
net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}")
@@ -523,6 +566,7 @@ def run_startup(
f" Model ID: {assigned_hf_repo}\n"
f" Shard: layers {assigned_shard_start}–{assigned_shard_end} "
f"({shard_count} of {assigned_num_layers})\n"
f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization)}\n"
f" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n"
f" Node ID: {tracker_node_id or 'unregistered'}\n"
@@ -538,6 +582,7 @@ def run_startup(
"model": model,
"device": device,
"vram_mb": vram_mb,
"ram_mb": ram_mb,
})
try:
assignment = _get_json(f"{tracker_url}/v1/nodes/assign?{assign_qs}")
@@ -616,12 +661,13 @@ def run_startup(
# Status summary
hw_str = device.upper()
if gpu_name:
hw_str += f" ({gpu_name}, {vram_mb} MB)"
hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)"
print(
f"\n{'=' * 32}\n"
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" Endpoint: {endpoint}\n"
f" Node ID: {node_id}\n"
f" Hardware: {hw_str}\n"