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
@@ -2,6 +2,7 @@
|
||||
|
||||
import http.server
|
||||
import json
|
||||
import time
|
||||
import threading
|
||||
import urllib.parse
|
||||
from pathlib import Path
|
||||
@@ -93,6 +94,8 @@ class _StubHTTPServer(http.server.HTTPServer):
|
||||
response_prefix: str,
|
||||
model: str,
|
||||
shard_path: Path | None,
|
||||
tracker_mode: bool = False,
|
||||
tracker_url: str | None = None,
|
||||
):
|
||||
super().__init__(addr, handler)
|
||||
self.shard_start = shard_start
|
||||
@@ -103,6 +106,9 @@ class _StubHTTPServer(http.server.HTTPServer):
|
||||
self.shard_path = shard_path
|
||||
self.received_activations: bool = False
|
||||
self.forward_chunk_count: int = 0
|
||||
self.tracker_mode: bool = tracker_mode
|
||||
self.tracker_url: str | None = tracker_url
|
||||
self.chat_completion_count: int = 0
|
||||
|
||||
|
||||
class _StubHandler(http.server.BaseHTTPRequestHandler):
|
||||
@@ -110,10 +116,13 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
|
||||
pass
|
||||
|
||||
def do_POST(self):
|
||||
server: _StubHTTPServer = self.server # type: ignore[assignment]
|
||||
if self.path == "/v1/infer":
|
||||
self._handle_infer()
|
||||
elif self.path == "/forward":
|
||||
self._handle_forward()
|
||||
elif self.path == "/v1/chat/completions" and server.tracker_mode:
|
||||
self._handle_chat_completions()
|
||||
else:
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
@@ -126,6 +135,82 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
|
||||
def _send_json(self, status: int, data: dict) -> None:
|
||||
payload = json.dumps(data).encode()
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(payload)))
|
||||
self.end_headers()
|
||||
self.wfile.write(payload)
|
||||
|
||||
def _handle_chat_completions(self) -> None:
|
||||
server: _StubHTTPServer = self.server # type: ignore[assignment]
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
try:
|
||||
body = json.loads(self.rfile.read(length) or b"{}")
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
self._send_json(400, {"error": "invalid JSON body"})
|
||||
return
|
||||
if not isinstance(body, dict):
|
||||
self._send_json(400, {"error": "JSON body must be an object"})
|
||||
return
|
||||
server.chat_completion_count += 1
|
||||
streaming = bool(body.get("stream", False))
|
||||
model = str(body.get("model", server.model))
|
||||
messages = body.get("messages", [])
|
||||
last_content = ""
|
||||
if isinstance(messages, list) and messages:
|
||||
last = messages[-1]
|
||||
if isinstance(last, dict):
|
||||
last_content = str(last.get("content", ""))
|
||||
text = f"{server.response_prefix} {last_content}"
|
||||
if streaming:
|
||||
self._send_sse_response(text, model)
|
||||
else:
|
||||
created = int(time.time())
|
||||
self._send_json(200, {
|
||||
"id": "chatcmpl-stub",
|
||||
"object": "chat.completion",
|
||||
"created": created,
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": text},
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
||||
})
|
||||
|
||||
def _send_sse_response(self, text: str, model: str) -> None:
|
||||
chunk_id = "chatcmpl-stub"
|
||||
created = int(time.time())
|
||||
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()
|
||||
|
||||
def _emit(data: str) -> None:
|
||||
self.wfile.write(f"data: {data}\n\n".encode())
|
||||
self.wfile.flush()
|
||||
|
||||
_emit(json.dumps({
|
||||
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
|
||||
}))
|
||||
_emit(json.dumps({
|
||||
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}],
|
||||
}))
|
||||
_emit(json.dumps({
|
||||
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
}))
|
||||
self.wfile.write(b"data: [DONE]\n\n")
|
||||
self.wfile.flush()
|
||||
|
||||
def _handle_shard_download(self, parsed: urllib.parse.ParseResult):
|
||||
server: _StubHTTPServer = self.server # type: ignore[assignment]
|
||||
params = urllib.parse.parse_qs(parsed.query)
|
||||
@@ -246,6 +331,7 @@ class StubNodeServer:
|
||||
shard_start / shard_end define which transformer layer range this node owns.
|
||||
is_last_shard controls whether the node returns a text response (True) or
|
||||
activation tensors (False) after processing its shard.
|
||||
tracker_mode enables the /v1/chat/completions endpoint for head-shard nodes.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -258,6 +344,8 @@ class StubNodeServer:
|
||||
response_prefix: str = "stub response to:",
|
||||
model: str = "stub-model",
|
||||
shard_path: Path | None = None,
|
||||
tracker_mode: bool = False,
|
||||
tracker_url: str | None = None,
|
||||
):
|
||||
self._host = host
|
||||
self._requested_port = port
|
||||
@@ -269,6 +357,8 @@ class StubNodeServer:
|
||||
self._response_prefix = response_prefix
|
||||
self._model = model
|
||||
self._shard_path = shard_path
|
||||
self._tracker_mode = tracker_mode
|
||||
self._tracker_url = tracker_url
|
||||
self._server: _StubHTTPServer | None = None
|
||||
self._thread: threading.Thread | None = None
|
||||
self.port: int | None = None
|
||||
@@ -283,6 +373,11 @@ class StubNodeServer:
|
||||
"""Number of binary /forward chunks handled since this node was started."""
|
||||
return self._server.forward_chunk_count if self._server is not None else 0
|
||||
|
||||
@property
|
||||
def chat_completion_count(self) -> int:
|
||||
"""Number of /v1/chat/completions requests handled since this node was started."""
|
||||
return self._server.chat_completion_count if self._server is not None else 0
|
||||
|
||||
def start(self) -> int:
|
||||
if self._server is not None:
|
||||
raise RuntimeError("StubNodeServer is already running")
|
||||
@@ -296,6 +391,8 @@ class StubNodeServer:
|
||||
self._response_prefix,
|
||||
self._model,
|
||||
self._shard_path,
|
||||
self._tracker_mode,
|
||||
self._tracker_url,
|
||||
)
|
||||
self.port = self._server.server_address[1]
|
||||
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
||||
|
||||
Reference in New Issue
Block a user