From b6272db93da5b99c82018eb0d4a8d7c8520e6e59 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Tue, 30 Jun 2026 23:08:34 +0300 Subject: [PATCH] feat: default quantization int8, GB display, shard heal cycle test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- packages/node/meshnet_node/cli.py | 4 +- packages/node/meshnet_node/startup.py | 66 ++++++++++-- tests/test_tracker_routing.py | 138 ++++++++++++++++++++++++++ 3 files changed, 196 insertions(+), 12 deletions(-) diff --git a/packages/node/meshnet_node/cli.py b/packages/node/meshnet_node/cli.py index 4250677..187792e 100644 --- a/packages/node/meshnet_node/cli.py +++ b/packages/node/meshnet_node/cli.py @@ -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") diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 753b251..5a02977 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -21,6 +21,40 @@ from .torch_server import TorchNodeServer from .wallet import load_or_create_wallet +_DEFAULT_BYTES_PER_LAYER = 30 * 1024 * 1024 + + +def _memory_budget(vram_mb: int, ram_mb: int) -> tuple[int, str]: + """Return the capacity budget in MB and whether it came from VRAM or RAM.""" + if vram_mb > 0: + return vram_mb, "VRAM" + return max(0, ram_mb), "RAM" + + +def _max_assignable_layers(memory_mb: int, total_layers: int | None) -> int: + if total_layers is None or total_layers <= 0 or memory_mb <= 0: + return 0 + budget_bytes = memory_mb * 1024 * 1024 + return min(total_layers, int((budget_bytes * 0.8) // _DEFAULT_BYTES_PER_LAYER)) + + +def _shard_budget_line(memory_mb: int, memory_source: str, total_layers: int | None, quantization: str) -> str: + memory_gb = memory_mb / 1024 + gb_str = f"{memory_gb:.1f} GB" + if total_layers is None or total_layers <= 0: + return f"Memory budget: {gb_str} {memory_source}; shard budget: unknown model layer count" + max_layers = _max_assignable_layers(memory_mb, total_layers) + # Remaining capacity after one full model load (rough estimate) + shard_bytes = max_layers * _DEFAULT_BYTES_PER_LAYER + remaining_gb = (memory_mb * 1024 * 1024 - shard_bytes) / (1024 ** 3) + remaining_str = f"; {remaining_gb:.1f} GB remaining after full load" if remaining_gb > 1 else "" + return ( + f"Memory budget: {gb_str} {memory_source}; " + f"Shard budget: up to {max_layers}/{total_layers} layers at {quantization}" + f"{remaining_str}" + ) + + def _post_json(url: str, payload: dict, timeout: float = 10.0) -> dict: data = json.dumps(payload).encode() req = urllib.request.Request( @@ -129,7 +163,11 @@ def _start_heartbeat( uptime = time.monotonic() - _start_time stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"} if node_ref is not None: - stats["total_requests"] = getattr(node_ref, "total_requests", 0) + stats["total_requests"] = getattr( + node_ref, + "total_requests", + getattr(node_ref, "chat_completion_count", 0), + ) stats["failed_requests"] = getattr(node_ref, "failed_requests", 0) stats["queue_depth"] = getattr(node_ref, "queue_depth", 0) return stats @@ -310,20 +348,24 @@ def run_startup( device: str = hw["device"] gpu_name: str | None = hw.get("gpu_name") vram_mb: int = hw.get("vram_mb", 0) + ram_mb: int = hw.get("ram_mb", 16 * 1024) if vram_mb_override is not None: vram_mb = vram_mb_override - print(f" Memory budget overridden to {vram_mb} MB via --memory", flush=True) + print(f" Memory budget overridden to {vram_mb / 1024:.1f} GB via --memory", flush=True) elif device == "cpu": - print(" WARNING: No CUDA GPU detected — running in CPU mode", flush=True) + print(f" WARNING: No CUDA GPU detected — running in CPU mode ({ram_mb / 1024:.1f} GB RAM)", flush=True) else: - print(f" GPU: {gpu_name} ({vram_mb} MB VRAM)", flush=True) + print(f" GPU: {gpu_name} ({vram_mb / 1024:.1f} GB VRAM, {ram_mb / 1024:.1f} GB RAM)", flush=True) + + memory_budget_mb, memory_budget_source = _memory_budget(vram_mb, ram_mb) + print(f" Memory budget: {memory_budget_mb} MB {memory_budget_source}", flush=True) + registration_capabilities = { + "vram_bytes": max(0, int(vram_mb)) * 1024 * 1024, + "ram_bytes": max(0, int(ram_mb)) * 1024 * 1024, "max_loaded_shards": max_loaded_shards, } - if vram_mb_override is not None or vram_mb > 0: - registration_capabilities["vram_bytes"] = max(0, int(vram_mb)) * 1024 * 1024 - # 2. Wallet print("Loading wallet...", flush=True) wallet_kwargs: dict = {} @@ -349,7 +391,7 @@ def run_startup( if shard_start is None and shard_end is None: try: qs = urllib.parse.urlencode({ - "device": device, "vram_mb": vram_mb, "hf_repo": model_id, + "device": device, "vram_mb": vram_mb, "ram_mb": ram_mb, "hf_repo": model_id, }) net_asgn = _get_json(f"{tracker_url}/v1/network/assign?{qs}", timeout=5.0) if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"): @@ -432,6 +474,7 @@ def run_startup( f" Wallet: {address}\n" f" Model ID: {model_id}\n" f" Shard: {shard_label}\n" + f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization)}\n" f" Quantization: {quantization}\n" f" Endpoint: {endpoint}\n" f" Node ID: {tracker_node_id or 'unregistered'}\n" @@ -445,7 +488,7 @@ def run_startup( # 3a. Auto-join: query tracker for network-wide HF model assignment. print("Querying tracker for network assignment...", flush=True) - assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": vram_mb}) + assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": vram_mb, "ram_mb": ram_mb}) net_assignment: dict = {} try: net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}") @@ -523,6 +566,7 @@ def run_startup( f" Model ID: {assigned_hf_repo}\n" f" Shard: layers {assigned_shard_start}–{assigned_shard_end} " f"({shard_count} of {assigned_num_layers})\n" + f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization)}\n" f" Quantization: {quantization}\n" f" Endpoint: {endpoint}\n" f" Node ID: {tracker_node_id or 'unregistered'}\n" @@ -538,6 +582,7 @@ def run_startup( "model": model, "device": device, "vram_mb": vram_mb, + "ram_mb": ram_mb, }) try: assignment = _get_json(f"{tracker_url}/v1/nodes/assign?{assign_qs}") @@ -616,12 +661,13 @@ def run_startup( # Status summary hw_str = device.upper() if gpu_name: - hw_str += f" ({gpu_name}, {vram_mb} MB)" + hw_str += f" ({gpu_name}, {vram_mb / 1024:.1f} GB)" print( f"\n{'=' * 32}\n" f"meshnet-node ready\n" f" Wallet: {address}\n" f" Shard: layers {shard_start}-{shard_end} ({assigned_model})\n" + f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assignment.get('model_layers_end', shard_end) + 1, quantization)}\n" f" Endpoint: {endpoint}\n" f" Node ID: {node_id}\n" f" Hardware: {hw_str}\n" diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index a6584fd..e4e08e6 100644 --- a/tests/test_tracker_routing.py +++ b/tests/test_tracker_routing.py @@ -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()