feat: add binary activation wire format

packages/node/meshnet_node/server.py

@@ -1,34 +1,85 @@
 """Stub node HTTP server for the Distributed Inference Network."""
 
-import base64
 import http.server
 import json
-import struct
 import threading
 import urllib.parse
 from pathlib import Path
 
 from .downloader import compute_shard_checksum, write_shard_archive
 
-# 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
+# Binary activation wire format (contract for all shard nodes):
+# POST /forward with raw tensor bytes in the body and tensor/session/chunk
+# metadata in X-Meshnet-* headers. Bodies may be zstd-compressed.
 
 _STUB_HIDDEN_DIM = 64
 _STUB_SEQ_LEN = 4
+_WIRE_VERSION = "2"
+_DTYPE_SIZES = {
+    "float16": 2,
+    "bfloat16": 2,
+    "float32": 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}}
+def _make_stub_binary_activation(shape: list[int], dtype: str) -> bytes:
+    """Return a zeroed stub tensor body for the requested shape and dtype."""
+    element_size = _DTYPE_SIZES[dtype]
+    element_count = 1
+    for dim in shape:
+        element_count *= dim
+    return b"\x00" * (element_count * element_size)
+
+
+def _parse_shape(value: str | None) -> list[int]:
+    if not value:
+        raise ValueError("missing X-Meshnet-Shape")
+    shape = [int(part) for part in value.split(",")]
+    if not shape or any(dim <= 0 for dim in shape):
+        raise ValueError("invalid X-Meshnet-Shape")
+    return shape
+
+
+def _decompress_body(body: bytes, encoding: str | None) -> bytes:
+    if not encoding:
+        return body
+    if encoding != "zstd":
+        raise ValueError("unsupported X-Meshnet-Encoding")
+    import zstandard as zstd
+
+    try:
+        return zstd.ZstdDecompressor().decompress(body)
+    except zstd.ZstdError as exc:
+        raise ValueError("invalid zstd activation body") from exc
+
+
+def _compress_body(body: bytes, encoding: str | None) -> bytes:
+    if not encoding:
+        return body
+    if encoding != "zstd":
+        raise ValueError("unsupported X-Meshnet-Encoding")
+    import zstandard as zstd
+
+    try:
+        return zstd.ZstdCompressor(level=1).compress(body)
+    except zstd.ZstdError as exc:
+        raise ValueError("could not zstd-compress activation body") from exc
+
+
+def _validate_activation_body(body: bytes, shape: list[int], dtype: str) -> None:
+    if dtype not in _DTYPE_SIZES:
+        raise ValueError("unsupported X-Meshnet-Dtype")
+    expected_len = len(_make_stub_binary_activation(shape, dtype))
+    if len(body) != expected_len:
+        raise ValueError("activation byte length does not match shape and dtype")
+    try:
+        import numpy as np
+    except ModuleNotFoundError:
+        return
+    if dtype == "bfloat16" and not hasattr(np, "bfloat16"):
+        return
+    np_dtype = getattr(np, "bfloat16") if dtype == "bfloat16" else np.dtype(dtype)
+    np.frombuffer(body, dtype=np_dtype).reshape(shape)
 
 
 class _StubHTTPServer(http.server.HTTPServer):
@@ -51,6 +102,7 @@ class _StubHTTPServer(http.server.HTTPServer):
         self.model = model
         self.shard_path = shard_path
         self.received_activations: bool = False
+        self.forward_chunk_count: int = 0
 
 
 class _StubHandler(http.server.BaseHTTPRequestHandler):
@@ -60,6 +112,8 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
     def do_POST(self):
         if self.path == "/v1/infer":
             self._handle_infer()
+        elif self.path == "/forward":
+            self._handle_forward()
         else:
             self.send_response(404)
             self.end_headers()
@@ -113,7 +167,12 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
             if server.is_last_shard:
                 payload = json.dumps({"text": f"{server.response_prefix} {prompt}"}).encode()
             else:
-                payload = json.dumps({"activations": _make_stub_activations(prompt)}).encode()
+                payload = json.dumps({
+                    "activations": {
+                        "shape": [1, _STUB_SEQ_LEN, _STUB_HIDDEN_DIM],
+                        "dtype": "bfloat16",
+                    }
+                }).encode()
         else:
             # Gateway-to-first-shard call: this node receives the original messages.
             messages = body.get("messages", [])
@@ -122,7 +181,12 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
                 # Single-node shortcut: no activation hop needed.
                 payload = json.dumps({"text": f"{server.response_prefix} {last_content}"}).encode()
             else:
-                payload = json.dumps({"activations": _make_stub_activations(last_content)}).encode()
+                payload = json.dumps({
+                    "activations": {
+                        "shape": [1, _STUB_SEQ_LEN, _STUB_HIDDEN_DIM],
+                        "dtype": "bfloat16",
+                    }
+                }).encode()
 
         self.send_response(200)
         self.send_header("Content-Type", "application/json")
@@ -130,6 +194,51 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
         self.end_headers()
         self.wfile.write(payload)
 
+    def _handle_forward(self):
+        server: _StubHTTPServer = self.server  # type: ignore[assignment]
+        try:
+            shape = _parse_shape(self.headers.get("X-Meshnet-Shape"))
+            dtype = self.headers.get("X-Meshnet-Dtype", "")
+            session = self.headers["X-Meshnet-Session"]
+            chunk_index = self.headers["X-Meshnet-Chunk-Index"]
+            chunk_total = self.headers["X-Meshnet-Chunk-Total"]
+            encoding = self.headers.get("X-Meshnet-Encoding")
+            length = int(self.headers.get("Content-Length", 0))
+            body = self.rfile.read(length)
+            raw_body = _decompress_body(body, encoding)
+            _validate_activation_body(raw_body, shape, dtype)
+            chunk_index_value = int(chunk_index)
+            chunk_total_value = int(chunk_total)
+            if chunk_total_value <= 0 or not 0 <= chunk_index_value < chunk_total_value:
+                raise ValueError("invalid chunk index/total")
+        except (KeyError, ValueError, TypeError):
+            self.send_response(400)
+            self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
+            self.end_headers()
+            return
+
+        server.forward_chunk_count += 1
+        if int(self.headers.get("X-Meshnet-Hop-Index", "0")) > 0:
+            server.received_activations = True
+
+        raw_payload = _make_stub_binary_activation(shape, dtype)
+        payload = _compress_body(raw_payload, encoding)
+        self.send_response(200)
+        self.send_header("Content-Type", "application/octet-stream")
+        self.send_header("Content-Length", str(len(payload)))
+        self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
+        self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in shape))
+        self.send_header("X-Meshnet-Dtype", dtype)
+        self.send_header("X-Meshnet-Session", session)
+        self.send_header("X-Meshnet-Chunk-Index", chunk_index)
+        self.send_header("X-Meshnet-Chunk-Total", chunk_total)
+        if encoding:
+            self.send_header("X-Meshnet-Encoding", encoding)
+        if server.is_last_shard:
+            self.send_header("X-Meshnet-Stub-Response-Prefix", server.response_prefix)
+        self.end_headers()
+        self.wfile.write(payload)
+
 
 class StubNodeServer:
     """Single-process stub node that returns fixed inference responses.
@@ -169,6 +278,11 @@ class StubNodeServer:
         """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
 
+    @property
+    def forward_chunk_count(self) -> int:
+        """Number of binary /forward chunks handled since this node was started."""
+        return self._server.forward_chunk_count if self._server is not None else 0
+
     def start(self) -> int:
         if self._server is not None:
             raise RuntimeError("StubNodeServer is already running")