Expose node shard capacity limits
This commit is contained in:
@@ -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={
|
||||
|
||||
Reference in New Issue
Block a user