diff --git a/.scratch/distributed-inference-network/issues/20-memory-budget-shard-slots-and-dropout-relocation.md b/.scratch/distributed-inference-network/issues/20-memory-budget-shard-slots-and-dropout-relocation.md index f128193..74f70cf 100644 --- a/.scratch/distributed-inference-network/issues/20-memory-budget-shard-slots-and-dropout-relocation.md +++ b/.scratch/distributed-inference-network/issues/20-memory-budget-shard-slots-and-dropout-relocation.md @@ -63,3 +63,4 @@ The current runtime still effectively has one active backend shard per node. A n ## Comments - 2026-06-30: Created after implementing the initial registration plumbing in commit `f1e4ed6` (`--memory`, `--max-shards`, tracker validation). This issue captures the remaining end-to-end behavior so it does not conflict with US-013. +- 2026-06-30: Implementation decision: `max_loaded_shards` is currently a validated and exposed capacity field, but multi-range assignment remains reserved because `TorchNodeServer` serves one active backend shard. The tracker therefore emits at most one active range per node while exposing `vram_bytes`, `ram_bytes`, `max_loaded_shards`, quantization, throughput, and computed `max_assignable_layers` in inspection endpoints. diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 8f5eef5..e15e773 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -438,6 +438,21 @@ def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int: return int((node.vram_bytes * 0.8) // layer_bytes) +def _node_capacity_summary(node: _NodeEntry, preset: dict | None = None) -> dict: + """Operator-facing capacity fields for inspection endpoints.""" + summary = { + "vram_bytes": node.vram_bytes, + "ram_bytes": node.ram_bytes, + "max_loaded_shards": node.max_loaded_shards, + "quantizations": list(node.quantizations), + "quantization": node.quantization, + "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec, + } + if preset is not None: + summary["max_assignable_layers"] = _node_layer_capacity(node, preset) + return summary + + def _coverage_map( nodes: list[_NodeEntry], required_start: int, @@ -480,6 +495,7 @@ def _node_health(node: "_NodeEntry", heartbeat_timeout: float) -> dict: "total_requests": node.total_requests, "heartbeat_success_rate": hb_rate, "inference_success_rate": inf_rate, + "capacity": _node_capacity_summary(node), } @@ -1114,6 +1130,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): with server.lock: self._purge_expired_nodes() nodes = list(server.registry.values()) + + def capacity_for(node: _NodeEntry) -> dict: + preset = None + if node.model: + preset = server.model_presets.get(node.model) + if preset is None and node.hf_repo and node.num_layers: + preset = _hf_rebalance_preset([node]) + return _node_capacity_summary(node, preset) + self._send_json(200, { "relay_url": server.relay_url, "nodes": [ @@ -1128,6 +1153,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "shard_end": node.shard_end, "tracker_mode": node.tracker_mode, "last_heartbeat": node.last_heartbeat, + "capacity": capacity_for(node), } for node in nodes ], diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index 4bb8b2c..a6584fd 100644 --- a/tests/test_tracker_routing.py +++ b/tests/test_tracker_routing.py @@ -864,6 +864,63 @@ def test_tracker_registration_rejects_invalid_max_loaded_shards(): tracker.stop() +def test_network_map_exposes_node_capacity_limits(): + tracker = TrackerServer(model_presets={ + "tiny-model": { + "total_layers": 12, + "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:9021", "model": "tiny-model", + "shard_start": 0, "shard_end": 7, + "vram_bytes": 10_000, "ram_bytes": 20_000, "max_loaded_shards": 2, + "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 4.0, + "hardware_profile": {}, "score": 1.0}, + ) + + network_map = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map") + + node = network_map["nodes"][0] + assert node["capacity"]["vram_bytes"] == 10_000 + assert node["capacity"]["ram_bytes"] == 20_000 + assert node["capacity"]["max_loaded_shards"] == 2 + assert node["capacity"]["max_assignable_layers"] == 8 + assert node["capacity"]["benchmark_tokens_per_sec"] == 4.0 + 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={ + "tiny-model": { + "total_layers": 8, + "bytes_per_layer": {"bfloat16": 1_000}, + }, + }) + tracker_port = tracker.start() + try: + reg = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9022", "model": "tiny-model", + "vram_bytes": 20_000, "ram_bytes": 20_000, "max_loaded_shards": 2, + "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0, + "hardware_profile": {}, "score": 1.0}, + ) + + assert reg["directive"]["action"] == "LOAD_SHARD" + assert reg["directive"]["start_layer"] == 0 + assert reg["directive"]["end_layer"] == 7 + hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{reg['node_id']}/heartbeat", {}) + assert hb == {} + finally: + tracker.stop() + + def test_tracker_routes_only_nodes_for_requested_model(): """A node registered for one model cannot satisfy another model's route.""" tracker = TrackerServer(model_presets={