diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 6c36a7c..3b42106 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -227,12 +227,15 @@ def run_startup( "score": 1.0, "tracker_mode": (shard_start == 0), } + tracker_node_id: str | None = None try: reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload) - node_id = reg_resp.get("node_id", "?") - print(f" Registered with tracker — node ID: {node_id}", flush=True) - _start_heartbeat(tracker_url, node_id, reg_payload) + tracker_node_id = str(reg_resp.get("node_id") or "?") + setattr(node, "tracker_node_id", tracker_node_id) + print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True) + _start_heartbeat(tracker_url, tracker_node_id, reg_payload) except Exception as exc: + setattr(node, "tracker_node_id", None) print(f" Warning: tracker registration failed: {exc}", flush=True) print( @@ -243,6 +246,7 @@ def run_startup( f" Shard: {shard_label}\n" f" Quantization: {quantization}\n" f" Endpoint: {endpoint}\n" + f" Node ID: {tracker_node_id or 'unregistered'}\n" f" Hardware: {device.upper()}\n" f"{'=' * 32}", flush=True, @@ -298,12 +302,15 @@ def run_startup( "score": 1.0, "tracker_mode": (assigned_shard_start == 0), } + tracker_node_id = None try: reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", auto_reg_payload) - node_id = reg_resp.get("node_id", "?") - print(f" Registered with tracker — node ID: {node_id}", flush=True) - _start_heartbeat(tracker_url, node_id, auto_reg_payload) + tracker_node_id = str(reg_resp.get("node_id") or "?") + setattr(node, "tracker_node_id", tracker_node_id) + print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True) + _start_heartbeat(tracker_url, tracker_node_id, auto_reg_payload) except Exception as exc: + setattr(node, "tracker_node_id", None) print(f" Warning: tracker registration failed: {exc}", flush=True) shard_count = assigned_shard_end - assigned_shard_start + 1 print( @@ -315,6 +322,7 @@ def run_startup( f"({shard_count} of {assigned_num_layers})\n" f" Quantization: {quantization}\n" f" Endpoint: {endpoint}\n" + f" Node ID: {tracker_node_id or 'unregistered'}\n" f" Hardware: {device.upper()}\n" f"{'=' * 32}", flush=True, @@ -385,7 +393,8 @@ def run_startup( "score": 1.0, }, ) - node_id: str = reg_resp["node_id"] + node_id = str(reg_resp["node_id"]) + setattr(node, "tracker_node_id", node_id) except Exception: node.stop() raise diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 7c201bb..6bf678f 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -22,11 +22,11 @@ HTTP API contract: """ import http.server +import hashlib import json import threading import time import urllib.parse -import uuid from typing import Any from .gossip import NodeGossip @@ -230,6 +230,27 @@ def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]: ] +def _nodes_and_bounds_for_model( + server: "_TrackerHTTPServer", + model: str, +) -> tuple[list[_NodeEntry], int, int] | None: + preset = server.model_presets.get(model) + if preset is not None: + required_start, required_end = _preset_layer_bounds(preset) + return [node for node in server.registry.values() if node.model == model], required_start, required_end + + nodes = [ + node for node in server.registry.values() + if (node.hf_repo == model or node.model == model) + and node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + ] + if not nodes: + return None + return nodes, 0, max(node.num_layers for node in nodes) - 1 + + def _load_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict: return { "action": "LOAD_SHARD", @@ -369,6 +390,27 @@ def _registration_ban_error(contracts: Any | None, wallet_address: str | None) - return None +def _node_id_for_registration( + endpoint: str, + model: str, + wallet_address: str | None, + shard_start: int | None, + shard_end: int | None, + hf_repo: str | None, +) -> str: + wallet_prefix = wallet_address[:8] if wallet_address else "anon" + stable_key = "|".join([ + wallet_address or "", + endpoint.rstrip("/"), + model, + hf_repo or "", + "" if shard_start is None else str(shard_start), + "" if shard_end is None else str(shard_end), + ]) + digest = hashlib.sha256(stable_key.encode()).hexdigest()[:12] + return f"{wallet_prefix}-{digest}" + + class _TrackerHTTPServer(http.server.HTTPServer): def __init__( self, @@ -404,7 +446,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(body))) self.end_headers() - self.wfile.write(body) + try: + self.wfile.write(body) + except BrokenPipeError: + pass def _read_json_body(self) -> dict | None: length = int(self.headers.get("Content-Length", 0)) @@ -484,6 +529,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned ] data = [] + seen_ids: set[str] = set() for name, preset in server.model_presets.items(): model_nodes = [node for node in alive if node.model == name] if not model_nodes: @@ -494,25 +540,73 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): required_start, required_end, ) + aliases = [name] + hf_repo = preset.get("hf_repo") + if hf_repo and hf_repo not in aliases: + aliases.append(hf_repo) data.append({ "id": name, "object": "model", "created": created, "owned_by": "meshnet", + "name": name, + "hf_repo": hf_repo, + "aliases": aliases, "shard_coverage_percentage": coverage, }) + seen_ids.add(name) + + hf_model_ids = sorted({ + node.hf_repo or node.model + for node in alive + if node.model is not None + and node.model not in server.model_presets + and node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + }) + for model_id in hf_model_ids: + if model_id is None or model_id in seen_ids: + continue + model_nodes = [ + node for node in alive + if node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + and (node.hf_repo == model_id or (node.hf_repo is None and node.model == model_id)) + ] + if not model_nodes: + continue + short_names = sorted({node.model for node in model_nodes if node.model}) + aliases = [model_id, *[name for name in short_names if name != model_id]] + required_start = 0 + required_end = max(node.num_layers for node in model_nodes) - 1 + data.append({ + "id": model_id, + "object": "model", + "created": created, + "owned_by": "meshnet", + "name": short_names[0] if short_names else model_id, + "hf_repo": model_id if any(node.hf_repo == model_id for node in model_nodes) else None, + "aliases": aliases, + "shard_coverage_percentage": _coverage_percentage( + model_nodes, + required_start, + required_end, + ), + }) + seen_ids.add(model_id) self._send_json(200, {"object": "list", "data": data}) def _handle_coverage(self, model: str): server: _TrackerHTTPServer = self.server # type: ignore[assignment] - preset = server.model_presets.get(model) - if preset is None: - self._send_json(404, {"error": f"unknown model preset: {model!r}"}) - return - required_start, required_end = _preset_layer_bounds(preset) with server.lock: self._purge_expired_nodes() - alive = [node for node in server.registry.values() if node.model == model] + resolved = _nodes_and_bounds_for_model(server, model) + if resolved is None: + self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) + return + alive, required_start, required_end = resolved if server.contracts is not None: alive = [ node for node in alive @@ -821,7 +915,14 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._send_json(400, {"error": "num_layers must be an integer"}) return - node_id = str(uuid.uuid4()) + node_id = _node_id_for_registration( + endpoint, + model, + wallet_address, + shard_start, + shard_end, + hf_repo, + ) entry = _NodeEntry( node_id=node_id, endpoint=endpoint.rstrip("/"), @@ -1191,6 +1292,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "route": [e.endpoint for e in route], "nodes": [ { + "node_id": e.node_id, "endpoint": e.endpoint, "wallet_address": e.wallet_address, "shard_start": e.shard_start, @@ -1250,6 +1352,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "route": [e.endpoint for e in route], "nodes": [ { + "node_id": e.node_id, "endpoint": e.endpoint, "wallet_address": e.wallet_address, "shard_start": e.shard_start, diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index b6f3127..e170dbe 100644 --- a/tests/test_node_startup.py +++ b/tests/test_node_startup.py @@ -370,6 +370,11 @@ def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, ca lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0}, ) monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer) + monkeypatch.setattr( + startup_mod, + "_post_json", + lambda _url, _payload, timeout=10.0: {"node_id": "node-test-123"}, + ) node = run_startup( tracker_url="http://127.0.0.1:8080", @@ -380,8 +385,10 @@ def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, ca ) assert node.backend.total_layers == 24 + assert node.tracker_node_id == "node-test-123" output = capsys.readouterr().out assert "Shard: layers 0–23; 24 of 24" in output + assert "Node ID: node-test-123" in output # --------------------------------------------------------------------------- diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index ab63a31..9508615 100644 --- a/tests/test_tracker_routing.py +++ b/tests/test_tracker_routing.py @@ -27,6 +27,69 @@ def _get_json(url: str) -> dict: return json.loads(r.read()) +def test_tracker_send_json_ignores_broken_pipe_after_client_disconnect(): + """A disconnected client must not dump a BrokenPipe traceback from the tracker.""" + from meshnet_tracker.server import _TrackerHandler + + class BrokenPipeWriter: + def write(self, _body): + raise BrokenPipeError + + class DummyHandler: + wfile = BrokenPipeWriter() + + def send_response(self, _status): + pass + + def send_header(self, _name, _value): + pass + + def end_headers(self): + pass + + _TrackerHandler._send_json(DummyHandler(), 200, {"ok": True}) + + +def test_tracker_registration_node_id_includes_wallet_prefix_and_stable_suffix(): + tracker = TrackerServer() + tracker_port = tracker.start() + wallet = "7j77FsPY1evV8tuf4Z73AVrWwxBEW1pvKwi4EvcRD3g" + try: + first = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9100", "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24, + "shard_start": 0, "shard_end": 21, "wallet_address": wallet, + "hardware_profile": {}, "score": 1.0}, + ) + second = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9100", "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24, + "shard_start": 0, "shard_end": 21, "wallet_address": wallet, + "hardware_profile": {}, "score": 1.0}, + ) + different_endpoint = _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": 20, "shard_end": 23, "wallet_address": wallet, + "hardware_profile": {}, "score": 1.0}, + ) + + assert first["node_id"].startswith("7j77FsPY-") + assert second["node_id"] == first["node_id"] + assert different_endpoint["node_id"].startswith("7j77FsPY-") + assert different_endpoint["node_id"] != first["node_id"] + + route_resp = _get_json( + f"http://127.0.0.1:{tracker_port}/v1/route?model=Qwen/Qwen2.5-0.5B-Instruct" + ) + assert route_resp["nodes"][0]["node_id"] == first["node_id"] + finally: + tracker.stop() + + def test_tracker_node_registration(): """A node can register with the tracker and receives a node_id.""" tracker = TrackerServer() @@ -131,6 +194,64 @@ def test_tracker_coverage_endpoint_reports_uncovered_ranges(): tracker.stop() +def test_tracker_coverage_endpoint_accepts_registered_hf_repo_or_short_name(): + """Coverage endpoint supports real HF models registered outside preset catalog.""" + tracker = TrackerServer() + tracker_port = tracker.start() + try: + _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, "hardware_profile": {}, "score": 1.0}, + ) + _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": 20, "shard_end": 23, "hardware_profile": {}, "score": 1.0}, + ) + + by_repo = _get_json( + f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen/Qwen2.5-0.5B-Instruct" + ) + by_short_name = _get_json( + f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen2.5-0.5B-Instruct" + ) + + assert by_repo["model"] == "Qwen/Qwen2.5-0.5B-Instruct" + assert by_repo["coverage"] == [ + {"start_layer": 0, "end_layer": 19, "node_count": 1}, + {"start_layer": 20, "end_layer": 21, "node_count": 2}, + {"start_layer": 22, "end_layer": 23, "node_count": 1}, + ] + assert by_short_name["coverage"] == by_repo["coverage"] + finally: + tracker.stop() + + +def test_tracker_models_endpoint_lists_registered_hf_repo_and_short_name_alias(): + tracker = TrackerServer() + tracker_port = tracker.start() + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9111", "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24, + "shard_start": 0, "shard_end": 23, "hardware_profile": {}, "score": 1.0}, + ) + + models_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/models") + + model = next(item for item in models_resp["data"] if item["id"] == "Qwen/Qwen2.5-0.5B-Instruct") + assert model["name"] == "Qwen2.5-0.5B-Instruct" + assert model["hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct" + assert model["aliases"] == ["Qwen/Qwen2.5-0.5B-Instruct", "Qwen2.5-0.5B-Instruct"] + assert model["shard_coverage_percentage"] == 100.0 + finally: + tracker.stop() + + def test_tracker_auto_assigns_new_node_to_uncovered_range_first(): """Capability-driven registration fills the first uncovered layer gap.""" tracker = TrackerServer(model_presets={