feat(us-022): X-Meshnet-Start-Layer pipeline protocol for overlapping shards
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:
@@ -145,6 +145,7 @@ class TorchModelShard:
|
|||||||
shape: list[int],
|
shape: list[int],
|
||||||
attention_mask_header: str | None,
|
attention_mask_header: str | None,
|
||||||
position_ids_header: str | None,
|
position_ids_header: str | None,
|
||||||
|
start_layer: int | None = None,
|
||||||
) -> TensorPayload | str:
|
) -> TensorPayload | str:
|
||||||
hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
|
hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
|
||||||
self.device
|
self.device
|
||||||
@@ -155,7 +156,9 @@ class TorchModelShard:
|
|||||||
position_ids = _tensor_from_int64_header(
|
position_ids = _tensor_from_int64_header(
|
||||||
position_ids_header, self.torch, self.device
|
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:
|
if self.is_tail:
|
||||||
return self.decode_tail(hidden_states)
|
return self.decode_tail(hidden_states)
|
||||||
return self._payload(hidden_states, attention_mask, position_ids)
|
return self._payload(hidden_states, attention_mask, position_ids)
|
||||||
@@ -278,7 +281,21 @@ class TorchModelShard:
|
|||||||
)
|
)
|
||||||
return dict(self.tokenizer(prompt, return_tensors="pt"))
|
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(
|
position_embeddings = _rotary_position_embeddings(
|
||||||
self.model,
|
self.model,
|
||||||
hidden_states,
|
hidden_states,
|
||||||
@@ -290,7 +307,7 @@ class TorchModelShard:
|
|||||||
self.torch,
|
self.torch,
|
||||||
)
|
)
|
||||||
with self.torch.inference_mode():
|
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(
|
hidden_states = _call_layer(
|
||||||
layer,
|
layer,
|
||||||
hidden_states,
|
hidden_states,
|
||||||
|
|||||||
@@ -141,12 +141,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
if int(self.headers.get("X-Meshnet-Hop-Index", "0")) > 0:
|
if int(self.headers.get("X-Meshnet-Hop-Index", "0")) > 0:
|
||||||
server.received_activations = True
|
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:
|
try:
|
||||||
result = server.backend.forward_bytes(
|
result = server.backend.forward_bytes(
|
||||||
raw_body,
|
raw_body,
|
||||||
shape,
|
shape,
|
||||||
self.headers.get("X-Meshnet-Attn-Mask"),
|
self.headers.get("X-Meshnet-Attn-Mask"),
|
||||||
self.headers.get("X-Meshnet-Position-Ids"),
|
self.headers.get("X-Meshnet-Position-Ids"),
|
||||||
|
start_layer=start_layer,
|
||||||
)
|
)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
self._send_json(500, {"error": str(exc)})
|
self._send_json(500, {"error": str(exc)})
|
||||||
@@ -300,17 +304,29 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
result_text = "".join(generated)
|
result_text = "".join(generated)
|
||||||
self._send_openai_response(result_text, model_name, stream, messages)
|
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.
|
# Fast path: tracker pre-resolved the downstream route and injected it as a header.
|
||||||
injected = self.headers.get("X-Meshnet-Route")
|
injected = self.headers.get("X-Meshnet-Route")
|
||||||
if injected:
|
if injected:
|
||||||
try:
|
try:
|
||||||
route = json.loads(injected)
|
route = json.loads(injected)
|
||||||
if isinstance(route, list):
|
if isinstance(route, list):
|
||||||
resolved = [str(ep) for ep in route]
|
hops: list[tuple[str, int]] = []
|
||||||
print(f" [node] using injected downstream route: {resolved}", flush=True)
|
for item in route:
|
||||||
return resolved
|
if isinstance(item, dict):
|
||||||
except (json.JSONDecodeError, TypeError):
|
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
|
pass
|
||||||
|
|
||||||
# Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject).
|
# 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)}"
|
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
|
||||||
with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
|
with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
|
||||||
route_resp = json.loads(r.read())
|
route_resp = json.loads(r.read())
|
||||||
route = route_resp.get("route", [])
|
|
||||||
own_port = server.server_address[1]
|
own_port = server.server_address[1]
|
||||||
resolved = [ep for ep in route if not ep.rstrip("/").endswith(f":{own_port}")]
|
nodes_info = route_resp.get("nodes", [])
|
||||||
print(f" [node] tracker downstream route: {resolved}", flush=True)
|
# nodes_info is ordered; find own node and compute start_layers post-hoc
|
||||||
return resolved
|
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:
|
except Exception as exc:
|
||||||
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
|
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
|
||||||
return []
|
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]
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||||
if not route:
|
if not route:
|
||||||
# Partial shard at tail: decode the activation from the previous node.
|
# 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_attn = attn_mask
|
||||||
current_pos = pos_ids
|
current_pos = pos_ids
|
||||||
|
|
||||||
for hop_index, node_url in enumerate(route):
|
for hop_index, (node_url, start_layer) in enumerate(route):
|
||||||
print(f" [node] pipeline hop {hop_index}: {node_url}", flush=True)
|
print(f" [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}", flush=True)
|
||||||
headers: dict[str, str] = {
|
headers: dict[str, str] = {
|
||||||
"Content-Type": "application/octet-stream",
|
"Content-Type": "application/octet-stream",
|
||||||
"X-Meshnet-Wire": _WIRE_VERSION,
|
"X-Meshnet-Wire": _WIRE_VERSION,
|
||||||
@@ -367,6 +396,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
"X-Meshnet-Chunk-Index": "0",
|
"X-Meshnet-Chunk-Index": "0",
|
||||||
"X-Meshnet-Chunk-Total": "1",
|
"X-Meshnet-Chunk-Total": "1",
|
||||||
"X-Meshnet-Hop-Index": str(hop_index),
|
"X-Meshnet-Hop-Index": str(hop_index),
|
||||||
|
"X-Meshnet-Start-Layer": str(start_layer),
|
||||||
}
|
}
|
||||||
if current_attn:
|
if current_attn:
|
||||||
headers["X-Meshnet-Attn-Mask"] = current_attn
|
headers["X-Meshnet-Attn-Mask"] = current_attn
|
||||||
|
|||||||
@@ -728,8 +728,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
]
|
]
|
||||||
rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1)
|
rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1)
|
||||||
route_nodes, _ = _select_route(all_nodes, rs, re)
|
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.
|
# 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(
|
route_debug = " -> ".join(
|
||||||
f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]"
|
f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]"
|
||||||
for n in route_nodes
|
for n in route_nodes
|
||||||
@@ -1334,12 +1342,19 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self._send_json(503, {"error": error})
|
self._send_json(503, {"error": error})
|
||||||
return
|
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, {
|
self._send_json(200, {
|
||||||
"route": [e.endpoint for e in route],
|
"route": [e.endpoint for e, _ in route_with_start],
|
||||||
"nodes": [
|
"nodes": [
|
||||||
{
|
{
|
||||||
"node_id": e.node_id,
|
"node_id": e.node_id,
|
||||||
"endpoint": e.endpoint,
|
"endpoint": e.endpoint,
|
||||||
|
"start_layer": start,
|
||||||
"wallet_address": e.wallet_address,
|
"wallet_address": e.wallet_address,
|
||||||
"shard_start": e.shard_start,
|
"shard_start": e.shard_start,
|
||||||
"shard_end": e.shard_end,
|
"shard_end": e.shard_end,
|
||||||
@@ -1347,7 +1362,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
"shard_checksum": e.shard_checksum,
|
"shard_checksum": e.shard_checksum,
|
||||||
"score": e.score,
|
"score": e.score,
|
||||||
}
|
}
|
||||||
for e in route
|
for e, start in route_with_start
|
||||||
],
|
],
|
||||||
})
|
})
|
||||||
|
|
||||||
|
|||||||
@@ -36,7 +36,7 @@ class _FakeBackend:
|
|||||||
position_ids_header=None,
|
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]
|
assert shape == [1, 6, 8]
|
||||||
return TensorPayload(
|
return TensorPayload(
|
||||||
body=body,
|
body=body,
|
||||||
@@ -50,7 +50,7 @@ class _FakeTailBackend(_FakeBackend):
|
|||||||
is_head = False
|
is_head = False
|
||||||
is_tail = True
|
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
|
assert len(body) == 1 * 6 * 8 * 2
|
||||||
return " Paris"
|
return " Paris"
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user