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

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

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

View File

@@ -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()