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

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