Expose node shard capacity limits
This commit is contained in:
@@ -63,3 +63,4 @@ The current runtime still effectively has one active backend shard per node. A n
|
|||||||
## Comments
|
## 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: 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.
|
||||||
|
|||||||
@@ -438,6 +438,21 @@ def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int:
|
|||||||
return int((node.vram_bytes * 0.8) // layer_bytes)
|
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(
|
def _coverage_map(
|
||||||
nodes: list[_NodeEntry],
|
nodes: list[_NodeEntry],
|
||||||
required_start: int,
|
required_start: int,
|
||||||
@@ -480,6 +495,7 @@ def _node_health(node: "_NodeEntry", heartbeat_timeout: float) -> dict:
|
|||||||
"total_requests": node.total_requests,
|
"total_requests": node.total_requests,
|
||||||
"heartbeat_success_rate": hb_rate,
|
"heartbeat_success_rate": hb_rate,
|
||||||
"inference_success_rate": inf_rate,
|
"inference_success_rate": inf_rate,
|
||||||
|
"capacity": _node_capacity_summary(node),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
@@ -1114,6 +1130,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
with server.lock:
|
with server.lock:
|
||||||
self._purge_expired_nodes()
|
self._purge_expired_nodes()
|
||||||
nodes = list(server.registry.values())
|
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, {
|
self._send_json(200, {
|
||||||
"relay_url": server.relay_url,
|
"relay_url": server.relay_url,
|
||||||
"nodes": [
|
"nodes": [
|
||||||
@@ -1128,6 +1153,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
"shard_end": node.shard_end,
|
"shard_end": node.shard_end,
|
||||||
"tracker_mode": node.tracker_mode,
|
"tracker_mode": node.tracker_mode,
|
||||||
"last_heartbeat": node.last_heartbeat,
|
"last_heartbeat": node.last_heartbeat,
|
||||||
|
"capacity": capacity_for(node),
|
||||||
}
|
}
|
||||||
for node in nodes
|
for node in nodes
|
||||||
],
|
],
|
||||||
|
|||||||
@@ -864,6 +864,63 @@ def test_tracker_registration_rejects_invalid_max_loaded_shards():
|
|||||||
tracker.stop()
|
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():
|
def test_tracker_routes_only_nodes_for_requested_model():
|
||||||
"""A node registered for one model cannot satisfy another model's route."""
|
"""A node registered for one model cannot satisfy another model's route."""
|
||||||
tracker = TrackerServer(model_presets={
|
tracker = TrackerServer(model_presets={
|
||||||
|
|||||||
Reference in New Issue
Block a user