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)
|
_purge_expired_nodes_locked(server)
|
||||||
|
|
||||||
def do_POST(self):
|
def do_POST(self):
|
||||||
|
if self.path == "/v1/chat/completions":
|
||||||
|
self._handle_proxy_chat()
|
||||||
|
return
|
||||||
if self.path == "/v1/nodes/register":
|
if self.path == "/v1/nodes/register":
|
||||||
self._handle_register()
|
self._handle_register()
|
||||||
return
|
return
|
||||||
@@ -463,6 +466,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self._handle_tracker_nodes(model)
|
self._handle_tracker_nodes(model)
|
||||||
elif parsed.path == "/v1/raft/status":
|
elif parsed.path == "/v1/raft/status":
|
||||||
self._handle_raft_status()
|
self._handle_raft_status()
|
||||||
|
elif parsed.path == "/v1/health":
|
||||||
|
self._send_json(200, {"status": "ok"})
|
||||||
else:
|
else:
|
||||||
self.send_response(404)
|
self.send_response(404)
|
||||||
self.end_headers()
|
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):
|
def _handle_register(self):
|
||||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||||
body = self._read_json_body()
|
body = self._read_json_body()
|
||||||
|
|||||||
Reference in New Issue
Block a user