554 lines
21 KiB
Python
554 lines
21 KiB
Python
"""HTTP server for real PyTorch-backed shard nodes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import http.server
|
|
import json
|
|
import sys
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.parse
|
|
import urllib.request
|
|
import uuid
|
|
from typing import Any
|
|
|
|
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,
|
|
tracker_mode: bool = False,
|
|
tracker_url: str | None = None,
|
|
):
|
|
super().__init__(addr, handler)
|
|
self.backend = backend
|
|
self.received_activations = False
|
|
self.forward_chunk_count = 0
|
|
self.tracker_mode = tracker_mode
|
|
self.tracker_url = tracker_url
|
|
|
|
|
|
class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
|
|
pass
|
|
|
|
def do_POST(self):
|
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
|
if self.path == "/forward":
|
|
self._handle_forward()
|
|
elif self.path == "/v1/infer":
|
|
self._handle_infer()
|
|
elif self.path == "/v1/chat/completions" and server.tracker_mode:
|
|
self._handle_chat_completions()
|
|
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)
|
|
|
|
def _handle_chat_completions(self) -> None:
|
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
|
body = self._read_json_body()
|
|
if body is None:
|
|
return
|
|
messages = body.get("messages", [])
|
|
if not isinstance(messages, list):
|
|
messages = []
|
|
stream = bool(body.get("stream", False))
|
|
model_name = str(body.get("model", ""))
|
|
max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 256)
|
|
temperature = float(body.get("temperature") or 1.0)
|
|
top_p = float(body.get("top_p") or 1.0)
|
|
|
|
# Fast path: this node owns the complete model — use HF generate() with KV cache.
|
|
# Avoids the single-token-per-forward-pass limitation of the distributed path.
|
|
if server.backend.is_head and server.backend.is_tail:
|
|
try:
|
|
if stream:
|
|
self._stream_openai_response(
|
|
server.backend.generate_text_streaming(messages, max_tokens, temperature, top_p),
|
|
model_name,
|
|
)
|
|
else:
|
|
text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
|
|
self._send_openai_response(text, model_name, False, messages)
|
|
except Exception as exc:
|
|
self._send_json(500, {"error": f"generation failed: {exc}"})
|
|
return
|
|
|
|
# Distributed path: encode prompt at the head, forward activations along the route.
|
|
prompt = " ".join(
|
|
str(m.get("content", ""))
|
|
for m in messages
|
|
if isinstance(m, dict) and m.get("role") == "user"
|
|
)
|
|
try:
|
|
payload = server.backend.encode_prompt(prompt)
|
|
except Exception as exc:
|
|
self._send_json(500, {"error": f"encode_prompt failed: {exc}"})
|
|
return
|
|
remaining_route = self._get_remaining_route(model_name)
|
|
result_text = self._run_downstream_pipeline(payload, remaining_route)
|
|
self._send_openai_response(result_text, model_name, stream, messages)
|
|
|
|
def _get_remaining_route(self, model: str) -> list[str]:
|
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
|
if server.tracker_url is None:
|
|
return []
|
|
try:
|
|
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(model)}"
|
|
with urllib.request.urlopen(url, timeout=5.0) as r:
|
|
route_resp = json.loads(r.read())
|
|
route = route_resp.get("route", [])
|
|
# Skip the first node in the route (self) since we're already the head
|
|
return list(route[1:])
|
|
except Exception:
|
|
return []
|
|
|
|
def _run_downstream_pipeline(self, payload: object, route: list[str]) -> str:
|
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
|
if not route:
|
|
# Partial shard at tail: decode the activation from the previous node.
|
|
# Full single-node (head+tail) is handled before entering this method.
|
|
if server.backend.is_tail:
|
|
try:
|
|
tensor = server.backend.torch.frombuffer(
|
|
bytearray(payload.body), # type: ignore[union-attr]
|
|
dtype=server.backend.torch.bfloat16,
|
|
).reshape(payload.shape).to(server.backend.device) # type: ignore[union-attr]
|
|
return server.backend.decode_tail(tensor)
|
|
except Exception as exc:
|
|
return f"decode error: {exc}"
|
|
return "no downstream route available for non-tail shard"
|
|
|
|
session = str(uuid.uuid4())
|
|
shape = payload.shape # type: ignore[union-attr]
|
|
attn_mask = payload.attention_mask_header # type: ignore[union-attr]
|
|
pos_ids = payload.position_ids_header # type: ignore[union-attr]
|
|
current_body = payload.body # type: ignore[union-attr]
|
|
current_shape = shape
|
|
current_attn = attn_mask
|
|
current_pos = pos_ids
|
|
|
|
for hop_index, node_url in enumerate(route):
|
|
headers: dict[str, str] = {
|
|
"Content-Type": "application/octet-stream",
|
|
"X-Meshnet-Wire": _WIRE_VERSION,
|
|
"X-Meshnet-Shape": ",".join(str(d) for d in current_shape),
|
|
"X-Meshnet-Dtype": "bfloat16",
|
|
"X-Meshnet-Session": session,
|
|
"X-Meshnet-Chunk-Index": "0",
|
|
"X-Meshnet-Chunk-Total": "1",
|
|
"X-Meshnet-Hop-Index": str(hop_index),
|
|
}
|
|
if current_attn:
|
|
headers["X-Meshnet-Attn-Mask"] = current_attn
|
|
if current_pos:
|
|
headers["X-Meshnet-Position-Ids"] = current_pos
|
|
req = urllib.request.Request(
|
|
f"{node_url}/forward",
|
|
data=current_body,
|
|
headers=headers,
|
|
method="POST",
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=10.0) as r:
|
|
resp_body = r.read()
|
|
resp_headers = {k.lower(): v for k, v in r.headers.items()}
|
|
except Exception as exc:
|
|
return f"pipeline error at {node_url}: {exc}"
|
|
content_type = resp_headers.get("content-type", "")
|
|
if "application/json" in content_type:
|
|
try:
|
|
data = json.loads(resp_body)
|
|
return str(data.get("text", ""))
|
|
except json.JSONDecodeError:
|
|
return resp_body.decode("utf-8", errors="replace")
|
|
# Binary activation — update and forward to next node
|
|
shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape))
|
|
current_shape = _parse_shape(shape_header)
|
|
current_body = resp_body
|
|
current_attn = resp_headers.get("x-meshnet-attn-mask")
|
|
current_pos = resp_headers.get("x-meshnet-position-ids")
|
|
return ""
|
|
|
|
def _stream_openai_response(self, token_iter, model: str) -> None:
|
|
"""Stream tokens from an iterator as SSE chunks."""
|
|
chunk_id = "chatcmpl-node"
|
|
created = int(time.time())
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
|
|
self.send_header("Cache-Control", "no-cache")
|
|
self.end_headers()
|
|
|
|
def _emit(data: str) -> None:
|
|
self.wfile.write(f"data: {data}\n\n".encode())
|
|
self.wfile.flush()
|
|
|
|
_emit(json.dumps({
|
|
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
|
|
}))
|
|
for token_text in token_iter:
|
|
if not token_text:
|
|
continue
|
|
_emit(json.dumps({
|
|
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {"content": token_text}, "finish_reason": None}],
|
|
}))
|
|
_emit(json.dumps({
|
|
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
|
}))
|
|
self.wfile.write(b"data: [DONE]\n\n")
|
|
self.wfile.flush()
|
|
|
|
def _send_openai_response(
|
|
self,
|
|
text: str,
|
|
model: str,
|
|
stream: bool,
|
|
messages: list[dict] | None = None,
|
|
) -> None:
|
|
chunk_id = "chatcmpl-node"
|
|
created = int(time.time())
|
|
if not stream:
|
|
usage = _usage_for_response(self.server.backend, messages or [], text) # type: ignore[attr-defined]
|
|
self._send_json(200, {
|
|
"id": chunk_id,
|
|
"object": "chat.completion",
|
|
"created": created,
|
|
"model": model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": text},
|
|
"finish_reason": "stop",
|
|
}],
|
|
"usage": usage,
|
|
})
|
|
return
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
|
|
self.send_header("Cache-Control", "no-cache")
|
|
self.end_headers()
|
|
|
|
def _emit(data: str) -> None:
|
|
self.wfile.write(f"data: {data}\n\n".encode())
|
|
self.wfile.flush()
|
|
|
|
_emit(json.dumps({
|
|
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
|
|
}))
|
|
_emit(json.dumps({
|
|
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}],
|
|
}))
|
|
_emit(json.dumps({
|
|
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
|
}))
|
|
self.wfile.write(b"data: [DONE]\n\n")
|
|
self.wfile.flush()
|
|
|
|
|
|
def _usage_for_response(backend: object, messages: list[dict], completion_text: str) -> dict[str, int]:
|
|
prompt_tokens = _backend_token_count(
|
|
backend,
|
|
"count_prompt_tokens",
|
|
messages,
|
|
fallback=_fallback_message_token_count(messages),
|
|
)
|
|
completion_tokens = _backend_token_count(
|
|
backend,
|
|
"count_text_tokens",
|
|
completion_text,
|
|
fallback=_fallback_text_token_count(completion_text),
|
|
)
|
|
return {
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": prompt_tokens + completion_tokens,
|
|
}
|
|
|
|
|
|
def _backend_token_count(backend: object, method_name: str, value: object, fallback: int) -> int:
|
|
method: Any = getattr(backend, method_name, None)
|
|
if callable(method):
|
|
try:
|
|
return max(0, int(method(value)))
|
|
except Exception:
|
|
pass
|
|
return max(0, int(fallback))
|
|
|
|
|
|
def _fallback_message_token_count(messages: list[dict]) -> int:
|
|
text = " ".join(
|
|
str(message.get("content", ""))
|
|
for message in messages
|
|
if isinstance(message, dict)
|
|
)
|
|
return _fallback_text_token_count(text)
|
|
|
|
|
|
def _fallback_text_token_count(text: str) -> int:
|
|
parts = text.split()
|
|
if parts:
|
|
return len(parts)
|
|
return 1 if text else 0
|
|
|
|
|
|
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,
|
|
tracker_mode: bool | None = None,
|
|
tracker_url: str | None = None,
|
|
) -> None:
|
|
self._host = host
|
|
self._requested_port = port
|
|
self._backend = backend or _load_backend(
|
|
model_id,
|
|
shard_start,
|
|
shard_end,
|
|
quantization,
|
|
)
|
|
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
|
|
self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
|
|
self._tracker_url = tracker_url
|
|
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._tracker_mode,
|
|
self._tracker_url,
|
|
)
|
|
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
|