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:
Dobromir Popov
2026-06-29 16:59:32 +03:00
parent 753f553766
commit dbf856f497
7 changed files with 580 additions and 7 deletions

View File

@@ -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.