popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(us-022): X-Meshnet-Start-Layer pipeline protocol for overlapping shards

Browse Source 
When _select_route picks two nodes with overlapping registrations (e.g.
A:0-22 and B:20-24), the tracker now injects start_layer per hop so B
executes only layers 23-24, not 20-24.

- model_backend: forward_bytes + _run_layers accept start_layer offset;
  clamped to shard_start to prevent out-of-bounds indexing
- torch_server: _handle_binary_forward reads X-Meshnet-Start-Layer header;
  _run_downstream_pipeline sends it per hop; route is now list[tuple[str,int]]
- server: proxy injects {endpoint, start_layer} objects in X-Meshnet-Route;
  /v1/route response includes start_layer per node in the nodes list
- test: fake backends accept start_layer=None kwarg

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-06-30 13:14:45 +03:00
parent e1ba120912
commit 96af82892f
4 changed files with 82 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
packages/node/meshnet_node/model_backend.py
View File

@@ -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,

54
packages/node/meshnet_node/torch_server.py
View File

@@ -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