fix inference
This commit is contained in:
@@ -24,6 +24,7 @@ HTTP API contract:
|
||||
import http.server
|
||||
import hashlib
|
||||
import json
|
||||
import socketserver
|
||||
import threading
|
||||
import time
|
||||
import urllib.parse
|
||||
@@ -291,7 +292,7 @@ def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]:
|
||||
for node_id in expired_ids:
|
||||
entry = server.registry.pop(node_id)
|
||||
print(
|
||||
f"[tracker] node expired: {node_id[:8]} {entry.endpoint} "
|
||||
f"[tracker] node expired: {node_id} {entry.endpoint} "
|
||||
f"(no heartbeat for >{server.heartbeat_timeout:.0f}s)",
|
||||
flush=True,
|
||||
)
|
||||
@@ -418,7 +419,9 @@ def _node_id_for_registration(
|
||||
return f"{wallet_prefix}-{digest}"
|
||||
|
||||
|
||||
class _TrackerHTTPServer(http.server.HTTPServer):
|
||||
class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
addr: tuple,
|
||||
@@ -706,6 +709,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
# 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"
|
||||
request_id = str(body.get("id") or f"req-{time.time_ns():x}")
|
||||
|
||||
# Pre-resolve the downstream route so the first-shard node skips its own
|
||||
# tracker query. We already hold the full registry picture — no need for
|
||||
@@ -726,6 +730,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
route_nodes, _ = _select_route(all_nodes, rs, re)
|
||||
# Strip the first-shard node we're about to proxy to — it's already handling the request.
|
||||
downstream_urls = json.dumps([n.endpoint for n in route_nodes if n.endpoint != node.endpoint])
|
||||
route_debug = " -> ".join(
|
||||
f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]"
|
||||
for n in route_nodes
|
||||
)
|
||||
print(
|
||||
f"[tracker] proxy route {request_id}: model={model!r} "
|
||||
f"head={node.node_id}@{node.endpoint} downstream={downstream_urls} "
|
||||
f"route={route_debug or '<empty>'}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
req = urllib.request.Request(
|
||||
target_url,
|
||||
@@ -743,6 +757,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
try:
|
||||
upstream = urllib.request.urlopen(req, timeout=300.0)
|
||||
print(f"[tracker] proxy connected {request_id}: {target_url}", flush=True)
|
||||
except urllib.error.HTTPError as exc:
|
||||
# Relay error status + body from node
|
||||
err_body = exc.read()
|
||||
@@ -756,6 +771,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
pass
|
||||
return
|
||||
except Exception as exc:
|
||||
print(f"[tracker] proxy failed {request_id}: {target_url}: {exc}", flush=True)
|
||||
self._send_json(503, {"error": {
|
||||
"message": f"upstream node unreachable: {exc}",
|
||||
"type": "service_unavailable",
|
||||
@@ -783,6 +799,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
else:
|
||||
# Non-streaming: buffer and relay
|
||||
resp_body = upstream.read()
|
||||
print(
|
||||
f"[tracker] proxy complete {request_id}: {target_url} bytes={len(resp_body)}",
|
||||
flush=True,
|
||||
)
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", content_type)
|
||||
self.send_header("Content-Length", str(len(resp_body)))
|
||||
@@ -966,7 +986,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
for eid in stale_ids:
|
||||
old = server.registry.pop(eid)
|
||||
print(
|
||||
f"[tracker] node re-registered: replaced {eid[:8]} with {node_id[:8]}"
|
||||
f"[tracker] node re-registered: replaced {eid} with {node_id}"
|
||||
f" {old.endpoint}",
|
||||
flush=True,
|
||||
)
|
||||
@@ -980,7 +1000,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded"
|
||||
repo_info = f" [{hf_repo}]" if hf_repo else ""
|
||||
print(
|
||||
f"[tracker] node registered: {node_id[:8]} {endpoint} {model}{repo_info} {shard_info}",
|
||||
f"[tracker] node registered: {node_id} {endpoint} {model}{repo_info} {shard_info}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user