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

feat: add two-node shard pipeline

Browse Source
This commit is contained in:
Dobromir Popov
2026-06-29 00:34:57 +03:00
parent 1141b51286
commit 7c8e85c571
4 changed files with 203 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
.scratch/distributed-inference-network/prd.json
View File

@@ -42,11 +42,12 @@
"Commit only this story's code changes with a focused conventional commit message"
],
"priority": 2,
"passes": false,
"passes": true,
"notes": "Source issue: .scratch/distributed-inference-network/issues/02-two-node-shard-pipeline.md",
"dependsOn": [
"US-001"
]
],
"completionNotes": "Completed by Ralph iteration 126384e5; verified by pytest two-node pipeline."
},
{
"id": "US-003",

61
packages/gateway/meshnet_gateway/server.py
View File

@@ -7,6 +7,12 @@ import time
import urllib.request
class _GatewayHTTPServer(http.server.HTTPServer):
def __init__(self, addr, handler, inference_route: list[str]):
super().__init__(addr, handler)
self.inference_route = inference_route
class _GatewayHandler(http.server.BaseHTTPRequestHandler):
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
pass
@@ -22,16 +28,15 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
length = int(self.headers.get("Content-Length", 0))
body = json.loads(self.rfile.read(length))
node_url = self.server.node_url # type: ignore[attr-defined]
payload = json.dumps({"messages": body.get("messages", [])}).encode()
req = urllib.request.Request(
f"{node_url}/v1/infer",
data=payload,
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(req) as r:
node_resp = json.loads(r.read())
route: list[str] = self.server.inference_route # type: ignore[attr-defined]
messages = body.get("messages", [])
# Send messages to the first node in the inference route.
node_resp = _post_json(f"{route[0]}/v1/infer", {"messages": messages})
# Forward activations through remaining nodes in the route.
for node_url in route[1:]:
node_resp = _post_json(f"{node_url}/v1/infer", {"activations": node_resp["activations"]})
response_body = json.dumps({
"id": "chatcmpl-stub",
@@ -60,17 +65,37 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
self.wfile.write(response_body)
class _GatewayHTTPServer(http.server.HTTPServer):
def __init__(self, addr, handler, node_url: str):
super().__init__(addr, handler)
self.node_url = node_url
def _post_json(url: str, payload: dict, timeout: float = 5.0) -> dict:
data = json.dumps(payload).encode()
req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
with urllib.request.urlopen(req, timeout=timeout) as r:
return json.loads(r.read())
class GatewayServer:
"""HTTP gateway that forwards /v1/chat/completions to a registered node."""
"""HTTP gateway that routes /v1/chat/completions through an ordered inference route.
def __init__(self, node_url: str, host: str = "127.0.0.1", port: int = 0):
self._node_url = node_url
Supply either ``node_url`` (single-node, US-001 compat) or ``inference_route``
(ordered list of node base URLs for multi-shard routing).
"""
def __init__(
self,
node_url: str | None = None,
*,
inference_route: list[str] | None = None,
host: str = "127.0.0.1",
port: int = 0,
):
if node_url is not None and inference_route is not None:
raise ValueError("Provide node_url or inference_route, not both")
if node_url is not None:
inference_route = [node_url]
elif inference_route is None:
raise ValueError("Provide either node_url or inference_route")
if not inference_route:
raise ValueError("inference_route must contain at least one node URL")
self._inference_route = inference_route
self._host = host
self._requested_port = port
self._server: _GatewayHTTPServer | None = None
@@ -84,7 +109,7 @@ class GatewayServer:
self._server = _GatewayHTTPServer(
(self._host, self._requested_port),
_GatewayHandler,
self._node_url,
self._inference_route,
)
self.port = self._server.server_address[1]
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)

107
packages/node/meshnet_node/server.py
View File

@@ -1,9 +1,40 @@
"""Stub node HTTP server for the Distributed Inference Network."""
import base64
import http.server
import json
import struct
import threading
# Activation wire format (contract for all shard nodes):
# JSON object: {"shape": [B, S, D], "dtype": "float32",
# "data": "<base64-encoded little-endian IEEE 754 float32 bytes>",
# "context": {"prompt": "<original user prompt>"}}
# shape: [batch_size, seq_len, hidden_dim] — all ints
# dtype: always "float32" in this network version
# data: base64(struct.pack(f"<{B*S*D}f", *values))
# context: opaque dict forwarded unchanged by every intermediate shard
_STUB_HIDDEN_DIM = 64
_STUB_SEQ_LEN = 4
def _make_stub_activations(prompt: str) -> dict:
"""Return zeroed stub activations carrying the original prompt as context."""
shape = [1, _STUB_SEQ_LEN, _STUB_HIDDEN_DIM]
n = shape[0] * shape[1] * shape[2]
data = base64.b64encode(struct.pack(f"<{n}f", *([0.0] * n))).decode()
return {"shape": shape, "dtype": "float32", "data": data, "context": {"prompt": prompt}}
class _StubHTTPServer(http.server.HTTPServer):
def __init__(self, addr, handler, shard_start: int, shard_end: int, is_last_shard: bool):
super().__init__(addr, handler)
self.shard_start = shard_start
self.shard_end = shard_end
self.is_last_shard = is_last_shard
self.received_activations: bool = False
class _StubHandler(http.server.BaseHTTPRequestHandler):
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
@@ -11,36 +42,84 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
def do_POST(self):
if self.path == "/v1/infer":
length = int(self.headers.get("Content-Length", 0))
body = json.loads(self.rfile.read(length))
messages = body.get("messages", [])
last_content = messages[-1]["content"] if messages else ""
payload = json.dumps({"text": f"stub response to: {last_content}"}).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
self._handle_infer()
else:
self.send_response(404)
self.end_headers()
def _handle_infer(self):
server: _StubHTTPServer = self.server # type: ignore[assignment]
length = int(self.headers.get("Content-Length", 0))
body = json.loads(self.rfile.read(length))
if "activations" in body:
# Shard-to-shard call: this node receives activations from the previous shard.
server.received_activations = True
prompt = body["activations"].get("context", {}).get("prompt", "")
if server.is_last_shard:
payload = json.dumps({"text": f"stub response to: {prompt}"}).encode()
else:
payload = json.dumps({"activations": _make_stub_activations(prompt)}).encode()
else:
# Gateway-to-first-shard call: this node receives the original messages.
messages = body.get("messages", [])
last_content = messages[-1]["content"] if messages else ""
if server.is_last_shard:
# Single-node shortcut: no activation hop needed.
payload = json.dumps({"text": f"stub response to: {last_content}"}).encode()
else:
payload = json.dumps({"activations": _make_stub_activations(last_content)}).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
class StubNodeServer:
"""Single-process stub node that returns fixed inference responses."""
"""Single-process stub node that returns fixed inference responses.
def __init__(self, host: str = "127.0.0.1", port: int = 0):
shard_start / shard_end define which transformer layer range this node owns.
is_last_shard controls whether the node returns a text response (True) or
activation tensors (False) after processing its shard.
"""
def __init__(
self,
host: str = "127.0.0.1",
port: int = 0,
shard_start: int = 0,
shard_end: int = 31,
is_last_shard: bool = True,
):
self._host = host
self._requested_port = port
self._server: http.server.HTTPServer | None = None
if shard_start > shard_end:
raise ValueError("shard_start must be less than or equal to shard_end")
self._shard_start = shard_start
self._shard_end = shard_end
self._is_last_shard = is_last_shard
self._server: _StubHTTPServer | None = None
self._thread: threading.Thread | None = None
self.port: int | None = None
@property
def received_activations(self) -> bool:
"""True if this node received an activation tensor since it was started."""
return self._server.received_activations if self._server is not None else False
def start(self) -> int:
if self._server is not None:
raise RuntimeError("StubNodeServer is already running")
self._server = http.server.HTTPServer((self._host, self._requested_port), _StubHandler)
self._server = _StubHTTPServer(
(self._host, self._requested_port),
_StubHandler,
self._shard_start,
self._shard_end,
self._is_last_shard,
)
self.port = self._server.server_address[1]
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
self._thread.start()

64
tests/test_two_node_pipeline.py Normal file
View File

@@ -0,0 +1,64 @@
"""US-002 integration test: two-node shard pipeline with activation tensors flowing A → B."""
import json
import urllib.request
from meshnet_node.server import StubNodeServer
from meshnet_gateway.server import GatewayServer
def test_two_node_activation_pipeline():
# Node A owns layers 0-15 (first shard, not last — returns activations).
node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
port_a = node_a.start()
# Node B owns layers 16-31 (last shard — returns text after receiving activations).
node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
port_b = node_b.start()
# Gateway hardcodes the two-node inference route for this story.
gateway = GatewayServer(
inference_route=[
f"http://127.0.0.1:{port_a}",
f"http://127.0.0.1:{port_b}",
]
)
gateway_port = gateway.start()
try:
payload = json.dumps({
"model": "stub-model",
"messages": [{"role": "user", "content": "hello"}],
}).encode()
req = urllib.request.Request(
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
data=payload,
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(req) as resp:
assert resp.status == 200
body = json.loads(resp.read())
# Gateway must assemble a valid OpenAI-format response.
assert body["id"] == "chatcmpl-stub"
assert body["object"] == "chat.completion"
assert body["model"] == "stub-model"
assert len(body["choices"]) == 1
choice = body["choices"][0]
assert choice["message"]["role"] == "assistant"
assert choice["message"]["content"] == "stub response to: hello"
assert choice["finish_reason"] == "stop"
# Node B must have received an activation tensor from node A — this is
# the observable proof that the A → B activation hop occurred.
assert node_b.received_activations, "node B did not receive activation tensors from node A"
# Node A must NOT have received activations (it is the first shard).
assert not node_a.received_activations, "node A should not receive activations — it is first in the route"
finally:
gateway.stop()
node_a.stop()
node_b.stop()