From 60ccf47cb4ebdd4195ee2377ff676b64caad1069 Mon Sep 17 00:00:00 2001
From: Dobromir Popov <dobromir.popov@gmail.com>
Date: Tue, 30 Jun 2026 01:07:38 +0300
Subject: [PATCH] feat(us-016): fix distributed inference route lookup and
 autoregressive generation

Route lookup was using the client-provided model name ("qwen2.5-0.5b") but
the tracker registers nodes under their full hf_repo ("Qwen/Qwen2.5-0.5B-Instruct").
This caused a 404 on /v1/route and the non-tail node fell back to the
"no downstream route available" error message.

Fix: _get_remaining_route now uses server.backend.model_id (the actual hf_repo)
for the tracker query. Skips self by port matching rather than blind route[0] drop.
Also prints a warning when route lookup fails so the cause is visible.

Distributed generation was also only producing 1 token (single greedy argmax
in decode_tail). Replaced with an autoregressive loop: head node encodes the
growing sequence and forwards to the downstream shard each step, collecting
one token per iteration up to max_tokens or EOS.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 packages/node/meshnet_node/torch_server.py | 74 +++++++++++++++++-----
 1 file changed, 58 insertions(+), 16 deletions(-)

diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py
index 3dd5c3c..fd41ed3 100644
--- a/packages/node/meshnet_node/torch_server.py
+++ b/packages/node/meshnet_node/torch_server.py
@@ -238,33 +238,75 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
                 self._send_json(500, {"error": f"generation failed: {exc}"})
             return
 
-        # Distributed path: encode prompt at the head, forward activations along the route.
-        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
+        # Distributed path: autoregressive generation across shards.
+        # We do N single-step forward passes (no cross-node KV cache), which is slow
+        # but correct. Each step: head encodes current sequence → forwards through route
+        # → tail returns the next token string → append → repeat.
         remaining_route = self._get_remaining_route(model_name)
-        result_text = self._run_downstream_pipeline(payload, remaining_route)
+        if not remaining_route:
+            self._send_openai_response(
+                "error: no downstream route — check tracker connectivity",
+                model_name, False, messages,
+            )
+            return
+
+        backend = server.backend
+        # Format with chat template so the model knows it's in assistant mode.
+        try:
+            if hasattr(backend.tokenizer, "apply_chat_template"):
+                prompt_text: str = backend.tokenizer.apply_chat_template(
+                    messages, add_generation_prompt=True, tokenize=False,
+                )
+            else:
+                raise AttributeError("no apply_chat_template")
+        except Exception:
+            prompt_text = " ".join(
+                str(m.get("content", ""))
+                for m in messages
+                if isinstance(m, dict) and m.get("role") == "user"
+            )
+
+        eos_token: str = getattr(backend.tokenizer, "eos_token", "") or ""
+        generated: list[str] = []
+        current_text = prompt_text
+
+        for _ in range(max_tokens):
+            try:
+                payload = backend.encode_prompt(current_text)
+            except Exception as exc:
+                print(f"  [node] distributed encode error: {exc}", flush=True)
+                break
+            token_str = self._run_downstream_pipeline(payload, remaining_route)
+            if not token_str:
+                break
+            # Stop on error responses or EOS.
+            if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")):
+                break
+            if eos_token and token_str == eos_token:
+                break
+            generated.append(token_str)
+            current_text = current_text + token_str
+
+        result_text = "".join(generated)
         self._send_openai_response(result_text, model_name, stream, messages)
 
     def _get_remaining_route(self, model: str) -> list[str]:
         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(model)}"
+            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 the first node in the route (self) since we're already the head
-            return list(route[1:])
-        except Exception:
+            # 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:
+            print(f"  [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
             return []
 
     def _run_downstream_pipeline(self, payload: object, route: list[str]) -> str: