Compare commits
2 Commits
ff4115f611
...
278be49539
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
278be49539 | ||
|
|
b6272db93d |
@@ -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")
|
||||
|
||||
@@ -94,6 +94,19 @@ def _make_bar(pct: float, width: int = 10) -> str:
|
||||
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:
|
||||
"""Start the live dashboard. Blocks until Ctrl-C. Returns cleanly."""
|
||||
if not is_interactive_tty():
|
||||
@@ -117,7 +130,8 @@ def _build_rich_renderable(
|
||||
from rich.text import Text # type: ignore[import]
|
||||
|
||||
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)
|
||||
delta_req = req_count - prev_req[0]
|
||||
@@ -163,6 +177,7 @@ def _build_rich_renderable(
|
||||
stats_lines = [
|
||||
f"Tokens/sec {tps_bar} {tps:.1f} t/s (EMA)",
|
||||
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"TAI earned 0.00 TAI (payments: US-006)",
|
||||
f"Uptime {_format_uptime(uptime)}",
|
||||
@@ -205,14 +220,17 @@ def _run_plain_loop(node, config: dict, start_time: float) -> None:
|
||||
try:
|
||||
while True:
|
||||
uptime = time.monotonic() - start_time
|
||||
req = getattr(node, "chat_completion_count", 0)
|
||||
stats = _node_stats(node)
|
||||
req = stats["total_requests"]
|
||||
gpu_stats = _gpu_stats()
|
||||
vram_str = ""
|
||||
if gpu_stats:
|
||||
g = gpu_stats[0]
|
||||
vram_str = f" VRAM{g['used_gb']:.1f}GB"
|
||||
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,
|
||||
)
|
||||
time.sleep(2)
|
||||
|
||||
@@ -1,10 +1,23 @@
|
||||
"""GPU hardware detection with graceful CPU fallback."""
|
||||
|
||||
import os
|
||||
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:
|
||||
"""Detect GPU model and available VRAM. Returns hardware profile dict."""
|
||||
ram_mb = _detect_ram_mb()
|
||||
try:
|
||||
import torch # type: ignore[import]
|
||||
if torch.cuda.is_available():
|
||||
@@ -12,7 +25,7 @@ def detect_hardware() -> dict:
|
||||
name = torch.cuda.get_device_name(idx)
|
||||
props = torch.cuda.get_device_properties(idx)
|
||||
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:
|
||||
pass
|
||||
|
||||
@@ -26,8 +39,54 @@ def detect_hardware() -> dict:
|
||||
parts = line.split(",", 1)
|
||||
gpu_name = parts[0].strip()
|
||||
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):
|
||||
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
|
||||
|
||||
@@ -14,13 +14,47 @@ from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
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 .server import StubNodeServer
|
||||
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,30 @@ 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)
|
||||
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
|
||||
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 / 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
|
||||
print("Loading wallet...", flush=True)
|
||||
wallet_kwargs: dict = {}
|
||||
@@ -349,7 +397,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,10 +480,12 @@ 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"
|
||||
f" Hardware: {device.upper()}\n"
|
||||
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
|
||||
f"{'=' * 32}",
|
||||
flush=True,
|
||||
)
|
||||
@@ -445,7 +495,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,10 +573,12 @@ 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"
|
||||
f" Hardware: {device.upper()}\n"
|
||||
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
|
||||
f"{'=' * 32}",
|
||||
flush=True,
|
||||
)
|
||||
@@ -538,6 +590,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,15 +669,17 @@ 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"
|
||||
f" Benchmark: {bench_tps:,.0f} (throughput index)\n"
|
||||
f"{'=' * 32}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
@@ -429,24 +429,38 @@ def _node_quantization(node: _NodeEntry, preset: dict) -> str:
|
||||
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:
|
||||
bytes_per_layer = _preset_bytes_per_layer(preset)
|
||||
quantization = _node_quantization(node, preset)
|
||||
layer_bytes = bytes_per_layer[quantization]
|
||||
if layer_bytes <= 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:
|
||||
"""Operator-facing capacity fields for inspection endpoints."""
|
||||
memory_budget_bytes, memory_budget_source = _node_memory_budget_bytes(node)
|
||||
summary = {
|
||||
"vram_bytes": node.vram_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,
|
||||
"quantizations": list(node.quantizations),
|
||||
"quantization": node.quantization,
|
||||
"benchmark_tokens_per_sec": node.benchmark_tokens_per_sec,
|
||||
"effective_throughput": round(_effective_throughput(node), 4),
|
||||
}
|
||||
if preset is not None:
|
||||
summary["max_assignable_layers"] = _node_layer_capacity(node, preset)
|
||||
@@ -1154,6 +1168,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
"tracker_mode": node.tracker_mode,
|
||||
"last_heartbeat": node.last_heartbeat,
|
||||
"capacity": capacity_for(node),
|
||||
"stats": _node_health(node, server.heartbeat_timeout),
|
||||
}
|
||||
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"
|
||||
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(
|
||||
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,
|
||||
)
|
||||
|
||||
@@ -1707,6 +1725,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
model — model preset name (default: first preset)
|
||||
device — "cuda" | "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
|
||||
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])
|
||||
except ValueError:
|
||||
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
|
||||
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)
|
||||
|
||||
# Collect covered intervals sorted by start layer.
|
||||
@@ -1798,6 +1825,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
Query params:
|
||||
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"
|
||||
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])
|
||||
except ValueError:
|
||||
vram_mb = 0
|
||||
try:
|
||||
ram_mb = int(params.get("ram_mb", ["0"])[0])
|
||||
except ValueError:
|
||||
ram_mb = 0
|
||||
device = params.get("device", ["cpu"])[0]
|
||||
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_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
|
||||
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
|
||||
else:
|
||||
max_layers = max(1, total_l // 2)
|
||||
|
||||
19
tests/conftest.py
Normal file
19
tests/conftest.py
Normal 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
|
||||
@@ -10,7 +10,7 @@ from pathlib import Path
|
||||
import pytest
|
||||
|
||||
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 (
|
||||
_infer_relay_url_from_tracker,
|
||||
_probationary_status_line,
|
||||
@@ -31,6 +31,8 @@ def test_detect_hardware_returns_valid_profile():
|
||||
hw = detect_hardware()
|
||||
assert hw["device"] in {"cuda", "cpu"}
|
||||
assert isinstance(hw.get("vram_mb"), int)
|
||||
assert isinstance(hw.get("ram_mb"), int)
|
||||
assert hw["ram_mb"] > 0
|
||||
if hw["device"] == "cpu":
|
||||
assert hw["gpu_name"] is None
|
||||
assert hw["vram_mb"] == 0
|
||||
@@ -39,6 +41,62 @@ def test_detect_hardware_returns_valid_profile():
|
||||
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):
|
||||
"""A new wallet is created when none exists, saved to disk."""
|
||||
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
|
||||
|
||||
|
||||
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(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
|
||||
@@ -894,6 +894,35 @@ def test_network_map_exposes_node_capacity_limits():
|
||||
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():
|
||||
"""max_loaded_shards is exposed but reserved until node runtime supports multi-range serving."""
|
||||
tracker = TrackerServer(model_presets={
|
||||
@@ -1555,3 +1584,112 @@ def test_torch_node_applies_tracker_load_shard_directive(monkeypatch):
|
||||
"tracker_mode": True,
|
||||
}
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user