feat: add binary activation wire format

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})
-        for node_url in route[1:]:
-            node_resp = _post_json(f"{node_url}/v1/infer", {"activations": node_resp["activations"]})
-
-        text = node_resp["text"]
+        prompt = _last_message_content(messages)
+        chunk_responses = _run_binary_pipeline(route, prompt)
+        _reassemble_binary_responses(chunk_responses)
+        server.last_binary_chunk_responses = chunk_responses
+        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"