popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add binary activation wire format

Browse Source
This commit is contained in:
Dobromir Popov
2026-06-29 14:58:55 +03:00
parent b17deb4d0e
commit 2a4383e353
6 changed files with 413 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
.scratch/distributed-inference-network/prd.json
View File

@@ -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",

228
packages/gateway/meshnet_gateway/server.py
View File

@@ -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")

4
packages/gateway/pyproject.toml
View File

@@ -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"

150
packages/node/meshnet_node/server.py
View File

@@ -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")

1
packages/node/pyproject.toml
View File

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

53
tests/test_two_node_pipeline.py
View File

@@ -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()