diff --git a/QUICKSTART.md b/QUICKSTART.md index 91411e2..a393d86 100644 --- a/QUICKSTART.md +++ b/QUICKSTART.md @@ -48,13 +48,10 @@ python3 -m venv .venv If `.venv/bin/meshnet-node` is missing, the editable install step did not finish successfully. Re-run the `.venv/bin/pip install -e ...` command above inside WSL. -WSL2 is still useful for local development, but do not rely on it for the -"another machine connects back to this node" LAN case. WSL2 commonly sits behind -Windows NAT/port-proxy behavior and may not accept inbound traffic from other LAN -machines without extra host networking setup. We intentionally leave that unfixed -because it is useful for testing NAT/relay scenarios. If you just want to bring up -a Windows node that other machines can reach directly, run the node in native -Windows PowerShell instead. +WSL2 sits behind Windows NAT and is **not directly reachable** from other LAN machines. +Direct cross-host hops time out. The relay path (see below) solves this: the WSL2 node +opens an outbound WebSocket to the relay server and all traffic flows through that tunnel. +No firewall rules, no `--advertise-host` needed — just point at the public tracker URL. ### Native Windows PowerShell node (not WSL) @@ -126,16 +123,21 @@ $env:HF_HOME = "D:\DEV\models" --port 8005 ``` -One-line variants: +One-line variants (direct LAN — node must be reachable by IP from other machines): ```powershell .\.venv\Scripts\meshnet-node.exe start --tracker http://192.168.0.179:8081 --model Qwen/Qwen2.5-0.5B-Instruct --advertise-host 192.168.0.20 -.\.venv\Scripts\meshnet-node.exe start --tracker http://ai.neuron.d-popov.com --model Qwen/Qwen2.5-0.5B-Instruct --advertise-host 192.168.0.20 +``` + +Via public hostname with relay (works from behind NAT, WSL2, 5G — no `--advertise-host` needed): + +```powershell +.\.venv\Scripts\meshnet-node.exe start --tracker https://ai.neuron.d-popov.com --model Qwen/Qwen2.5-0.5B-Instruct ``` `--host 0.0.0.0` binds the node to all Windows interfaces. `--advertise-host` -is what the tracker gives to other nodes, so it must be the Windows LAN IP that -the tracker and peer nodes can actually reach. +is what the tracker gives to other nodes for direct connections; omit it when using +the relay path since all traffic flows through the relay tunnel instead. If you want verbose per-hop pipeline logs while debugging a split model, add `--debug`. Leave it off for normal runs; otherwise every generated token logs @@ -144,6 +146,7 @@ lines like: ```text [node] pipeline hop 0: http://127.0.0.1:8005 start_layer=22 [node] pipeline hop 0 returned text=' token' + [node] pipeline hop 1: wss://ai.neuron.d-popov.com/rpc/abc123 relay start_layer=12 ``` 6. From the tracker machine, verify Windows is reachable: @@ -152,32 +155,47 @@ lines like: curl http://192.168.0.42:8005/v1/health ``` -If that endpoint returns 404, that is okay: it still proves the TCP connection -reached the node process. If it times out or connection-refuses, check the -Windows Firewall rule, `--host 0.0.0.0`, the selected LAN IP, and that the node is -still running. +If that endpoint returns 404 or 501, that is okay: it still proves the TCP +connection reached the node process. If it times out or connection-refuses, check +the Windows Firewall rule, `--host 0.0.0.0`, the selected LAN IP, and that the +node is still running. -### Public tracker + WSS relay +--- -For internet nodes, expose one public HTTPS host and proxy these paths: +## Public tracker + relay (internet / NAT nodes) -```text -/v1/* -> meshnet-tracker, for registration, heartbeats, routing, and OpenAI requests -/ws -> meshnet-relay, for outbound node gossip/bridge connections -/rpc/* -> meshnet-relay, for tracker-to-node relay requests +This setup lets nodes connect from anywhere — behind home NAT, 5G, WSL2, or +on a different continent — without opening firewall ports. + +### Architecture + +``` +Client → HTTPS → ai.neuron.d-popov.com (nginx) + ├─ /v1/* → meshnet-tracker :8081 + ├─ /ws → meshnet-relay :8765 (node persistent outbound WS) + └─ /rpc/* → meshnet-relay :8765 (caller opens WS per hop) ``` -Start the tracker with the public relay URL it should advertise: +### Start the relay and tracker (server side) ```bash +# Terminal 1 — relay (WebSocket hub) .venv/bin/meshnet-relay --host 0.0.0.0 --port 8765 + +# Terminal 2 — tracker (advertises relay URL to nodes) .venv/bin/meshnet-tracker start \ --host 0.0.0.0 \ --port 8081 \ --relay-url wss://ai.neuron.d-popov.com/ws ``` -Then a node only needs the public tracker address: +The `--relay-url` flag embeds the relay address in `/v1/network/map`. Every node +queries that endpoint on startup and auto-connects if a relay URL is present. + +### Start a node (any machine, any network) + +No `--advertise-host` needed. The node discovers the relay URL from the tracker +and opens a persistent outbound WebSocket: ```bash .venv/bin/meshnet-node start \ @@ -201,6 +219,84 @@ curl -s "https://ai.neuron.d-popov.com/v1/network/map" | python3 -m json.tool curl -s "https://ai.neuron.d-popov.com/v1/route?model=qwen2.5-0.5b" | python3 -m json.tool ``` +Expected startup output (relay path): + +``` + Auto-detected 24 layers → shard 0–23 + Relay connected — wss://ai.neuron.d-popov.com/rpc/abc1def2ef3f4567 +================================ +meshnet-node ready + Wallet:
+ Model ID: Qwen/Qwen2.5-0.5B-Instruct + Shard: layers 0–23; 24 of 24 + Quantization: bfloat16 + Endpoint: http://172.29.104.23:7001 + Node ID: + Hardware: CPU +================================ +``` + +The `Endpoint` shown is the local IP (unreachable from outside). Other nodes reach +this one via `wss://ai.neuron.d-popov.com/rpc/` instead. + +### How relay hops work + +When node A needs to forward activations to node B (behind NAT): + +1. Tracker injects `X-Meshnet-Route` with `relay_addr` for each behind-NAT hop. +2. Node A opens a WebSocket to `wss://relay/rpc/{peer_id_B}`. +3. Relay forwards the `relay-http-request` envelope to Node B's persistent connection. +4. Node B processes `/forward` locally, returns `relay-http-response`. +5. Relay sends the response back to Node A over the same WebSocket. +6. Node A closes the WebSocket and continues the pipeline. + +Binary activation tensors (bfloat16) are Base64-encoded through the relay JSON +protocol and decoded on both sides — no precision loss. + +If the relay hop fails (relay down, peer disconnected), the node logs a warning and +falls back to a direct HTTP attempt before returning an error. + +### Test from WSL2 using the public tracker + +In WSL2 (which gets a `172.x.x.x` virtual IP — unreachable from other machines): + +```bash +# WSL2 Terminal 1 — head node (layers 0–11, handles chat requests) +.venv/bin/meshnet-node start \ + --tracker https://ai.neuron.d-popov.com \ + --model Qwen/Qwen2.5-0.5B-Instruct \ + --shard-start 0 --shard-end 11 + +# WSL2 Terminal 2 — tail node (layers 12–23) +.venv/bin/meshnet-node start \ + --tracker https://ai.neuron.d-popov.com \ + --model Qwen/Qwen2.5-0.5B-Instruct \ + --shard-start 12 --shard-end 23 +``` + +Both nodes connect to the relay automatically. When a chat request arrives at Node A, +it forwards activations to Node B via `wss://ai.neuron.d-popov.com/rpc/{peer_id_B}`. + +Send inference through the tracker (which picks the head node and injects the route): + +```bash +curl -s https://ai.neuron.d-popov.com/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen2.5-0.5B-Instruct", + "messages": [{"role": "user", "content": "What is 7 times 8?"}], + "stream": false + }' | python3 -m json.tool +``` + +Or send directly to Node A's local port (within WSL): + +```bash +curl -s http://localhost:7001/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{"model": "Qwen/Qwen2.5-0.5B-Instruct", "messages": [{"role": "user", "content": "Hi"}]}' +``` + --- ## Step 1 — Start the tracker (Terminal 1) @@ -338,7 +434,7 @@ HF_HOME=/run/media/popov/d/DEV/models \ --port 8002 ``` -Send the request to Node A — it tokenizes, runs layers 0–13, passes binary +Send the request to Node A — it tokenizes, runs layers 0–11, passes binary activations to Node B, and streams the final response back. --- @@ -382,18 +478,6 @@ tracker with a relay URL so the node registers a `relay_addr`. --- -## Start the relay node (for NAT traversal) - -```bash -.venv/bin/pip install -e packages/relay -.venv/bin/meshnet-relay --port 8765 -``` - -Nodes behind NAT connect to the relay and advertise their relay address to the -tracker. See `docs/adr/0010-p2p-gossip-and-nat-relay.md`. - ---- - ## Run all tests ```bash diff --git a/packages/node/meshnet_node/relay_bridge.py b/packages/node/meshnet_node/relay_bridge.py index d7b63ea..4127160 100644 --- a/packages/node/meshnet_node/relay_bridge.py +++ b/packages/node/meshnet_node/relay_bridge.py @@ -2,6 +2,7 @@ from __future__ import annotations +import base64 import json import logging import threading @@ -114,20 +115,35 @@ class RelayHttpBridge: request_id = str(payload.get("request_id") or "") method = str(payload.get("method") or "POST").upper() path = str(payload.get("path") or "/") - body_text = payload.get("body") or "" headers = payload.get("headers") if isinstance(payload.get("headers"), dict) else {} + # body_base64 carries binary data (e.g. bfloat16 activation tensors) safely. + # Fallback to text "body" for backward-compat with non-binary requests. + body_b64 = payload.get("body_base64") + if body_b64: + data = base64.b64decode(body_b64) + else: + body_text = payload.get("body") or "" + data = body_text.encode() if isinstance(body_text, str) else bytes(body_text) + url = f"{self.local_base_url}{path}" - data = body_text.encode() if isinstance(body_text, str) else bytes(body_text) req = urllib.request.Request(url, data=data, headers=headers, method=method) try: with urllib.request.urlopen(req, timeout=300.0) as resp: - return { + resp_bytes = resp.read() + resp_headers = dict(resp.headers) + # Forward all X-Meshnet-* headers so the caller can reconstruct the activation. + is_binary = "octet-stream" in resp.headers.get("Content-Type", "") + result: dict = { "request_id": request_id, "status": resp.status, - "headers": {"Content-Type": resp.headers.get("Content-Type", "application/json")}, - "body": resp.read().decode(errors="replace"), + "headers": resp_headers, } + if is_binary: + result["body_base64"] = base64.b64encode(resp_bytes).decode() + else: + result["body"] = resp_bytes.decode(errors="replace") + return result except urllib.error.HTTPError as exc: return { "request_id": request_id, diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 0c4cedb..e2c5c0c 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -123,6 +123,26 @@ def _start_heartbeat( except Exception: return False + def _apply_directives(directives: list[dict]) -> None: + if not directives: + return + if node_ref is None or not hasattr(node_ref, "apply_tracker_directives"): + print(f" [node] tracker directives received: {directives}", flush=True) + return + try: + applied = node_ref.apply_tracker_directives(directives) + except Exception as exc: + print(f" [node] WARNING: failed to apply tracker directives: {exc}", flush=True) + return + if applied: + model_id = applied.get("model", register_payload.get("hf_repo") or register_payload.get("model")) + register_payload["model"] = str(model_id).split("/")[-1] + register_payload["hf_repo"] = model_id + register_payload["shard_start"] = applied["shard_start"] + register_payload["shard_end"] = applied["shard_end"] + register_payload["quantization"] = applied.get("quantization", register_payload.get("quantization")) + register_payload["tracker_mode"] = bool(applied.get("tracker_mode", False)) + def _loop() -> None: nonlocal node_id hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat" @@ -150,6 +170,7 @@ def _start_heartbeat( try: resp = _post_json(hb_url, _get_stats()) + _apply_directives(resp.get("directives", [])) new_asgn = resp.get("new_assignment") if new_asgn: print( @@ -282,6 +303,7 @@ def run_startup( print(f" {probationary_line}", flush=True) if model_id: # treat "" the same as None — no explicit model given + user_pinned_shard = shard_start is not None or shard_end is not None # Auto-detect shard range from model config if not explicitly provided if shard_start is None or shard_end is None: detected = _detect_num_layers(model_id) @@ -355,6 +377,7 @@ def run_startup( "quantization": quantization, "score": 1.0, "tracker_mode": (shard_start == 0), + "managed_assignment": not user_pinned_shard, **relay_fields, } tracker_node_id: str | None = None @@ -442,6 +465,7 @@ def run_startup( "quantization": quantization, "score": 1.0, "tracker_mode": (assigned_shard_start == 0), + "managed_assignment": True, **relay_fields, } tracker_node_id = None diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 036f17b..1f64c2c 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -2,6 +2,7 @@ from __future__ import annotations +import base64 import http.server import json import sys @@ -29,6 +30,40 @@ from .server import ( ) +def _relay_hop( + relay_addr: str, + path: str, + body: bytes, + headers: dict[str, str], + timeout: float = 120.0, +) -> tuple[int, dict[str, str], bytes]: + """Send a single HTTP-shaped request through a relay RPC WebSocket. + + relay_addr is the wss://relay.../rpc/{peer_id} URL. + Returns (status, response_headers_lower, response_body). + Raises on connection failure so callers can fall back to direct. + """ + import websockets.sync.client as wsc # type: ignore[import] + + request_id = f"{time.time_ns():x}" + payload = json.dumps({ + "request_id": request_id, + "method": "POST", + "path": path, + "headers": headers, + "body_base64": base64.b64encode(body).decode(), + }) + with wsc.connect(relay_addr, open_timeout=timeout) as ws: + ws.send(payload) + raw = ws.recv(timeout=timeout) + resp = json.loads(raw) + status = int(resp.get("status", 503)) + resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()} + body_b64 = resp.get("body_base64") + resp_body = base64.b64decode(body_b64) if body_b64 else (resp.get("body") or "").encode() + return status, resp_headers, resp_body + + class _TorchHTTPServer(http.server.HTTPServer): def __init__( self, @@ -326,8 +361,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): result_text = "".join(generated) self._send_openai_response(result_text, model_name, stream, messages) - def _get_remaining_route(self, model: str) -> list[tuple[str, int]]: - """Return downstream hops as (endpoint, start_layer) pairs. + def _get_remaining_route(self, model: str) -> list[dict]: + """Return downstream hops as dicts with endpoint, start_layer, and optional relay_addr. Fast path reads X-Meshnet-Route header injected by the tracker. Slow path queries the tracker's /v1/route endpoint as a fallback. @@ -340,13 +375,19 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): try: route = json.loads(injected) if isinstance(route, list): - hops: list[tuple[str, int]] = [] + hops: list[dict] = [] for item in route: if isinstance(item, dict): - hops.append((str(item["endpoint"]), int(item.get("start_layer", 0)))) + hop = { + "endpoint": str(item["endpoint"]), + "start_layer": int(item.get("start_layer", 0)), + } + if item.get("relay_addr"): + hop["relay_addr"] = str(item["relay_addr"]) + hops.append(hop) elif isinstance(item, str): - hops.append((item, 0)) # backward-compat: plain string, no start_layer - print(f" [node] using injected downstream route: {[ep for ep, _ in hops]}", flush=True) + hops.append({"endpoint": item, "start_layer": 0}) + print(f" [node] using injected downstream route: {[h['endpoint'] for h in hops]}", flush=True) return hops except (json.JSONDecodeError, TypeError, KeyError): pass @@ -362,7 +403,6 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): route_resp = json.loads(r.read()) own_port = server.server_address[1] nodes_info = route_resp.get("nodes", []) - # nodes_info is ordered; find own node and compute start_layers post-hoc hops = [] covered_up_to: int | None = None for node_info in nodes_info: @@ -371,18 +411,19 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): covered_up_to = node_info.get("shard_end") continue if covered_up_to is None: - # Own node not found yet; use node's shard_start as fallback covered_up_to = (node_info.get("shard_start") or 1) - 1 - start_l = covered_up_to + 1 - hops.append((ep, start_l)) + hop = {"endpoint": ep, "start_layer": covered_up_to + 1} + if node_info.get("relay_addr"): + hop["relay_addr"] = str(node_info["relay_addr"]) + hops.append(hop) covered_up_to = node_info.get("shard_end", covered_up_to) - print(f" [node] tracker downstream route: {[ep for ep, _ in hops]}", flush=True) + print(f" [node] tracker downstream route: {[h['endpoint'] for h in hops]}", flush=True) return hops except Exception as exc: print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True) return [] - def _run_downstream_pipeline(self, payload: object, route: list[tuple[str, int]]) -> str: + def _run_downstream_pipeline(self, payload: object, route: list[dict]) -> str: server: _TorchHTTPServer = self.server # type: ignore[assignment] if not route: # Partial shard at tail: decode the activation from the previous node. @@ -407,9 +448,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): current_attn = attn_mask current_pos = pos_ids - for hop_index, (node_url, start_layer) in enumerate(route): + for hop_index, hop in enumerate(route): + node_url = hop["endpoint"] + start_layer = hop.get("start_layer", 0) + relay_addr = hop.get("relay_addr") if server.debug: - print(f" [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}", flush=True) + print( + f" [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}" + + (f" relay={relay_addr}" if relay_addr else ""), + flush=True, + ) headers: dict[str, str] = { "Content-Type": "application/octet-stream", "X-Meshnet-Wire": _WIRE_VERSION, @@ -425,19 +473,38 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): headers["X-Meshnet-Attn-Mask"] = current_attn if current_pos: headers["X-Meshnet-Position-Ids"] = current_pos - req = urllib.request.Request( - f"{node_url}/forward", - data=current_body, - headers=headers, - method="POST", - ) - try: - with urllib.request.urlopen(req, timeout=120.0) as r: - resp_body = r.read() - resp_headers = {k.lower(): v for k, v in r.headers.items()} - except Exception as exc: - print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True) - return f"pipeline error at {node_url}: {exc}" + if relay_addr: + try: + status, resp_headers, resp_body = _relay_hop( + relay_addr, "/forward", current_body, headers, timeout=120.0, + ) + if status >= 400: + print( + f" [node] relay hop {hop_index} returned {status} from {relay_addr}", + flush=True, + ) + return f"pipeline error at {node_url} via relay: status {status}" + except Exception as exc: + print( + f" [node] relay hop {hop_index} failed at {relay_addr}: {exc}; " + f"falling back to direct {node_url}", + flush=True, + ) + relay_addr = None # fall through to direct + if not relay_addr: + req = urllib.request.Request( + f"{node_url}/forward", + data=current_body, + headers=headers, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=120.0) as r: + resp_body = r.read() + resp_headers = {k.lower(): v for k, v in r.headers.items()} + except Exception as exc: + print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True) + return f"pipeline error at {node_url}: {exc}" content_type = resp_headers.get("content-type", "") if "application/json" in content_type: try: @@ -662,6 +729,40 @@ class TorchNodeServer: def queue_depth(self) -> int: return self._server.queue_depth if self._server is not None else 0 + def apply_tracker_directives(self, directives: list[dict]) -> dict | None: + """Apply tracker LOAD_SHARD directives by hot-swapping the loaded backend.""" + load_directive = next( + (directive for directive in reversed(directives) if directive.get("action") == "LOAD_SHARD"), + None, + ) + if load_directive is None: + return None + shard_start = int(load_directive["shard_start"]) + shard_end = int(load_directive["shard_end"]) + quantization = str(load_directive.get("quantization") or self._backend.quantization) + model_id = str(load_directive.get("model") or self._backend.model_id) + print( + f" [node] loading reassigned shard: {model_id} layers {shard_start}-{shard_end}", + flush=True, + ) + new_backend = _load_backend(model_id, shard_start, shard_end, quantization) + self._backend = new_backend + self._tracker_mode = shard_start == 0 + if self._server is not None: + self._server.backend = new_backend + self._server.tracker_mode = self._tracker_mode + print( + f" [node] loaded reassigned shard: {model_id} layers {shard_start}-{shard_end}", + flush=True, + ) + return { + "model": model_id, + "shard_start": shard_start, + "shard_end": shard_end, + "quantization": quantization, + "tracker_mode": self._tracker_mode, + } + def start(self) -> int: if self._server is not None: raise RuntimeError("TorchNodeServer is already running") diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 97eaf0d..927b8a6 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -706,9 +706,109 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None: ) +def _hf_rebalance_preset(nodes: list[_NodeEntry]) -> dict: + total_layers = max(node.num_layers or 0 for node in nodes) + return { + "layers_start": 0, + "layers_end": total_layers - 1, + "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024}, + } + + +def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> None: + model_nodes = [ + node for node in server.registry.values() + if node.hf_repo == hf_repo + and node.num_layers is not None + ] + managed_nodes = [node for node in model_nodes if node.managed_assignment] + if not model_nodes or not managed_nodes: + return + + preset = _hf_rebalance_preset(model_nodes) + required_start, required_end = _preset_layer_bounds(preset) + total_layers = required_end - required_start + 1 + if total_layers <= 0: + return + + previous_ranges = { + node.node_id: (node.shard_start, node.shard_end, node.quantization) + for node in managed_nodes + } + for node in managed_nodes: + node.shard_start = None + node.shard_end = None + + managed_nodes.sort( + key=lambda node: ( + -node.benchmark_tokens_per_sec, + -_node_layer_capacity(node, preset), + node.node_id, + ) + ) + base_nodes = [node for node in model_nodes if not node.managed_assignment] + coverage = _coverage_map(base_nodes, required_start, required_end) + gaps = _coverage_gaps(coverage) + if not gaps: + gaps = [(required_start, required_end)] + + eligible_nodes = [ + node for node in managed_nodes + if _node_layer_capacity(node, preset) > 0 + ] + node_index = 0 + for gap_start, gap_end in gaps: + cursor = gap_start + while cursor <= gap_end and node_index < len(eligible_nodes): + node = eligible_nodes[node_index] + remaining_layers = gap_end - cursor + 1 + remaining_nodes_after = len(eligible_nodes) - node_index - 1 + capacity = min( + _node_layer_capacity(node, preset), + total_layers, + max(1, remaining_layers - remaining_nodes_after), + ) + if capacity <= 0: + node_index += 1 + continue + quantization = _node_quantization(node, preset) + node.quantization = quantization + node.shard_start = cursor + node.shard_end = min(gap_end, cursor + capacity - 1) + cursor = node.shard_end + 1 + node_index += 1 + + for node in managed_nodes: + previous_start, previous_end, previous_quantization = previous_ranges[node.node_id] + current_range = (node.shard_start, node.shard_end, node.quantization) + if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]: + continue + if previous_start is not None and previous_end is not None: + node.pending_directives.append( + _drop_directive( + node, + hf_repo, + previous_start, + previous_end, + previous_quantization or _node_quantization(node, preset), + ) + ) + node.pending_directives.append( + _load_directive( + node, + hf_repo, + node.shard_start, + node.shard_end, + node.quantization or _node_quantization(node, preset), + ) + ) + + def _rebalance_all_locked(server: "_TrackerHTTPServer") -> None: for model in list(server.model_presets): _rebalance_model_locked(server, model) + for hf_repo in sorted({node.hf_repo for node in server.registry.values() if node.hf_repo}): + _rebalance_hf_model_locked(server, hf_repo) def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None: @@ -1095,7 +1195,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): covered_up_to = rs - 1 route_hops: list[dict] = [] for rn in route_nodes: - route_hops.append({"endpoint": rn.endpoint, "start_layer": covered_up_to + 1}) + hop: dict = {"endpoint": rn.endpoint, "start_layer": covered_up_to + 1} + if rn.relay_addr: + hop["relay_addr"] = rn.relay_addr + route_hops.append(hop) covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to # Strip the first-shard node we're about to proxy to — it's already handling the request. downstream_hops = [h for h in route_hops if h["endpoint"].rstrip("/") != node.endpoint.rstrip("/")] @@ -1332,6 +1435,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): return tracker_mode = bool(body.get("tracker_mode", False)) + managed_assignment = bool(body.get("managed_assignment", False)) hf_repo = body.get("hf_repo") if hf_repo is not None and not isinstance(hf_repo, str): self._send_json(400, {"error": "hf_repo must be a string"}) @@ -1371,7 +1475,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): quantizations=quantizations, benchmark_tokens_per_sec=benchmark_tokens_per_sec, quantization=quantization, - managed_assignment=not explicit_shard, + managed_assignment=managed_assignment or not explicit_shard, tracker_mode=tracker_mode, hf_repo=hf_repo, num_layers=num_layers, @@ -1395,7 +1499,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): ) server.registry[node_id] = entry if entry.managed_assignment: - _rebalance_model_locked(server, model) + if entry.hf_repo: + _rebalance_hf_model_locked(server, entry.hf_repo) + else: + _rebalance_model_locked(server, model) assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None if assignment_directive is not None: entry.pending_directives.clear() @@ -1459,7 +1566,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): entry.uptime_seconds = float(body["uptime_seconds"]) if "status" in body and body["status"] in ("ready", "loading"): entry.status = body["status"] - _rebalance_model_locked(server, entry.model or "stub-model") + if entry.hf_repo: + _rebalance_hf_model_locked(server, entry.hf_repo) + else: + _rebalance_model_locked(server, entry.model or "stub-model") directives = list(entry.pending_directives) entry.pending_directives.clear() new_assignment = entry.pending_new_assignment diff --git a/tests/test_mining_cli.py b/tests/test_mining_cli.py index e26331f..f52f498 100644 --- a/tests/test_mining_cli.py +++ b/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 diff --git a/tests/test_real_model_backend.py b/tests/test_real_model_backend.py index 2757670..de3b740 100644 --- a/tests/test_real_model_backend.py +++ b/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") diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index 6f3c4c4..bd1d3fe 100644 --- a/tests/test_tracker_routing.py +++ b/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