feat: tracker-as-first-layer-node inference entry point (US-014)
- Tracker: add GET /v1/tracker-nodes/<model> returning nodes registered with tracker_mode=true whose shard_start matches the model's first layer - Node: StubNodeServer and TorchNodeServer accept tracker_mode/tracker_url; when tracker_mode=True (or auto-detected via shard_start==0 for Torch), /v1/chat/completions is served alongside /forward - TorchNodeServer: full pipeline implementation — encode_prompt → route selection via tracker → binary forward through remaining hops → decode - Gateway: _handle_chat_completions checks _get_tracker_nodes() first and proxies round-robin to tracker-nodes; falls back to existing direct pipeline when none found (preserves all US-005 backward compat) - CLI: --tracker-mode and --tracker-url flags added to meshnet-node start - Test: two stub tracker-nodes + two mid-shard nodes for gpt2; 10 requests; round-robin 5/5 split verified; all OpenAI-format responses validated - All 78 tests pass Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -56,6 +56,7 @@ class _GatewayHTTPServer(http.server.HTTPServer):
|
||||
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):
|
||||
@@ -243,11 +244,28 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
def _handle_chat_completions(self):
|
||||
server: _GatewayHTTPServer = self.server # type: ignore[assignment]
|
||||
body = self._read_json_body()
|
||||
if body is None:
|
||||
# Read raw bytes first so we can proxy them if tracker-nodes 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
|
||||
streaming = bool(body.get("stream", False))
|
||||
|
||||
model = str(body.get("model", "stub-model"))
|
||||
tracker_nodes = _get_tracker_nodes(server, model)
|
||||
if tracker_nodes:
|
||||
# Proxy to a tracker-node (round-robin by request count)
|
||||
target = tracker_nodes[server.request_count % len(tracker_nodes)]
|
||||
server.request_count += 1
|
||||
return self._proxy_to_tracker_node(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:
|
||||
@@ -266,6 +284,37 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler):
|
||||
else:
|
||||
self._send_json(200, completion)
|
||||
|
||||
def _proxy_to_tracker_node(self, url: str, body_bytes: bytes) -> None:
|
||||
"""Forward a raw request body to a tracker-node and stream the response back."""
|
||||
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")
|
||||
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"tracker-node 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:
|
||||
@@ -767,6 +816,17 @@ def _get_json(url: str, timeout: float = 5.0) -> dict:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def _get_tracker_nodes(server: _GatewayHTTPServer, model: str) -> list[str]:
|
||||
"""Return tracker-node endpoint URLs for this model, or empty list on any failure."""
|
||||
if server.tracker_url is None:
|
||||
return []
|
||||
try:
|
||||
resp = _get_json(f"{server.tracker_url}/v1/tracker-nodes/{urllib.parse.quote(model)}")
|
||||
return [n["endpoint"] for n in resp.get("tracker_nodes", [])]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
class GatewayServer:
|
||||
"""HTTP gateway that routes /v1/chat/completions through an ordered inference route.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user