1687 lines
69 KiB
Python
1687 lines
69 KiB
Python
"""HTTP server for real PyTorch-backed shard nodes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import http.client
|
|
import http.server
|
|
import json
|
|
import sys
|
|
import threading
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any, Mapping
|
|
|
|
from .model_backend import (
|
|
InsufficientVRAMError,
|
|
KVCacheMiss,
|
|
MissingModelDependencyError,
|
|
Quantization,
|
|
TailTokenResult,
|
|
TorchModelShard,
|
|
_tensor_from_bfloat16_bytes,
|
|
validate_quantization,
|
|
)
|
|
from .seam_telemetry import GenerationTelemetry
|
|
from .activation_compression import (
|
|
CompressionPolicies,
|
|
CompressionPolicy,
|
|
compress_activation,
|
|
decompress_activation,
|
|
)
|
|
|
|
|
|
class _PipelineCacheMiss(Exception):
|
|
"""A downstream hop reported 409 cache_miss — head must re-prefill."""
|
|
|
|
|
|
class _RelayRequestUncertainError(ConnectionError):
|
|
"""A relay request may have reached the peer but produced no response."""
|
|
|
|
|
|
class _DirectRequestUncertainError(ConnectionError):
|
|
"""A direct request may have reached the downstream node but did not finish."""
|
|
|
|
|
|
from .server import (
|
|
_WIRE_VERSION,
|
|
_parse_shape,
|
|
_validate_activation_body,
|
|
)
|
|
|
|
|
|
def _endpoint_key(url: str) -> str:
|
|
"""Normalize http(s) endpoints for host:port comparison."""
|
|
parsed = urllib.parse.urlparse(url.rstrip("/"))
|
|
host = (parsed.hostname or "").lower()
|
|
if not host:
|
|
return url.rstrip("/").lower()
|
|
port = parsed.port
|
|
if port is None:
|
|
port = 443 if parsed.scheme == "https" else 80
|
|
return f"{host}:{port}"
|
|
|
|
|
|
def _own_endpoint_key(server: _TorchHTTPServer) -> str:
|
|
advertised = getattr(server, "advertised_endpoint", None)
|
|
if advertised:
|
|
return _endpoint_key(advertised)
|
|
host, port = server.server_address
|
|
return _endpoint_key(f"http://{host}:{port}")
|
|
|
|
|
|
def _clamp_downstream_hops(
|
|
hops: list[dict],
|
|
backend: TorchModelShard | None,
|
|
) -> list[dict]:
|
|
"""Ensure downstream start_layer continues after this shard's layers."""
|
|
if not hops or backend is None:
|
|
return hops
|
|
shard_end = getattr(backend, "shard_end", None)
|
|
if shard_end is None:
|
|
return hops
|
|
min_start = int(shard_end) + 1
|
|
clamped: list[dict] = []
|
|
for hop in hops:
|
|
adjusted = dict(hop)
|
|
if int(adjusted.get("start_layer", 0)) < min_start:
|
|
adjusted["start_layer"] = min_start
|
|
clamped.append(adjusted)
|
|
return clamped
|
|
|
|
|
|
def _format_downstream_route(hops: list[dict]) -> str:
|
|
return ", ".join(
|
|
f"{h['endpoint']}@{h.get('start_layer', 0)}" for h in hops
|
|
)
|
|
|
|
|
|
def _write_progress_line(state: list[bool], message: str, *, final: bool = False) -> None:
|
|
"""Rewrite one in-place progress line (\\r) or finish with a newline."""
|
|
if final:
|
|
if state[0]:
|
|
sys.stdout.write("\r" + message + "\n")
|
|
state[0] = False
|
|
else:
|
|
print(message, flush=True)
|
|
return
|
|
if state[0]:
|
|
sys.stdout.write("\r" + message)
|
|
else:
|
|
sys.stdout.write(message)
|
|
state[0] = True
|
|
sys.stdout.flush()
|
|
|
|
|
|
class _RelayHopClient:
|
|
"""Persistent relay connection scoped to one generation handler."""
|
|
|
|
def __init__(self, relay_addr: str, timeout: float = 120.0) -> None:
|
|
self.relay_addr = relay_addr
|
|
self.timeout = timeout
|
|
self._ws = None
|
|
self._lock = threading.RLock()
|
|
|
|
def request(
|
|
self,
|
|
path: str,
|
|
body: bytes,
|
|
headers: dict[str, str],
|
|
) -> tuple[int, dict[str, str], bytes]:
|
|
import websockets.sync.client as wsc # type: ignore[import]
|
|
|
|
from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size
|
|
|
|
with self._lock:
|
|
if self._ws is None:
|
|
self._ws = wsc.connect(
|
|
self.relay_addr,
|
|
open_timeout=self.timeout,
|
|
max_size=ws_max_size(),
|
|
compression=None,
|
|
)
|
|
request_id = headers.get("X-Meshnet-Activation-Id") or uuid.uuid4().hex
|
|
frame = encode_binary_frame({
|
|
"request_id": request_id,
|
|
"method": "POST",
|
|
"path": path,
|
|
"headers": headers,
|
|
}, body)
|
|
try:
|
|
# A send failure is uncertain too: bytes may already have been
|
|
# accepted by the kernel or peer before the exception surfaced.
|
|
self._ws.send(frame)
|
|
raw = self._ws.recv(timeout=self.timeout)
|
|
if isinstance(raw, (bytes, bytearray)):
|
|
resp_header, resp_body = decode_binary_frame(bytes(raw))
|
|
response_id = str(resp_header.get("request_id") or "")
|
|
status = int(resp_header.get("status", 503))
|
|
resp_headers = {
|
|
k.lower(): v for k, v in (resp_header.get("headers") or {}).items()
|
|
}
|
|
else:
|
|
resp = json.loads(raw)
|
|
response_id = str(resp.get("request_id") or "")
|
|
status = int(resp.get("status", 503))
|
|
resp_headers = {
|
|
k.lower(): v for k, v in (resp.get("headers") or {}).items()
|
|
}
|
|
body_b64 = resp.get("body_base64")
|
|
resp_body = (
|
|
base64.b64decode(body_b64)
|
|
if body_b64 else (resp.get("body") or "").encode()
|
|
)
|
|
if response_id and response_id != request_id:
|
|
raise ValueError("relay response request_id did not match request")
|
|
return status, resp_headers, resp_body
|
|
except Exception as exc:
|
|
self.close()
|
|
raise _RelayRequestUncertainError(
|
|
"relay connection failed after forwarding request; refusing replay"
|
|
) from exc
|
|
|
|
def close(self) -> None:
|
|
with self._lock:
|
|
if self._ws is not None:
|
|
try:
|
|
self._ws.close()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
self._ws = None
|
|
|
|
|
|
class _DirectHopClient:
|
|
"""One serialized HTTP/1.1 connection to one downstream hop.
|
|
|
|
Generation handlers own these clients, so a cached Route Session reuses a
|
|
TCP connection without sharing it between concurrent Route Sessions.
|
|
"""
|
|
|
|
def __init__(self, endpoint: str, timeout: float = 120.0) -> None:
|
|
parsed = urllib.parse.urlsplit(endpoint.rstrip("/"))
|
|
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
|
|
raise ValueError(f"invalid downstream endpoint: {endpoint!r}")
|
|
self.timeout = timeout
|
|
self._scheme = parsed.scheme
|
|
self._host = parsed.hostname
|
|
self._port = parsed.port
|
|
self._base_path = parsed.path.rstrip("/")
|
|
self._connection: http.client.HTTPConnection | None = None
|
|
self._lock = threading.RLock()
|
|
|
|
def _connect(self) -> http.client.HTTPConnection:
|
|
connection_type = (
|
|
http.client.HTTPSConnection if self._scheme == "https" else http.client.HTTPConnection
|
|
)
|
|
return connection_type(self._host, self._port, timeout=self.timeout)
|
|
|
|
def request(
|
|
self, path: str, body: bytes, headers: Mapping[str, str],
|
|
) -> tuple[int, dict[str, str], bytes]:
|
|
request_path = f"{self._base_path}{path if path.startswith('/') else '/' + path}"
|
|
with self._lock:
|
|
if self._connection is None:
|
|
self._connection = self._connect()
|
|
try:
|
|
self._connection.request("POST", request_path, body=body, headers=dict(headers))
|
|
response = self._connection.getresponse()
|
|
response_body = response.read()
|
|
response_headers = {key.lower(): value for key, value in response.headers.items()}
|
|
return response.status, response_headers, response_body
|
|
except Exception as exc:
|
|
self.close()
|
|
raise _DirectRequestUncertainError(
|
|
"direct connection failed after forwarding request; refusing replay"
|
|
) from exc
|
|
|
|
def close(self) -> None:
|
|
with self._lock:
|
|
if self._connection is not None:
|
|
try:
|
|
self._connection.close()
|
|
finally:
|
|
self._connection = None
|
|
|
|
|
|
def _relay_hop(
|
|
relay_addr: str,
|
|
path: str,
|
|
body: bytes,
|
|
headers: dict[str, str],
|
|
timeout: float = 120.0,
|
|
) -> tuple[int, dict[str, str], bytes]:
|
|
"""Send one request through a short-lived relay client (compatibility API).
|
|
|
|
relay_addr is the wss://relay.../rpc/{peer_id} URL. The request and any
|
|
binary response travel as binary frames (JSON header + raw body); relay
|
|
error responses and legacy peers still answer with JSON text frames.
|
|
Returns (status, response_headers_lower, response_body).
|
|
Raises on connection failure so callers can fall back to direct.
|
|
"""
|
|
client = _RelayHopClient(relay_addr, timeout)
|
|
try:
|
|
return client.request(path, body, headers)
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
_COMPRESSION_POLICIES = CompressionPolicies()
|
|
|
|
|
|
def _maybe_compress_activation(
|
|
body: bytes, policy: CompressionPolicy | None = None,
|
|
) -> tuple[bytes, str | None]:
|
|
"""Compatibility wrapper for callers that only need wire body and encoding."""
|
|
result = compress_activation(body, policy or _COMPRESSION_POLICIES.for_condition("lan", "prefill"))
|
|
return result.body, result.encoding
|
|
|
|
|
|
def _is_cache_miss_body(body: bytes) -> bool:
|
|
try:
|
|
return json.loads(body).get("error") == "cache_miss"
|
|
except (json.JSONDecodeError, AttributeError, UnicodeDecodeError):
|
|
return False
|
|
|
|
|
|
def _response_error_snippet(body: bytes, limit: int = 500) -> str:
|
|
"""Return a compact error string from a downstream JSON/text response body."""
|
|
try:
|
|
payload = json.loads(body)
|
|
if isinstance(payload, dict):
|
|
message = payload.get("error") or payload.get("detail") or payload
|
|
return str(message)[:limit]
|
|
except (json.JSONDecodeError, TypeError, UnicodeDecodeError):
|
|
pass
|
|
return body.decode("utf-8", errors="replace")[:limit]
|
|
|
|
|
|
class _TorchHTTPServer(http.server.HTTPServer):
|
|
def __init__(
|
|
self,
|
|
addr,
|
|
handler,
|
|
backend: TorchModelShard,
|
|
tracker_mode: bool = False,
|
|
tracker_url: str | None = None,
|
|
route_timeout: float = 30.0,
|
|
debug: bool = False,
|
|
max_loaded_shards: int = 1,
|
|
):
|
|
super().__init__(addr, handler)
|
|
self.backend = backend
|
|
self.backends: dict[str, TorchModelShard] = {backend.model_id: backend}
|
|
self.received_activations = False
|
|
self.forward_chunk_count = 0
|
|
self.tracker_mode = tracker_mode
|
|
self.tracker_url = tracker_url
|
|
self.route_timeout = route_timeout
|
|
self.debug = debug
|
|
self.max_loaded_shards = max(1, max_loaded_shards)
|
|
self.advertised_endpoint: str | None = None
|
|
self.total_requests: int = 0
|
|
self.failed_requests: int = 0
|
|
self.queue_depth: int = 0
|
|
self._stats_lock = threading.Lock()
|
|
self._active_requests: dict[str, dict[str, Any]] = {}
|
|
self._decode_log: dict[str, dict[str, float]] = {}
|
|
|
|
def note_decode_step(
|
|
self, session: str, now: float | None = None,
|
|
) -> int | None:
|
|
"""Count one decode forward; return the cumulative step count when a
|
|
log line is due (first step of a session, then every 5s), else None."""
|
|
if now is None:
|
|
now = time.monotonic()
|
|
with self._stats_lock:
|
|
rec = self._decode_log.get(session)
|
|
if rec is None:
|
|
if len(self._decode_log) >= 64:
|
|
stale = [
|
|
sid for sid, r in self._decode_log.items()
|
|
if now - r["seen"] > 600.0
|
|
]
|
|
for sid in stale:
|
|
del self._decode_log[sid]
|
|
while len(self._decode_log) >= 64:
|
|
self._decode_log.pop(next(iter(self._decode_log)))
|
|
self._decode_log[session] = {"steps": 1.0, "logged": now, "seen": now}
|
|
return 1
|
|
rec["steps"] += 1
|
|
rec["seen"] = now
|
|
if now - rec["logged"] >= 5.0:
|
|
rec["logged"] = now
|
|
return int(rec["steps"])
|
|
return None
|
|
|
|
def snapshot_current_requests(self) -> list[dict[str, Any]]:
|
|
"""In-flight request snapshots for tracker heartbeats."""
|
|
now = time.monotonic()
|
|
with self._stats_lock:
|
|
out: list[dict[str, Any]] = []
|
|
for rec in self._active_requests.values():
|
|
elapsed = max(now - float(rec["started"]), 1e-6)
|
|
tokens = int(rec.get("tokens") or 0)
|
|
out.append({
|
|
"request_id": str(rec["request_id"]),
|
|
"model": str(rec.get("model") or ""),
|
|
"kind": str(rec.get("kind") or "chat"),
|
|
"tokens": tokens,
|
|
"elapsed_seconds": round(elapsed, 1),
|
|
"tokens_per_sec": round(tokens / elapsed, 2) if tokens > 0 else 0.0,
|
|
"routing_complete": bool(rec.get("routing_complete")),
|
|
"telemetry": rec["telemetry"].snapshot(now=now) if rec.get("telemetry") else None,
|
|
})
|
|
return out
|
|
|
|
def resolve_backend(self, model_name: str | None) -> TorchModelShard | None:
|
|
if not model_name:
|
|
return self.backend
|
|
wanted = model_name.strip().lower()
|
|
for key, shard_backend in self.backends.items():
|
|
key_l = key.lower()
|
|
if key_l == wanted or key_l.rsplit("/", 1)[-1] == wanted:
|
|
return shard_backend
|
|
return self.backend
|
|
|
|
def chat_enabled(self) -> bool:
|
|
return any(
|
|
shard_backend.is_head
|
|
for shard_backend in self.backends.values()
|
|
)
|
|
|
|
|
|
class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|
# HTTP/1.1 is required for Route Session-owned downstream connections.
|
|
# Finite responses below provide Content-Length; streams are chunked.
|
|
protocol_version = "HTTP/1.1"
|
|
|
|
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
|
|
pass
|
|
|
|
def _request_id(self) -> str:
|
|
return (
|
|
self.headers.get("X-Meshnet-Request-Id")
|
|
or self.headers.get("X-Request-Id")
|
|
or f"local-{time.time_ns():x}"
|
|
)
|
|
|
|
def _request_log_suffix(self) -> str:
|
|
activation_id = self.headers.get("X-Meshnet-Activation-Id")
|
|
if activation_id:
|
|
return f" activation_id={activation_id}"
|
|
req_id = self.headers.get("X-Meshnet-Request-Id") or self.headers.get("X-Request-Id")
|
|
return f" request_id={req_id}" if req_id else ""
|
|
|
|
def _track_request_begin(
|
|
self,
|
|
server: "_TorchHTTPServer",
|
|
request_id: str,
|
|
model: str,
|
|
) -> None:
|
|
with server._stats_lock:
|
|
server._active_requests[request_id] = {
|
|
"request_id": request_id,
|
|
"model": model,
|
|
"kind": "chat",
|
|
"started": time.monotonic(),
|
|
"tokens": 0,
|
|
"routing_complete": False,
|
|
"telemetry": None,
|
|
}
|
|
|
|
def _track_request_progress(
|
|
self,
|
|
server: "_TorchHTTPServer",
|
|
request_id: str,
|
|
*,
|
|
tokens: int,
|
|
routing_complete: bool = False,
|
|
) -> None:
|
|
with server._stats_lock:
|
|
rec = server._active_requests.get(request_id)
|
|
if rec is None:
|
|
return
|
|
rec["tokens"] = tokens
|
|
if routing_complete:
|
|
rec["routing_complete"] = True
|
|
|
|
def _track_request_end(self, server: "_TorchHTTPServer", request_id: str) -> None:
|
|
with server._stats_lock:
|
|
server._active_requests.pop(request_id, None)
|
|
|
|
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.chat_enabled():
|
|
self._handle_chat_completions()
|
|
else:
|
|
self.send_response(404)
|
|
self.send_header("Content-Length", "0")
|
|
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_activation(body, encoding).body
|
|
_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.send_header("Content-Length", "0")
|
|
self.end_headers()
|
|
return
|
|
|
|
server.forward_chunk_count += 1
|
|
hop_index = int(self.headers.get("X-Meshnet-Hop-Index", "0"))
|
|
if hop_index > 0:
|
|
server.received_activations = True
|
|
# Session KV-cache protocol: prefill establishes per-session state on
|
|
# this node's layer range; decode reuses it. Absent header = legacy
|
|
# stateless call (also the signature fake backends implement).
|
|
cache_mode = self.headers.get("X-Meshnet-Cache")
|
|
if chunk_index_value == 0:
|
|
shard_start = getattr(server.backend, "shard_start", "?")
|
|
shard_end = getattr(server.backend, "shard_end", "?")
|
|
if cache_mode == "decode":
|
|
# One decode forward arrives per generated token — log a
|
|
# periodic per-session summary instead of one line per token.
|
|
steps = server.note_decode_step(session)
|
|
if steps is not None:
|
|
print(
|
|
f" [node] decoding layers={shard_start}-{shard_end} "
|
|
f"session={session[:8]} steps={steps}"
|
|
f"{self._request_log_suffix()}",
|
|
flush=True,
|
|
)
|
|
else:
|
|
print(
|
|
f" [node] forward hop={hop_index} "
|
|
f"layers={shard_start}-{shard_end} "
|
|
f"session={session[:8]}{self._request_log_suffix()}",
|
|
flush=True,
|
|
)
|
|
|
|
start_layer_header = self.headers.get("X-Meshnet-Start-Layer")
|
|
start_layer = int(start_layer_header) if start_layer_header else None
|
|
forward_kwargs: dict[str, object] = {}
|
|
if cache_mode in ("prefill", "decode"):
|
|
past_len_header = self.headers.get("X-Meshnet-Past-Len")
|
|
forward_kwargs = {
|
|
"session_id": session,
|
|
"cache_mode": cache_mode,
|
|
"past_len": int(past_len_header) if past_len_header else None,
|
|
}
|
|
|
|
try:
|
|
result = server.backend.forward_bytes(
|
|
raw_body,
|
|
shape,
|
|
self.headers.get("X-Meshnet-Attn-Mask"),
|
|
self.headers.get("X-Meshnet-Position-Ids"),
|
|
start_layer=start_layer,
|
|
**forward_kwargs,
|
|
)
|
|
except KVCacheMiss as exc:
|
|
self._send_json(409, {"error": "cache_miss", "detail": str(exc)})
|
|
return
|
|
except Exception as exc:
|
|
print(
|
|
f" [node] forward failed layers={getattr(server.backend, 'shard_start', '?')}-"
|
|
f"{getattr(server.backend, 'shard_end', '?')} session={session[:8]}: {exc}"
|
|
f"{self._request_log_suffix()}",
|
|
flush=True,
|
|
)
|
|
self._send_json(500, {"error": str(exc)})
|
|
return
|
|
|
|
if isinstance(result, TailTokenResult):
|
|
self._send_json(200, {"text": result.text, "token_id": result.token_id})
|
|
return
|
|
if isinstance(result, str):
|
|
self._send_json(200, {"text": result})
|
|
return
|
|
|
|
route_condition = self.headers.get("X-Meshnet-Compression-Route", "lan")
|
|
phase_condition = cache_mode if cache_mode in {"prefill", "decode"} else "prefill"
|
|
response_compression = compress_activation(
|
|
result.body, _COMPRESSION_POLICIES.for_condition(route_condition, phase_condition),
|
|
)
|
|
response_body = response_compression.body
|
|
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 response_compression.encoding:
|
|
self.send_header("X-Meshnet-Encoding", response_compression.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()
|
|
try:
|
|
self.wfile.write(payload)
|
|
except BrokenPipeError:
|
|
pass # client disconnected before we could respond — not an error
|
|
|
|
def _handle_chat_completions(self) -> None:
|
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
|
request_id = self._request_id()
|
|
with server._stats_lock:
|
|
server.total_requests += 1
|
|
server.queue_depth += 1
|
|
try:
|
|
self._do_chat_completions(server, request_id)
|
|
finally:
|
|
self._track_request_end(server, request_id)
|
|
with server._stats_lock:
|
|
server.queue_depth -= 1
|
|
|
|
def _record_failed_request(self) -> None:
|
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
|
with server._stats_lock:
|
|
server.failed_requests += 1
|
|
|
|
def _do_chat_completions(self, server: "_TorchHTTPServer", request_id: str) -> None:
|
|
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", ""))
|
|
backend = server.resolve_backend(model_name)
|
|
if backend is None or not backend.is_head:
|
|
self._send_json(400, {"error": "model not loaded on this node"})
|
|
return
|
|
max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 5120)
|
|
temperature = float(body.get("temperature") or 1.0)
|
|
top_p = float(body.get("top_p") or 1.0)
|
|
|
|
self._track_request_begin(server, request_id, model_name)
|
|
print(
|
|
f" [node] processing chat model={model_name!r} stream={stream} "
|
|
f"max_tokens={max_tokens}{self._request_log_suffix()}",
|
|
flush=True,
|
|
)
|
|
|
|
# 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 backend.is_head and backend.is_tail:
|
|
gen_started = time.monotonic()
|
|
progress_line = [False]
|
|
try:
|
|
if stream:
|
|
token_count = 0
|
|
|
|
def _counting_stream():
|
|
nonlocal token_count
|
|
for token_text in backend.generate_text_streaming(
|
|
messages, max_tokens, temperature, top_p,
|
|
):
|
|
if token_text:
|
|
token_count += 1
|
|
self._track_request_progress(
|
|
server, request_id, tokens=token_count, routing_complete=True,
|
|
)
|
|
yield token_text
|
|
|
|
self._stream_openai_response(_counting_stream(), model_name)
|
|
elapsed = time.monotonic() - gen_started
|
|
tps = token_count / max(elapsed, 1e-6)
|
|
_write_progress_line(
|
|
progress_line,
|
|
f" [node] chat complete (stream) tokens={token_count} "
|
|
f"elapsed_s={elapsed:.1f} tps={tps:.2f}{self._request_log_suffix()}",
|
|
final=True,
|
|
)
|
|
else:
|
|
text = backend.generate_text(messages, max_tokens, temperature, top_p)
|
|
completion_tokens = _backend_token_count(
|
|
backend, "count_text_tokens", text, fallback=len(text.split()) or 1,
|
|
)
|
|
print(
|
|
f" [node] chat complete tokens={completion_tokens} "
|
|
f"elapsed_s={time.monotonic() - gen_started:.1f}{self._request_log_suffix()}",
|
|
flush=True,
|
|
)
|
|
self._send_openai_response(text, model_name, False, messages, backend=backend)
|
|
except Exception as exc:
|
|
self._record_failed_request()
|
|
print(
|
|
f" [node] chat failed after {time.monotonic() - gen_started:.1f}s: {exc}"
|
|
f"{self._request_log_suffix()}",
|
|
flush=True,
|
|
)
|
|
self._send_json(500, {"error": f"generation failed: {exc}"})
|
|
return
|
|
|
|
# Distributed path: autoregressive generation across shards with a
|
|
# sharded per-node KV cache. Step 0 prefills the full prompt through the
|
|
# route (each node caches state for its own layer range, keyed by a
|
|
# per-generation session id); steps 1+ send only the newest token's
|
|
# hidden state. A 409 cache_miss from any hop (eviction/restart/route
|
|
# change) falls back to a full re-prefill — the old stateless behavior.
|
|
remaining_route = self._get_remaining_route(model_name, backend=backend)
|
|
print(
|
|
f" [node] chat route model={model_name!r} max_tokens={max_tokens} "
|
|
f"downstream={remaining_route}",
|
|
flush=True,
|
|
)
|
|
if not remaining_route:
|
|
self._send_openai_response(
|
|
"error: no downstream route — check tracker connectivity",
|
|
model_name, False, messages, backend=backend,
|
|
)
|
|
return
|
|
|
|
# Format with chat template so the model knows it's in assistant mode.
|
|
try:
|
|
if hasattr(backend.tokenizer, "apply_chat_template"):
|
|
prompt_text: str = backend.tokenizer.apply_chat_template(
|
|
messages, add_generation_prompt=True, tokenize=False,
|
|
)
|
|
else:
|
|
raise AttributeError("no apply_chat_template")
|
|
except Exception:
|
|
prompt_text = " ".join(
|
|
str(m.get("content", ""))
|
|
for m in messages
|
|
if isinstance(m, dict) and m.get("role") == "user"
|
|
)
|
|
|
|
eos_token: str = getattr(backend.tokenizer, "eos_token", "") or ""
|
|
generated: list[str] = []
|
|
current_text = prompt_text
|
|
|
|
session_id = str(uuid.uuid4())
|
|
telemetry = GenerationTelemetry(session_id)
|
|
with server._stats_lock:
|
|
current = server._active_requests.get(request_id)
|
|
if current is not None:
|
|
current["telemetry"] = telemetry
|
|
use_kv = bool(getattr(backend, "supports_kv_cache", False))
|
|
# EOS detection by id must work on the stateless path too: the tail
|
|
# returns token_id regardless of caching, and EOS usually decodes to
|
|
# "" (skip_special_tokens), so the text comparison never fires.
|
|
eos_ids: set[int] = set()
|
|
try:
|
|
eos_ids = set(backend.eos_token_ids())
|
|
except Exception:
|
|
eos_ids = set()
|
|
|
|
stream_emit = None
|
|
if stream:
|
|
stream_emit = self._start_openai_stream(model_name)
|
|
self._track_request_progress(server, request_id, tokens=0, routing_complete=True)
|
|
|
|
_GENERATION_LOG_INTERVAL = 5.0
|
|
gen_started = time.monotonic()
|
|
last_gen_log = gen_started
|
|
progress_line = [False]
|
|
last_token_id: int | None = None
|
|
failure_reason: str | None = None
|
|
relay_clients: dict[str, _RelayHopClient] = {}
|
|
direct_clients: dict[str, _DirectHopClient] = {}
|
|
|
|
def _prefill_step() -> tuple[str, int | None]:
|
|
"""Full-sequence prefill: initial step and cache-miss recovery."""
|
|
payload = (
|
|
backend.encode_prompt(current_text, session_id=session_id)
|
|
if use_kv
|
|
else backend.encode_prompt(current_text)
|
|
)
|
|
return self._run_downstream_pipeline(
|
|
payload, remaining_route, backend=backend,
|
|
session=session_id, cache_mode="prefill" if use_kv else None,
|
|
relay_clients=relay_clients,
|
|
direct_clients=direct_clients if use_kv else None,
|
|
telemetry=telemetry,
|
|
)
|
|
|
|
for step in range(max_tokens):
|
|
try:
|
|
if use_kv and step > 0 and last_token_id is not None:
|
|
try:
|
|
payload = backend.encode_next_token(last_token_id, session_id)
|
|
token_str, token_id = self._run_downstream_pipeline(
|
|
payload, remaining_route, backend=backend,
|
|
session=session_id, cache_mode="decode",
|
|
relay_clients=relay_clients,
|
|
direct_clients=direct_clients,
|
|
telemetry=telemetry,
|
|
)
|
|
except (KVCacheMiss, _PipelineCacheMiss) as miss:
|
|
# Evicted/restarted node or head lost its own session:
|
|
# re-prefill the whole sequence once and continue cached.
|
|
print(
|
|
f" [node] kv cache miss at step {step} ({miss}); "
|
|
f"re-prefilling {len(current_text)} chars",
|
|
flush=True,
|
|
)
|
|
token_str, token_id = _prefill_step()
|
|
else:
|
|
token_str, token_id = _prefill_step()
|
|
except _PipelineCacheMiss as exc:
|
|
print(f" [node] unexpected cache miss on prefill session={session_id[:8]}: {exc}", flush=True)
|
|
failure_reason = f"cache miss on prefill: {exc}"
|
|
break
|
|
except Exception as exc:
|
|
print(f" [node] distributed encode error session={session_id[:8]}: {exc}", flush=True)
|
|
failure_reason = f"distributed encode error: {exc}"
|
|
break
|
|
# Stop on error responses or EOS.
|
|
if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")):
|
|
failure_reason = token_str
|
|
break
|
|
if token_id is not None and token_id in eos_ids:
|
|
break
|
|
if eos_token and token_str == eos_token:
|
|
break
|
|
if not token_str and token_id is None:
|
|
break
|
|
last_token_id = token_id
|
|
# token_str can be empty for a skipped special token that is not
|
|
# EOS — keep generating from its token_id without emitting text.
|
|
if token_str:
|
|
generated.append(token_str)
|
|
if stream_emit is not None:
|
|
stream_emit(token_str)
|
|
current_text = current_text + token_str
|
|
self._track_request_progress(
|
|
server,
|
|
request_id,
|
|
tokens=len(generated),
|
|
routing_complete=True,
|
|
)
|
|
telemetry.note_tokens(len(generated))
|
|
now = time.monotonic()
|
|
if telemetry.report_due:
|
|
summary = telemetry.snapshot(now=now)
|
|
seams = summary["seams"]
|
|
if seams:
|
|
latest = seams[-1]
|
|
print(
|
|
f" [node] seam telemetry session={session_id[:8]} "
|
|
f"phase={latest['phase']} hop={latest['hop']} "
|
|
f"activations={latest['activations']} "
|
|
f"avg_ms={latest['avg_latency_ms']:.2f} "
|
|
f"wire_bytes={latest['wire_bytes']} "
|
|
f"response_bytes={latest['response_bytes']} "
|
|
f"tps={summary['tokens_per_sec']:.2f} "
|
|
f"activation_id={latest['last_activation_id']}",
|
|
flush=True,
|
|
)
|
|
telemetry.mark_reported(now=now)
|
|
if step == 0 or now - last_gen_log >= _GENERATION_LOG_INTERVAL:
|
|
elapsed = now - gen_started
|
|
token_count = len(generated)
|
|
tps = token_count / max(elapsed, 1e-6)
|
|
_write_progress_line(
|
|
progress_line,
|
|
f" [node] generating step={step + 1}/{max_tokens} session={session_id[:8]} "
|
|
f"tokens={token_count} elapsed_s={elapsed:.1f} tps={tps:.2f}",
|
|
)
|
|
last_gen_log = now
|
|
|
|
if use_kv:
|
|
try:
|
|
backend.release_session(session_id)
|
|
except Exception:
|
|
pass
|
|
for relay_client in relay_clients.values():
|
|
relay_client.close()
|
|
for direct_client in direct_clients.values():
|
|
direct_client.close()
|
|
|
|
if generated:
|
|
elapsed = time.monotonic() - gen_started
|
|
token_count = len(generated)
|
|
tps = token_count / max(elapsed, 1e-6)
|
|
_write_progress_line(
|
|
progress_line,
|
|
f" [node] generation complete session={session_id[:8]} tokens={token_count} "
|
|
f"elapsed_s={elapsed:.1f} tps={tps:.2f}",
|
|
final=True,
|
|
)
|
|
telemetry.close()
|
|
|
|
result_text = "".join(generated)
|
|
# A failure before the first token is an upstream error, not an empty
|
|
# completion — tell the client instead of returning a blank 200.
|
|
if failure_reason and not generated:
|
|
if stream_emit is not None:
|
|
stream_emit(None, error=failure_reason)
|
|
return
|
|
self._send_json(502, {
|
|
"error": {"message": failure_reason, "type": "upstream_error"},
|
|
})
|
|
return
|
|
if stream_emit is not None:
|
|
stream_emit(None)
|
|
return
|
|
self._send_openai_response(result_text, model_name, stream, messages, backend=backend)
|
|
|
|
def _get_remaining_route(self, model: str, *, backend: TorchModelShard | None = None) -> list[dict]:
|
|
"""Return downstream hops as dicts with endpoint, start_layer, and optional relay_addr.
|
|
|
|
Fast path reads X-Meshnet-Route header injected by the tracker.
|
|
Slow path queries the tracker's /v1/route endpoint as a fallback.
|
|
start_layer tells each downstream node which layer to begin from,
|
|
enabling correct execution when shard ranges overlap.
|
|
"""
|
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
|
active_backend = backend or server.backend
|
|
|
|
# Fast path: tracker pre-resolved the downstream route and injected it as a header.
|
|
injected = self.headers.get("X-Meshnet-Route")
|
|
if injected:
|
|
try:
|
|
route = json.loads(injected)
|
|
if isinstance(route, list):
|
|
hops: list[dict] = []
|
|
for item in route:
|
|
if isinstance(item, dict):
|
|
hop = {
|
|
"endpoint": str(item["endpoint"]),
|
|
"start_layer": int(item.get("start_layer", 0)),
|
|
}
|
|
if item.get("relay_addr"):
|
|
hop["relay_addr"] = str(item["relay_addr"])
|
|
hops.append(hop)
|
|
elif isinstance(item, str):
|
|
hops.append({"endpoint": item, "start_layer": 0})
|
|
hops = _clamp_downstream_hops(hops, active_backend)
|
|
print(
|
|
f" [node] using injected downstream route: {_format_downstream_route(hops)}",
|
|
flush=True,
|
|
)
|
|
return hops
|
|
except (json.JSONDecodeError, TypeError, KeyError):
|
|
pass
|
|
|
|
# Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject).
|
|
if server.tracker_url is None:
|
|
return []
|
|
route_model = getattr(active_backend, "model_id", None) or model
|
|
try:
|
|
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
|
|
with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
|
|
route_resp = json.loads(r.read())
|
|
own_key = _own_endpoint_key(server)
|
|
nodes_info = route_resp.get("nodes", [])
|
|
hops: list[dict] = []
|
|
passed_self = False
|
|
for node_info in nodes_info:
|
|
ep = node_info.get("endpoint", "")
|
|
if not ep:
|
|
continue
|
|
if _endpoint_key(ep) == own_key:
|
|
passed_self = True
|
|
continue
|
|
if not passed_self:
|
|
continue
|
|
hop = {
|
|
"endpoint": ep,
|
|
"start_layer": int(node_info.get("start_layer", 0)),
|
|
}
|
|
if node_info.get("relay_addr"):
|
|
hop["relay_addr"] = str(node_info["relay_addr"])
|
|
hops.append(hop)
|
|
hops = _clamp_downstream_hops(hops, active_backend)
|
|
print(
|
|
f" [node] tracker downstream route: {_format_downstream_route(hops)}",
|
|
flush=True,
|
|
)
|
|
return hops
|
|
except Exception as exc:
|
|
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
|
|
return []
|
|
|
|
def _run_downstream_pipeline(
|
|
self,
|
|
payload: object,
|
|
route: list[dict],
|
|
*,
|
|
backend: TorchModelShard | None = None,
|
|
session: str | None = None,
|
|
cache_mode: str | None = None,
|
|
relay_clients: dict[str, _RelayHopClient] | None = None,
|
|
direct_clients: dict[str, _DirectHopClient] | None = None,
|
|
telemetry: GenerationTelemetry | None = None,
|
|
) -> tuple[str, int | None]:
|
|
"""Forward an activation through the downstream route.
|
|
|
|
Returns (token_text, token_id) — token_id is None when a hop predates
|
|
the KV-cache protocol. Raises _PipelineCacheMiss when a hop responds
|
|
409 cache_miss (evicted/restarted node) so the caller can re-prefill.
|
|
"""
|
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
|
active_backend = backend or server.backend
|
|
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 active_backend.is_tail:
|
|
try:
|
|
tensor = _tensor_from_bfloat16_bytes(
|
|
payload.body, payload.shape, active_backend.torch, # type: ignore[union-attr]
|
|
).to(active_backend.device)
|
|
if hasattr(active_backend, "decode_tail_token"):
|
|
tail = active_backend.decode_tail_token(tensor)
|
|
return tail.text, tail.token_id
|
|
return active_backend.decode_tail(tensor), None
|
|
except Exception as exc:
|
|
return f"decode error: {exc}", None
|
|
return "no downstream route available for non-tail shard", None
|
|
|
|
# Session is stable across all steps of one generation when the caller
|
|
# provides it (KV-cache protocol); fresh per call otherwise (legacy).
|
|
session = session or 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, hop in enumerate(route):
|
|
node_url = hop["endpoint"]
|
|
start_layer = hop.get("start_layer", 0)
|
|
relay_addr = hop.get("relay_addr")
|
|
if server.debug:
|
|
print(
|
|
f" [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}"
|
|
+ (f" relay={relay_addr}" if relay_addr else ""),
|
|
flush=True,
|
|
)
|
|
phase = cache_mode if cache_mode in {"prefill", "decode"} else "prefill"
|
|
route_condition = "relay" if relay_addr else "lan"
|
|
compression = compress_activation(
|
|
current_body, _COMPRESSION_POLICIES.for_condition(route_condition, phase),
|
|
)
|
|
wire_body, wire_encoding = compression.body, compression.encoding
|
|
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),
|
|
"X-Meshnet-Start-Layer": str(start_layer),
|
|
"X-Meshnet-Activation-Id": uuid.uuid4().hex,
|
|
"X-Meshnet-Compression-Route": route_condition,
|
|
}
|
|
if wire_encoding:
|
|
headers["X-Meshnet-Encoding"] = wire_encoding
|
|
if telemetry is not None:
|
|
telemetry.record_compression(
|
|
phase=phase, hop=hop_index, node=node_url,
|
|
input_bytes=compression.input_bytes, output_bytes=compression.output_bytes,
|
|
elapsed_seconds=compression.elapsed_seconds,
|
|
)
|
|
if cache_mode:
|
|
headers["X-Meshnet-Cache"] = cache_mode
|
|
past_len = getattr(payload, "past_len", None)
|
|
if cache_mode == "decode" and past_len is not None:
|
|
headers["X-Meshnet-Past-Len"] = str(past_len)
|
|
if current_attn:
|
|
headers["X-Meshnet-Attn-Mask"] = current_attn
|
|
if current_pos:
|
|
headers["X-Meshnet-Position-Ids"] = current_pos
|
|
if relay_addr:
|
|
connection_reused = bool(
|
|
relay_clients is not None and relay_addr in relay_clients
|
|
)
|
|
seam_started = time.monotonic()
|
|
try:
|
|
if relay_clients is None:
|
|
status, resp_headers, resp_body = _relay_hop(
|
|
relay_addr, "/forward", wire_body, headers, timeout=120.0,
|
|
)
|
|
else:
|
|
relay_client = relay_clients.setdefault(
|
|
relay_addr, _RelayHopClient(relay_addr, timeout=120.0),
|
|
)
|
|
status, resp_headers, resp_body = relay_client.request(
|
|
"/forward", wire_body, headers,
|
|
)
|
|
if telemetry is not None:
|
|
telemetry.record_seam(
|
|
activation_id=headers["X-Meshnet-Activation-Id"],
|
|
phase=cache_mode or "prefill",
|
|
hop=hop_index,
|
|
node=node_url,
|
|
latency_seconds=time.monotonic() - seam_started,
|
|
wire_bytes=len(wire_body),
|
|
response_bytes=len(resp_body),
|
|
connection_reused=connection_reused,
|
|
)
|
|
if status == 409 and _is_cache_miss_body(resp_body):
|
|
raise _PipelineCacheMiss(node_url)
|
|
if status >= 400:
|
|
detail = _response_error_snippet(resp_body)
|
|
print(
|
|
f" [node] relay hop {hop_index} session={session[:8]} returned "
|
|
f"{status} from {relay_addr}: {detail}",
|
|
flush=True,
|
|
)
|
|
return f"pipeline error at {node_url} via relay: status {status}: {detail}", None
|
|
except _PipelineCacheMiss:
|
|
raise
|
|
except _RelayRequestUncertainError as exc:
|
|
# The activation may already have mutated the downstream
|
|
# KV cache. Do not replay it on a direct connection.
|
|
print(
|
|
f" [node] relay hop {hop_index} session={session[:8]} outcome is uncertain at "
|
|
f"{relay_addr}: {exc}",
|
|
flush=True,
|
|
)
|
|
return f"pipeline relay outcome uncertain at {node_url}: {exc}", None
|
|
except Exception as exc:
|
|
print(
|
|
f" [node] relay hop {hop_index} session={session[:8]} failed at {relay_addr}: {exc}; "
|
|
f"falling back to direct {node_url}",
|
|
flush=True,
|
|
)
|
|
relay_addr = None # fall through to direct
|
|
if not relay_addr:
|
|
connection_reused = bool(
|
|
direct_clients is not None and node_url in direct_clients
|
|
)
|
|
seam_started = time.monotonic()
|
|
try:
|
|
if direct_clients is None:
|
|
direct_client = _DirectHopClient(node_url, timeout=120.0)
|
|
try:
|
|
status, resp_headers, resp_body = direct_client.request(
|
|
"/forward", wire_body, headers,
|
|
)
|
|
finally:
|
|
direct_client.close()
|
|
else:
|
|
direct_client = direct_clients.setdefault(
|
|
node_url, _DirectHopClient(node_url, timeout=120.0),
|
|
)
|
|
status, resp_headers, resp_body = direct_client.request(
|
|
"/forward", wire_body, headers,
|
|
)
|
|
if telemetry is not None:
|
|
telemetry.record_seam(
|
|
activation_id=headers["X-Meshnet-Activation-Id"],
|
|
phase=cache_mode or "prefill",
|
|
hop=hop_index,
|
|
node=node_url,
|
|
latency_seconds=time.monotonic() - seam_started,
|
|
wire_bytes=len(wire_body),
|
|
response_bytes=len(resp_body),
|
|
connection_reused=connection_reused,
|
|
)
|
|
if status == 409 and _is_cache_miss_body(resp_body):
|
|
raise _PipelineCacheMiss(node_url)
|
|
if status >= 400:
|
|
detail = _response_error_snippet(resp_body)
|
|
print(
|
|
f" [node] pipeline hop {hop_index} session={session[:8]} failed at {node_url}: "
|
|
f"status {status}: {detail}", flush=True,
|
|
)
|
|
return f"pipeline error at {node_url}: status {status}: {detail}", None
|
|
except _PipelineCacheMiss:
|
|
raise
|
|
except _DirectRequestUncertainError as exc:
|
|
print(
|
|
f" [node] pipeline hop {hop_index} session={session[:8]} outcome is uncertain "
|
|
f"at {node_url}: {exc}",
|
|
flush=True,
|
|
)
|
|
return f"pipeline direct outcome uncertain at {node_url}: {exc}", None
|
|
except Exception as exc:
|
|
print(
|
|
f" [node] pipeline hop {hop_index} session={session[:8]} "
|
|
f"failed at {node_url}: {exc}", flush=True,
|
|
)
|
|
return f"pipeline error at {node_url}: {exc}", None
|
|
content_type = resp_headers.get("content-type", "")
|
|
if "application/json" in content_type:
|
|
try:
|
|
data = json.loads(resp_body)
|
|
text = str(data.get("text", ""))
|
|
token_id = data.get("token_id")
|
|
if server.debug:
|
|
print(f" [node] pipeline hop {hop_index} returned text={text!r}", flush=True)
|
|
return text, int(token_id) if token_id is not None else None
|
|
except json.JSONDecodeError:
|
|
return resp_body.decode("utf-8", errors="replace"), None
|
|
# 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)
|
|
try:
|
|
decompression = decompress_activation(
|
|
resp_body, resp_headers.get("x-meshnet-encoding")
|
|
)
|
|
current_body = decompression.body
|
|
if telemetry is not None:
|
|
telemetry.record_compression(
|
|
phase=phase, hop=hop_index, node=node_url,
|
|
input_bytes=decompression.input_bytes, output_bytes=decompression.output_bytes,
|
|
elapsed_seconds=decompression.elapsed_seconds, decompression=True,
|
|
)
|
|
except ValueError as exc:
|
|
print(
|
|
f" [node] pipeline hop {hop_index} session={session[:8]} "
|
|
f"bad response encoding: {exc}", flush=True,
|
|
)
|
|
return f"pipeline error at {node_url}: {exc}", None
|
|
current_attn = resp_headers.get("x-meshnet-attn-mask")
|
|
current_pos = resp_headers.get("x-meshnet-position-ids")
|
|
return "", None
|
|
|
|
def _stream_openai_response(self, token_iter, model: str) -> None:
|
|
"""Stream tokens from an iterator as SSE chunks."""
|
|
emit = self._start_openai_stream(model)
|
|
for token_text in token_iter:
|
|
if not token_text:
|
|
continue
|
|
emit(token_text)
|
|
emit(None)
|
|
|
|
def _start_openai_stream(self, model: str):
|
|
"""Open an OpenAI-compatible SSE response and return a token emitter."""
|
|
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.send_header("Transfer-Encoding", "chunked")
|
|
self.end_headers()
|
|
|
|
def _emit(data: str) -> bool:
|
|
try:
|
|
body = f"data: {data}\n\n".encode()
|
|
self.wfile.write(f"{len(body):X}\r\n".encode() + body + b"\r\n")
|
|
self.wfile.flush()
|
|
return True
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
return False
|
|
|
|
_emit(json.dumps({
|
|
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
|
|
}))
|
|
|
|
def emit_token(token_text: str | None, *, error: str | None = None) -> None:
|
|
if error is not None:
|
|
# OpenAI-style mid-stream error frame; clients surface it
|
|
# instead of showing an empty completion.
|
|
_emit(json.dumps({"error": {"message": error, "type": "upstream_error"}}))
|
|
try:
|
|
self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n")
|
|
self.wfile.flush()
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
pass
|
|
return
|
|
if token_text is None:
|
|
_emit(json.dumps({
|
|
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
|
}))
|
|
try:
|
|
self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n")
|
|
self.wfile.flush()
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
pass
|
|
return
|
|
_emit(json.dumps({
|
|
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {"content": token_text}, "finish_reason": None}],
|
|
}))
|
|
|
|
return emit_token
|
|
|
|
def _send_openai_response(
|
|
self,
|
|
text: str,
|
|
model: str,
|
|
stream: bool,
|
|
messages: list[dict] | None = None,
|
|
backend: TorchModelShard | None = None,
|
|
) -> None:
|
|
chunk_id = "chatcmpl-node"
|
|
created = int(time.time())
|
|
active_backend = backend or self.server.backend # type: ignore[attr-defined]
|
|
if not stream:
|
|
usage = _usage_for_response(active_backend, messages or [], text)
|
|
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.send_header("Transfer-Encoding", "chunked")
|
|
self.end_headers()
|
|
|
|
def _emit(data: str) -> bool:
|
|
try:
|
|
body = f"data: {data}\n\n".encode()
|
|
self.wfile.write(f"{len(body):X}\r\n".encode() + body + b"\r\n")
|
|
self.wfile.flush()
|
|
return True
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
return False
|
|
|
|
_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"}],
|
|
}))
|
|
try:
|
|
self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n")
|
|
self.wfile.flush()
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
pass
|
|
|
|
|
|
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,
|
|
route_timeout: float = 30.0,
|
|
cache_dir: Path | None = None,
|
|
debug: bool = False,
|
|
max_loaded_shards: int = 1,
|
|
force_cpu: bool = False,
|
|
recipe_params: Mapping[str, Any] | None = None,
|
|
) -> None:
|
|
self._host = host
|
|
self._requested_port = port
|
|
self._max_loaded_shards = max(1, max_loaded_shards)
|
|
self._backend = backend or _load_backend(
|
|
model_id,
|
|
shard_start,
|
|
shard_end,
|
|
quantization,
|
|
cache_dir,
|
|
force_cpu=force_cpu,
|
|
recipe_params=recipe_params,
|
|
)
|
|
self._backends: dict[str, TorchModelShard] = {self._backend.model_id: self._backend}
|
|
# 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._route_timeout = route_timeout
|
|
self._cache_dir = cache_dir
|
|
self._debug = debug
|
|
self._server: _TorchHTTPServer | None = None
|
|
self._thread: threading.Thread | None = None
|
|
self.port: int | None = None
|
|
|
|
@property
|
|
def route_timeout(self) -> float:
|
|
return self._route_timeout
|
|
|
|
@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
|
|
|
|
@property
|
|
def total_requests(self) -> int:
|
|
return self._server.total_requests if self._server is not None else 0
|
|
|
|
@property
|
|
def failed_requests(self) -> int:
|
|
return self._server.failed_requests if self._server is not None else 0
|
|
|
|
@property
|
|
def queue_depth(self) -> int:
|
|
return self._server.queue_depth if self._server is not None else 0
|
|
|
|
@property
|
|
def current_requests(self) -> list[dict[str, Any]]:
|
|
if self._server is None:
|
|
return []
|
|
return self._server.snapshot_current_requests()
|
|
|
|
@property
|
|
def loaded_model_ids(self) -> list[str]:
|
|
return list(self._backends.keys())
|
|
|
|
def apply_tracker_directives(self, directives: list[dict]) -> dict | None:
|
|
"""Apply tracker shard directives (LOAD_SHARD replace, ADD_SHARD load-more)."""
|
|
drop_directive = next(
|
|
(directive for directive in reversed(directives) if directive.get("action") == "DROP_SHARD"),
|
|
None,
|
|
)
|
|
if drop_directive is not None:
|
|
model_id = str(drop_directive.get("model") or "")
|
|
removed = self._backends.pop(model_id, None)
|
|
if removed is None:
|
|
return None
|
|
if self._backends:
|
|
self._backend = next(iter(self._backends.values()))
|
|
self._tracker_mode = self._backend.shard_start == 0
|
|
else:
|
|
self._backend = None
|
|
self._tracker_mode = False
|
|
if self._server is not None:
|
|
self._server.backends = dict(self._backends)
|
|
self._server.backend = self._backend
|
|
self._server.tracker_mode = self._tracker_mode
|
|
return {"action": "DROP_SHARD", "model": model_id}
|
|
add_directive = next(
|
|
(directive for directive in reversed(directives) if directive.get("action") == "ADD_SHARD"),
|
|
None,
|
|
)
|
|
load_directive = next(
|
|
(directive for directive in reversed(directives) if directive.get("action") == "LOAD_SHARD"),
|
|
None,
|
|
)
|
|
directive = add_directive or load_directive
|
|
if directive is None:
|
|
return None
|
|
shard_start = int(directive["shard_start"])
|
|
shard_end = int(directive["shard_end"])
|
|
quantization = str(directive.get("quantization") or self._backend.quantization)
|
|
model_id = str(directive.get("model") or self._backend.model_id)
|
|
replacing = directive.get("action") == "LOAD_SHARD"
|
|
if not replacing and len(self._backends) >= self._max_loaded_shards:
|
|
print(
|
|
f" [node] WARNING: ignoring ADD_SHARD for {model_id!r} — "
|
|
f"loaded {len(self._backends)}/{self._max_loaded_shards} slots full",
|
|
flush=True,
|
|
)
|
|
return None
|
|
action_label = "reassigned" if replacing else "additional"
|
|
print(
|
|
f" [node] loading {action_label} shard: {model_id} layers {shard_start}-{shard_end}",
|
|
flush=True,
|
|
)
|
|
try:
|
|
if replacing:
|
|
self._backends.clear()
|
|
new_backend = _load_backend(model_id, shard_start, shard_end, quantization, self._cache_dir)
|
|
except TypeError:
|
|
new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
|
|
self._backends[model_id] = new_backend
|
|
if replacing or shard_start == 0:
|
|
self._backend = new_backend
|
|
self._tracker_mode = shard_start == 0
|
|
print(
|
|
f" [node] loaded {action_label} shard: {model_id} layers {shard_start}-{shard_end}",
|
|
flush=True,
|
|
)
|
|
if self._server is not None:
|
|
self._server.backends = dict(self._backends)
|
|
if replacing or shard_start == 0:
|
|
self._server.backend = new_backend
|
|
self._server.tracker_mode = self._tracker_mode
|
|
return {
|
|
"action": directive.get("action"),
|
|
"model": model_id,
|
|
"shard_start": shard_start,
|
|
"shard_end": shard_end,
|
|
"quantization": quantization,
|
|
"tracker_mode": shard_start == 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._route_timeout,
|
|
self._debug,
|
|
self._max_loaded_shards,
|
|
)
|
|
self._server.backends = dict(self._backends)
|
|
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 set_advertised_endpoint(self, endpoint: str) -> None:
|
|
"""Set the LAN-facing endpoint used for route self-detection."""
|
|
if self._server is not None:
|
|
self._server.advertised_endpoint = endpoint
|
|
|
|
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,
|
|
cache_dir: Path | None = None,
|
|
force_cpu: bool = False,
|
|
recipe_params: Mapping[str, Any] | None = None,
|
|
) -> TorchModelShard:
|
|
from .model_backend import load_torch_shard
|
|
|
|
quant = validate_quantization(quantization)
|
|
try:
|
|
return load_torch_shard(
|
|
model_id,
|
|
shard_start,
|
|
shard_end,
|
|
quant,
|
|
cache_dir,
|
|
force_cpu=force_cpu,
|
|
recipe_params=recipe_params,
|
|
)
|
|
except MissingModelDependencyError:
|
|
raise
|
|
except InsufficientVRAMError as exc:
|
|
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
|
|
raise
|