Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

tests/test_tracker_routing.py

@@ -671,6 +671,42 @@ def test_tracker_rebalances_after_middle_range_node_timeout():
         tracker.stop()
 
 
+def test_tracker_rebalances_managed_hf_node_after_peer_timeout():
+    """HF nodes auto-assigned by the tracker receive LOAD_SHARD after a peer dies."""
+    tracker = TrackerServer(heartbeat_timeout=0.15, rebalance_interval=10.0)
+    tracker_port = tracker.start()
+    try:
+        managed = _post_json(
+            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+            {"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct",
+             "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
+             "shard_start": 0, "shard_end": 21, "managed_assignment": True,
+             "vram_bytes": 1_000_000_000, "hardware_profile": {}, "score": 1.0},
+        )
+        expired = _post_json(
+            f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+            {"endpoint": "http://127.0.0.1:9102", "model": "Qwen2.5-0.5B-Instruct",
+             "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
+             "shard_start": 22, "shard_end": 23,
+             "vram_bytes": 1_000_000_000, "hardware_profile": {}, "score": 1.0},
+        )
+
+        time.sleep(0.10)
+        _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{managed['node_id']}/heartbeat", {})
+        time.sleep(0.10)
+
+        hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{managed['node_id']}/heartbeat", {})
+
+        assert expired["node_id"] not in tracker._registry
+        load_directives = [d for d in hb.get("directives", []) if d["action"] == "LOAD_SHARD"]
+        assert load_directives
+        assert load_directives[-1]["model"] == "Qwen/Qwen2.5-0.5B-Instruct"
+        assert load_directives[-1]["start_layer"] == 0
+        assert load_directives[-1]["end_layer"] == 23
+    finally:
+        tracker.stop()
+
+
 def test_tracker_route_error_no_nodes():
     """Tracker returns 503 with clear error when the registry is empty."""
     tracker = TrackerServer()
@@ -1399,3 +1435,47 @@ def test_route_timeout_config_is_exposed_on_server():
 
     node = TorchNodeServer(backend=_MinimalBackend(), route_timeout=45.0)
     assert node.route_timeout == 45.0
+
+
+def test_torch_node_applies_tracker_load_shard_directive(monkeypatch):
+    from meshnet_node import torch_server
+    from meshnet_node.torch_server import TorchNodeServer
+
+    class _MinimalBackend:
+        def __init__(self, model_id="Qwen/Qwen2.5-0.5B-Instruct", shard_start=0, shard_end=21, quantization="bfloat16"):
+            self.model_id = model_id
+            self.shard_start = shard_start
+            self.shard_end = shard_end
+            self.quantization = quantization
+            self.total_layers = 24
+            self.is_head = shard_start == 0
+            self.is_tail = shard_end == 23
+
+        def generate_text(self, *a, **kw): return ""
+        def count_prompt_tokens(self, *a): return 0
+        def count_text_tokens(self, *a): return 0
+
+    loaded = []
+
+    def fake_load_backend(model_id, shard_start, shard_end, quantization):
+        loaded.append((model_id, shard_start, shard_end, quantization))
+        return _MinimalBackend(model_id, shard_start, shard_end, quantization)
+
+    monkeypatch.setattr(torch_server, "_load_backend", fake_load_backend)
+    node = TorchNodeServer(backend=_MinimalBackend())
+
+    applied = node.apply_tracker_directives([
+        {"action": "DROP_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct", "shard_start": 0, "shard_end": 21},
+        {"action": "LOAD_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct", "shard_start": 0, "shard_end": 23,
+         "quantization": "bfloat16"},
+    ])
+
+    assert loaded == [("Qwen/Qwen2.5-0.5B-Instruct", 0, 23, "bfloat16")]
+    assert applied == {
+        "model": "Qwen/Qwen2.5-0.5B-Instruct",
+        "shard_start": 0,
+        "shard_end": 23,
+        "quantization": "bfloat16",
+        "tracker_mode": True,
+    }
+    assert node.backend.shard_end == 23