130 lines
4.4 KiB
Python
130 lines
4.4 KiB
Python
"""Gateway HTTP server — accepts OpenAI-format requests and routes to inference nodes."""
|
|
|
|
import http.server
|
|
import json
|
|
import threading
|
|
import time
|
|
import urllib.request
|
|
|
|
|
|
class _GatewayHTTPServer(http.server.HTTPServer):
|
|
def __init__(self, addr, handler, inference_route: list[str]):
|
|
super().__init__(addr, handler)
|
|
self.inference_route = inference_route
|
|
|
|
|
|
class _GatewayHandler(http.server.BaseHTTPRequestHandler):
|
|
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
|
|
pass
|
|
|
|
def do_POST(self):
|
|
if self.path == "/v1/chat/completions":
|
|
self._handle_chat_completions()
|
|
else:
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
|
|
def _handle_chat_completions(self):
|
|
length = int(self.headers.get("Content-Length", 0))
|
|
body = json.loads(self.rfile.read(length))
|
|
|
|
route: list[str] = self.server.inference_route # type: ignore[attr-defined]
|
|
messages = body.get("messages", [])
|
|
|
|
# Send messages to the first node in the inference route.
|
|
node_resp = _post_json(f"{route[0]}/v1/infer", {"messages": messages})
|
|
|
|
# Forward activations through remaining nodes in the route.
|
|
for node_url in route[1:]:
|
|
node_resp = _post_json(f"{node_url}/v1/infer", {"activations": node_resp["activations"]})
|
|
|
|
response_body = json.dumps({
|
|
"id": "chatcmpl-stub",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": body.get("model", "stub-model"),
|
|
"choices": [{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": node_resp["text"],
|
|
},
|
|
"finish_reason": "stop",
|
|
}],
|
|
"usage": {
|
|
"prompt_tokens": 0,
|
|
"completion_tokens": 0,
|
|
"total_tokens": 0,
|
|
},
|
|
}).encode()
|
|
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(response_body)))
|
|
self.end_headers()
|
|
self.wfile.write(response_body)
|
|
|
|
|
|
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())
|
|
|
|
|
|
class GatewayServer:
|
|
"""HTTP gateway that routes /v1/chat/completions through an ordered inference route.
|
|
|
|
Supply either ``node_url`` (single-node, US-001 compat) or ``inference_route``
|
|
(ordered list of node base URLs for multi-shard routing).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
node_url: str | None = None,
|
|
*,
|
|
inference_route: list[str] | None = None,
|
|
host: str = "127.0.0.1",
|
|
port: int = 0,
|
|
):
|
|
if node_url is not None and inference_route is not None:
|
|
raise ValueError("Provide node_url or inference_route, not both")
|
|
if node_url is not None:
|
|
inference_route = [node_url]
|
|
elif inference_route is None:
|
|
raise ValueError("Provide either node_url or inference_route")
|
|
if not inference_route:
|
|
raise ValueError("inference_route must contain at least one node URL")
|
|
self._inference_route = inference_route
|
|
self._host = host
|
|
self._requested_port = port
|
|
self._server: _GatewayHTTPServer | None = None
|
|
self._thread: threading.Thread | None = None
|
|
self.port: int | None = None
|
|
|
|
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.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
|