feat: add binary activation wire format

.scratch/distributed-inference-network/prd.json

@@ -259,11 +259,12 @@
         "Commit only this story's changes"
       ],
       "priority": 11,
-      "passes": false,
+      "passes": true,
       "notes": "Source issue: .scratch/distributed-inference-network/issues/11-binary-wire-format.md",
       "dependsOn": [
         "US-002"
-      ]
+      ],
+      "completionNotes": "Completed by fresh Ralph/Codex session 3f3bed75; verified by pytest, compileall, and diff check."
     },
     {
       "id": "US-012",

packages/gateway/meshnet_gateway/server.py

@@ -3,7 +3,9 @@
 import http.server
 import hashlib
 import json
+import os
 from collections import Counter
+from dataclasses import dataclass
 import threading
 import time
 import urllib.error
@@ -12,6 +14,27 @@ import urllib.request
 import uuid
 from typing import Any
 
+_STUB_HIDDEN_DIM = 64
+_STUB_DTYPE = "bfloat16"
+_WIRE_VERSION = "2"
+_DTYPE_SIZES = {
+    "float16": 2,
+    "bfloat16": 2,
+    "float32": 4,
+}
+
+
+@dataclass
+class _BinaryActivation:
+    body: bytes
+    shape: list[int]
+    dtype: str
+    session: str
+    chunk_index: int
+    chunk_total: int
+    encoding: str | None
+    headers: dict[str, str]
+
 
 class _GatewayHTTPServer(http.server.HTTPServer):
     def __init__(
@@ -32,6 +55,7 @@ class _GatewayHTTPServer(http.server.HTTPServer):
         self.contracts = contracts
         self.minimum_stake = minimum_stake
         self.cost_per_layer_token_lamport = cost_per_layer_token_lamport
+        self.last_binary_chunk_responses: list[_BinaryActivation] = []
 
 
 class _GatewayHandler(http.server.BaseHTTPRequestHandler):
@@ -232,7 +256,7 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
         except _TrackerUnavailable:
             self._send_json_error(503, "tracker unavailable")
             return
-        except (urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError) as exc:
+        except (urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError, ValueError) as exc:
             self._send_openai_error(503, f"inference pipeline error: {exc}", "pipeline_error")
             return
 
@@ -264,7 +288,7 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
         except _TrackerUnavailable:
             self._send_json_error(503, "tracker unavailable")
             return
-        except (urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError) as exc:
+        except (urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError, ValueError) as exc:
             self._send_openai_error(503, f"inference pipeline error: {exc}", "pipeline_error")
             return
 
@@ -283,11 +307,12 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
         route_nodes = route_payload.get("nodes", [])
         messages = body.get("messages", [])
 
-        node_resp = _post_json(f"{route[0]}/v1/infer", {"messages": messages})
+        prompt = _last_message_content(messages)
-        for node_url in route[1:]:
+        chunk_responses = _run_binary_pipeline(route, prompt)
-            node_resp = _post_json(f"{node_url}/v1/infer", {"activations": node_resp["activations"]})
+        _reassemble_binary_responses(chunk_responses)
-
+        server.last_binary_chunk_responses = chunk_responses
-        text = node_resp["text"]
+        response_prefix = _stub_response_prefix(chunk_responses)
+        text = f"{response_prefix} {prompt}"
         model = body.get("model", "stub-model")
         self._record_inference(model, messages, text, route_nodes)
         return _completion_response(str(model), text)
@@ -553,6 +578,190 @@ def _post_json(url: str, payload: dict, timeout: float = 5.0) -> dict:
         return json.loads(r.read())
 
 
+def _last_message_content(messages: object) -> str:
+    if not isinstance(messages, list) or not messages:
+        return ""
+    last_message = messages[-1]
+    if not isinstance(last_message, dict):
+        return ""
+    content = last_message.get("content", "")
+    return content if isinstance(content, str) else str(content)
+
+
+def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) -> list[_BinaryActivation]:
+    session = str(uuid.uuid4())
+    chunk_token_count = _chunk_token_count()
+    total_tokens = max(1, _prompt_token_count(prompt))
+    chunk_total = max(1, (total_tokens + chunk_token_count - 1) // chunk_token_count)
+    responses: list[_BinaryActivation] = []
+
+    for chunk_index in range(chunk_total):
+        remaining_tokens = total_tokens - (chunk_index * chunk_token_count)
+        seq_len = min(chunk_token_count, remaining_tokens)
+        activation = _BinaryActivation(
+            body=_make_stub_binary_activation([1, seq_len, _STUB_HIDDEN_DIM], _STUB_DTYPE),
+            shape=[1, seq_len, _STUB_HIDDEN_DIM],
+            dtype=_STUB_DTYPE,
+            session=session,
+            chunk_index=chunk_index,
+            chunk_total=chunk_total,
+            encoding=_preferred_binary_encoding(),
+            headers={},
+        )
+        for hop_index, node_url in enumerate(route):
+            activation = _post_binary_forward(
+                f"{node_url}/forward",
+                activation,
+                hop_index=hop_index,
+                timeout=timeout,
+            )
+        responses.append(activation)
+    return responses
+
+
+def _chunk_token_count() -> int:
+    raw_value = os.environ.get("MESHNET_CHUNK_TOKENS", "128")
+    try:
+        value = int(raw_value)
+    except ValueError:
+        return 128
+    return value if value > 0 else 128
+
+
+def _prompt_token_count(prompt: str) -> int:
+    return len(prompt.split()) if prompt.split() else max(1, len(prompt))
+
+
+def _preferred_binary_encoding() -> str | None:
+    try:
+        import zstandard  # noqa: F401
+    except ModuleNotFoundError:
+        return None
+    return "zstd"
+
+
+def _make_stub_binary_activation(shape: list[int], dtype: str) -> bytes:
+    element_count = 1
+    for dim in shape:
+        element_count *= dim
+    return b"\x00" * (element_count * _DTYPE_SIZES[dtype])
+
+
+def _post_binary_forward(
+    url: str,
+    activation: _BinaryActivation,
+    *,
+    hop_index: int,
+    timeout: float,
+) -> _BinaryActivation:
+    request_body = _compress_body(activation.body, activation.encoding)
+    headers = {
+        "Content-Type": "application/octet-stream",
+        "X-Meshnet-Wire": _WIRE_VERSION,
+        "X-Meshnet-Shape": ",".join(str(dim) for dim in activation.shape),
+        "X-Meshnet-Dtype": activation.dtype,
+        "X-Meshnet-Session": activation.session,
+        "X-Meshnet-Chunk-Index": str(activation.chunk_index),
+        "X-Meshnet-Chunk-Total": str(activation.chunk_total),
+        "X-Meshnet-Hop-Index": str(hop_index),
+    }
+    if activation.encoding:
+        headers["X-Meshnet-Encoding"] = activation.encoding
+    req = urllib.request.Request(url, data=request_body, headers=headers, method="POST")
+    with urllib.request.urlopen(req, timeout=timeout) as response:
+        response_body = response.read()
+        response_headers = {key.lower(): value for key, value in response.headers.items()}
+
+    encoding = response_headers.get("x-meshnet-encoding")
+    raw_body = _decompress_body(response_body, encoding)
+    shape = _parse_shape(response_headers["x-meshnet-shape"])
+    dtype = response_headers["x-meshnet-dtype"]
+    session = response_headers["x-meshnet-session"]
+    chunk_index = int(response_headers["x-meshnet-chunk-index"])
+    chunk_total = int(response_headers["x-meshnet-chunk-total"])
+    if session != activation.session:
+        raise ValueError("binary activation response changed session")
+    if chunk_index != activation.chunk_index or chunk_total != activation.chunk_total:
+        raise ValueError("binary activation response changed chunk metadata")
+    if dtype != activation.dtype or shape != activation.shape:
+        raise ValueError("binary activation response shape/dtype mismatch")
+    _validate_activation_body(raw_body, shape, dtype)
+    return _BinaryActivation(
+        body=raw_body,
+        shape=shape,
+        dtype=dtype,
+        session=session,
+        chunk_index=chunk_index,
+        chunk_total=chunk_total,
+        encoding=encoding,
+        headers=response_headers,
+    )
+
+
+def _reassemble_binary_responses(chunks: list[_BinaryActivation]) -> bytes:
+    """Validate and concatenate final route chunk responses into one activation body."""
+    if not chunks:
+        raise ValueError("binary pipeline returned no chunks")
+    first = chunks[0]
+    expected_total = first.chunk_total
+    if len(chunks) != expected_total:
+        raise ValueError("binary pipeline returned incomplete chunk set")
+    if [chunk.chunk_index for chunk in chunks] != list(range(expected_total)):
+        raise ValueError("binary pipeline returned out-of-order chunks")
+    for chunk in chunks:
+        if chunk.session != first.session or chunk.dtype != first.dtype:
+            raise ValueError("binary chunks do not belong to the same activation")
+        if len(chunk.shape) != 3 or chunk.shape[0] != first.shape[0] or chunk.shape[2] != first.shape[2]:
+            raise ValueError("binary chunks have incompatible shapes")
+    return b"".join(chunk.body for chunk in chunks)
+
+
+def _compress_body(body: bytes, encoding: str | None) -> bytes:
+    if not encoding:
+        return body
+    if encoding != "zstd":
+        raise ValueError("unsupported binary encoding")
+    import zstandard as zstd
+
+    try:
+        return zstd.ZstdCompressor(level=1).compress(body)
+    except zstd.ZstdError as exc:
+        raise ValueError("could not zstd-compress binary activation") from exc
+
+
+def _decompress_body(body: bytes, encoding: str | None) -> bytes:
+    if not encoding:
+        return body
+    if encoding != "zstd":
+        raise ValueError("unsupported binary encoding")
+    import zstandard as zstd
+
+    try:
+        return zstd.ZstdDecompressor().decompress(body)
+    except zstd.ZstdError as exc:
+        raise ValueError("invalid zstd binary activation") from exc
+
+
+def _parse_shape(value: str) -> list[int]:
+    shape = [int(part) for part in value.split(",")]
+    if not shape or any(dim <= 0 for dim in shape):
+        raise ValueError("invalid binary activation shape")
+    return shape
+
+
+def _validate_activation_body(body: bytes, shape: list[int], dtype: str) -> None:
+    if dtype not in _DTYPE_SIZES:
+        raise ValueError("unsupported binary activation dtype")
+    if len(body) != len(_make_stub_binary_activation(shape, dtype)):
+        raise ValueError("binary activation byte length does not match headers")
+
+
+def _stub_response_prefix(chunk_responses: list[_BinaryActivation]) -> str:
+    if not chunk_responses:
+        return "stub response to:"
+    return chunk_responses[-1].headers.get("x-meshnet-stub-response-prefix", "stub response to:")
+
+
 def _get_json(url: str, timeout: float = 5.0) -> dict:
     with urllib.request.urlopen(url, timeout=timeout) as r:
         return json.loads(r.read())
@@ -606,6 +815,11 @@ class GatewayServer:
         self._thread: threading.Thread | None = None
         self.port: int | None = None
 
+    @property
+    def last_binary_chunk_responses(self) -> list[_BinaryActivation]:
+        """Most recent binary chunk responses returned by the final route node."""
+        return self._server.last_binary_chunk_responses if self._server is not None else []
+
     def start(self) -> int:
         if self._server is not None:
             raise RuntimeError("GatewayServer is already running")

packages/gateway/pyproject.toml

@@ -8,6 +8,10 @@ version = "0.1.0"
 description = "Distributed Inference Network HTTP gateway"
 requires-python = ">=3.10"
 
+dependencies = [
+    "zstandard>=0.22",
+]
+
 [project.scripts]
 meshnet-gateway = "meshnet_gateway.cli:main"
 

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):
+# Binary activation wire format (contract for all shard nodes):
-# JSON object: {"shape": [B, S, D], "dtype": "float32",
+# POST /forward with raw tensor bytes in the body and tensor/session/chunk
-#               "data": "<base64-encoded little-endian IEEE 754 float32 bytes>",
+# metadata in X-Meshnet-* headers. Bodies may be zstd-compressed.
-#               "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
+_WIRE_VERSION = "2"
+_DTYPE_SIZES = {
+    "float16": 2,
+    "bfloat16": 2,
+    "float32": 4,
+}
 
 
-def _make_stub_activations(prompt: str) -> dict:
+def _make_stub_binary_activation(shape: list[int], dtype: str) -> bytes:
-    """Return zeroed stub activations carrying the original prompt as context."""
+    """Return a zeroed stub tensor body for the requested shape and dtype."""
-    shape = [1, _STUB_SEQ_LEN, _STUB_HIDDEN_DIM]
+    element_size = _DTYPE_SIZES[dtype]
-    n = shape[0] * shape[1] * shape[2]
+    element_count = 1
-    data = base64.b64encode(struct.pack(f"<{n}f", *([0.0] * n))).decode()
+    for dim in shape:
-    return {"shape": shape, "dtype": "float32", "data": data, "context": {"prompt": prompt}}
+        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")

packages/node/pyproject.toml

@@ -11,6 +11,7 @@ requires-python = ">=3.10"
 dependencies = [
     "cryptography>=41",
     "huggingface-hub>=0.20",
+    "zstandard>=0.22",
 ]
 
 [project.scripts]

tests/test_two_node_pipeline.py

@@ -1,4 +1,4 @@
-"""US-002 integration test: two-node shard pipeline with activation tensors flowing A → B."""
+"""Integration tests for two-node shard pipelines with activation tensors flowing A to B."""
 
 import json
 import urllib.request
@@ -62,3 +62,54 @@ def test_two_node_activation_pipeline():
         gateway.stop()
         node_a.stop()
         node_b.stop()
+
+
+def test_binary_activation_pipeline_chunks_512_token_prefill():
+    node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
+    port_a = node_a.start()
+    node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
+    port_b = node_b.start()
+    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:
+        prompt = " ".join(f"tok{i}" for i in range(512))
+        payload = json.dumps({
+            "model": "stub-model",
+            "messages": [{"role": "user", "content": prompt}],
+        }).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())
+
+        assert body["choices"][0]["message"]["content"] == f"stub response to: {prompt}"
+        assert node_a.forward_chunk_count == 4
+        assert node_b.forward_chunk_count == 4
+        assert not node_a.received_activations
+        assert node_b.received_activations
+
+        chunk_responses = gateway.last_binary_chunk_responses
+        assert len(chunk_responses) == 4
+        for index, chunk in enumerate(chunk_responses):
+            assert chunk.shape == [1, 128, 64]
+            assert chunk.dtype == "bfloat16"
+            assert chunk.chunk_index == index
+            assert chunk.chunk_total == 4
+            assert chunk.headers["x-meshnet-wire"] == "2"
+            assert len(chunk.body) == 1 * 128 * 64 * 2
+    finally:
+        gateway.stop()
+        node_a.stop()
+        node_b.stop()