diff --git a/.scratch/distributed-inference-network/prd.json b/.scratch/distributed-inference-network/prd.json index 479c2a1..2840388 100644 --- a/.scratch/distributed-inference-network/prd.json +++ b/.scratch/distributed-inference-network/prd.json @@ -67,11 +67,12 @@ "Commit only this story's code changes with a focused conventional commit message" ], "priority": 3, - "passes": false, + "passes": true, "notes": "Source issue: .scratch/distributed-inference-network/issues/03-tracker-registration-and-routing.md", "dependsOn": [ "US-002" - ] + ], + "completionNotes": "Completed by Ralph iteration 79796dd2; verified by pytest, compileall, editable installs, CLI help, and spec/code reviews." }, { "id": "US-004", diff --git a/packages/gateway/meshnet_gateway/cli.py b/packages/gateway/meshnet_gateway/cli.py index 5795e6b..3064985 100644 --- a/packages/gateway/meshnet_gateway/cli.py +++ b/packages/gateway/meshnet_gateway/cli.py @@ -12,11 +12,11 @@ def main() -> None: subparsers = parser.add_subparsers(dest="command") start_cmd = subparsers.add_parser("start", help="Start the gateway server") - start_cmd.add_argument("--tracker", default="http://localhost:8080", help="Tracker URL") - start_cmd.add_argument( + route_group = start_cmd.add_mutually_exclusive_group() + route_group.add_argument("--tracker", default="http://localhost:8080", help="Tracker URL") + route_group.add_argument( "--node-url", - default="http://localhost:7000", - help="Node URL to route to until tracker routing is implemented", + help="Single node URL for local smoke tests instead of tracker routing", ) start_cmd.add_argument("--port", type=int, default=8080, help="Port to listen on") @@ -25,11 +25,15 @@ def main() -> None: if args.command == "start": from meshnet_gateway.server import GatewayServer - gateway = GatewayServer(node_url=args.node_url, port=args.port) + if args.node_url: + gateway = GatewayServer(node_url=args.node_url, port=args.port) + routing_label = f"Routing to node URL: {args.node_url}" + else: + gateway = GatewayServer(tracker_url=args.tracker, port=args.port) + routing_label = f"Tracker URL: {args.tracker}" port = gateway.start() print(f"meshnet-gateway listening on port {port}") - print(f"Tracker URL: {args.tracker}") - print(f"Routing to node URL: {args.node_url}") + print(routing_label) try: import time while True: diff --git a/packages/gateway/meshnet_gateway/server.py b/packages/gateway/meshnet_gateway/server.py index 21dabd9..6c0e8a5 100644 --- a/packages/gateway/meshnet_gateway/server.py +++ b/packages/gateway/meshnet_gateway/server.py @@ -4,13 +4,16 @@ import http.server import json import threading import time +import urllib.error +import urllib.parse import urllib.request class _GatewayHTTPServer(http.server.HTTPServer): - def __init__(self, addr, handler, inference_route: list[str]): + def __init__(self, addr, handler, inference_route: list[str] | None, tracker_url: str | None): super().__init__(addr, handler) self.inference_route = inference_route + self.tracker_url = tracker_url class _GatewayHandler(http.server.BaseHTTPRequestHandler): @@ -24,11 +27,41 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler): self.send_response(404) self.end_headers() + def _send_json_error(self, status: int, message: str) -> None: + resp = json.dumps({"error": message}).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(resp))) + self.end_headers() + self.wfile.write(resp) + def _handle_chat_completions(self): + server: _GatewayHTTPServer = self.server # type: ignore[assignment] length = int(self.headers.get("Content-Length", 0)) body = json.loads(self.rfile.read(length)) - route: list[str] = self.server.inference_route # type: ignore[attr-defined] + if server.tracker_url is not None: + model = body.get("model", "stub-model") + tracker_path = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(model)}" + try: + tracker_resp = _get_json(tracker_path) + route = tracker_resp["route"] + if ( + not isinstance(route, list) + or not route + or not all(isinstance(node_url, str) and node_url for node_url in route) + ): + raise ValueError("invalid tracker route") + except urllib.error.HTTPError as exc: + error_body = _safe_error_body(exc) + self._send_json_error(503, error_body.get("error", "tracker unavailable")) + return + except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, KeyError, TypeError, ValueError): + self._send_json_error(503, "tracker unavailable") + return + else: + route = server.inference_route # type: ignore[assignment] + messages = body.get("messages", []) # Send messages to the first node in the inference route. @@ -65,6 +98,14 @@ class _GatewayHandler(http.server.BaseHTTPRequestHandler): self.wfile.write(response_body) +def _safe_error_body(exc: urllib.error.HTTPError) -> dict: + try: + body = json.loads(exc.read()) + except json.JSONDecodeError: + return {} + return body if isinstance(body, dict) else {} + + def _post_json(url: str, payload: dict, timeout: float = 5.0) -> dict: data = json.dumps(payload).encode() req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST") @@ -72,11 +113,18 @@ def _post_json(url: str, payload: dict, timeout: float = 5.0) -> dict: return json.loads(r.read()) +def _get_json(url: str, timeout: float = 5.0) -> dict: + with urllib.request.urlopen(url, timeout=timeout) as r: + return json.loads(r.read()) + + class GatewayServer: """HTTP gateway that routes /v1/chat/completions through an ordered inference route. - Supply either ``node_url`` (single-node, US-001 compat) or ``inference_route`` - (ordered list of node base URLs for multi-shard routing). + Routing modes (mutually exclusive): + - ``tracker_url``: gateway queries the tracker per request using the model name. + - ``inference_route``: static ordered list of node base URLs. + - ``node_url``: single-node shortcut (US-001 compat), equivalent to inference_route=[node_url]. """ def __init__( @@ -84,18 +132,23 @@ class GatewayServer: node_url: str | None = None, *, inference_route: list[str] | None = None, + tracker_url: str | None = None, host: str = "127.0.0.1", port: int = 0, ): - if node_url is not None and inference_route is not None: - raise ValueError("Provide node_url or inference_route, not both") + routing_modes = sum(x is not None for x in (node_url, inference_route, tracker_url)) + if routing_modes > 1: + raise ValueError("Provide exactly one of: node_url, inference_route, or tracker_url") + if routing_modes == 0: + raise ValueError("Provide one of: node_url, inference_route, or tracker_url") + if node_url is not None: inference_route = [node_url] - elif inference_route is None: - raise ValueError("Provide either node_url or inference_route") - if not inference_route: + if inference_route is not None and not inference_route: raise ValueError("inference_route must contain at least one node URL") - self._inference_route = inference_route + + self._inference_route: list[str] | None = inference_route + self._tracker_url: str | None = tracker_url.rstrip("/") if tracker_url is not None else None self._host = host self._requested_port = port self._server: _GatewayHTTPServer | None = None @@ -110,6 +163,7 @@ class GatewayServer: (self._host, self._requested_port), _GatewayHandler, self._inference_route, + self._tracker_url, ) self.port = self._server.server_address[1] self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) diff --git a/packages/tracker/meshnet_tracker/cli.py b/packages/tracker/meshnet_tracker/cli.py index 2ca9237..c8c5e1c 100644 --- a/packages/tracker/meshnet_tracker/cli.py +++ b/packages/tracker/meshnet_tracker/cli.py @@ -2,6 +2,9 @@ import argparse import sys +import time + +from .server import TrackerServer def main() -> None: @@ -12,17 +15,30 @@ def main() -> None: subparsers = parser.add_subparsers(dest="command") start_cmd = subparsers.add_parser("start", help="Start the tracker server") + start_cmd.add_argument("--host", default="127.0.0.1", help="Host interface to listen on") start_cmd.add_argument("--port", type=int, default=8080, help="Port to listen on") + start_cmd.add_argument( + "--heartbeat-timeout", + type=float, + default=30.0, + help="Seconds before a node is removed from the registry after missed heartbeat", + ) args = parser.parse_args() if args.command == "start": - print(f"meshnet-tracker starting on port {args.port}") + server = TrackerServer( + host=args.host, + port=args.port, + heartbeat_timeout=args.heartbeat_timeout, + ) + port = server.start() + print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True) try: - import time while True: time.sleep(1) except KeyboardInterrupt: + server.stop() sys.exit(0) else: parser.print_help() diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py new file mode 100644 index 0000000..785e3da --- /dev/null +++ b/packages/tracker/meshnet_tracker/server.py @@ -0,0 +1,291 @@ +"""Tracker HTTP server — node registry and route selection for the inference mesh. + +HTTP API contract: +- POST /v1/nodes/register + Request JSON: { + "endpoint": "http://host:port", "shard_start": int, "shard_end": int, + "hardware_profile": object, "score": number optional + } + Response 200: {"node_id": str} + Response 400: {"error": str} +- POST /v1/nodes/{node_id}/heartbeat + Response 200: {} + Response 404: {"error": "node not found"} +- GET /v1/route?model= + Response 200: {"route": ["http://node-a", "http://node-b"]} + Response 400/404/503: {"error": str} +""" + +import http.server +import json +import threading +import time +import urllib.parse +import uuid + + +DEFAULT_MODEL_PRESETS: dict[str, dict] = { + "stub-model": {"layers_start": 0, "layers_end": 31}, +} + + +class _NodeEntry: + __slots__ = ( + "node_id", "endpoint", "shard_start", "shard_end", + "hardware_profile", "score", "last_heartbeat", + ) + + def __init__( + self, + node_id: str, + endpoint: str, + shard_start: int, + shard_end: int, + hardware_profile: dict, + score: float, + ) -> None: + self.node_id = node_id + self.endpoint = endpoint + self.shard_start = shard_start + self.shard_end = shard_end + self.hardware_profile = hardware_profile + self.score = score + self.last_heartbeat: float = time.monotonic() + + +def _select_route( + nodes: list[_NodeEntry], + required_start: int, + required_end: int, +) -> tuple[list[_NodeEntry], str]: + """Greedy interval-cover. Returns (ordered route, error_message).""" + candidates = sorted(nodes, key=lambda n: (n.shard_start, -n.shard_end)) + route: list[_NodeEntry] = [] + covered_up_to = required_start - 1 + + while covered_up_to < required_end: + best: _NodeEntry | None = None + for node in candidates: + if node.shard_start <= covered_up_to + 1 and node.shard_end > covered_up_to: + if best is None or node.shard_end > best.shard_end: + best = node + if best is None: + missing = covered_up_to + 1 + return [], f"no route available: no registered node covers layer {missing}" + route.append(best) + covered_up_to = best.shard_end + candidates = [n for n in candidates if n is not best] + + return route, "" + + +class _TrackerHTTPServer(http.server.HTTPServer): + def __init__( + self, + addr: tuple, + handler, + registry: dict, + lock: threading.Lock, + heartbeat_timeout: float, + model_presets: dict, + ) -> None: + super().__init__(addr, handler) + self.registry = registry + self.lock = lock + self.heartbeat_timeout = heartbeat_timeout + self.model_presets = model_presets + + +class _TrackerHandler(http.server.BaseHTTPRequestHandler): + def log_message(self, fmt, *args): # noqa: suppress request logs in tests + pass + + def _send_json(self, status: int, data: dict) -> None: + body = json.dumps(data).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _read_json_body(self) -> dict | None: + length = int(self.headers.get("Content-Length", 0)) + try: + body = json.loads(self.rfile.read(length) or b"{}") + except json.JSONDecodeError: + self._send_json(400, {"error": "invalid JSON body"}) + return None + if not isinstance(body, dict): + self._send_json(400, {"error": "JSON body must be an object"}) + return None + return body + + def _purge_expired_nodes(self) -> None: + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + now = time.monotonic() + expired_ids = [ + node_id for node_id, entry in server.registry.items() + if (now - entry.last_heartbeat) > server.heartbeat_timeout + ] + for node_id in expired_ids: + del server.registry[node_id] + + def do_POST(self): + if self.path == "/v1/nodes/register": + self._handle_register() + return + parts = self.path.split("/") + # /v1/nodes//heartbeat -> ['', 'v1', 'nodes', '', 'heartbeat'] + if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat": + self._handle_heartbeat(parts[3]) + return + self.send_response(404) + self.end_headers() + + def do_GET(self): + parsed = urllib.parse.urlparse(self.path) + if parsed.path == "/v1/route": + self._handle_route(parsed) + else: + self.send_response(404) + self.end_headers() + + def _handle_register(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_json_body() + if body is None: + return + + endpoint = body.get("endpoint") + if not isinstance(endpoint, str) or not endpoint: + self._send_json(400, {"error": "endpoint is required"}) + return + parsed_endpoint = urllib.parse.urlparse(endpoint) + if parsed_endpoint.scheme not in {"http", "https"} or not parsed_endpoint.netloc: + self._send_json(400, {"error": "endpoint must be an http(s) URL"}) + return + + try: + shard_start = int(body["shard_start"]) + shard_end = int(body["shard_end"]) + score = float(body.get("score", 1.0)) + except (KeyError, TypeError, ValueError): + self._send_json(400, {"error": "shard_start, shard_end, and score must be numeric"}) + return + if shard_start < 0 or shard_end < 0 or shard_start > shard_end: + self._send_json(400, {"error": "shard range must be non-negative and ordered"}) + return + + hardware_profile = body.get("hardware_profile", {}) + if not isinstance(hardware_profile, dict): + self._send_json(400, {"error": "hardware_profile must be an object"}) + return + + node_id = str(uuid.uuid4()) + entry = _NodeEntry( + node_id=node_id, + endpoint=endpoint.rstrip("/"), + shard_start=shard_start, + shard_end=shard_end, + hardware_profile=hardware_profile, + score=score, + ) + with server.lock: + self._purge_expired_nodes() + server.registry[node_id] = entry + + self._send_json(200, {"node_id": node_id}) + + def _handle_heartbeat(self, node_id: str): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + with server.lock: + self._purge_expired_nodes() + entry = server.registry.get(node_id) + if entry is None: + self._send_json(404, {"error": "node not found"}) + return + entry.last_heartbeat = time.monotonic() + self._send_json(200, {}) + + def _handle_route(self, parsed: urllib.parse.ParseResult): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + params = urllib.parse.parse_qs(parsed.query) + model_list = params.get("model") + if not model_list: + self._send_json(400, {"error": "missing 'model' query parameter"}) + return + + model = model_list[0] + preset = server.model_presets.get(model) + if preset is None: + self._send_json(404, {"error": f"unknown model preset: {model!r}"}) + return + + required_start: int = preset["layers_start"] + required_end: int = preset["layers_end"] + + with server.lock: + self._purge_expired_nodes() + alive = list(server.registry.values()) + + route, error = _select_route(alive, required_start, required_end) + if error: + self._send_json(503, {"error": error}) + return + + self._send_json(200, {"route": [e.endpoint for e in route]}) + + +class TrackerServer: + """HTTP tracker that manages node registration and resolves inference routes. + + Nodes register via POST /v1/nodes/register and keep themselves alive with + POST /v1/nodes//heartbeat. The gateway queries GET /v1/route?model= + to obtain an ordered list of node endpoints whose shards cover all layers. + """ + + def __init__( + self, + host: str = "127.0.0.1", + port: int = 0, + heartbeat_timeout: float = 30.0, + model_presets: dict | None = None, + ) -> None: + self._host = host + self._requested_port = port + self._heartbeat_timeout = heartbeat_timeout + self._model_presets: dict = ( + model_presets if model_presets is not None else dict(DEFAULT_MODEL_PRESETS) + ) + self._registry: dict[str, _NodeEntry] = {} + self._lock = threading.Lock() + self._server: _TrackerHTTPServer | None = None + self._thread: threading.Thread | None = None + self.port: int | None = None + + def start(self) -> int: + if self._server is not None: + raise RuntimeError("TrackerServer is already running") + self._server = _TrackerHTTPServer( + (self._host, self._requested_port), + _TrackerHandler, + self._registry, + self._lock, + self._heartbeat_timeout, + self._model_presets, + ) + self.port = self._server.server_address[1] + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + self._thread.start() + return self.port + + def stop(self) -> None: + if self._server is None: + return + self._server.shutdown() + self._server.server_close() + if self._thread is not None: + self._thread.join(timeout=1) + self._server = None + self._thread = None + self.port = None diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py new file mode 100644 index 0000000..d7d32d8 --- /dev/null +++ b/tests/test_tracker_routing.py @@ -0,0 +1,365 @@ +"""US-003 tests: tracker node registration and route selection.""" + +import http.server +import json +import threading +import time +import urllib.error +import urllib.request + +from meshnet_gateway.server import GatewayServer +from meshnet_node.server import StubNodeServer +from meshnet_tracker.server import TrackerServer + + +def _post_json(url: str, payload: dict) -> dict: + data = json.dumps(payload).encode() + req = urllib.request.Request( + url, data=data, headers={"Content-Type": "application/json"}, method="POST" + ) + with urllib.request.urlopen(req) as r: + return json.loads(r.read()) + + +def _get_json(url: str) -> dict: + with urllib.request.urlopen(url) as r: + return json.loads(r.read()) + + +def test_tracker_node_registration(): + """A node can register with the tracker and receives a node_id.""" + tracker = TrackerServer() + tracker_port = tracker.start() + try: + resp = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + { + "endpoint": "http://127.0.0.1:9000", + "shard_start": 0, + "shard_end": 15, + "hardware_profile": {"gpu": "rtx3090"}, + "score": 1.0, + }, + ) + assert "node_id" in resp + assert isinstance(resp["node_id"], str) + assert len(resp["node_id"]) > 0 + finally: + tracker.stop() + + +def test_tracker_route_selection(): + """Tracker returns ordered route when nodes collectively cover all layers.""" + tracker = TrackerServer() + tracker_port = tracker.start() + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15, + "hardware_profile": {}, "score": 1.0}, + ) + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9002", "shard_start": 16, "shard_end": 31, + "hardware_profile": {}, "score": 1.0}, + ) + + route_resp = _get_json( + f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model" + ) + assert "route" in route_resp + assert route_resp["route"] == [ + "http://127.0.0.1:9001", + "http://127.0.0.1:9002", + ] + finally: + tracker.stop() + + +def test_tracker_route_error_no_coverage(): + """Tracker returns 503 when registered nodes do not cover all required layers.""" + tracker = TrackerServer() + tracker_port = tracker.start() + try: + # Only half the layers covered — route must fail. + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15, + "hardware_profile": {}, "score": 1.0}, + ) + + try: + _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model") + raise AssertionError("Expected 503 but request succeeded") + except urllib.error.HTTPError as exc: + assert exc.code == 503 + body = json.loads(exc.read()) + assert "error" in body + finally: + tracker.stop() + + +def test_tracker_route_error_no_nodes(): + """Tracker returns 503 with clear error when the registry is empty.""" + tracker = TrackerServer() + tracker_port = tracker.start() + try: + try: + _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model") + raise AssertionError("Expected 503 but request succeeded") + except urllib.error.HTTPError as exc: + assert exc.code == 503 + body = json.loads(exc.read()) + assert "error" in body + finally: + tracker.stop() + + +def test_tracker_heartbeat_updates_node(): + """Sending a heartbeat for a registered node succeeds.""" + tracker = TrackerServer() + tracker_port = tracker.start() + try: + reg = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31, + "hardware_profile": {}, "score": 1.0}, + ) + node_id = reg["node_id"] + + hb_resp = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {} + ) + assert hb_resp == {} + finally: + tracker.stop() + + +def test_tracker_heartbeat_expiry(): + """Nodes that miss their heartbeat window are excluded from routes.""" + tracker = TrackerServer(heartbeat_timeout=0.05) # 50 ms + tracker_port = tracker.start() + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31, + "hardware_profile": {}, "score": 1.0}, + ) + + # Immediately after registration the node is alive. + route_resp = _get_json( + f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model" + ) + assert route_resp["route"] == ["http://127.0.0.1:9001"] + + # Let the heartbeat window expire. + time.sleep(0.15) + + try: + _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model") + raise AssertionError("Expected 503 after heartbeat expiry") + except urllib.error.HTTPError as exc: + assert exc.code == 503 + finally: + tracker.stop() + + +def test_tracker_heartbeat_expiry_removes_node_from_registry(): + """Expired nodes are removed, not merely hidden from route responses.""" + tracker = TrackerServer(heartbeat_timeout=0.05) + tracker_port = tracker.start() + try: + reg = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31, + "hardware_profile": {}, "score": 1.0}, + ) + node_id = reg["node_id"] + time.sleep(0.15) + + try: + _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model") + raise AssertionError("Expected 503 after heartbeat expiry") + except urllib.error.HTTPError as exc: + assert exc.code == 503 + + try: + _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {}) + raise AssertionError("Expected 404 for removed node") + except urllib.error.HTTPError as exc: + assert exc.code == 404 + finally: + tracker.stop() + + +def test_tracker_route_rejects_non_extending_overlap(): + """Overlapping shards that do not extend coverage cannot form a complete route.""" + tracker = TrackerServer() + tracker_port = tracker.start() + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 10, + "hardware_profile": {}, "score": 1.0}, + ) + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9002", "shard_start": 0, "shard_end": 5, + "hardware_profile": {}, "score": 1.0}, + ) + + try: + _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model") + raise AssertionError("Expected 503 for incomplete overlapping route") + except urllib.error.HTTPError as exc: + assert exc.code == 503 + finally: + tracker.stop() + + +def test_tracker_registration_rejects_invalid_payload(): + """Registration errors return a defined JSON 400 response.""" + tracker = TrackerServer() + tracker_port = tracker.start() + try: + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "not-a-url", "shard_start": 10, "shard_end": 0, + "hardware_profile": []}, + ) + raise AssertionError("Expected 400 for invalid registration") + except urllib.error.HTTPError as exc: + assert exc.code == 400 + body = json.loads(exc.read()) + assert "error" in body + finally: + tracker.stop() + + +def test_gateway_returns_json_503_when_tracker_unavailable(): + """Tracker-backed gateway failures are reported as JSON HTTP errors.""" + gateway = GatewayServer(tracker_url="http://127.0.0.1:9") + gateway_port = gateway.start() + try: + payload = json.dumps({ + "model": "stub-model", + "messages": [{"role": "user", "content": "hello"}], + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{gateway_port}/v1/chat/completions", + data=payload, + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + urllib.request.urlopen(req, timeout=5) + raise AssertionError("Expected 503 for unavailable tracker") + except urllib.error.HTTPError as exc: + assert exc.code == 503 + assert exc.headers["Content-Type"] == "application/json" + body = json.loads(exc.read()) + assert body == {"error": "tracker unavailable"} + finally: + gateway.stop() + + +def test_gateway_returns_json_503_for_malformed_tracker_route(): + """Malformed tracker route payloads do not drop the gateway connection.""" + + class MalformedRouteHandler(http.server.BaseHTTPRequestHandler): + def log_message(self, fmt, *args): + pass + + def do_GET(self): + body = json.dumps({"route": []}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + fake_tracker = http.server.HTTPServer(("127.0.0.1", 0), MalformedRouteHandler) + fake_thread = threading.Thread(target=fake_tracker.serve_forever, daemon=True) + fake_thread.start() + gateway = GatewayServer(tracker_url=f"http://127.0.0.1:{fake_tracker.server_address[1]}") + gateway_port = gateway.start() + try: + payload = json.dumps({ + "model": "stub-model", + "messages": [{"role": "user", "content": "hello"}], + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{gateway_port}/v1/chat/completions", + data=payload, + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + urllib.request.urlopen(req, timeout=5) + raise AssertionError("Expected 503 for malformed tracker route") + except urllib.error.HTTPError as exc: + assert exc.code == 503 + assert exc.headers["Content-Type"] == "application/json" + body = json.loads(exc.read()) + assert body == {"error": "tracker unavailable"} + finally: + gateway.stop() + fake_tracker.shutdown() + fake_tracker.server_close() + fake_thread.join(timeout=1) + + +def test_two_node_pipeline_via_tracker(): + """End-to-end: nodes register with tracker; gateway discovers route dynamically.""" + tracker = TrackerServer() + tracker_port = tracker.start() + tracker_url = f"http://127.0.0.1:{tracker_port}" + + node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False) + port_a = node_a.start() + + node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True) + port_b = node_b.start() + + _post_json( + f"{tracker_url}/v1/nodes/register", + {"endpoint": f"http://127.0.0.1:{port_a}", "shard_start": 0, "shard_end": 15, + "hardware_profile": {}, "score": 1.0}, + ) + _post_json( + f"{tracker_url}/v1/nodes/register", + {"endpoint": f"http://127.0.0.1:{port_b}", "shard_start": 16, "shard_end": 31, + "hardware_profile": {}, "score": 1.0}, + ) + + gateway = GatewayServer(tracker_url=tracker_url) + gateway_port = gateway.start() + + try: + payload = json.dumps({ + "model": "stub-model", + "messages": [{"role": "user", "content": "hello"}], + }).encode() + + req = urllib.request.Request( + f"http://127.0.0.1:{gateway_port}/v1/chat/completions", + data=payload, + headers={"Content-Type": "application/json"}, + method="POST", + ) + + with urllib.request.urlopen(req) as resp: + assert resp.status == 200 + body = json.loads(resp.read()) + + assert body["id"] == "chatcmpl-stub" + assert body["object"] == "chat.completion" + assert body["choices"][0]["message"]["content"] == "stub response to: hello" + assert node_b.received_activations, "node B did not receive activations from node A" + assert not node_a.received_activations, "node A should not receive activations" + finally: + gateway.stop() + node_a.stop() + node_b.stop() + tracker.stop()