feat: add two-node shard pipeline

.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",

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]
+        route: list[str] = self.server.inference_route  # type: ignore[attr-defined]
-        payload = json.dumps({"messages": body.get("messages", [])}).encode()
+        messages = body.get("messages", [])
-        req = urllib.request.Request(
+
-            f"{node_url}/v1/infer",
+        # Send messages to the first node in the inference route.
-            data=payload,
+        node_resp = _post_json(f"{route[0]}/v1/infer", {"messages": messages})
-            headers={"Content-Type": "application/json"},
+
-            method="POST",
+        # Forward activations through remaining nodes in the route.
-        )
+        for node_url in route[1:]:
-        with urllib.request.urlopen(req) as r:
+            node_resp = _post_json(f"{node_url}/v1/infer", {"activations": node_resp["activations"]})
-            node_resp = json.loads(r.read())
 
         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 _post_json(url: str, payload: dict, timeout: float = 5.0) -> dict:
-    def __init__(self, addr, handler, node_url: str):
+    data = json.dumps(payload).encode()
-        super().__init__(addr, handler)
+    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
-        self.node_url = node_url
+    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):
+    Supply either ``node_url`` (single-node, US-001 compat) or ``inference_route``
-        self._node_url = node_url
+    (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)

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": "<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":
+            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)
-        else:
-            self.send_response(404)
-            self.end_headers()
 
 
 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()

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()