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

tests/test_mining_cli.py

@@ -452,6 +452,7 @@ def test_default_cli_passes_advertise_host(monkeypatch):
         "meshnet-node",
         "--tracker", "http://192.168.0.179:8081",
         "--advertise-host", "192.168.0.42",
+        "--debug",
         "--no-tui",
     ])
 
@@ -465,3 +466,4 @@ def test_default_cli_passes_advertise_host(monkeypatch):
 
     assert captured["tracker_url"] == "http://192.168.0.179:8081"
     assert captured["advertise_host"] == "192.168.0.42"
+    assert captured["debug"] is True

tests/test_real_model_backend.py

@@ -81,6 +81,38 @@ class _FakeFullBackend(_FakeBackend):
         return 1
 
 
+class _FakeChatTokenizer:
+    eos_token = ""
+
+    def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
+        assert add_generation_prompt is True
+        assert tokenize is False
+        return "debug prompt"
+
+
+class _FakePipelineHeadBackend(_FakeBackend):
+    tokenizer = _FakeChatTokenizer()
+
+    def encode_prompt(self, prompt: str) -> TensorPayload:
+        assert prompt == "debug prompt"
+        return TensorPayload(
+            body=b"\x00" * (1 * 6 * 8 * 2),
+            shape=[1, 6, 8],
+            attention_mask_header=None,
+            position_ids_header=None,
+        )
+
+
+class _FakePipelineTailBackend(_FakeTailBackend):
+    def __init__(self) -> None:
+        self.start_layers: list[int | None] = []
+
+    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
+        self.start_layers.append(start_layer)
+        assert len(body) == 1 * 6 * 8 * 2
+        return " token"
+
+
 def test_quantization_flag_validation():
     assert validate_quantization("bfloat16") == "bfloat16"
     assert validate_quantization("int8") == "int8"
@@ -198,6 +230,75 @@ def test_full_model_chat_completion_uses_generation_not_single_token_decode():
         node.stop()
 
 
+def test_pipeline_hop_logs_are_suppressed_without_debug(capsys):
+    tail_backend = _FakePipelineTailBackend()
+    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
+    tail = TorchNodeServer(backend=tail_backend)
+    head_port = head.start()
+    tail_port = tail.start()
+    try:
+        payload = json.dumps({
+            "model": "fake-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "max_tokens": 1,
+        }).encode()
+        req = urllib.request.Request(
+            f"http://127.0.0.1:{head_port}/v1/chat/completions",
+            data=payload,
+            headers={
+                "Content-Type": "application/json",
+                "X-Meshnet-Route": json.dumps([
+                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
+                ]),
+            },
+            method="POST",
+        )
+        with urllib.request.urlopen(req, timeout=5) as resp:
+            body = json.loads(resp.read())
+    finally:
+        head.stop()
+        tail.stop()
+
+    out = capsys.readouterr().out
+    assert body["choices"][0]["message"]["content"] == " token"
+    assert tail_backend.start_layers == [22]
+    assert "pipeline hop 0:" not in out
+    assert "pipeline hop 0 returned text" not in out
+
+
+def test_pipeline_hop_logs_are_enabled_with_debug(capsys):
+    head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True, debug=True)
+    tail = TorchNodeServer(backend=_FakePipelineTailBackend())
+    head_port = head.start()
+    tail_port = tail.start()
+    try:
+        payload = json.dumps({
+            "model": "fake-model",
+            "messages": [{"role": "user", "content": "hello"}],
+            "max_tokens": 1,
+        }).encode()
+        req = urllib.request.Request(
+            f"http://127.0.0.1:{head_port}/v1/chat/completions",
+            data=payload,
+            headers={
+                "Content-Type": "application/json",
+                "X-Meshnet-Route": json.dumps([
+                    {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
+                ]),
+            },
+            method="POST",
+        )
+        with urllib.request.urlopen(req, timeout=5) as resp:
+            json.loads(resp.read())
+    finally:
+        head.stop()
+        tail.stop()
+
+    out = capsys.readouterr().out
+    assert f"  [node] pipeline hop 0: http://127.0.0.1:{tail_port} start_layer=22" in out
+    assert "  [node] pipeline hop 0 returned text=' token'" in out
+
+
 def test_int_tensor_header_serializes_torch_tensors():
     torch = pytest.importorskip("torch")
 

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