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

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:       <address>
+  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:      <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/<peer_id>` 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

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,

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

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")

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

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