feat: tracker exposes OpenAI-compatible /v1/chat/completions proxy
The tracker is now the single entrypoint for inference. Clients POST to
the tracker's /v1/chat/completions and it forwards to a live tracker-mode
(first-shard) node for the requested model, applying round-robin load
balancing across multiple first-shard nodes when available.
Streaming (text/event-stream) is relayed chunk-by-chunk. Non-streaming
responses are buffered and forwarded. BrokenPipeError on client disconnect
is silenced. Upstream errors relay the original HTTP status and body.
Also adds GET /v1/health → {"status": "ok"} on the tracker.
Usage:
curl -s http://tracker:8081/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{"model":"Qwen/Qwen2.5-0.5B-Instruct","messages":[...]}'
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -423,6 +423,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
_purge_expired_nodes_locked(server)
|
||||
|
||||
def do_POST(self):
|
||||
if self.path == "/v1/chat/completions":
|
||||
self._handle_proxy_chat()
|
||||
return
|
||||
if self.path == "/v1/nodes/register":
|
||||
self._handle_register()
|
||||
return
|
||||
@@ -463,6 +466,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
self._handle_tracker_nodes(model)
|
||||
elif parsed.path == "/v1/raft/status":
|
||||
self._handle_raft_status()
|
||||
elif parsed.path == "/v1/health":
|
||||
self._send_json(200, {"status": "ok"})
|
||||
else:
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
@@ -550,6 +555,119 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
],
|
||||
})
|
||||
|
||||
# ---------------------------------------------------------------- OpenAI proxy
|
||||
|
||||
def _handle_proxy_chat(self) -> None:
|
||||
"""Proxy POST /v1/chat/completions to a tracker-mode (first-shard) node.
|
||||
|
||||
Picks a live tracker-mode node for the requested model using round-robin,
|
||||
then forwards the request verbatim and relays the response (including
|
||||
streaming SSE chunks) back to the caller.
|
||||
"""
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
raw_body = self.rfile.read(length) if length else b"{}"
|
||||
|
||||
try:
|
||||
body = json.loads(raw_body)
|
||||
except json.JSONDecodeError:
|
||||
self._send_json(400, {"error": {"message": "invalid JSON", "type": "invalid_request_error", "code": "invalid_request"}})
|
||||
return
|
||||
|
||||
model: str = body.get("model", "")
|
||||
is_stream: bool = bool(body.get("stream", False))
|
||||
|
||||
# Find a live tracker-mode node for this model
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
candidates = [
|
||||
n for n in server.registry.values()
|
||||
if n.tracker_mode and (n.model == model or n.hf_repo == model)
|
||||
]
|
||||
|
||||
if not candidates:
|
||||
# Fall back: any node serving shard_start=0 for this model
|
||||
with server.lock:
|
||||
candidates = [
|
||||
n for n in server.registry.values()
|
||||
if n.shard_start == 0 and (n.model == model or n.hf_repo == model)
|
||||
]
|
||||
|
||||
if not candidates:
|
||||
self._send_json(503, {"error": {
|
||||
"message": f"no nodes available for model {model!r}",
|
||||
"type": "service_unavailable",
|
||||
"code": "model_not_available",
|
||||
}})
|
||||
return
|
||||
|
||||
# Simple round-robin via list length modulo (stateless, good enough)
|
||||
node = candidates[int(time.time() * 1000) % len(candidates)]
|
||||
target_url = f"{node.endpoint}/v1/chat/completions"
|
||||
|
||||
req = urllib.request.Request(
|
||||
target_url,
|
||||
data=raw_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
# Copy Authorization header from client if present
|
||||
auth = self.headers.get("Authorization")
|
||||
if auth:
|
||||
req.add_header("Authorization", auth)
|
||||
|
||||
try:
|
||||
upstream = urllib.request.urlopen(req, timeout=300.0)
|
||||
except urllib.error.HTTPError as exc:
|
||||
# Relay error status + body from node
|
||||
err_body = exc.read()
|
||||
self.send_response(exc.code)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(err_body)))
|
||||
self.end_headers()
|
||||
try:
|
||||
self.wfile.write(err_body)
|
||||
except BrokenPipeError:
|
||||
pass
|
||||
return
|
||||
except Exception as exc:
|
||||
self._send_json(503, {"error": {
|
||||
"message": f"upstream node unreachable: {exc}",
|
||||
"type": "service_unavailable",
|
||||
"code": "upstream_error",
|
||||
}})
|
||||
return
|
||||
|
||||
with upstream:
|
||||
content_type = upstream.headers.get("Content-Type", "application/json")
|
||||
if is_stream or "text/event-stream" in content_type:
|
||||
# Relay SSE stream chunk-by-chunk
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
|
||||
self.send_header("Cache-Control", "no-cache")
|
||||
self.end_headers()
|
||||
try:
|
||||
while True:
|
||||
line = upstream.readline()
|
||||
if not line:
|
||||
break
|
||||
self.wfile.write(line)
|
||||
self.wfile.flush()
|
||||
except BrokenPipeError:
|
||||
pass
|
||||
else:
|
||||
# Non-streaming: buffer and relay
|
||||
resp_body = upstream.read()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", content_type)
|
||||
self.send_header("Content-Length", str(len(resp_body)))
|
||||
self.end_headers()
|
||||
try:
|
||||
self.wfile.write(resp_body)
|
||||
except BrokenPipeError:
|
||||
pass
|
||||
|
||||
def _handle_register(self):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
body = self._read_json_body()
|
||||
|
||||
Reference in New Issue
Block a user