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

@@ -39,6 +39,16 @@ def main() -> None:
"--advertise-host",
help="Reachable host/IP to advertise to the tracker (defaults to FQDN when binding 0.0.0.0)",
)
start_cmd.add_argument(
"--tracker-mode",
action="store_true",
help="Enable client-facing /v1/chat/completions (auto-enabled when shard-start=0)",
)
start_cmd.add_argument(
"--tracker-url",
default=None,
help="Tracker URL for route selection (used in tracker mode)",
)
args = parser.parse_args()

View File

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

View File

@@ -6,6 +6,11 @@ import http.server
import json
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from .model_backend import (
InsufficientVRAMError,
@@ -24,11 +29,20 @@ from .server import (
class _TorchHTTPServer(http.server.HTTPServer):
def __init__(self, addr, handler, backend: TorchModelShard):
def __init__(
self,
addr,
handler,
backend: TorchModelShard,
tracker_mode: bool = False,
tracker_url: str | None = None,
):
super().__init__(addr, handler)
self.backend = backend
self.received_activations = False
self.forward_chunk_count = 0
self.tracker_mode = tracker_mode
self.tracker_url = tracker_url
class _TorchHandler(http.server.BaseHTTPRequestHandler):
@@ -36,10 +50,13 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
pass
def do_POST(self):
server: _TorchHTTPServer = self.server # type: ignore[assignment]
if self.path == "/forward":
self._handle_forward()
elif self.path == "/v1/infer":
self._handle_infer()
elif self.path == "/v1/chat/completions" and server.tracker_mode:
self._handle_chat_completions()
else:
self.send_response(404)
self.end_headers()
@@ -190,6 +207,152 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
self.end_headers()
self.wfile.write(payload)
def _handle_chat_completions(self) -> None:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
body = self._read_json_body()
if body is None:
return
messages = body.get("messages", [])
stream = bool(body.get("stream", False))
model = str(body.get("model", ""))
prompt = " ".join(
str(m.get("content", ""))
for m in messages
if isinstance(m, dict) and m.get("role") == "user"
)
try:
payload = server.backend.encode_prompt(prompt)
except Exception as exc:
self._send_json(500, {"error": f"encode_prompt failed: {exc}"})
return
remaining_route = self._get_remaining_route(model)
result_text = self._run_downstream_pipeline(payload, remaining_route)
self._send_openai_response(result_text, model, stream)
def _get_remaining_route(self, model: str) -> list[str]:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
if server.tracker_url is None:
return []
try:
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(model)}"
with urllib.request.urlopen(url, timeout=5.0) as r:
route_resp = json.loads(r.read())
route = route_resp.get("route", [])
# Skip the first node in the route (self) since we're already the head
return list(route[1:])
except Exception:
return []
def _run_downstream_pipeline(self, payload: object, route: list[str]) -> str:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
if not route:
# Single-node mode: decode tail locally if we're the tail
if server.backend.is_tail:
try:
tensor = server.backend.torch.frombuffer(
bytearray(payload.body), # type: ignore[union-attr]
dtype=server.backend.torch.bfloat16,
).reshape(payload.shape).to(server.backend.device) # type: ignore[union-attr]
return server.backend.decode_tail(tensor)
except Exception as exc:
return f"decode error: {exc}"
return ""
session = str(uuid.uuid4())
shape = payload.shape # type: ignore[union-attr]
attn_mask = payload.attention_mask_header # type: ignore[union-attr]
pos_ids = payload.position_ids_header # type: ignore[union-attr]
current_body = payload.body # type: ignore[union-attr]
current_shape = shape
current_attn = attn_mask
current_pos = pos_ids
for hop_index, node_url in enumerate(route):
headers: dict[str, str] = {
"Content-Type": "application/octet-stream",
"X-Meshnet-Wire": _WIRE_VERSION,
"X-Meshnet-Shape": ",".join(str(d) for d in current_shape),
"X-Meshnet-Dtype": "bfloat16",
"X-Meshnet-Session": session,
"X-Meshnet-Chunk-Index": "0",
"X-Meshnet-Chunk-Total": "1",
"X-Meshnet-Hop-Index": str(hop_index),
}
if current_attn:
headers["X-Meshnet-Attn-Mask"] = current_attn
if current_pos:
headers["X-Meshnet-Position-Ids"] = current_pos
req = urllib.request.Request(
f"{node_url}/forward",
data=current_body,
headers=headers,
method="POST",
)
try:
with urllib.request.urlopen(req, timeout=10.0) as r:
resp_body = r.read()
resp_headers = {k.lower(): v for k, v in r.headers.items()}
except Exception as exc:
return f"pipeline error at {node_url}: {exc}"
content_type = resp_headers.get("content-type", "")
if "application/json" in content_type:
try:
data = json.loads(resp_body)
return str(data.get("text", ""))
except json.JSONDecodeError:
return resp_body.decode("utf-8", errors="replace")
# Binary activation — update and forward to next node
shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape))
current_shape = _parse_shape(shape_header)
current_body = resp_body
current_attn = resp_headers.get("x-meshnet-attn-mask")
current_pos = resp_headers.get("x-meshnet-position-ids")
return ""
def _send_openai_response(self, text: str, model: str, stream: bool) -> None:
chunk_id = "chatcmpl-node"
created = int(time.time())
if not stream:
self._send_json(200, {
"id": chunk_id,
"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},
})
return
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()
class TorchNodeServer:
"""HTTP server backed by a HuggingFace causal language model shard."""
@@ -203,6 +366,8 @@ class TorchNodeServer:
shard_end: int = 6,
quantization: str = "bfloat16",
backend: TorchModelShard | None = None,
tracker_mode: bool | None = None,
tracker_url: str | None = None,
) -> None:
self._host = host
self._requested_port = port
@@ -212,6 +377,9 @@ class TorchNodeServer:
shard_end,
quantization,
)
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
self._tracker_url = tracker_url
self._server: _TorchHTTPServer | None = None
self._thread: threading.Thread | None = None
self.port: int | None = None
@@ -235,6 +403,8 @@ class TorchNodeServer:
(self._host, self._requested_port),
_TorchHandler,
self._backend,
self._tracker_mode,
self._tracker_url,
)
self.port = self._server.server_address[1]
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)