diff --git a/.scratch/distributed-inference-network/prd.json b/.scratch/distributed-inference-network/prd.json index dc8fcb6..95ebb54 100644 --- a/.scratch/distributed-inference-network/prd.json +++ b/.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", diff --git a/packages/gateway/meshnet_gateway/server.py b/packages/gateway/meshnet_gateway/server.py index 77961f1..3b3b629 100644 --- a/packages/gateway/meshnet_gateway/server.py +++ b/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") diff --git a/packages/gateway/pyproject.toml b/packages/gateway/pyproject.toml index 0523889..694e0aa 100644 --- a/packages/gateway/pyproject.toml +++ b/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" diff --git a/packages/node/meshnet_node/server.py b/packages/node/meshnet_node/server.py index d514109..bad4ae8 100644 --- a/packages/node/meshnet_node/server.py +++ b/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": "", -# "context": {"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") diff --git a/packages/node/pyproject.toml b/packages/node/pyproject.toml index 46225c8..1b8d4f3 100644 --- a/packages/node/pyproject.toml +++ b/packages/node/pyproject.toml @@ -11,6 +11,7 @@ requires-python = ">=3.10" dependencies = [ "cryptography>=41", "huggingface-hub>=0.20", + "zstandard>=0.22", ] [project.scripts] diff --git a/tests/test_two_node_pipeline.py b/tests/test_two_node_pipeline.py index 514e673..7e8ecf9 100644 --- a/tests/test_two_node_pipeline.py +++ b/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()