ram pool map

tests/test_tracker_routing.py

@@ -13,7 +13,13 @@ from meshnet_gateway.server import GatewayServer, _banned_route_wallet
 from meshnet_node.server import StubNodeServer
 from meshnet_contracts import LocalSolanaContracts
 from meshnet_tracker.auth import sign_hive_request
-from meshnet_tracker.server import TrackerServer, _NodeEntry, _registration_ban_error
+from meshnet_tracker.server import (
+    TrackerServer,
+    _NodeEntry,
+    _memory_pool_map,
+    _registration_ban_error,
+    _scale_demanded_models_locked,
+)
 
 _TEST_HIVE_SECRET = "test-hive-secret"
 
@@ -2340,6 +2346,7 @@ def test_torch_node_applies_tracker_load_shard_directive(monkeypatch):
 
     assert loaded == [("Qwen/Qwen2.5-0.5B-Instruct", 0, 23, "bfloat16")]
     assert applied == {
+        "action": "LOAD_SHARD",
         "model": "Qwen/Qwen2.5-0.5B-Instruct",
         "shard_start": 0,
         "shard_end": 23,
@@ -2456,3 +2463,105 @@ def test_shard_heal_cycle_surviving_node_covers_dead_peers_gap(monkeypatch):
         )
     finally:
         tracker.stop()
+
+
+def test_network_map_exposes_memory_pool():
+    tracker = TrackerServer(model_presets={
+        "tiny-model": {
+            "total_layers": 8,
+            "bytes_per_layer": {"bfloat16": 1_000},
+            "hf_repo": "org/TinyModel",
+        },
+    })
+    tracker_port = tracker.start()
+    try:
+        _post_json(
+            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+            {"endpoint": "http://127.0.0.1:9050", "model": "tiny-model",
+             "hf_repo": "org/TinyModel", "num_layers": 8,
+             "shard_start": 0, "shard_end": 7, "max_loaded_shards": 2,
+             "vram_bytes": 20_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
+             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
+        )
+        net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
+        pool = net["memory_pool"]
+        assert pool["total_spare_slots"] == 1
+        assert pool["hosts"][0]["loaded_slots"] == 1
+        assert pool["hosts"][0]["max_loaded_shards"] == 2
+        assert pool["hosts"][0]["memory_spare_bytes"] > 0
+    finally:
+        tracker.stop()
+
+
+def test_same_endpoint_can_register_multiple_models():
+    tracker = TrackerServer()
+    tracker_port = tracker.start()
+    try:
+        base = {
+            "endpoint": "http://127.0.0.1:9055",
+            "num_layers": 24,
+            "shard_start": 0,
+            "shard_end": 23,
+            "hardware_profile": {},
+            "score": 1.0,
+            "max_loaded_shards": 2,
+            "vram_bytes": 50_000_000,
+            "ram_bytes": 50_000_000,
+        }
+        first = _post_json(
+            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+            {**base, "model": "Qwen2.5-0.5B-Instruct", "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct"},
+        )
+        second = _post_json(
+            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+            {**base, "model": "OtherModel", "hf_repo": "org/OtherModel"},
+        )
+        assert first["node_id"] != second["node_id"]
+        assert len(tracker._registry) == 2
+    finally:
+        tracker.stop()
+
+
+def test_scale_demanded_models_queues_add_shard_on_spare_host():
+    tracker = TrackerServer(model_presets={
+        "model-a": {
+            "total_layers": 4,
+            "bytes_per_layer": {"bfloat16": 1_000},
+            "hf_repo": "org/ModelA",
+        },
+        "model-b": {
+            "total_layers": 4,
+            "bytes_per_layer": {"bfloat16": 1_000},
+            "hf_repo": "org/ModelB",
+        },
+    })
+    tracker_port = tracker.start()
+    try:
+        reg_b = _post_json(
+            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+            {"endpoint": "http://127.0.0.1:9061", "model": "model-b",
+             "hf_repo": "org/ModelB", "num_layers": 4,
+             "shard_start": 0, "shard_end": 3, "max_loaded_shards": 2,
+             "vram_bytes": 20_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
+             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
+        )
+        _post_json(
+            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+            {"endpoint": "http://127.0.0.1:9060", "model": "model-a",
+             "hf_repo": "org/ModelA", "num_layers": 4,
+             "shard_start": 0, "shard_end": 3, "max_loaded_shards": 1,
+             "vram_bytes": 20_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
+             "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
+        )
+        assert tracker._stats is not None
+        for _ in range(400):
+            tracker._stats.record_request("org/ModelA")
+        with tracker._lock:
+            _scale_demanded_models_locked(tracker._server)  # type: ignore[arg-type]
+        node_b = tracker._registry[reg_b["node_id"]]
+        assignment = node_b.pending_new_assignment
+        assert assignment is not None
+        assert assignment["action"] == "ADD_SHARD"
+        assert assignment["model"] == "org/ModelA"
+    finally:
+        tracker.stop()