From 3dff7098075f4d2866f055554cc38661515199c3 Mon Sep 17 00:00:00 2001 From: CrimsonDump Date: Thu, 13 Aug 2026 15:14:33 +0800 Subject: [PATCH 1/2] fix(pd): role chunk before decode errors; count only consumed MTP steps --- tilert/pd_vllm/pd_router.py | 2 ++ tilert/pd_vllm/profiles/mla_nsa.py | 28 +++++++++++++++++----------- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/tilert/pd_vllm/pd_router.py b/tilert/pd_vllm/pd_router.py index 896a98d..768c643 100644 --- a/tilert/pd_vllm/pd_router.py +++ b/tilert/pd_vllm/pd_router.py @@ -384,6 +384,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: 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, From a706bec2c9f236729cc699258099810fa2f70eee Mon Sep 17 00:00:00 2001 From: CrimsonDump Date: Fri, 14 Aug 2026 18:08:31 +0800 Subject: [PATCH 2/2] feat(pd_router): queue for a free decode node instead of failing fast MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 摘要 一个 decode engine 一次只服务一条序列，所以 router 给每个请求预留一个节点，池子占满 时直接返回 429。对**单客户端扇出**的负载这个语义太硬：一个 agentic 会话派生出并发的 子对话时，突发宽度大于池子但持续很短，本可以稍等即得的负载被判了 429。 `--queue-timeout <秒>` 让请求等待而不是失败；默认 `0` 保持原来的 fail-fast 行为， 所以这项改动在默认配置下是惰性的。 # 改了什么 ## `Pool` —— `Lock` 换 `Condition` `acquire` 在 `queue_timeout > 0` 时阻塞等待，`release` 归还后 `notify` 一个等待者。 等待者在 `while` 里重扫全部节点，所以不会丢唤醒。两个调用点本来就在线程池里 （非流式的 `_handle` 整个被 `run_in_threadpool` 调起；流式那条显式包了一层）， 阻塞不会卡住事件循环。 ## 修一处会让节点永久 busy 的取消泄漏 客户端断连会取消这个任务，而 `CancelledError` 继承 `BaseException`、 不被原来的 `except Exception` 捕获；流式尚未开始，生成器的 `finally` （平时的释放路径）也还不存在。于是节点的 `busy` 永不复位 —— 单节点部署下 router 之后会永远返回 429。 两条路径：**prefill 期间断连**（这条在加排队之前就存在）、**排队期间断连** （这条是排队新引入的取消点，窗口宽度就是 `--queue-timeout`）。 修法是把预留放进 `try`、`node` 先置 `None`、两个 except 都释放； `acquire` 那次 await 另包一层 `anyio.CancelScope(shield=True)`， 使断连不能打断等待本身、留下一个谁也拿不到的预留。 ⚠️ 只加 shield 不够：anyio 把挂起的取消投递在 shield 作用域的**出口**， 赋值已完成但 `if node is None` 那行还没执行到，所以预留必须落在 `try` 里面。 ## 可观测性与 429 语义 * 排队超过 0.1 s 打一行 `queued %.1fs for decode node %s`，否则这段等待在响应里 完全看不出来； * 429 的 body 区分两种情况：`all decode nodes busy`（fail-fast）与 `no decode node free after waiting %.1fs`（等到超时）。 ## README Topology A 的 router 段落补一段说明 `--queue-timeout` 的适用场景与默认值。 # 测试 * `Pool` 语义 16 条断言：fail-fast 兼容、超时触发、6 路 fan-out 全部成功且串行化、 多节点并行度 = 节点数、反复争用无饥饿； * 取消泄漏 6 条断言（真 anyio + 真 `Pool`，三种结构 × 两种断连时机）： 旧结构两种都漏、只加 shield 仍漏、本改动都不漏； * `pre-commit run` 全绿（isort / black / flake8+插件 / mypy / bandit / pyupgrade / codespell）。 Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 2 + tilert/pd_vllm/pd_router.py | 108 ++++++++++++++++++++++++++++++------ 2 files changed, 92 insertions(+), 18 deletions(-) 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 768c643..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") @@ -468,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 = [] @@ -483,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",