942 lines
36 KiB
Python
942 lines
36 KiB
Python
"""Gateway HTTP server — accepts OpenAI-format requests and routes to inference nodes."""
|
|
|
|
import http.server
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from collections import Counter
|
|
from dataclasses import dataclass
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.parse
|
|
import urllib.request
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from .prefill_backpressure import BoundedPrefillSender, PrefillTransferLimits
|
|
|
|
_STUB_HIDDEN_DIM = 64
|
|
_STUB_DTYPE = "bfloat16"
|
|
_WIRE_VERSION = "2"
|
|
_DTYPE_SIZES = {
|
|
"float16": 2,
|
|
"bfloat16": 2,
|
|
"float32": 4,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class _BinaryActivation:
|
|
body: bytes
|
|
shape: list[int]
|
|
dtype: str
|
|
session: str
|
|
chunk_index: int
|
|
chunk_total: int
|
|
encoding: str | None
|
|
headers: dict[str, str]
|
|
|
|
|
|
class _GatewayHTTPServer(http.server.HTTPServer):
|
|
def __init__(
|
|
self,
|
|
addr,
|
|
handler,
|
|
inference_route: list[str] | None,
|
|
tracker_url: str | None,
|
|
model_presets: list[str] | None,
|
|
contracts: Any | None,
|
|
minimum_stake: int,
|
|
cost_per_layer_token_lamport: int,
|
|
):
|
|
super().__init__(addr, handler)
|
|
self.inference_route = inference_route
|
|
self.tracker_url = tracker_url
|
|
self.model_presets = model_presets
|
|
self.contracts = contracts
|
|
self.minimum_stake = minimum_stake
|
|
self.cost_per_layer_token_lamport = cost_per_layer_token_lamport
|
|
self.last_binary_chunk_responses: list[_BinaryActivation] = []
|
|
self.request_count: int = 0
|
|
|
|
|
|
class _GatewayHandler(http.server.BaseHTTPRequestHandler):
|
|
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
|
|
pass
|
|
|
|
def do_GET(self):
|
|
parsed = urllib.parse.urlparse(self.path)
|
|
if parsed.path == "/v1/health":
|
|
self._handle_health()
|
|
elif parsed.path == "/v1/models":
|
|
self._handle_models()
|
|
elif parsed.path == "/v1/models/available":
|
|
self._handle_available_models()
|
|
elif parsed.path == "/v1/wallet/balance":
|
|
self._handle_wallet_balance()
|
|
elif parsed.path == "/v1/estimate_cost":
|
|
self._handle_estimate_cost(parsed)
|
|
else:
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
|
|
def do_POST(self):
|
|
parsed = urllib.parse.urlparse(self.path)
|
|
if parsed.path == "/v1/chat/completions":
|
|
self._handle_chat_completions()
|
|
elif parsed.path == "/v1/wallet/top_up":
|
|
self._handle_wallet_top_up()
|
|
elif parsed.path == "/v1/request":
|
|
self._handle_meshnet_request()
|
|
else:
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
|
|
def _send_json(self, status: int, data: dict) -> None:
|
|
body = json.dumps(data).encode()
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def _send_json_error(self, status: int, message: str) -> None:
|
|
self._send_json(status, {"error": message})
|
|
|
|
def _send_openai_error(self, status: int, message: str, code: str) -> None:
|
|
self._send_json(status, {
|
|
"error": {
|
|
"message": message,
|
|
"type": "invalid_request_error",
|
|
"param": "model",
|
|
"code": code,
|
|
}
|
|
})
|
|
|
|
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_error(400, "invalid JSON body")
|
|
return None
|
|
if not isinstance(body, dict):
|
|
self._send_json_error(400, "JSON body must be an object")
|
|
return None
|
|
return body
|
|
|
|
def _handle_health(self) -> None:
|
|
self._send_json(200, {"status": "ok"})
|
|
|
|
def _handle_models(self) -> None:
|
|
server: _GatewayHTTPServer = self.server # type: ignore[assignment]
|
|
created = int(time.time())
|
|
if server.tracker_url is not None:
|
|
try:
|
|
models_resp = _get_json(f"{server.tracker_url}/v1/models")
|
|
self._send_json(200, models_resp)
|
|
return
|
|
except Exception:
|
|
self._send_json_error(503, "tracker unavailable")
|
|
return
|
|
presets = server.model_presets or ["stub-model"]
|
|
data = [
|
|
{"id": p, "object": "model", "created": created, "owned_by": "meshnet"}
|
|
for p in presets
|
|
]
|
|
self._send_json(200, {"object": "list", "data": data})
|
|
|
|
def _handle_available_models(self) -> None:
|
|
server: _GatewayHTTPServer = self.server # type: ignore[assignment]
|
|
if server.tracker_url is not None:
|
|
try:
|
|
models_resp = _get_json(f"{server.tracker_url}/v1/models")
|
|
except Exception:
|
|
self._send_json_error(503, "tracker unavailable")
|
|
return
|
|
data = [
|
|
{
|
|
"id": model["id"],
|
|
"shard_coverage_percentage": float(model.get("shard_coverage_percentage", 0.0)),
|
|
}
|
|
for model in models_resp.get("data", [])
|
|
if isinstance(model, dict)
|
|
and isinstance(model.get("id"), str)
|
|
and float(model.get("shard_coverage_percentage", 0.0)) > 0.0
|
|
]
|
|
self._send_json(200, {"data": data})
|
|
return
|
|
presets = server.model_presets or ["stub-model"]
|
|
self._send_json(200, {
|
|
"data": [
|
|
{"id": model, "shard_coverage_percentage": 100.0}
|
|
for model in presets
|
|
]
|
|
})
|
|
|
|
def _handle_wallet_balance(self) -> None:
|
|
server: _GatewayHTTPServer = self.server # type: ignore[assignment]
|
|
api_key = _api_key_from_authorization(self.headers.get("Authorization"))
|
|
if not api_key:
|
|
self._send_json_error(401, "missing API key")
|
|
return
|
|
if server.contracts is None:
|
|
self._send_json(200, {"sol": 0.0, "usdc": 0.0})
|
|
return
|
|
balance = server.contracts.payment.get_balance(api_key)
|
|
self._send_json(200, {
|
|
"sol": _lamports_to_sol(balance.lamports),
|
|
"usdc": balance.usdc_micro / 1_000_000,
|
|
})
|
|
|
|
def _handle_wallet_top_up(self) -> None:
|
|
api_key = _api_key_from_authorization(self.headers.get("Authorization"))
|
|
if not api_key:
|
|
self._send_json_error(401, "missing API key")
|
|
return
|
|
body = self._read_json_body()
|
|
if body is None:
|
|
return
|
|
try:
|
|
amount_sol = float(body.get("amount_sol", 0.0))
|
|
amount_usdc = float(body.get("amount_usdc", 0.0))
|
|
except (TypeError, ValueError):
|
|
self._send_json_error(400, "top-up amount must be numeric")
|
|
return
|
|
if amount_sol <= 0 and amount_usdc <= 0:
|
|
self._send_json_error(400, "top-up amount must be positive")
|
|
return
|
|
payment_address = _payment_address_for_api_key(api_key)
|
|
query = urllib.parse.urlencode({
|
|
"amount": amount_sol,
|
|
"spl-token": "USDC" if amount_usdc > 0 else "SOL",
|
|
})
|
|
self._send_json(200, {
|
|
"payment_address": payment_address,
|
|
"qr": f"solana:{payment_address}?{query}",
|
|
"amount_sol": amount_sol,
|
|
"amount_usdc": amount_usdc,
|
|
})
|
|
|
|
def _handle_estimate_cost(self, parsed: urllib.parse.ParseResult) -> None:
|
|
server: _GatewayHTTPServer = self.server # type: ignore[assignment]
|
|
params = urllib.parse.parse_qs(parsed.query)
|
|
model = params.get("model", [""])[0]
|
|
try:
|
|
tokens = int(params.get("tokens", ["0"])[0])
|
|
except ValueError:
|
|
self._send_json_error(400, "tokens must be an integer")
|
|
return
|
|
if not model:
|
|
self._send_json_error(400, "model is required")
|
|
return
|
|
if tokens <= 0:
|
|
self._send_json_error(400, "tokens must be positive")
|
|
return
|
|
layer_count = 32
|
|
if server.tracker_url is None and server.model_presets is not None and model not in server.model_presets:
|
|
self._send_json_error(404, "model not available")
|
|
return
|
|
lamports = layer_count * tokens * server.cost_per_layer_token_lamport
|
|
self._send_json(200, {
|
|
"model": model,
|
|
"tokens": tokens,
|
|
"lamports": lamports,
|
|
"sol": _lamports_to_sol(lamports),
|
|
})
|
|
|
|
def _handle_chat_completions(self):
|
|
server: _GatewayHTTPServer = self.server # type: ignore[assignment]
|
|
# Read raw bytes first so we can proxy them if head workers are available.
|
|
length = int(self.headers.get("Content-Length", 0))
|
|
raw_body = self.rfile.read(length)
|
|
try:
|
|
body = json.loads(raw_body or b"{}")
|
|
except (json.JSONDecodeError, ValueError):
|
|
self._send_json_error(400, "invalid JSON body")
|
|
return
|
|
if not isinstance(body, dict):
|
|
self._send_json_error(400, "JSON body must be an object")
|
|
return
|
|
|
|
model = str(body.get("model", "stub-model"))
|
|
head_workers = _get_head_workers(server, model)
|
|
if head_workers:
|
|
target = head_workers[server.request_count % len(head_workers)]
|
|
server.request_count += 1
|
|
return self._proxy_to_head_worker(target, raw_body)
|
|
|
|
# Fallback: use existing direct pipeline (backward compat)
|
|
streaming = bool(body.get("stream", False))
|
|
try:
|
|
completion = self._build_completion(body)
|
|
except _ModelUnavailable as exc:
|
|
self._send_openai_error(503, exc.message, "model_not_available")
|
|
return
|
|
except _TrackerUnavailable:
|
|
self._send_json_error(503, "tracker unavailable")
|
|
return
|
|
except (urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError, ValueError) as exc:
|
|
self._send_openai_error(503, f"inference pipeline error: {exc}", "pipeline_error")
|
|
return
|
|
|
|
if streaming:
|
|
text = completion["choices"][0]["message"]["content"]
|
|
self._send_sse_response(text, completion["model"])
|
|
else:
|
|
self._send_json(200, completion)
|
|
|
|
def _proxy_to_head_worker(self, url: str, body_bytes: bytes) -> None:
|
|
"""Forward a raw request body to a head worker and relay SSE without buffering."""
|
|
target_url = f"{url}/v1/chat/completions"
|
|
req = urllib.request.Request(
|
|
target_url,
|
|
data=body_bytes,
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=30.0) as r:
|
|
content_type = r.headers.get("Content-Type", "application/json")
|
|
if "text/event-stream" in content_type:
|
|
self.send_response(r.status)
|
|
self.send_header("Content-Type", content_type)
|
|
self.send_header("Cache-Control", "no-cache")
|
|
self.send_header("X-Accel-Buffering", "no")
|
|
self.end_headers()
|
|
while True:
|
|
line = r.readline()
|
|
if not line:
|
|
break
|
|
self.wfile.write(line)
|
|
self.wfile.flush()
|
|
return
|
|
resp_body = r.read()
|
|
status = r.status
|
|
except urllib.error.HTTPError as exc:
|
|
resp_body = exc.read()
|
|
self.send_response(exc.code)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(resp_body)))
|
|
self.end_headers()
|
|
self.wfile.write(resp_body)
|
|
return
|
|
except Exception as exc:
|
|
self._send_json_error(503, f"head worker unreachable: {exc}")
|
|
return
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", content_type)
|
|
self.send_header("Content-Length", str(len(resp_body)))
|
|
self.end_headers()
|
|
self.wfile.write(resp_body)
|
|
|
|
def _handle_meshnet_request(self) -> None:
|
|
body = self._read_json_body()
|
|
if body is None:
|
|
return
|
|
try:
|
|
redundancy = int(body.get("redundancy", 1))
|
|
except (TypeError, ValueError):
|
|
self._send_json_error(400, "redundancy must be an integer")
|
|
return
|
|
if redundancy < 1:
|
|
self._send_json_error(400, "redundancy must be at least 1")
|
|
return
|
|
|
|
try:
|
|
routes = self._resolve_routes(str(body.get("model", "stub-model")), redundancy)
|
|
responses = [self._build_completion(body, route_payload=route_payload) for route_payload in routes]
|
|
except _ModelUnavailable as exc:
|
|
self._send_openai_error(503, exc.message, "model_not_available")
|
|
return
|
|
except _TrackerUnavailable:
|
|
self._send_json_error(503, "tracker unavailable")
|
|
return
|
|
except (urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError, ValueError) as exc:
|
|
self._send_openai_error(503, f"inference pipeline error: {exc}", "pipeline_error")
|
|
return
|
|
|
|
majority_response = _majority_response(responses)
|
|
self._send_json(200, {
|
|
"redundancy": redundancy,
|
|
"majority_response": majority_response,
|
|
"responses": responses,
|
|
})
|
|
|
|
def _build_completion(self, body: dict, route_payload: dict | None = None) -> dict:
|
|
server: _GatewayHTTPServer = self.server # type: ignore[assignment]
|
|
if route_payload is None:
|
|
route_payload = self._resolve_route(str(body.get("model", "stub-model")))
|
|
route = route_payload["route"]
|
|
route_nodes = route_payload.get("nodes", [])
|
|
messages = body.get("messages", [])
|
|
|
|
prompt = _last_message_content(messages)
|
|
chunk_responses = _run_binary_pipeline(route, prompt)
|
|
_reassemble_binary_responses(chunk_responses)
|
|
server.last_binary_chunk_responses = chunk_responses
|
|
response_prefix = _stub_response_prefix(chunk_responses)
|
|
text = f"{response_prefix} {prompt}"
|
|
model = body.get("model", "stub-model")
|
|
self._record_inference(model, messages, text, route_nodes)
|
|
return _completion_response(str(model), text)
|
|
|
|
def _resolve_route(self, model: str) -> dict:
|
|
server: _GatewayHTTPServer = self.server # type: ignore[assignment]
|
|
if server.tracker_url is not None:
|
|
tracker_path = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(model)}"
|
|
try:
|
|
tracker_resp = _get_json(tracker_path)
|
|
_validate_route_payload(tracker_resp)
|
|
banned_wallet = _banned_route_wallet(server.contracts, tracker_resp.get("nodes", []))
|
|
if banned_wallet is not None:
|
|
raise _ModelUnavailable(f"model not available: route includes banned wallet {banned_wallet}")
|
|
return tracker_resp
|
|
except urllib.error.HTTPError as exc:
|
|
error_body = _safe_error_body(exc)
|
|
raise _ModelUnavailable(error_body.get("error", "model not available")) from exc
|
|
except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, KeyError, TypeError, ValueError) as exc:
|
|
raise _TrackerUnavailable() from exc
|
|
return {"route": server.inference_route, "nodes": []}
|
|
|
|
def _resolve_routes(self, model: str, redundancy: int) -> list[dict]:
|
|
server: _GatewayHTTPServer = self.server # type: ignore[assignment]
|
|
if server.tracker_url is None:
|
|
return [self._resolve_route(model) for _ in range(redundancy)]
|
|
tracker_path = (
|
|
f"{server.tracker_url}/v1/routes?"
|
|
f"model={urllib.parse.quote(model)}&redundancy={redundancy}"
|
|
)
|
|
try:
|
|
tracker_resp = _get_json(tracker_path)
|
|
routes = tracker_resp["routes"]
|
|
if not isinstance(routes, list) or len(routes) != redundancy:
|
|
raise ValueError("invalid tracker routes")
|
|
for route_payload in routes:
|
|
_validate_route_payload(route_payload)
|
|
banned_wallet = _banned_route_wallet(server.contracts, route_payload.get("nodes", []))
|
|
if banned_wallet is not None:
|
|
raise _ModelUnavailable(f"model not available: route includes banned wallet {banned_wallet}")
|
|
return routes
|
|
except urllib.error.HTTPError as exc:
|
|
error_body = _safe_error_body(exc)
|
|
raise _ModelUnavailable(error_body.get("error", "model not available")) from exc
|
|
except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, KeyError, TypeError, ValueError) as exc:
|
|
raise _TrackerUnavailable() from exc
|
|
|
|
def _record_inference(
|
|
self,
|
|
model: object,
|
|
messages: object,
|
|
text: str,
|
|
route_nodes: object,
|
|
) -> None:
|
|
server: _GatewayHTTPServer = self.server # type: ignore[assignment]
|
|
session_id = str(uuid.uuid4())
|
|
if server.contracts is not None and isinstance(route_nodes, list):
|
|
api_key = _api_key_from_authorization(self.headers.get("Authorization"))
|
|
if api_key:
|
|
tokens = max(1, _estimate_token_count(messages, text))
|
|
for node in route_nodes:
|
|
wallet_address = node.get("wallet_address") if isinstance(node, dict) else None
|
|
if not wallet_address:
|
|
continue
|
|
server.contracts.payment.record_attribution(
|
|
session_id=session_id,
|
|
api_key=api_key,
|
|
node_wallet=wallet_address,
|
|
layer_start=node["shard_start"],
|
|
layer_end=node["shard_end"],
|
|
tokens=tokens,
|
|
speed_score=float(node.get("score", 1.0)),
|
|
)
|
|
if server.contracts is not None:
|
|
server.contracts.validation.record_completed_inference(
|
|
session_id=session_id,
|
|
model=str(model),
|
|
messages=messages if isinstance(messages, list) else [],
|
|
observed_output=text,
|
|
route_nodes=route_nodes if isinstance(route_nodes, list) else [],
|
|
)
|
|
|
|
def _send_sse_response(self, text: str, model: str) -> None:
|
|
chunk_id = "chatcmpl-stub"
|
|
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("X-Accel-Buffering", "no")
|
|
self.end_headers()
|
|
|
|
def _emit(data: str) -> None:
|
|
self.wfile.write(f"data: {data}\n\n".encode())
|
|
self.wfile.flush()
|
|
|
|
# Role delta
|
|
_emit(json.dumps({
|
|
"id": chunk_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
|
|
}))
|
|
|
|
# Content delta
|
|
_emit(json.dumps({
|
|
"id": chunk_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}],
|
|
}))
|
|
|
|
# Stop delta
|
|
_emit(json.dumps({
|
|
"id": chunk_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": model,
|
|
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
|
}))
|
|
|
|
# SSE sentinel
|
|
self.wfile.write(b"data: [DONE]\n\n")
|
|
self.wfile.flush()
|
|
|
|
|
|
def _safe_error_body(exc: urllib.error.HTTPError) -> dict:
|
|
try:
|
|
body = json.loads(exc.read())
|
|
except json.JSONDecodeError:
|
|
return {}
|
|
return body if isinstance(body, dict) else {}
|
|
|
|
|
|
def _valid_route_nodes(route_nodes: object) -> bool:
|
|
if not isinstance(route_nodes, list):
|
|
return False
|
|
for node in route_nodes:
|
|
if not isinstance(node, dict):
|
|
return False
|
|
if not isinstance(node.get("endpoint"), str) or not node["endpoint"]:
|
|
return False
|
|
wallet_address = node.get("wallet_address")
|
|
if wallet_address is not None and not isinstance(wallet_address, str):
|
|
return False
|
|
try:
|
|
shard_start = int(node["shard_start"])
|
|
shard_end = int(node["shard_end"])
|
|
except (KeyError, TypeError, ValueError):
|
|
return False
|
|
if shard_start < 0 or shard_end < shard_start:
|
|
return False
|
|
node["shard_start"] = shard_start
|
|
node["shard_end"] = shard_end
|
|
return True
|
|
|
|
|
|
def _validate_route_payload(route_payload: dict) -> None:
|
|
route = route_payload["route"]
|
|
if (
|
|
not isinstance(route, list)
|
|
or not route
|
|
or not all(isinstance(node_url, str) and node_url for node_url in route)
|
|
):
|
|
raise ValueError("invalid tracker route")
|
|
route_nodes = route_payload.get("nodes", [])
|
|
if route_nodes and not _valid_route_nodes(route_nodes):
|
|
raise ValueError("invalid tracker route metadata")
|
|
|
|
|
|
def _banned_route_wallet(contracts: Any | None, route_nodes: object) -> str | None:
|
|
if contracts is None or not isinstance(route_nodes, list):
|
|
return None
|
|
for node in route_nodes:
|
|
if not isinstance(node, dict):
|
|
continue
|
|
wallet_address = node.get("wallet_address")
|
|
if isinstance(wallet_address, str) and contracts.registry.get_wallet(wallet_address).banned:
|
|
return wallet_address
|
|
return None
|
|
|
|
|
|
def _api_key_from_authorization(header: str | None) -> str | None:
|
|
if not header:
|
|
return None
|
|
scheme, _, token = header.partition(" ")
|
|
if scheme.lower() != "bearer" or not token.strip():
|
|
return None
|
|
return token.strip()
|
|
|
|
|
|
def _estimate_token_count(messages: object, completion_text: str) -> int:
|
|
"""Request-token placeholder until shard nodes return tokenizer usage."""
|
|
return 1
|
|
|
|
|
|
def _lamports_to_sol(lamports: int) -> float:
|
|
return lamports / 1_000_000_000
|
|
|
|
|
|
def _payment_address_for_api_key(api_key: str) -> str:
|
|
alphabet = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"
|
|
value = int.from_bytes(hashlib.sha256(f"meshnet:{api_key}".encode()).digest(), "big")
|
|
chars = []
|
|
while value:
|
|
value, remainder = divmod(value, 58)
|
|
chars.append(alphabet[remainder])
|
|
return "".join(reversed(chars))[:44].rjust(32, "1")
|
|
|
|
|
|
def _completion_response(model: str, text: str) -> dict:
|
|
return {
|
|
"id": "chatcmpl-stub",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": text,
|
|
},
|
|
"finish_reason": "stop",
|
|
}],
|
|
"usage": {
|
|
"prompt_tokens": 0,
|
|
"completion_tokens": 0,
|
|
"total_tokens": 0,
|
|
},
|
|
}
|
|
|
|
|
|
def _majority_response(responses: list[dict]) -> dict:
|
|
if not responses:
|
|
raise ValueError("responses must not be empty")
|
|
contents = [
|
|
response["choices"][0]["message"]["content"]
|
|
for response in responses
|
|
]
|
|
majority_content, _ = Counter(contents).most_common(1)[0]
|
|
for response in responses:
|
|
if response["choices"][0]["message"]["content"] == majority_content:
|
|
return response
|
|
return responses[0]
|
|
|
|
|
|
class _ModelUnavailable(Exception):
|
|
def __init__(self, message: str) -> None:
|
|
super().__init__(message)
|
|
self.message = message
|
|
|
|
|
|
class _TrackerUnavailable(Exception):
|
|
pass
|
|
|
|
|
|
def _post_json(url: str, payload: dict, timeout: float = 5.0) -> dict:
|
|
data = json.dumps(payload).encode()
|
|
req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
|
|
with urllib.request.urlopen(req, timeout=timeout) as r:
|
|
return json.loads(r.read())
|
|
|
|
|
|
def _last_message_content(messages: object) -> str:
|
|
if not isinstance(messages, list) or not messages:
|
|
return ""
|
|
last_message = messages[-1]
|
|
if not isinstance(last_message, dict):
|
|
return ""
|
|
content = last_message.get("content", "")
|
|
return content if isinstance(content, str) else str(content)
|
|
|
|
|
|
def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) -> list[_BinaryActivation]:
|
|
session = str(uuid.uuid4())
|
|
limits = PrefillTransferLimits.from_env()
|
|
chunk_token_count = limits.chunk_tokens
|
|
total_tokens = max(1, _prompt_token_count(prompt))
|
|
chunk_total = max(1, (total_tokens + chunk_token_count - 1) // chunk_token_count)
|
|
|
|
def chunks():
|
|
for chunk_index in range(chunk_total):
|
|
remaining_tokens = total_tokens - (chunk_index * chunk_token_count)
|
|
seq_len = min(chunk_token_count, remaining_tokens)
|
|
yield _BinaryActivation(
|
|
body=_make_stub_binary_activation([1, seq_len, _STUB_HIDDEN_DIM], _STUB_DTYPE),
|
|
shape=[1, seq_len, _STUB_HIDDEN_DIM],
|
|
dtype=_STUB_DTYPE,
|
|
session=session,
|
|
chunk_index=chunk_index,
|
|
chunk_total=chunk_total,
|
|
encoding=_preferred_binary_encoding(),
|
|
headers={},
|
|
)
|
|
|
|
def forward(activation: _BinaryActivation) -> _BinaryActivation:
|
|
for hop_index, node_url in enumerate(route):
|
|
activation = _post_binary_forward(
|
|
f"{node_url}/forward",
|
|
activation,
|
|
hop_index=hop_index,
|
|
timeout=timeout,
|
|
)
|
|
return activation
|
|
|
|
# Each completed response is retained only for the gateway's existing test
|
|
# diagnostic surface. At every hop the sender owns one chunk at a time.
|
|
return BoundedPrefillSender(limits).send(chunks(), body_size=lambda item: len(item.body), forward=forward)
|
|
|
|
|
|
def _chunk_token_count() -> int:
|
|
return PrefillTransferLimits.from_env().chunk_tokens
|
|
|
|
|
|
def _prompt_token_count(prompt: str) -> int:
|
|
return len(prompt.split()) if prompt.split() else max(1, len(prompt))
|
|
|
|
|
|
def _preferred_binary_encoding() -> str | None:
|
|
try:
|
|
import zstandard # noqa: F401
|
|
except ModuleNotFoundError:
|
|
return None
|
|
return "zstd"
|
|
|
|
|
|
def _make_stub_binary_activation(shape: list[int], dtype: str) -> bytes:
|
|
element_count = 1
|
|
for dim in shape:
|
|
element_count *= dim
|
|
return b"\x00" * (element_count * _DTYPE_SIZES[dtype])
|
|
|
|
|
|
def _post_binary_forward(
|
|
url: str,
|
|
activation: _BinaryActivation,
|
|
*,
|
|
hop_index: int,
|
|
timeout: float,
|
|
) -> _BinaryActivation:
|
|
request_body = _compress_body(activation.body, activation.encoding)
|
|
headers = {
|
|
"Content-Type": "application/octet-stream",
|
|
"X-Meshnet-Wire": _WIRE_VERSION,
|
|
"X-Meshnet-Shape": ",".join(str(dim) for dim in activation.shape),
|
|
"X-Meshnet-Dtype": activation.dtype,
|
|
"X-Meshnet-Session": activation.session,
|
|
"X-Meshnet-Chunk-Index": str(activation.chunk_index),
|
|
"X-Meshnet-Chunk-Total": str(activation.chunk_total),
|
|
"X-Meshnet-Hop-Index": str(hop_index),
|
|
}
|
|
if activation.encoding:
|
|
headers["X-Meshnet-Encoding"] = activation.encoding
|
|
req = urllib.request.Request(url, data=request_body, headers=headers, method="POST")
|
|
with urllib.request.urlopen(req, timeout=timeout) as response:
|
|
response_body = response.read()
|
|
response_headers = {key.lower(): value for key, value in response.headers.items()}
|
|
|
|
encoding = response_headers.get("x-meshnet-encoding")
|
|
raw_body = _decompress_body(response_body, encoding)
|
|
# Legacy single-chunk peers returned only the body before chunk metadata
|
|
# was added. Treat absent metadata as the caller's single chunk, while a
|
|
# peer which partially implements the new protocol still fails closed.
|
|
if not response_headers.get("x-meshnet-shape"):
|
|
if activation.chunk_total != 1:
|
|
raise ValueError("legacy peer cannot acknowledge multi-chunk prefill")
|
|
shape = activation.shape
|
|
dtype = activation.dtype
|
|
session = activation.session
|
|
chunk_index = activation.chunk_index
|
|
chunk_total = activation.chunk_total
|
|
else:
|
|
shape = _parse_shape(response_headers["x-meshnet-shape"])
|
|
dtype = response_headers["x-meshnet-dtype"]
|
|
session = response_headers["x-meshnet-session"]
|
|
chunk_index = int(response_headers["x-meshnet-chunk-index"])
|
|
chunk_total = int(response_headers["x-meshnet-chunk-total"])
|
|
if session != activation.session:
|
|
raise ValueError("binary activation response changed session")
|
|
if chunk_index != activation.chunk_index or chunk_total != activation.chunk_total:
|
|
raise ValueError("binary activation response changed chunk metadata")
|
|
if dtype != activation.dtype or shape != activation.shape:
|
|
raise ValueError("binary activation response shape/dtype mismatch")
|
|
_validate_activation_body(raw_body, shape, dtype)
|
|
return _BinaryActivation(
|
|
body=raw_body,
|
|
shape=shape,
|
|
dtype=dtype,
|
|
session=session,
|
|
chunk_index=chunk_index,
|
|
chunk_total=chunk_total,
|
|
encoding=encoding,
|
|
headers=response_headers,
|
|
)
|
|
|
|
|
|
def _reassemble_binary_responses(chunks: list[_BinaryActivation]) -> bytes:
|
|
"""Validate and concatenate final route chunk responses into one activation body."""
|
|
if not chunks:
|
|
raise ValueError("binary pipeline returned no chunks")
|
|
first = chunks[0]
|
|
expected_total = first.chunk_total
|
|
if len(chunks) != expected_total:
|
|
raise ValueError("binary pipeline returned incomplete chunk set")
|
|
if [chunk.chunk_index for chunk in chunks] != list(range(expected_total)):
|
|
raise ValueError("binary pipeline returned out-of-order chunks")
|
|
for chunk in chunks:
|
|
if chunk.session != first.session or chunk.dtype != first.dtype:
|
|
raise ValueError("binary chunks do not belong to the same activation")
|
|
if len(chunk.shape) != 3 or chunk.shape[0] != first.shape[0] or chunk.shape[2] != first.shape[2]:
|
|
raise ValueError("binary chunks have incompatible shapes")
|
|
return b"".join(chunk.body for chunk in chunks)
|
|
|
|
|
|
def _compress_body(body: bytes, encoding: str | None) -> bytes:
|
|
if not encoding:
|
|
return body
|
|
if encoding != "zstd":
|
|
raise ValueError("unsupported binary encoding")
|
|
import zstandard as zstd
|
|
|
|
try:
|
|
return zstd.ZstdCompressor(level=1).compress(body)
|
|
except zstd.ZstdError as exc:
|
|
raise ValueError("could not zstd-compress binary activation") from exc
|
|
|
|
|
|
def _decompress_body(body: bytes, encoding: str | None) -> bytes:
|
|
if not encoding:
|
|
return body
|
|
if encoding != "zstd":
|
|
raise ValueError("unsupported binary encoding")
|
|
import zstandard as zstd
|
|
|
|
try:
|
|
return zstd.ZstdDecompressor().decompress(body)
|
|
except zstd.ZstdError as exc:
|
|
raise ValueError("invalid zstd binary activation") from exc
|
|
|
|
|
|
def _parse_shape(value: str) -> list[int]:
|
|
shape = [int(part) for part in value.split(",")]
|
|
if not shape or any(dim <= 0 for dim in shape):
|
|
raise ValueError("invalid binary activation shape")
|
|
return shape
|
|
|
|
|
|
def _validate_activation_body(body: bytes, shape: list[int], dtype: str) -> None:
|
|
if dtype not in _DTYPE_SIZES:
|
|
raise ValueError("unsupported binary activation dtype")
|
|
if len(body) != len(_make_stub_binary_activation(shape, dtype)):
|
|
raise ValueError("binary activation byte length does not match headers")
|
|
|
|
|
|
def _stub_response_prefix(chunk_responses: list[_BinaryActivation]) -> str:
|
|
if not chunk_responses:
|
|
return "stub response to:"
|
|
return chunk_responses[-1].headers.get("x-meshnet-stub-response-prefix", "stub response to:")
|
|
|
|
|
|
def _get_json(url: str, timeout: float = 5.0) -> dict:
|
|
with urllib.request.urlopen(url, timeout=timeout) as r:
|
|
return json.loads(r.read())
|
|
|
|
|
|
def _get_head_workers(server: _GatewayHTTPServer, model: str) -> list[str]:
|
|
"""Return head-worker endpoint URLs for this model, or empty list on failure."""
|
|
if server.tracker_url is None:
|
|
return []
|
|
try:
|
|
resp = _get_json(f"{server.tracker_url}/v1/head-workers/{urllib.parse.quote(model)}")
|
|
return [n["endpoint"] for n in resp.get("head_workers", [])]
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
class GatewayServer:
|
|
"""HTTP gateway that routes /v1/chat/completions through an ordered inference route.
|
|
|
|
Routing modes (mutually exclusive):
|
|
- ``tracker_url``: gateway queries the tracker per request using the model name.
|
|
- ``inference_route``: static ordered list of node base URLs.
|
|
- ``node_url``: single-node shortcut (US-001 compat), equivalent to inference_route=[node_url].
|
|
|
|
Optional ``model_presets`` lists the model names returned by GET /v1/models when using
|
|
a static inference route. Ignored when using tracker_url (tracker provides the list).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
node_url: str | None = None,
|
|
*,
|
|
inference_route: list[str] | None = None,
|
|
tracker_url: str | None = None,
|
|
model_presets: list[str] | None = None,
|
|
contracts: Any | None = None,
|
|
minimum_stake: int = 0,
|
|
cost_per_layer_token_lamport: int = 1,
|
|
host: str = "127.0.0.1",
|
|
port: int = 0,
|
|
):
|
|
routing_modes = sum(x is not None for x in (node_url, inference_route, tracker_url))
|
|
if routing_modes > 1:
|
|
raise ValueError("Provide exactly one of: node_url, inference_route, or tracker_url")
|
|
if routing_modes == 0:
|
|
raise ValueError("Provide one of: node_url, inference_route, or tracker_url")
|
|
|
|
if node_url is not None:
|
|
inference_route = [node_url]
|
|
if inference_route is not None and not inference_route:
|
|
raise ValueError("inference_route must contain at least one node URL")
|
|
|
|
self._inference_route: list[str] | None = inference_route
|
|
self._tracker_url: str | None = tracker_url.rstrip("/") if tracker_url is not None else None
|
|
self._model_presets: list[str] | None = model_presets
|
|
self._contracts = contracts
|
|
self._minimum_stake = minimum_stake
|
|
self._cost_per_layer_token_lamport = cost_per_layer_token_lamport
|
|
self._host = host
|
|
self._requested_port = port
|
|
self._server: _GatewayHTTPServer | None = None
|
|
self._thread: threading.Thread | None = None
|
|
self.port: int | None = None
|
|
|
|
@property
|
|
def last_binary_chunk_responses(self) -> list[_BinaryActivation]:
|
|
"""Most recent binary chunk responses returned by the final route node."""
|
|
return self._server.last_binary_chunk_responses if self._server is not None else []
|
|
|
|
def start(self) -> int:
|
|
if self._server is not None:
|
|
raise RuntimeError("GatewayServer is already running")
|
|
|
|
self._server = _GatewayHTTPServer(
|
|
(self._host, self._requested_port),
|
|
_GatewayHandler,
|
|
self._inference_route,
|
|
self._tracker_url,
|
|
self._model_presets,
|
|
self._contracts,
|
|
self._minimum_stake,
|
|
self._cost_per_layer_token_lamport,
|
|
)
|
|
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
|