diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 2d975e2..23bbc94 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -294,18 +294,26 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self._send_openai_response(result_text, model_name, stream, messages) def _get_remaining_route(self, model: str) -> list[str]: + # Fast path: tracker pre-resolved the downstream route and injected it as a header. + injected = self.headers.get("X-Meshnet-Route") + if injected: + try: + route = json.loads(injected) + if isinstance(route, list): + return [str(ep) for ep in route] + except (json.JSONDecodeError, TypeError): + pass + + # Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject). server: _TorchHTTPServer = self.server # type: ignore[assignment] if server.tracker_url is None: return [] - # Use the backend's actual hf_repo, not the client-provided model name (which may be - # a lowercased or abbreviated alias that doesn't match what the tracker registered). route_model = getattr(server.backend, "model_id", None) or model try: url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}" with urllib.request.urlopen(url, timeout=5.0) as r: route_resp = json.loads(r.read()) route = route_resp.get("route", []) - # Skip our own endpoint from the route (match by port so host aliases don't matter). own_port = server.server_address[1] return [ep for ep in route if not ep.rstrip("/").endswith(f":{own_port}")] except Exception as exc: diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 377264a..7c201bb 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -606,10 +606,33 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): node = candidates[int(time.time() * 1000) % len(candidates)] target_url = f"{node.endpoint}/v1/chat/completions" + # Pre-resolve the downstream route so the first-shard node skips its own + # tracker query. We already hold the full registry picture — no need for + # a second round-trip. + route_model = node.hf_repo or node.model or model + with server.lock: + if route_model in server.model_presets: + preset = server.model_presets[route_model] + rs, re = _preset_layer_bounds(preset) + all_nodes: list = [n for n in server.registry.values() if n.model == route_model] + else: + all_nodes = [ + n for n in server.registry.values() + if (n.hf_repo == route_model or n.model == route_model) + and n.shard_start is not None and n.num_layers is not None + ] + rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1) + route_nodes, _ = _select_route(all_nodes, rs, re) + # Strip the first-shard node we're about to proxy to — it's already handling the request. + downstream_urls = json.dumps([n.endpoint for n in route_nodes if n.endpoint != node.endpoint]) + req = urllib.request.Request( target_url, data=raw_body, - headers={"Content-Type": "application/json"}, + headers={ + "Content-Type": "application/json", + "X-Meshnet-Route": downstream_urls, + }, method="POST", ) # Copy Authorization header from client if present