Expose node shard capacity limits

This commit is contained in:
Dobromir Popov
2026-06-30 21:29:30 +02:00
parent 94d2216a28
commit b06b86a6f0
3 changed files with 84 additions and 0 deletions

View File

@@ -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={