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,
|
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")
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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
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
|
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,
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
Reference in New Issue
Block a user