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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user