feat: add real PyTorch model backend
This commit is contained in:
Dobromir Popov
271
packages/node/meshnet_node/torch_server.py
Normal file
271
packages/node/meshnet_node/torch_server.py
Normal file
@@ -0,0 +1,271 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user