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:
@@ -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