272 lines
9.6 KiB
Python
272 lines
9.6 KiB
Python
"""HTTP server for real PyTorch-backed shard nodes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import http.server
|
|
import json
|
|
import sys
|
|
import threading
|
|
|
|
from .model_backend import (
|
|
InsufficientVRAMError,
|
|
MissingModelDependencyError,
|
|
Quantization,
|
|
TorchModelShard,
|
|
validate_quantization,
|
|
)
|
|
from .server import (
|
|
_WIRE_VERSION,
|
|
_compress_body,
|
|
_decompress_body,
|
|
_parse_shape,
|
|
_validate_activation_body,
|
|
)
|
|
|
|
|
|
class _TorchHTTPServer(http.server.HTTPServer):
|
|
def __init__(self, addr, handler, backend: TorchModelShard):
|
|
super().__init__(addr, handler)
|
|
self.backend = backend
|
|
self.received_activations = False
|
|
self.forward_chunk_count = 0
|
|
|
|
|
|
class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
|
|
pass
|
|
|
|
def do_POST(self):
|
|
if self.path == "/forward":
|
|
self._handle_forward()
|
|
elif self.path == "/v1/infer":
|
|
self._handle_infer()
|
|
else:
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
|
|
def _handle_infer(self) -> None:
|
|
body = self._read_json_body()
|
|
if body is None:
|
|
return
|
|
messages = body.get("messages", [])
|
|
prompt = ""
|
|
if isinstance(messages, list) and messages:
|
|
last = messages[-1]
|
|
if isinstance(last, dict):
|
|
prompt = str(last.get("content", ""))
|
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
|
try:
|
|
payload = server.backend.encode_prompt(prompt)
|
|
if server.backend.is_tail:
|
|
text = server.backend.decode_tail(
|
|
server.backend.torch.frombuffer(
|
|
bytearray(payload.body),
|
|
dtype=server.backend.torch.bfloat16,
|
|
)
|
|
.reshape(payload.shape)
|
|
.to(server.backend.device)
|
|
)
|
|
self._send_json(200, {"text": text})
|
|
return
|
|
self._send_json(200, {"activations": {"shape": payload.shape, "dtype": "bfloat16"}})
|
|
except Exception as exc:
|
|
self._send_json(500, {"error": str(exc)})
|
|
|
|
def _handle_forward(self) -> None:
|
|
content_type = self.headers.get("Content-Type", "")
|
|
if content_type.startswith("application/json"):
|
|
self._handle_prompt_forward()
|
|
return
|
|
self._handle_binary_forward()
|
|
|
|
def _handle_prompt_forward(self) -> None:
|
|
body = self._read_json_body()
|
|
if body is None:
|
|
return
|
|
prompt = str(body.get("prompt", ""))
|
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
|
try:
|
|
payload = server.backend.encode_prompt(prompt)
|
|
except Exception as exc:
|
|
self._send_json(400, {"error": str(exc)})
|
|
return
|
|
self._send_activation(payload)
|
|
|
|
def _handle_binary_forward(self) -> None:
|
|
server: _TorchHTTPServer = 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)
|
|
if dtype != "bfloat16":
|
|
raise ValueError("real model backend requires bfloat16 activation input")
|
|
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
|
|
|
|
try:
|
|
result = server.backend.forward_bytes(
|
|
raw_body,
|
|
shape,
|
|
self.headers.get("X-Meshnet-Attn-Mask"),
|
|
self.headers.get("X-Meshnet-Position-Ids"),
|
|
)
|
|
except Exception as exc:
|
|
self._send_json(500, {"error": str(exc)})
|
|
return
|
|
|
|
if isinstance(result, str):
|
|
self._send_json(200, {"text": result})
|
|
return
|
|
|
|
response_body = _compress_body(result.body, encoding)
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/octet-stream")
|
|
self.send_header("Content-Length", str(len(response_body)))
|
|
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
|
|
self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in result.shape))
|
|
self.send_header("X-Meshnet-Dtype", "bfloat16")
|
|
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 result.attention_mask_header:
|
|
self.send_header("X-Meshnet-Attn-Mask", result.attention_mask_header)
|
|
if result.position_ids_header:
|
|
self.send_header("X-Meshnet-Position-Ids", result.position_ids_header)
|
|
self.end_headers()
|
|
self.wfile.write(response_body)
|
|
|
|
def _send_activation(self, payload) -> None:
|
|
body = payload.body
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/octet-stream")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
|
|
self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in payload.shape))
|
|
self.send_header("X-Meshnet-Dtype", "bfloat16")
|
|
if payload.attention_mask_header:
|
|
self.send_header("X-Meshnet-Attn-Mask", payload.attention_mask_header)
|
|
if payload.position_ids_header:
|
|
self.send_header("X-Meshnet-Position-Ids", payload.position_ids_header)
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def _read_json_body(self) -> dict | None:
|
|
length = int(self.headers.get("Content-Length", 0))
|
|
try:
|
|
body = json.loads(self.rfile.read(length) or b"{}")
|
|
except (json.JSONDecodeError, ValueError):
|
|
self._send_json(400, {"error": "invalid JSON body"})
|
|
return None
|
|
if not isinstance(body, dict):
|
|
self._send_json(400, {"error": "JSON body must be an object"})
|
|
return None
|
|
return body
|
|
|
|
def _send_json(self, status: int, data: dict) -> None:
|
|
payload = json.dumps(data).encode()
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(payload)))
|
|
self.end_headers()
|
|
self.wfile.write(payload)
|
|
|
|
|
|
class TorchNodeServer:
|
|
"""HTTP server backed by a HuggingFace causal language model shard."""
|
|
|
|
def __init__(
|
|
self,
|
|
host: str = "127.0.0.1",
|
|
port: int = 0,
|
|
model_id: str = "openai-community/gpt2",
|
|
shard_start: int = 0,
|
|
shard_end: int = 6,
|
|
quantization: str = "bfloat16",
|
|
backend: TorchModelShard | None = None,
|
|
) -> None:
|
|
self._host = host
|
|
self._requested_port = port
|
|
self._backend = backend or _load_backend(
|
|
model_id,
|
|
shard_start,
|
|
shard_end,
|
|
quantization,
|
|
)
|
|
self._server: _TorchHTTPServer | None = None
|
|
self._thread: threading.Thread | None = None
|
|
self.port: int | None = None
|
|
|
|
@property
|
|
def backend(self) -> TorchModelShard:
|
|
return self._backend
|
|
|
|
@property
|
|
def received_activations(self) -> bool:
|
|
return self._server.received_activations if self._server is not None else False
|
|
|
|
@property
|
|
def forward_chunk_count(self) -> int:
|
|
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("TorchNodeServer is already running")
|
|
self._server = _TorchHTTPServer(
|
|
(self._host, self._requested_port),
|
|
_TorchHandler,
|
|
self._backend,
|
|
)
|
|
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
|
|
|
|
|
|
def _load_backend(
|
|
model_id: str,
|
|
shard_start: int,
|
|
shard_end: int,
|
|
quantization: str,
|
|
) -> TorchModelShard:
|
|
from .model_backend import load_torch_shard
|
|
|
|
quant = validate_quantization(quantization)
|
|
try:
|
|
return load_torch_shard(model_id, shard_start, shard_end, quant)
|
|
except MissingModelDependencyError:
|
|
raise
|
|
except InsufficientVRAMError as exc:
|
|
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
|
|
raise
|