feat: add two-node shard pipeline

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