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>
This commit is contained in:
Dobromir Popov
2026-06-30 01:07:38 +03:00
parent c75e9708ae
commit 60ccf47cb4

View File

@@ -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(
# 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)
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 = server.backend.encode_prompt(prompt)
payload = backend.encode_prompt(current_text)
except Exception as exc:
self._send_json(500, {"error": f"encode_prompt failed: {exc}"})
return
remaining_route = self._get_remaining_route(model_name)
result_text = self._run_downstream_pipeline(payload, remaining_route)
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: