diff --git a/.scratch/distributed-inference-network/prd.json b/.scratch/distributed-inference-network/prd.json index aee07ca..479c2a1 100644 --- a/.scratch/distributed-inference-network/prd.json +++ b/.scratch/distributed-inference-network/prd.json @@ -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", diff --git a/packages/gateway/meshnet_gateway/server.py b/packages/gateway/meshnet_gateway/server.py index 7256f93..21dabd9 100644 --- a/packages/gateway/meshnet_gateway/server.py +++ b/packages/gateway/meshnet_gateway/server.py @@ -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) diff --git a/packages/node/meshnet_node/server.py b/packages/node/meshnet_node/server.py index 3f647fd..d01512c 100644 --- a/packages/node/meshnet_node/server.py +++ b/packages/node/meshnet_node/server.py @@ -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": "", +# "context": {"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() diff --git a/tests/test_two_node_pipeline.py b/tests/test_two_node_pipeline.py new file mode 100644 index 0000000..514e673 --- /dev/null +++ b/tests/test_two_node_pipeline.py @@ -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()