diff --git a/README.md b/README.md index 463cb6c..f3d2db7 100644 --- a/README.md +++ b/README.md @@ -363,6 +363,8 @@ python -m tilert.pd_vllm.pd_router \ Send OpenAI requests to `http://:23333/v1/chat/completions`. The router runs the prefill on vLLM (first token), hands the attention state to the TileRT decode node over RDMA, and streams the completion back. +A decode engine serves one sequence at a time, so the router reserves a node per request and answers `429` while they are all busy. Add `--queue-timeout ` to make a request wait for a free node instead of failing: useful when a single client fans out into concurrent sub-conversations — an agentic session spawning sub-agents, say — and the burst is wider than the pool but short-lived. Waits longer than 0.1 s are logged. The default, `0`, keeps the fail-fast behaviour. + ### Topology B: shared prefill → TileRT decode **and** native vLLM decode One prefill pool feeds two decode pools side by side, composed under vLLM's `MultiConnector`. Each request is claimed by exactly one connector — the TileRT connector claims requests marked with `tilert_host`, and vLLM's native connector handles the rest — so latency-critical traffic goes to TileRT while general traffic stays on native vLLM decode, behind the same OpenAI surface. diff --git a/tilert/pd_vllm/pd_router.py b/tilert/pd_vllm/pd_router.py index 896a98d..b3b75b4 100644 --- a/tilert/pd_vllm/pd_router.py +++ b/tilert/pd_vllm/pd_router.py @@ -4,7 +4,8 @@ non-streaming. Flow per request (phase-1 hybrid, see design doc): - 1. pick a free decode node (in-memory busy tracking; all busy -> 429) + 1. pick a free decode node (in-memory busy tracking; all busy -> wait up to + --queue-timeout, then 429) 2. forward to vLLM with max_tokens=1 + logprobs and inject kv_transfer_params {tilert_host, tilert_ctrl_port} — the connector claims the request and RDMA-sends state to the decode node @@ -40,6 +41,9 @@ logger = logging.getLogger("pd_vllm.router") +# Only log a queue wait once it is long enough to explain a latency bump. +QUEUE_LOG_SECONDS = 0.1 + class DecodeNode: def __init__(self, host: str, ctrl_port: int, http_port: int): @@ -54,21 +58,43 @@ def http_base(self) -> str: class Pool: - def __init__(self, nodes: list[DecodeNode]): + """Decode-node reservation. + + ``queue_timeout`` > 0 makes ``acquire`` wait for a node instead of failing + fast. A decode engine serves one sequence at a time, so a client that puts + more than one request in flight per node — a multi-turn agentic session + fanning out into concurrent sub-conversations, for instance — otherwise gets + 429s for load the pool can serve a moment later. 0 keeps the fail-fast + behaviour. + """ + + def __init__(self, nodes: list[DecodeNode], queue_timeout: float = 0.0): self.nodes = nodes - self._lock = threading.Lock() + self.queue_timeout = queue_timeout + self._cv = threading.Condition() def acquire(self) -> DecodeNode | None: - with self._lock: - for n in self.nodes: - if not n.busy: - n.busy = True - return n - return None + """Reserve a node, or None once ``queue_timeout`` elapses. + + Blocks while waiting; both call sites already hop off the event loop via + ``run_in_threadpool``, so other streams keep being served. + """ + deadline = time.monotonic() + self.queue_timeout + with self._cv: + while True: + for n in self.nodes: + if not n.busy: + n.busy = True + return n + remaining = deadline - time.monotonic() + if remaining <= 0: + return None + self._cv.wait(remaining) def release(self, node: DecodeNode) -> None: - with self._lock: + with self._cv: node.busy = False + self._cv.notify() def first_token_from_logprobs(resp: dict, is_chat: bool) -> int: @@ -154,6 +180,29 @@ def build_app(ctx: RouterCtx) -> FastAPI: app = FastAPI() pool = ctx.pool + def _acquire_node() -> tuple[DecodeNode | None, float]: + """Reserve a node, and report how long the caller had to queue for it. + + A decode engine serves one sequence at a time, so a client that fans out + — an agentic session spawning sub-conversations, say — serialises here. + That wait is invisible in the response, hence the log line. + """ + t0 = time.monotonic() + node = pool.acquire() + waited = time.monotonic() - t0 + if node is not None and waited >= QUEUE_LOG_SECONDS: + logger.info("queued %.1fs for decode node %s", waited, node.host) + return node, waited + + def _busy_response(waited: float) -> JSONResponse: + """429 body: a pool that is full reads differently from one we waited on.""" + detail = ( + f"no decode node free after waiting {waited:.1f}s" + if pool.queue_timeout > 0 + else "all decode nodes busy" + ) + return JSONResponse({"error": detail}, status_code=429) + @app.get("/health") def health(): return {"status": "ok", "decode_free": sum(1 for n in pool.nodes if not n.busy)} @@ -178,9 +227,9 @@ def _max_tokens_of(body): # ── non-streaming ──────────────────────────────────────────────────── def _handle(path: str, body: dict): is_chat = path.endswith("chat/completions") - node = pool.acquire() + node, waited = _acquire_node() if node is None: - return JSONResponse({"error": "all decode nodes busy"}, status_code=429) + return _busy_response(waited) t0 = time.time() try: prefill = _prefill(path, body, node) @@ -255,20 +304,37 @@ def _handle(path: str, body: dict): # ── streaming (chat only) ──────────────────────────────────────────── async def _handle_stream(path: str, body: dict, request: Request): + import anyio from starlette.concurrency import run_in_threadpool - node = pool.acquire() - if node is None: - return JSONResponse({"error": "all decode nodes busy"}, status_code=429) - + # The reservation must be inside the try: a client disconnect cancels this + # task, and until streaming starts the generator's finally — the usual + # release path — does not exist yet. Both handlers below release. + node = None try: + # Shielded so a disconnect cannot interrupt the queue wait itself and + # strand a reservation the worker thread already made. Bounded by + # --queue-timeout, which is what we were waiting for anyway. + with anyio.CancelScope(shield=True): + node, waited = await run_in_threadpool(_acquire_node) + if node is None: + return _busy_response(waited) prefill = await run_in_threadpool(_prefill, path, body, node) rid = derive_rid(prefill["id"]) first_token_id = first_token_from_logprobs(prefill, True) except Exception as e: - pool.release(node) + if node is not None: + pool.release(node) logger.exception("pd stream request failed before streaming") return JSONResponse({"error": str(e)}, status_code=502) + except BaseException: + # CancelledError is not an Exception. anyio delivers a pending + # cancellation when the shielded scope exits — after the reservation + # is made, before streaming starts — so this clause is what keeps the + # node from staying busy forever. See drafts/mock_router_cancel.py. + if node is not None: + pool.release(node) + raise chunk_id = prefill["id"] model = prefill.get("model") @@ -384,6 +450,8 @@ def _role_once(): if finish_reason == "cancelled": finish_reason = "stop" elif "error" in msg: + if not role_sent: + yield _role_once() yield _chunk({"content": f"\n[decode error: {msg['error']}]"}) finish_reason = "stop" if client_gone: @@ -466,6 +534,12 @@ def main() -> None: default="glm47", help="output parser (reasoning + tool calls)", ) + ap.add_argument( + "--queue-timeout", + type=float, + default=0.0, + help="seconds to wait for a free decode node before answering 429 (0: fail fast)", + ) args = ap.parse_args() nodes = [] @@ -481,7 +555,7 @@ def main() -> None: args.model_path, trust_remote_code=True ) # nosec B615 - ctx = RouterCtx(args.vllm_url, Pool(nodes), tokenizer, args.parser) + ctx = RouterCtx(args.vllm_url, Pool(nodes, args.queue_timeout), tokenizer, args.parser) app = build_app(ctx) logger.info( "router on :%d -> vllm=%s, %d decode node(s), parser=%s", diff --git a/tilert/pd_vllm/profiles/mla_nsa.py b/tilert/pd_vllm/profiles/mla_nsa.py index f8747bd..29c1b35 100644 --- a/tilert/pd_vllm/profiles/mla_nsa.py +++ b/tilert/pd_vllm/profiles/mla_nsa.py @@ -441,19 +441,25 @@ def _decode_mtp(self, first_token_id, budget, on_token, cancel_event): if fwd == 0: fwd += 1 continue - accepted.extend(per_step) fwd += 1 - for tok in emitted: - if len(tokens) >= budget: - break - tok = int(tok) - if tok in stop_ids: - finished = True - finish = "stop" + offset = 0 + for na in per_step: + step_emit = emitted[offset : offset + na] + offset += na + for tok in step_emit: + if len(tokens) >= budget: + break + tok = int(tok) + if tok in stop_ids: + finished = True + finish = "stop" + break + tokens.append(tok) + if on_token: + on_token(tok) + accepted.append(na) + if finished or len(tokens) >= budget: break - tokens.append(tok) - if on_token: - on_token(tok) dl.reset_sequence() self.last_stats = { "finish_reason": finish,