2 Commits

Author SHA1 Message Date
Dobromir Popov
278be49539 node stats and benchmark, dynamic realocation working 2026-07-01 10:02:17 +03:00
Dobromir Popov
b6272db93d feat: default quantization int8, GB display, shard heal cycle test
- 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>
2026-06-30 23:08:34 +03:00
8 changed files with 476 additions and 27 deletions

View File

@@ -23,7 +23,7 @@ def _run_node(cfg: dict) -> None:
model_id=cfg.get("model_hf_repo") or None, model_id=cfg.get("model_hf_repo") or None,
shard_start=cfg.get("shard_start"), shard_start=cfg.get("shard_start"),
shard_end=cfg.get("shard_end"), 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, 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, cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
host=cfg.get("host", "0.0.0.0"), 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("--model-id", help="HuggingFace repo ID")
start_cmd.add_argument("--shard-start", type=int) start_cmd.add_argument("--shard-start", type=int)
start_cmd.add_argument("--shard-end", 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("--host", default="0.0.0.0")
start_cmd.add_argument("--advertise-host") start_cmd.add_argument("--advertise-host")
start_cmd.add_argument("--tracker-mode", action="store_true") start_cmd.add_argument("--tracker-mode", action="store_true")

View File

@@ -94,6 +94,19 @@ def _make_bar(pct: float, width: int = 10) -> str:
return "" * filled + "" * (width - filled) return "" * filled + "" * (width - filled)
def _node_stats(node) -> dict:
total = int(getattr(node, "total_requests", getattr(node, "chat_completion_count", 0)) or 0)
failed = int(getattr(node, "failed_requests", 0) or 0)
queue_depth = int(getattr(node, "queue_depth", 0) or 0)
success_rate = ((total - failed) / total * 100.0) if total else 100.0
return {
"total_requests": total,
"failed_requests": failed,
"queue_depth": queue_depth,
"success_rate": success_rate,
}
def run_dashboard(node, config: dict, start_time: float) -> None: def run_dashboard(node, config: dict, start_time: float) -> None:
"""Start the live dashboard. Blocks until Ctrl-C. Returns cleanly.""" """Start the live dashboard. Blocks until Ctrl-C. Returns cleanly."""
if not is_interactive_tty(): if not is_interactive_tty():
@@ -117,7 +130,8 @@ def _build_rich_renderable(
from rich.text import Text # type: ignore[import] from rich.text import Text # type: ignore[import]
uptime = time.monotonic() - start_time uptime = time.monotonic() - start_time
req_count = getattr(node, "chat_completion_count", 0) stats = _node_stats(node)
req_count = stats["total_requests"]
# Tokens/sec EMA (approximate: 20 tokens per request heuristic when no real counter) # Tokens/sec EMA (approximate: 20 tokens per request heuristic when no real counter)
delta_req = req_count - prev_req[0] delta_req = req_count - prev_req[0]
@@ -163,6 +177,7 @@ def _build_rich_renderable(
stats_lines = [ stats_lines = [
f"Tokens/sec {tps_bar} {tps:.1f} t/s (EMA)", f"Tokens/sec {tps_bar} {tps:.1f} t/s (EMA)",
f"Requests {req_count:,} served", f"Requests {req_count:,} served",
f"Success {stats['success_rate']:.1f}% failed {stats['failed_requests']:,} queue {stats['queue_depth']}",
f"Peers 0 connected (gossip: US-017)", f"Peers 0 connected (gossip: US-017)",
f"TAI earned 0.00 TAI (payments: US-006)", f"TAI earned 0.00 TAI (payments: US-006)",
f"Uptime {_format_uptime(uptime)}", f"Uptime {_format_uptime(uptime)}",
@@ -205,14 +220,17 @@ def _run_plain_loop(node, config: dict, start_time: float) -> None:
try: try:
while True: while True:
uptime = time.monotonic() - start_time uptime = time.monotonic() - start_time
req = getattr(node, "chat_completion_count", 0) stats = _node_stats(node)
req = stats["total_requests"]
gpu_stats = _gpu_stats() gpu_stats = _gpu_stats()
vram_str = "" vram_str = ""
if gpu_stats: if gpu_stats:
g = gpu_stats[0] g = gpu_stats[0]
vram_str = f" VRAM{g['used_gb']:.1f}GB" vram_str = f" VRAM{g['used_gb']:.1f}GB"
print( print(
f"[{model_name} req{req}{vram_str} up{_format_uptime(uptime)}]", f"[{model_name} req{req} ok{stats['success_rate']:.1f}% "
f"fail{stats['failed_requests']} q{stats['queue_depth']}"
f"{vram_str} up{_format_uptime(uptime)}]",
flush=True, flush=True,
) )
time.sleep(2) time.sleep(2)

View File

@@ -1,10 +1,23 @@
"""GPU hardware detection with graceful CPU fallback.""" """GPU hardware detection with graceful CPU fallback."""
import os
import subprocess import subprocess
import time
def _detect_ram_mb() -> int:
"""Return host physical RAM in MB, or 0 if unavailable."""
try:
pages = os.sysconf("SC_PHYS_PAGES")
page_size = os.sysconf("SC_PAGE_SIZE")
return int((pages * page_size) // (1024 * 1024))
except (AttributeError, OSError, ValueError):
return 0
def detect_hardware() -> dict: def detect_hardware() -> dict:
"""Detect GPU model and available VRAM. Returns hardware profile dict.""" """Detect GPU model and available VRAM. Returns hardware profile dict."""
ram_mb = _detect_ram_mb()
try: try:
import torch # type: ignore[import] import torch # type: ignore[import]
if torch.cuda.is_available(): if torch.cuda.is_available():
@@ -12,7 +25,7 @@ def detect_hardware() -> dict:
name = torch.cuda.get_device_name(idx) name = torch.cuda.get_device_name(idx)
props = torch.cuda.get_device_properties(idx) props = torch.cuda.get_device_properties(idx)
vram_mb = props.total_memory // (1024 * 1024) vram_mb = props.total_memory // (1024 * 1024)
return {"device": "cuda", "gpu_name": name, "vram_mb": vram_mb} return {"device": "cuda", "gpu_name": name, "vram_mb": vram_mb, "ram_mb": ram_mb}
except ImportError: except ImportError:
pass pass
@@ -26,8 +39,54 @@ def detect_hardware() -> dict:
parts = line.split(",", 1) parts = line.split(",", 1)
gpu_name = parts[0].strip() gpu_name = parts[0].strip()
vram_mb = int(parts[1].strip()) if len(parts) > 1 else 0 vram_mb = int(parts[1].strip()) if len(parts) > 1 else 0
return {"device": "cuda", "gpu_name": gpu_name, "vram_mb": vram_mb} return {"device": "cuda", "gpu_name": gpu_name, "vram_mb": vram_mb, "ram_mb": ram_mb}
except (FileNotFoundError, subprocess.TimeoutExpired, ValueError, IndexError): except (FileNotFoundError, subprocess.TimeoutExpired, ValueError, IndexError):
pass pass
return {"device": "cpu", "gpu_name": None, "vram_mb": 0} return {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": ram_mb}
def benchmark_throughput(device_str: str = "cpu") -> float:
"""
Estimate compute throughput via a synthetic transformer GEMM benchmark.
Runs hidden_size × (hidden_size*4) matmul — the dominant op in FFN layers —
and returns iterations/second as a relative speed index. Higher = faster.
The value is used as benchmark_tokens_per_sec in tracker registration for
routing tiebreaks; it is not an absolute token rate.
Falls back to 1.0 if torch is unavailable.
"""
try:
import torch # type: ignore[import]
device = torch.device(device_str)
# bfloat16 on CUDA matches real inference dtype; float32 on CPU avoids
# precision-downcast surprises on older hardware without bfloat16 support.
dtype = torch.bfloat16 if device_str == "cuda" else torch.float32
# hidden_size=2048 is representative of a mid-sized model; large enough
# that BLAS finds an efficient kernel on both GPU and CPU.
hidden_size = 2048
a = torch.randn(1, hidden_size, dtype=dtype, device=device)
b = torch.randn(hidden_size, hidden_size * 4, dtype=dtype, device=device)
def _sync() -> None:
if device_str == "cuda":
torch.cuda.synchronize()
# Warmup: prime caches and JIT compilation.
for _ in range(10):
torch.matmul(a, b)
_sync()
n_iters = 50
t0 = time.perf_counter()
for _ in range(n_iters):
torch.matmul(a, b)
_sync()
elapsed = time.perf_counter() - t0
return round(n_iters / max(elapsed, 1e-9), 2)
except Exception:
return 1.0

View File

@@ -14,13 +14,47 @@ from pathlib import Path
from typing import Any from typing import Any
from .downloader import compute_shard_checksum, download_shard from .downloader import compute_shard_checksum, download_shard
from .hardware import detect_hardware from .hardware import detect_hardware, benchmark_throughput
from .relay_bridge import RelayHttpBridge, peer_id_from_wallet from .relay_bridge import RelayHttpBridge, peer_id_from_wallet
from .server import StubNodeServer from .server import StubNodeServer
from .torch_server import TorchNodeServer from .torch_server import TorchNodeServer
from .wallet import load_or_create_wallet 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: def _post_json(url: str, payload: dict, timeout: float = 10.0) -> dict:
data = json.dumps(payload).encode() data = json.dumps(payload).encode()
req = urllib.request.Request( req = urllib.request.Request(
@@ -129,7 +163,11 @@ def _start_heartbeat(
uptime = time.monotonic() - _start_time uptime = time.monotonic() - _start_time
stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"} stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"}
if node_ref is not None: 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["failed_requests"] = getattr(node_ref, "failed_requests", 0)
stats["queue_depth"] = getattr(node_ref, "queue_depth", 0) stats["queue_depth"] = getattr(node_ref, "queue_depth", 0)
return stats return stats
@@ -310,20 +348,30 @@ def run_startup(
device: str = hw["device"] device: str = hw["device"]
gpu_name: str | None = hw.get("gpu_name") gpu_name: str | None = hw.get("gpu_name")
vram_mb: int = hw.get("vram_mb", 0) vram_mb: int = hw.get("vram_mb", 0)
ram_mb: int = hw.get("ram_mb", 16 * 1024)
if vram_mb_override is not None: if vram_mb_override is not None:
vram_mb = vram_mb_override 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": 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: 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)
registration_capabilities = {
"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
memory_budget_mb, memory_budget_source = _memory_budget(vram_mb, ram_mb)
print(f" Memory budget: {memory_budget_mb / 1024:.1f} GB {memory_budget_source}", flush=True)
print("Benchmarking compute...", flush=True)
bench_tps = benchmark_throughput(device)
device_label = "GPU" if device == "cuda" else "CPU"
print(f" {device_label} throughput index: {bench_tps:,.0f}", 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,
"benchmark_tokens_per_sec": bench_tps,
}
# 2. Wallet # 2. Wallet
print("Loading wallet...", flush=True) print("Loading wallet...", flush=True)
wallet_kwargs: dict = {} wallet_kwargs: dict = {}
@@ -349,7 +397,7 @@ def run_startup(
if shard_start is None and shard_end is None: if shard_start is None and shard_end is None:
try: try:
qs = urllib.parse.urlencode({ 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) 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"): if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
@@ -432,10 +480,12 @@ def run_startup(
f" Wallet: {address}\n" f" Wallet: {address}\n"
f" Model ID: {model_id}\n" f" Model ID: {model_id}\n"
f" Shard: {shard_label}\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" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n" f" Endpoint: {endpoint}\n"
f" Node ID: {tracker_node_id or 'unregistered'}\n" f" Node ID: {tracker_node_id or 'unregistered'}\n"
f" Hardware: {device.upper()}\n" f" Hardware: {device.upper()}\n"
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
f"{'=' * 32}", f"{'=' * 32}",
flush=True, flush=True,
) )
@@ -445,7 +495,7 @@ def run_startup(
# 3a. Auto-join: query tracker for network-wide HF model assignment. # 3a. Auto-join: query tracker for network-wide HF model assignment.
print("Querying tracker for network assignment...", flush=True) 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 = {} net_assignment: dict = {}
try: try:
net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}") net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}")
@@ -523,10 +573,12 @@ def run_startup(
f" Model ID: {assigned_hf_repo}\n" f" Model ID: {assigned_hf_repo}\n"
f" Shard: layers {assigned_shard_start}{assigned_shard_end} " f" Shard: layers {assigned_shard_start}{assigned_shard_end} "
f"({shard_count} of {assigned_num_layers})\n" 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" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n" f" Endpoint: {endpoint}\n"
f" Node ID: {tracker_node_id or 'unregistered'}\n" f" Node ID: {tracker_node_id or 'unregistered'}\n"
f" Hardware: {device.upper()}\n" f" Hardware: {device.upper()}\n"
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
f"{'=' * 32}", f"{'=' * 32}",
flush=True, flush=True,
) )
@@ -538,6 +590,7 @@ def run_startup(
"model": model, "model": model,
"device": device, "device": device,
"vram_mb": vram_mb, "vram_mb": vram_mb,
"ram_mb": ram_mb,
}) })
try: try:
assignment = _get_json(f"{tracker_url}/v1/nodes/assign?{assign_qs}") assignment = _get_json(f"{tracker_url}/v1/nodes/assign?{assign_qs}")
@@ -616,15 +669,17 @@ def run_startup(
# Status summary # Status summary
hw_str = device.upper() hw_str = device.upper()
if gpu_name: if gpu_name:
hw_str += f" ({gpu_name}, {vram_mb} MB)" hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)"
print( print(
f"\n{'=' * 32}\n" f"\n{'=' * 32}\n"
f"meshnet-node ready\n" f"meshnet-node ready\n"
f" Wallet: {address}\n" f" Wallet: {address}\n"
f" Shard: layers {shard_start}-{shard_end} ({assigned_model})\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" Endpoint: {endpoint}\n"
f" Node ID: {node_id}\n" f" Node ID: {node_id}\n"
f" Hardware: {hw_str}\n" f" Hardware: {hw_str}\n"
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
f"{'=' * 32}", f"{'=' * 32}",
flush=True, flush=True,
) )

View File

@@ -429,24 +429,38 @@ def _node_quantization(node: _NodeEntry, preset: dict) -> str:
return next(iter(bytes_per_layer)) return next(iter(bytes_per_layer))
def _node_memory_budget_bytes(node: _NodeEntry) -> tuple[int, str]:
"""Return the memory pool used for shard-capacity planning."""
if node.vram_bytes > 0:
return node.vram_bytes, "vram"
if node.ram_bytes > 0:
return node.ram_bytes, "ram"
return DEFAULT_RAM_BYTES, "ram-default"
def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int: def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int:
bytes_per_layer = _preset_bytes_per_layer(preset) bytes_per_layer = _preset_bytes_per_layer(preset)
quantization = _node_quantization(node, preset) quantization = _node_quantization(node, preset)
layer_bytes = bytes_per_layer[quantization] layer_bytes = bytes_per_layer[quantization]
if layer_bytes <= 0: if layer_bytes <= 0:
return 0 return 0
return int((node.vram_bytes * 0.8) // layer_bytes) memory_budget_bytes, _ = _node_memory_budget_bytes(node)
return int((memory_budget_bytes * 0.8) // layer_bytes)
def _node_capacity_summary(node: _NodeEntry, preset: dict | None = None) -> dict: def _node_capacity_summary(node: _NodeEntry, preset: dict | None = None) -> dict:
"""Operator-facing capacity fields for inspection endpoints.""" """Operator-facing capacity fields for inspection endpoints."""
memory_budget_bytes, memory_budget_source = _node_memory_budget_bytes(node)
summary = { summary = {
"vram_bytes": node.vram_bytes, "vram_bytes": node.vram_bytes,
"ram_bytes": node.ram_bytes, "ram_bytes": node.ram_bytes,
"memory_budget_bytes": memory_budget_bytes,
"memory_budget_source": memory_budget_source,
"max_loaded_shards": node.max_loaded_shards, "max_loaded_shards": node.max_loaded_shards,
"quantizations": list(node.quantizations), "quantizations": list(node.quantizations),
"quantization": node.quantization, "quantization": node.quantization,
"benchmark_tokens_per_sec": node.benchmark_tokens_per_sec, "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec,
"effective_throughput": round(_effective_throughput(node), 4),
} }
if preset is not None: if preset is not None:
summary["max_assignable_layers"] = _node_layer_capacity(node, preset) summary["max_assignable_layers"] = _node_layer_capacity(node, preset)
@@ -1154,6 +1168,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
"tracker_mode": node.tracker_mode, "tracker_mode": node.tracker_mode,
"last_heartbeat": node.last_heartbeat, "last_heartbeat": node.last_heartbeat,
"capacity": capacity_for(node), "capacity": capacity_for(node),
"stats": _node_health(node, server.heartbeat_timeout),
} }
for node in nodes for node in nodes
], ],
@@ -1567,8 +1582,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded" shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded"
repo_info = f" [{hf_repo}]" if hf_repo else "" repo_info = f" [{hf_repo}]" if hf_repo else ""
budget_bytes, budget_source = _node_memory_budget_bytes(entry)
budget_gb = budget_bytes / (1024 ** 3)
print( print(
f"[tracker] node registered: {node_id} {endpoint} {model}{repo_info} {shard_info}", f"[tracker] node registered: {node_id} {endpoint} {model}{repo_info} {shard_info} "
f"capacity={budget_gb:.1f}GB {budget_source} slots={max_loaded_shards}",
flush=True, flush=True,
) )
@@ -1707,6 +1725,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
model — model preset name (default: first preset) model — model preset name (default: first preset)
device — "cuda" | "cpu" device — "cuda" | "cpu"
vram_mb — integer VRAM in MB (0 for CPU) vram_mb — integer VRAM in MB (0 for CPU)
ram_mb — integer system RAM in MB, used when vram_mb=0
The greedy strategy: find the first gap in current layer coverage The greedy strategy: find the first gap in current layer coverage
and assign it. If no gap exists, assign the full model range so the and assign it. If no gap exists, assign the full model range so the
@@ -1745,8 +1764,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
vram_mb = int(params.get("vram_mb", ["0"])[0]) vram_mb = int(params.get("vram_mb", ["0"])[0])
except ValueError: except ValueError:
vram_mb = 0 vram_mb = 0
try:
ram_mb = int(params.get("ram_mb", ["0"])[0])
except ValueError:
ram_mb = 0
max_layers = required_end - required_start + 1 max_layers = required_end - required_start + 1
if device != "cuda" or vram_mb < 8192: memory_mb = vram_mb if vram_mb > 0 else ram_mb
if memory_mb > 0:
layer_bytes = _preset_bytes_per_layer(preset).get("bfloat16", 30 * 1024 * 1024)
max_layers = min(max_layers, max(1, int(((memory_mb * 1024 * 1024) * 0.8) // layer_bytes)))
elif device != "cuda" or vram_mb < 8192:
max_layers = min(max_layers, 16) max_layers = min(max_layers, 16)
# Collect covered intervals sorted by start layer. # Collect covered intervals sorted by start layer.
@@ -1798,6 +1825,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
Query params: Query params:
vram_mb — integer VRAM in MB (0 = CPU-only node) vram_mb — integer VRAM in MB (0 = CPU-only node)
ram_mb — integer system RAM in MB, used when vram_mb=0
device — "cuda" | "cpu" device — "cuda" | "cpu"
hf_repo — optional; if set, restrict search to this repo only hf_repo — optional; if set, restrict search to this repo only
@@ -1811,6 +1839,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
vram_mb = int(params.get("vram_mb", ["0"])[0]) vram_mb = int(params.get("vram_mb", ["0"])[0])
except ValueError: except ValueError:
vram_mb = 0 vram_mb = 0
try:
ram_mb = int(params.get("ram_mb", ["0"])[0])
except ValueError:
ram_mb = 0
device = params.get("device", ["cpu"])[0] device = params.get("device", ["cpu"])[0]
filter_repo = params.get("hf_repo", [None])[0] # optional repo filter filter_repo = params.get("hf_repo", [None])[0] # optional repo filter
@@ -1894,9 +1926,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
best_gap_start = 0 best_gap_start = 0
best_num_layers = repo_layers[best_repo] best_num_layers = repo_layers[best_repo]
# Capacity: CPU nodes get at most half the layers; CUDA nodes based on VRAM. # Capacity: use the same 80%-of-memory rule as registered node planning.
total_l = best_num_layers total_l = best_num_layers
if device == "cuda" and vram_mb >= 8192: memory_mb = vram_mb if vram_mb > 0 else ram_mb
if memory_mb > 0:
max_layers = min(
total_l,
max(1, int(((memory_mb * 1024 * 1024) * 0.8) // (30 * 1024 * 1024))),
)
elif device == "cuda" and vram_mb >= 8192:
max_layers = total_l max_layers = total_l
else: else:
max_layers = max(1, total_l // 2) max_layers = max(1, total_l // 2)

19
tests/conftest.py Normal file
View File

@@ -0,0 +1,19 @@
"""Shared pytest fixtures for the meshnet test suite."""
import pytest
@pytest.fixture(autouse=True)
def _stub_benchmark_throughput(monkeypatch):
"""Replace the GEMM benchmark with a fixed value in all tests.
The benchmark runs 60 matmuls (warmup + measure) which adds ~100ms per test
on CPU. Tests verify registration flow, not hardware speed — stub it out.
Tests that specifically exercise benchmark_throughput import it directly from
meshnet_node.hardware and are not affected by this patch.
"""
try:
import meshnet_node.startup as startup_mod
monkeypatch.setattr(startup_mod, "benchmark_throughput", lambda _device: 999.0)
except ImportError:
pass

View File

@@ -10,7 +10,7 @@ from pathlib import Path
import pytest import pytest
from meshnet_node.downloader import download_shard, write_shard_archive from meshnet_node.downloader import download_shard, write_shard_archive
from meshnet_node.hardware import detect_hardware from meshnet_node.hardware import detect_hardware, benchmark_throughput
from meshnet_node.startup import ( from meshnet_node.startup import (
_infer_relay_url_from_tracker, _infer_relay_url_from_tracker,
_probationary_status_line, _probationary_status_line,
@@ -31,6 +31,8 @@ def test_detect_hardware_returns_valid_profile():
hw = detect_hardware() hw = detect_hardware()
assert hw["device"] in {"cuda", "cpu"} assert hw["device"] in {"cuda", "cpu"}
assert isinstance(hw.get("vram_mb"), int) assert isinstance(hw.get("vram_mb"), int)
assert isinstance(hw.get("ram_mb"), int)
assert hw["ram_mb"] > 0
if hw["device"] == "cpu": if hw["device"] == "cpu":
assert hw["gpu_name"] is None assert hw["gpu_name"] is None
assert hw["vram_mb"] == 0 assert hw["vram_mb"] == 0
@@ -39,6 +41,62 @@ def test_detect_hardware_returns_valid_profile():
assert hw["vram_mb"] > 0 assert hw["vram_mb"] > 0
def test_benchmark_throughput_cpu_returns_positive():
"""CPU benchmark returns a positive float greater than the 1.0 error fallback."""
result = benchmark_throughput("cpu")
assert isinstance(result, float)
assert result > 1.0, f"expected benchmark > 1.0, got {result}"
def test_benchmark_throughput_fallback_on_bad_device():
"""benchmark_throughput returns 1.0 (not raises) when device is invalid."""
result = benchmark_throughput("invalid_device_xyz")
assert result == 1.0
def test_benchmark_throughput_is_registered_in_payload(monkeypatch, tmp_path):
"""benchmark_tokens_per_sec from the benchmark is included in the tracker registration."""
import meshnet_node.startup as startup_mod
captured: dict = {}
class FakeNode:
backend = None
tracker_node_id = None
def start(self):
return 7099
def stop(self):
pass
def apply_tracker_directives(self, directives):
return None
monkeypatch.setattr(startup_mod, "detect_hardware",
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16384})
monkeypatch.setattr(startup_mod, "benchmark_throughput", lambda _device: 42.5)
monkeypatch.setattr(startup_mod, "TorchNodeServer", lambda **_kw: FakeNode())
monkeypatch.setattr(startup_mod, "_detect_num_layers", lambda _model_id: 24)
monkeypatch.setattr(startup_mod, "RelayHttpBridge", None)
monkeypatch.setattr(startup_mod, "_get_json",
lambda _url, timeout=10.0: {"relay_url": None, "nodes": []})
monkeypatch.setattr(startup_mod, "_post_json",
lambda _url, payload, timeout=10.0: (captured.update(payload) or {"node_id": "x"}))
monkeypatch.setattr(startup_mod, "_start_heartbeat", lambda *a, **kw: None)
node = run_startup(
tracker_url="http://localhost:8080",
model_id="Qwen/Qwen2.5-0.5B-Instruct",
shard_start=0,
shard_end=23,
wallet_path=tmp_path / "wallet.json",
)
node.stop()
assert captured.get("benchmark_tokens_per_sec") == 42.5
def test_wallet_generates_new_keypair(tmp_path): def test_wallet_generates_new_keypair(tmp_path):
"""A new wallet is created when none exists, saved to disk.""" """A new wallet is created when none exists, saved to disk."""
wallet_file = tmp_path / "wallet.json" wallet_file = tmp_path / "wallet.json"
@@ -490,6 +548,70 @@ def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, ca
assert "Node ID: node-test-123" in output assert "Node ID: node-test-123" in output
def test_real_model_startup_autodetects_cpu_memory_budget_and_logs_shard_budget(
tmp_path,
monkeypatch,
capsys,
):
"""Without --memory, startup reports RAM-backed capacity to the tracker and operator."""
import meshnet_node.startup as startup_mod
captured_registration = {}
class FakeBackend:
total_layers = 24
class FakeTorchNodeServer:
def __init__(self, **kwargs):
self.kwargs = kwargs
self.backend = FakeBackend()
self.port = None
self.total_requests = 0
self.failed_requests = 0
self.queue_depth = 0
def start(self):
self.port = 8001
return self.port
def stop(self):
pass
monkeypatch.setattr(
startup_mod,
"detect_hardware",
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16384},
)
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
monkeypatch.setattr(
startup_mod,
"_post_json",
lambda _url, _payload, timeout=10.0: (
captured_registration.update(_payload) or {"node_id": "node-auto-mem"}
),
)
node = run_startup(
tracker_url="http://127.0.0.1:8080",
model_id="Qwen/Qwen2.5-0.5B-Instruct",
shard_start=0,
shard_end=23,
wallet_path=tmp_path / "wallet.json",
)
try:
pass
finally:
node.stop()
assert captured_registration["vram_bytes"] == 0
assert captured_registration["ram_bytes"] == 16384 * 1024 * 1024
assert captured_registration["max_loaded_shards"] == 1
output = capsys.readouterr().out
assert "Memory budget: 16.0 GB RAM" in output
assert "Shard budget: up to 24/24 layers at bfloat16" in output
assert "GB remaining after full load" in output
assert "Node ID: node-auto-mem" in output
def test_public_tracker_model_node_registers_relay_metadata_from_tracker_url_only( def test_public_tracker_model_node_registers_relay_metadata_from_tracker_url_only(
tmp_path, tmp_path,
monkeypatch, monkeypatch,

View File

@@ -894,6 +894,35 @@ def test_network_map_exposes_node_capacity_limits():
tracker.stop() tracker.stop()
def test_tracker_capacity_uses_ram_when_node_has_no_vram():
"""CPU-only nodes should expose RAM-backed shard capacity, not default GPU capacity."""
tracker = TrackerServer(model_presets={
"tiny-model": {
"total_layers": 20,
"bytes_per_layer": {"bfloat16": 1_000},
},
})
tracker_port = tracker.start()
try:
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{"endpoint": "http://127.0.0.1:9020", "model": "tiny-model",
"vram_bytes": 0, "ram_bytes": 16_000, "quantizations": ["bfloat16"],
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
)
network_map = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
capacity = network_map["nodes"][0]["capacity"]
assert capacity["vram_bytes"] == 0
assert capacity["ram_bytes"] == 16_000
assert capacity["memory_budget_bytes"] == 16_000
assert capacity["memory_budget_source"] == "ram"
assert capacity["max_assignable_layers"] == 12
finally:
tracker.stop()
def test_rebalance_keeps_one_active_range_even_when_multiple_slots_advertised(): def test_rebalance_keeps_one_active_range_even_when_multiple_slots_advertised():
"""max_loaded_shards is exposed but reserved until node runtime supports multi-range serving.""" """max_loaded_shards is exposed but reserved until node runtime supports multi-range serving."""
tracker = TrackerServer(model_presets={ tracker = TrackerServer(model_presets={
@@ -1555,3 +1584,112 @@ def test_torch_node_applies_tracker_load_shard_directive(monkeypatch):
"tracker_mode": True, "tracker_mode": True,
} }
assert node.backend.shard_end == 23 assert node.backend.shard_end == 23
def test_shard_heal_cycle_surviving_node_covers_dead_peers_gap(monkeypatch):
"""End-to-end heal: kill one managed node, surviving node receives LOAD_SHARD and hot-swaps.
Cycle:
1. Two managed nodes (A: 0-11, B: 12-23) register with tracker.
2. Node A stops heartbeating; tracker expires it and triggers rebalance.
3. Node B's next heartbeat response contains LOAD_SHARD(0, 23).
4. Node B (TorchNodeServer) applies the directive — backend hot-swapped.
5. Coverage endpoint confirms full model is covered by Node B alone.
"""
from meshnet_node import torch_server
from meshnet_node.torch_server import TorchNodeServer
# --- minimal fake backend (no GPU / PyTorch needed) ---
class _FakeBackend:
def __init__(self, model_id="Qwen/Qwen2.5-0.5B-Instruct", shard_start=0, shard_end=23, quantization="int8"):
self.model_id = model_id
self.shard_start = shard_start
self.shard_end = shard_end
self.quantization = quantization
self.total_layers = 24
self.is_head = shard_start == 0
self.is_tail = shard_end == 23
def generate_text(self, *a, **kw): return ""
def count_prompt_tokens(self, *a): return 0
def count_text_tokens(self, *a): return 0
loaded_shards: list[tuple] = []
def fake_load(model_id, shard_start, shard_end, quantization):
loaded_shards.append((model_id, shard_start, shard_end))
return _FakeBackend(model_id, shard_start, shard_end, quantization)
monkeypatch.setattr(torch_server, "_load_backend", fake_load)
# Use a very short timeout so Node A expires quickly.
tracker = TrackerServer(heartbeat_timeout=0.15, rebalance_interval=10.0)
tracker_port = tracker.start()
node_b = TorchNodeServer(backend=_FakeBackend(shard_start=12, shard_end=23))
base_reg = {
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
"model": "Qwen2.5-0.5B-Instruct",
"num_layers": 24,
"vram_bytes": 2_000_000_000,
"ram_bytes": 0,
"hardware_profile": {},
"score": 1.0,
"managed_assignment": True,
}
try:
# Step 1: register both nodes as managed.
reg_a = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base_reg, "endpoint": "http://127.0.0.1:19001", "shard_start": 0, "shard_end": 11},
)
reg_b = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base_reg, "endpoint": "http://127.0.0.1:19002", "shard_start": 12, "shard_end": 23},
)
node_a_id = reg_a["node_id"]
node_b_id = reg_b["node_id"]
# Initial heartbeat to mark both alive.
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_a_id}/heartbeat", {})
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_b_id}/heartbeat", {})
# Step 2: let Node A's heartbeat expire (only Node B keeps heartbeating).
time.sleep(0.10)
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_b_id}/heartbeat", {})
time.sleep(0.10)
# Step 3: Node B's heartbeat triggers purge of A and gets LOAD_SHARD.
hb_resp = _post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_b_id}/heartbeat", {}
)
directives = hb_resp.get("directives", [])
load_dirs = [d for d in directives if d["action"] == "LOAD_SHARD"]
assert load_dirs, f"Expected LOAD_SHARD directive, got: {directives}"
assert load_dirs[-1]["shard_start"] == 0
assert load_dirs[-1]["shard_end"] == 23
assert node_a_id not in tracker._registry
# Step 4: Node B applies the directive — backend hot-swapped.
applied = node_b.apply_tracker_directives(directives)
assert applied is not None
assert applied["shard_start"] == 0
assert applied["shard_end"] == 23
assert applied["tracker_mode"] is True
assert node_b.backend.shard_start == 0
assert node_b.backend.shard_end == 23
assert loaded_shards == [("Qwen/Qwen2.5-0.5B-Instruct", 0, 23)]
# Step 5: re-register Node B with its new shard so tracker reflects healed state.
_post_json(
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
{**base_reg, "endpoint": "http://127.0.0.1:19002", "shard_start": 0, "shard_end": 23},
)
coverage_resp = _get_json(
f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen%2FQwen2.5-0.5B-Instruct"
)
assert all(seg["node_count"] >= 1 for seg in coverage_resp["coverage"]), (
f"Coverage gap after heal: {coverage_resp['coverage']}"
)
finally:
tracker.stop()