diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index 68eb51b..73584b8 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/packages/node/meshnet_node/model_backend.py @@ -145,6 +145,7 @@ class TorchModelShard: shape: list[int], attention_mask_header: str | None, position_ids_header: str | None, + start_layer: int | None = None, ) -> TensorPayload | str: hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to( self.device @@ -155,7 +156,9 @@ class TorchModelShard: position_ids = _tensor_from_int64_header( position_ids_header, self.torch, self.device ) - hidden_states = self._run_layers(hidden_states, attention_mask, position_ids) + hidden_states = self._run_layers( + hidden_states, attention_mask, position_ids, start_layer=start_layer + ) if self.is_tail: return self.decode_tail(hidden_states) return self._payload(hidden_states, attention_mask, position_ids) @@ -278,7 +281,21 @@ class TorchModelShard: ) return dict(self.tokenizer(prompt, return_tensors="pt")) - def _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any: + def _run_layers( + self, + hidden_states: Any, + attention_mask: Any, + position_ids: Any, + start_layer: int | None = None, + ) -> Any: + # start_layer overrides shard_start for overlapping-shard routing + # (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent + # indexing outside the loaded weights. + effective_start = ( + max(self.shard_start, start_layer) + if start_layer is not None + else self.shard_start + ) position_embeddings = _rotary_position_embeddings( self.model, hidden_states, @@ -290,7 +307,7 @@ class TorchModelShard: self.torch, ) with self.torch.inference_mode(): - for layer in self.layers[self.shard_start:self.shard_end + 1]: + for layer in self.layers[effective_start:self.shard_end + 1]: hidden_states = _call_layer( layer, hidden_states, diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 0285ba7..d26991f 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -141,12 +141,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): if int(self.headers.get("X-Meshnet-Hop-Index", "0")) > 0: server.received_activations = True + start_layer_header = self.headers.get("X-Meshnet-Start-Layer") + start_layer = int(start_layer_header) if start_layer_header else None + try: result = server.backend.forward_bytes( raw_body, shape, self.headers.get("X-Meshnet-Attn-Mask"), self.headers.get("X-Meshnet-Position-Ids"), + start_layer=start_layer, ) except Exception as exc: self._send_json(500, {"error": str(exc)}) @@ -300,17 +304,29 @@ 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[str]: + def _get_remaining_route(self, model: str) -> list[tuple[str, int]]: + """Return downstream hops as (endpoint, start_layer) pairs. + + Fast path reads X-Meshnet-Route header injected by the tracker. + Slow path queries the tracker's /v1/route endpoint as a fallback. + start_layer tells each downstream node which layer to begin from, + enabling correct execution when shard ranges overlap. + """ # Fast path: tracker pre-resolved the downstream route and injected it as a header. injected = self.headers.get("X-Meshnet-Route") if injected: try: route = json.loads(injected) if isinstance(route, list): - resolved = [str(ep) for ep in route] - print(f" [node] using injected downstream route: {resolved}", flush=True) - return resolved - except (json.JSONDecodeError, TypeError): + hops: list[tuple[str, int]] = [] + for item in route: + if isinstance(item, dict): + hops.append((str(item["endpoint"]), int(item.get("start_layer", 0)))) + 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) + return hops + except (json.JSONDecodeError, TypeError, KeyError): pass # Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject). @@ -322,16 +338,29 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}" with urllib.request.urlopen(url, timeout=server.route_timeout) as r: route_resp = json.loads(r.read()) - route = route_resp.get("route", []) own_port = server.server_address[1] - resolved = [ep for ep in route if not ep.rstrip("/").endswith(f":{own_port}")] - print(f" [node] tracker downstream route: {resolved}", flush=True) - return resolved + 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: + ep = node_info.get("endpoint", "") + if ep.rstrip("/").endswith(f":{own_port}"): + 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)) + covered_up_to = node_info.get("shard_end", covered_up_to) + print(f" [node] tracker downstream route: {[ep for ep, _ 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[str]) -> str: + def _run_downstream_pipeline(self, payload: object, route: list[tuple[str, int]]) -> str: server: _TorchHTTPServer = self.server # type: ignore[assignment] if not route: # Partial shard at tail: decode the activation from the previous node. @@ -356,8 +385,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): current_attn = attn_mask current_pos = pos_ids - for hop_index, node_url in enumerate(route): - print(f" [node] pipeline hop {hop_index}: {node_url}", flush=True) + for hop_index, (node_url, start_layer) in enumerate(route): + print(f" [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}", flush=True) headers: dict[str, str] = { "Content-Type": "application/octet-stream", "X-Meshnet-Wire": _WIRE_VERSION, @@ -367,6 +396,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): "X-Meshnet-Chunk-Index": "0", "X-Meshnet-Chunk-Total": "1", "X-Meshnet-Hop-Index": str(hop_index), + "X-Meshnet-Start-Layer": str(start_layer), } if current_attn: headers["X-Meshnet-Attn-Mask"] = current_attn diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index d79af49..89bced5 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -728,8 +728,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): ] rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1) route_nodes, _ = _select_route(all_nodes, rs, re) + # Compute start_layer for each hop: each node begins where the previous ended + 1. + # This allows overlapping shard registrations without double-computation. + 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}) + 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_urls = json.dumps([n.endpoint for n in route_nodes if n.endpoint != node.endpoint]) + downstream_hops = [h for h in route_hops if h["endpoint"].rstrip("/") != node.endpoint.rstrip("/")] + downstream_urls = json.dumps(downstream_hops) route_debug = " -> ".join( f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]" for n in route_nodes @@ -1334,12 +1342,19 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._send_json(503, {"error": error}) return + covered_up_to = required_start - 1 + route_with_start: list[tuple] = [] + for rn in route: + route_with_start.append((rn, covered_up_to + 1)) + covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to + self._send_json(200, { - "route": [e.endpoint for e in route], + "route": [e.endpoint for e, _ in route_with_start], "nodes": [ { "node_id": e.node_id, "endpoint": e.endpoint, + "start_layer": start, "wallet_address": e.wallet_address, "shard_start": e.shard_start, "shard_end": e.shard_end, @@ -1347,7 +1362,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "shard_checksum": e.shard_checksum, "score": e.score, } - for e in route + for e, start in route_with_start ], }) diff --git a/tests/test_real_model_backend.py b/tests/test_real_model_backend.py index fe9159b..2757670 100644 --- a/tests/test_real_model_backend.py +++ b/tests/test_real_model_backend.py @@ -36,7 +36,7 @@ class _FakeBackend: position_ids_header=None, ) - def forward_bytes(self, body, shape, attention_mask_header, position_ids_header): + def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): assert shape == [1, 6, 8] return TensorPayload( body=body, @@ -50,7 +50,7 @@ class _FakeTailBackend(_FakeBackend): is_head = False is_tail = True - def forward_bytes(self, body, shape, attention_mask_header, position_ids_header): + def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): assert len(body) == 1 * 6 * 8 * 2 return " Paris"