feat: add two-node shard pipeline

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)