316 lines
11 KiB
Python
316 lines
11 KiB
Python
"""Stub node HTTP server for the Distributed Inference Network."""
|
|
|
|
import http.server
|
|
import json
|
|
import threading
|
|
import urllib.parse
|
|
from pathlib import Path
|
|
|
|
from .downloader import compute_shard_checksum, write_shard_archive
|
|
|
|
# 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_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):
|
|
def __init__(
|
|
self,
|
|
addr,
|
|
handler,
|
|
shard_start: int,
|
|
shard_end: int,
|
|
is_last_shard: bool,
|
|
response_prefix: str,
|
|
model: str,
|
|
shard_path: Path | None,
|
|
):
|
|
super().__init__(addr, handler)
|
|
self.shard_start = shard_start
|
|
self.shard_end = shard_end
|
|
self.is_last_shard = is_last_shard
|
|
self.response_prefix = response_prefix
|
|
self.model = model
|
|
self.shard_path = shard_path
|
|
self.received_activations: bool = False
|
|
self.forward_chunk_count: int = 0
|
|
|
|
|
|
class _StubHandler(http.server.BaseHTTPRequestHandler):
|
|
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
|
|
pass
|
|
|
|
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()
|
|
|
|
def do_GET(self):
|
|
parsed = urllib.parse.urlparse(self.path)
|
|
if parsed.path == "/v1/shards/download":
|
|
self._handle_shard_download(parsed)
|
|
else:
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
|
|
def _handle_shard_download(self, parsed: urllib.parse.ParseResult):
|
|
server: _StubHTTPServer = self.server # type: ignore[assignment]
|
|
params = urllib.parse.parse_qs(parsed.query)
|
|
model = params.get("model", [""])[0]
|
|
try:
|
|
shard_start = int(params.get("shard_start", [""])[0])
|
|
shard_end = int(params.get("shard_end", [""])[0])
|
|
except ValueError:
|
|
self.send_response(400)
|
|
self.end_headers()
|
|
return
|
|
|
|
if (
|
|
server.shard_path is None
|
|
or not server.shard_path.exists()
|
|
or model != server.model
|
|
or shard_start != server.shard_start
|
|
or shard_end != server.shard_end
|
|
):
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
return
|
|
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/x-tar")
|
|
self.send_header("X-Meshnet-Shard-Checksum", compute_shard_checksum(server.shard_path))
|
|
self.end_headers()
|
|
write_shard_archive(server.shard_path, self.wfile)
|
|
|
|
def _handle_infer(self):
|
|
server: _StubHTTPServer = self.server # type: ignore[assignment]
|
|
length = int(self.headers.get("Content-Length", 0))
|
|
body = json.loads(self.rfile.read(length))
|
|
|
|
if "activations" in body:
|
|
# Shard-to-shard call: this node receives activations from the previous shard.
|
|
server.received_activations = True
|
|
prompt = body["activations"].get("context", {}).get("prompt", "")
|
|
if server.is_last_shard:
|
|
payload = json.dumps({"text": f"{server.response_prefix} {prompt}"}).encode()
|
|
else:
|
|
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", [])
|
|
last_content = messages[-1]["content"] if messages else ""
|
|
if server.is_last_shard:
|
|
# Single-node shortcut: no activation hop needed.
|
|
payload = json.dumps({"text": f"{server.response_prefix} {last_content}"}).encode()
|
|
else:
|
|
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")
|
|
self.send_header("Content-Length", str(len(payload)))
|
|
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.
|
|
|
|
shard_start / shard_end define which transformer layer range this node owns.
|
|
is_last_shard controls whether the node returns a text response (True) or
|
|
activation tensors (False) after processing its shard.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
host: str = "127.0.0.1",
|
|
port: int = 0,
|
|
shard_start: int = 0,
|
|
shard_end: int = 31,
|
|
is_last_shard: bool = True,
|
|
response_prefix: str = "stub response to:",
|
|
model: str = "stub-model",
|
|
shard_path: Path | None = None,
|
|
):
|
|
self._host = host
|
|
self._requested_port = port
|
|
if shard_start > shard_end:
|
|
raise ValueError("shard_start must be less than or equal to shard_end")
|
|
self._shard_start = shard_start
|
|
self._shard_end = shard_end
|
|
self._is_last_shard = is_last_shard
|
|
self._response_prefix = response_prefix
|
|
self._model = model
|
|
self._shard_path = shard_path
|
|
self._server: _StubHTTPServer | None = None
|
|
self._thread: threading.Thread | None = None
|
|
self.port: int | None = None
|
|
|
|
@property
|
|
def received_activations(self) -> bool:
|
|
"""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")
|
|
|
|
self._server = _StubHTTPServer(
|
|
(self._host, self._requested_port),
|
|
_StubHandler,
|
|
self._shard_start,
|
|
self._shard_end,
|
|
self._is_last_shard,
|
|
self._response_prefix,
|
|
self._model,
|
|
self._shard_path,
|
|
)
|
|
self.port = self._server.server_address[1]
|
|
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
|
self._thread.start()
|
|
return self.port
|
|
|
|
def stop(self) -> None:
|
|
if self._server is None:
|
|
return
|
|
|
|
self._server.shutdown()
|
|
self._server.server_close()
|
|
if self._thread is not None:
|
|
self._thread.join(timeout=1)
|
|
self._server = None
|
|
self._thread = None
|
|
self.port = None
|