From 2713ff6bcae9122ee25eaefebc3d4e4e39072255 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 8 Jul 2026 02:25:36 -0700 Subject: [PATCH 01/18] [None][feat] add disaggregated coordinator serving Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- docs/source/features/disagg-serving.md | 107 +++- tensorrt_llm/commands/serve.py | 306 +++++++++- tensorrt_llm/llmapi/disagg_utils.py | 81 ++- tensorrt_llm/llmapi/llm_args.py | 16 + tensorrt_llm/serve/coordinator_server.py | 163 +++++ tensorrt_llm/serve/disagg_coordinator.py | 517 ++++++++++++++++ tensorrt_llm/serve/openai_client.py | 6 +- tensorrt_llm/serve/openai_disagg_server.py | 130 ++-- tensorrt_llm/serve/openai_disagg_service.py | 202 +++--- tensorrt_llm/serve/openai_server.py | 27 + tensorrt_llm/serve/perf_metrics.py | 8 + tensorrt_llm/serve/responses_utils.py | 108 ++++ tensorrt_llm/serve/router.py | 578 +++++++++++++++--- tensorrt_llm/serve/router_utils.py | 377 ++++++++++++ .../disaggregated/test_coordinator_e2e.py | 430 +++++++++++++ .../disaggregated/test_coordinator_worker.py | 287 +++++++++ .../test_disagg_openai_client.py | 6 +- .../disaggregated/test_disagg_utils.py | 45 +- .../test_openai_disagg_service.py | 24 +- 19 files changed, 3106 insertions(+), 312 deletions(-) create mode 100644 tensorrt_llm/serve/coordinator_server.py create mode 100644 tensorrt_llm/serve/disagg_coordinator.py create mode 100644 tensorrt_llm/serve/router_utils.py create mode 100644 tests/unittest/disaggregated/test_coordinator_e2e.py create mode 100644 tests/unittest/disaggregated/test_coordinator_worker.py diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index f38e3c818f57..9fb7926a0adb 100644 --- a/docs/source/features/disagg-serving.md +++ b/docs/source/features/disagg-serving.md @@ -94,11 +94,19 @@ The optimizations required for KV cache transmission vary depending on whether i ### Unique Global Request ID -A disaggregated-serving request can provide a unique global request ID via `DisaggregatedParams.disagg_request_id`. -When this field is a positive integer, the context and generation requests share that value as their internal request ID, which enables end-to-end tracking. -To avoid collisions with worker-local or warm-up requests, it is recommended to use a value larger than `1 << 42 = 4398046511104`. -If the field is unset or non-positive, the context and generation requests instead receive separate local sequence IDs, rotating within the range `(0, 1<<42)`, assigned by the respective workers. When `disagg_request_id` is specified, do not route the context and generation requests to the same worker. -This field is optional at present; however, some forthcoming features will depend on this unique identifier. +The context and generation phases of one request must share a single request ID: the ctx↔gen KV-cache transfer is keyed by it, so a collision (two in-flight requests with the same ID) corrupts the transfer. This shared ID is carried on `DisaggregatedParams.disagg_request_id`. + +The disaggregated server generates this ID itself as a **snowflake** — a self-contained 64-bit positive integer that is unique without any cross-process coordination. The bit layout is: + +``` +[ 0 (1 bit) | timestamp_ms (39 bits) | node_id (8 bits) | process_id (6 bits) | counter (10 bits) ] +``` + +- `node_id` (0–255) identifies the node (defaults to a hash of the MAC address; overridable via `node_id` in the disaggregated config). +- `process_id` (0–63) identifies the orchestrator process on that node. In a [coordinator + worker fleet](#coordinator-and-worker-fleet) each fleet worker receives a distinct value, so co-located workers never emit the same ID in the same millisecond. It is set from the `TRTLLM_DISAGG_WORKER_PROCESS_ID` environment variable (assigned automatically per worker by the launcher). +- The `(node_id, process_id)` pair therefore makes the ID unique across all orchestrator processes without a shared counter or an extra network round trip — each worker mints its own IDs locally. + +Global disaggregated IDs occupy the range `[1 << 40, 2**63)`; worker-local and warm-up request IDs occupy the disjoint range `[0, 1 << 40)`, so the two never collide. If a client supplies its own positive `disagg_request_id`, that value is used verbatim and must be globally unique; when unset, the server mints a snowflake ID as above. ## Usage @@ -233,6 +241,95 @@ Example (two-node deployment): - **Client entrypoint** - Send requests or use a load balancer forwarding to `node-a:8000` and `node-b:8000` +### Coordinator and Worker Fleet + +A single disaggregated server process is itself a single-threaded orchestrator and can become a throughput bottleneck (it terminates every client connection, runs routing, and proxies the ctx→gen hop). To scale the orchestrator on one node without standing up multiple independent instances, `trtllm-serve disaggregated` can run a **fleet** of stateless disaggregated-server worker processes behind a shared **coordinator**. + +The two roles split as follows: + +- **Coordinator** — a single process that owns all cluster state: the ctx/gen routers, worker readiness, and (for the KV-cache-aware router) the single ZMQ event-ingest endpoint. It exposes an internal coordination API (`/select`, `/finish`, `/cluster_info`, `/health`). +- **Fleet workers** — `num_workers` stateless disaggregated servers that share the public port via `SO_REUSEPORT` (each worker is its own process binding the same port, so the kernel load-balances incoming connections across them by 4-tuple hash). Each holds a lightweight delegating client: it computes the routing key locally (e.g. block hashes) and delegates the placement decision to the coordinator over HTTP. Workers own no routing state, so routing stays globally consistent no matter which worker terminates a connection. Each worker also gets a distinct `process_id` for the [global request ID](#unique-global-request-id). + +This is controlled by two fields in the disaggregated config: + +- `num_workers` (int, default `1`) — number of disaggregated-server worker processes to run on the public port. +- `disagg_coordinator_url` (str, optional) — URL of an already-running coordinator. When set, this process starts **no** coordinator and its fleet delegates to that external one. + +The three resulting topologies: + +| `num_workers` | `disagg_coordinator_url` | Behavior | +|---------------|--------------------------|----------| +| `1` | unset | Single self-contained server with an in-process coordinator (the default; unchanged from earlier examples). | +| `> 1` | unset | An **implicit** coordinator starts in this process (on `port - 1`) and a fleet of `num_workers` delegating servers runs on the public port. | +| any | set | **No** coordinator starts here; a fleet of `num_workers` delegating servers points at the external `disagg_coordinator_url`. | + +```{note} +The fleet is most useful with a *stateful* router (`kv_cache_aware`, `conversation`) where placement must be globally consistent — that decision is delegated to the coordinator. With a *stateless* router (`round_robin`, `load_balancing`) each worker simply places locally and no coordinator round-trip occurs. +``` + +#### Example: implicit coordinator + 4-worker fleet + +Extend the `disagg_config.yaml` from the [trtllm-serve](#trtllm-serve) example with `num_workers` and a router type: + +```yaml +hostname: localhost +port: 8000 +backend: pytorch +# Run 4 stateless disaggregated-server workers on port 8000, with an implicit +# coordinator started in-process on port 7999 (port - 1). +num_workers: 4 +context_servers: + num_instances: 2 + urls: + - "localhost:8001" + - "localhost:8002" + router: + type: kv_cache_aware +generation_servers: + num_instances: 1 + urls: + - "localhost:8003" + router: + type: kv_cache_aware +``` + +Launch it exactly as before — the coordinator and fleet are started for you: + +``` +trtllm-serve disaggregated -c disagg_config.yaml +``` + +Clients still send requests to the public endpoint (`localhost:8000`); the fleet transparently delegates routing to the coordinator. + +#### Example: external coordinator + +To point a fleet at a coordinator already running elsewhere (for example, one shared across nodes), set `disagg_coordinator_url` and omit the coordinator from this process: + +```yaml +hostname: localhost +port: 8000 +backend: pytorch +num_workers: 4 +disagg_coordinator_url: "http://coordinator-host:7999" +context_servers: + num_instances: 2 + urls: + - "localhost:8001" + - "localhost:8002" + router: + type: kv_cache_aware +generation_servers: + num_instances: 1 + urls: + - "localhost:8003" + router: + type: kv_cache_aware +``` + +```{note} +A fleet worker fails fast if its coordinator is unreachable: on startup it probes the coordinator's `/cluster_info` with bounded retry (up to `--server_start_timeout` seconds) and exits with an error rather than coming up and returning `Cluster is not ready` for every request. +``` + ## Environment Variables TRT-LLM uses some environment variables to control the behavior of disaggregated service. diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 2f39353c521c..f96b0448e44a 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1,4 +1,5 @@ import asyncio +import atexit import gc import importlib import inspect @@ -9,6 +10,7 @@ import socket import subprocess # nosec B404 import sys +import time import uuid from pathlib import Path from typing import Any, Dict, Optional, Sequence, Set @@ -1643,6 +1645,35 @@ def disaggregated( logger.info(f"Reserving disaggregated server address " f"{disagg_cfg.hostname}:{disagg_cfg.port} (pid={os.getpid()})") + metadata_server_cfg = parse_metadata_server_config_file( + metadata_server_config_file) + + # Topology from explicit config (num_workers + disagg_coordinator_url): + # (a) url set -> delegate to that external coordinator; (b) url absent, + # num_workers>1 -> implicit in-process coordinator + delegating fleet; + # (c) url absent, num_workers==1 -> single self-contained server. + num_workers = disagg_cfg.num_workers + coordinator_url = disagg_cfg.disagg_coordinator_url + + if coordinator_url: + # (a) External coordinator: fork a fleet of delegating servers (or a + # single one) pointed at it; never start a coordinator in this process. + _serve_disagg_fleet(disagg_cfg, config_file, + metadata_server_config_file, request_timeout, + server_start_timeout, num_workers, coordinator_url) + return + + if num_workers > 1: + # (b) Implicit coordinator in this process (on port-1) + a delegating + # uvicorn fleet (workers=N) on the public port. See below. + _serve_coordinator_and_fleet(disagg_cfg, config_file, + metadata_server_config_file, + metadata_server_cfg, request_timeout, + server_start_timeout, num_workers) + return + + # (c) num_workers==1, no external coordinator: a single disagg server with an + # in-process (local) coordinator. Pre-bind the socket (validates port), serve. with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: try: s.bind((disagg_cfg.hostname, disagg_cfg.port)) @@ -1656,9 +1687,6 @@ def disaggregated( f"Failed to bind socket to {disagg_cfg.hostname}:{disagg_cfg.port}: {e}. " f"Port holder(s): {holder}") - metadata_server_cfg = parse_metadata_server_config_file( - metadata_server_config_file) - server = OpenAIDisaggServer( config=disagg_cfg, req_timeout_secs=request_timeout, @@ -1682,6 +1710,256 @@ def disaggregated( uvloop.run(server(disagg_cfg.hostname, disagg_cfg.port, sockets=[s])) +def _launch_disagg_fleet(disagg_cfg, config_file, metadata_server_config_file, + request_timeout, server_start_timeout, num_workers, + coordinator_url): + """Launch ``num_workers`` delegating disagg servers (each its own ``Popen`` + running ``_run_fleet_worker``) on the shared public port, pointed at + ``coordinator_url``. Separate Popens (not ``uvicorn --workers N``) so each gets + an explicit per-worker TLLM_DISAGG_WORKER_PROCESS_ID and binds its own + SO_REUSEPORT socket for kernel-balanced connections. MPI/PMIX/SLURM env is + stripped (workers are plain HTTP processes). Returns the list of Popen handles. + """ + from tensorrt_llm.llmapi.disagg_utils import disagg_process_id_space + public_host, public_port = disagg_cfg.hostname, disagg_cfg.port + base_env = { + k: v + for k, v in os.environ.items() + if not k.startswith(("SLURM_", "PMIX_", "PMI_", "OMPI_", "UCX_", + "I_MPI_", "HYDRA_", "MPI_")) + } + # num_workers is explicit config now; ensure no stale WEB_CONCURRENCY leaks in + # and re-forks each plain-HTTP worker into a nested fleet. + base_env.pop("WEB_CONCURRENCY", None) + base_env[DisaggWorkerEnvs.TLLM_DISAGG_COORDINATOR_URL] = coordinator_url + base_env[DisaggWorkerEnvs.TLLM_DISAGG_CONFIG_FILE] = os.path.abspath( + config_file) + if metadata_server_config_file: + base_env[DisaggWorkerEnvs.TLLM_DISAGG_METADATA_CONFIG_FILE] = \ + os.path.abspath(metadata_server_config_file) + base_env[DisaggWorkerEnvs.TLLM_DISAGG_REQUEST_TIMEOUT] = str( + request_timeout) + base_env[DisaggWorkerEnvs.TLLM_DISAGG_SERVER_START_TIMEOUT] = str( + server_start_timeout) + # Propagate the parent's log level so fleet workers' INFO logs (e.g. the + # per-request [ttft_split] / [coord_api] breakdowns) are not dropped. + base_env[DisaggWorkerEnvs.TLLM_DISAGG_LOG_LEVEL] = logger.level + + cmd = [ + sys.executable, "-c", + "from tensorrt_llm.commands.serve import _run_fleet_worker; " + "_run_fleet_worker()" + ] + logger.info( + f"Launching disagg fleet: {num_workers} SO_REUSEPORT workers on " + f"{public_host}:{public_port}, coordinator={coordinator_url}") + + procid_space = disagg_process_id_space() + fleet = [] + for i in range(num_workers): + worker_env = dict(base_env) + # Explicit per-worker process index (no shared counter file). Wrap into + # the snowflake process_id space so ids stay in the 6-bit field. + worker_env[DisaggWorkerEnvs.TLLM_DISAGG_WORKER_PROCESS_ID] = str( + i % procid_space) + p = subprocess.Popen(cmd, + env=worker_env, + stdout=sys.stdout, + stderr=sys.stderr, + start_new_session=True) + logger.info(f"Disagg fleet worker {i} launched (pid={p.pid})") + fleet.append(p) + + def _cleanup(): + for p in fleet: + if p.poll() is None: + p.terminate() + for p in fleet: + try: + p.wait(timeout=10) + except Exception: + p.kill() + + atexit.register(_cleanup) + return fleet + + +def _serve_disagg_fleet(disagg_cfg, config_file, metadata_server_config_file, + request_timeout, server_start_timeout, num_workers, + coordinator_url): + """External coordinator: fork the delegating fleet and block on it. + + No coordinator is started in this process -- the fleet delegates to the + already-running coordinator at ``coordinator_url``. + """ + fleet = _launch_disagg_fleet(disagg_cfg, config_file, + metadata_server_config_file, request_timeout, + server_start_timeout, num_workers, + coordinator_url) + # Block until any worker exits; a nonzero exit from any worker is a failure. + while True: + for i, p in enumerate(fleet): + rc = p.poll() + if rc is not None: + if rc != 0: + raise RuntimeError( + f"Disagg fleet worker {i} (pid={p.pid}) exited with " + f"code {rc}") + # A clean exit of one worker ends the fleet. + return + time.sleep(1) + + +def _serve_coordinator_and_fleet(disagg_cfg, config_file, + metadata_server_config_file, + metadata_server_cfg, request_timeout, + server_start_timeout, num_workers): + """workers>1, no external URL: coordinator server here + a delegating fleet. + + This process runs the coordinator server (owns the ctx/gen routers, cluster + state, and centralized ZMQ ingest) on ``port-1``; a uvicorn fleet on the + public port delegates to it. + """ + from tensorrt_llm.serve.coordinator_server import CoordinatorServer + from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService + + public_host, public_port = disagg_cfg.hostname, disagg_cfg.port + coord_port = int( + os.environ.get(DisaggWorkerEnvs.TLLM_DISAGG_COORDINATOR_PORT, + public_port - 1)) + # The fleet is co-located with the implicit coordinator, so route its hot + # /select,/finish over a Unix domain socket (avoids the TCP loopback stack + # that dominated per-request latency). Opt-out via the UDS env = "0". + use_uds = os.environ.get(DisaggWorkerEnvs.TLLM_DISAGG_COORDINATOR_UDS, + "1") == "1" + coord_uds = ( + f"/tmp/trtllm_disagg_coord_{public_port}.sock" # nosec B108 + if use_uds else None) + # Fleet points at the UDS when enabled; the TCP port stays up for health. + coord_url = f"unix:{coord_uds}" if coord_uds else \ + f"http://{public_host}:{coord_port}" + + # 1. Launch the delegating fleet pointed at the implicit coordinator we start + # below (port-1 for TCP; UDS for the hot path). Workers hold + # CoordinatorClients (no core), so they can't race the ZMQ ingest bind. + _launch_disagg_fleet(disagg_cfg, config_file, metadata_server_config_file, + request_timeout, server_start_timeout, num_workers, + coord_url) + + # 2. Build + serve the coordinator in this process. It OWNS routing state and + # builds the owner routers itself (single shared namespace-aware core + ONE + # ZMQ ingest server, started once here). + def _client_factory(router, role, max_retries=1): + from tensorrt_llm.serve.openai_client import OpenAIHttpClient + return OpenAIHttpClient(router, role, request_timeout, max_retries) + + coordinator = DisaggCoordinatorService( + disagg_cfg, + _client_factory, + metadata_config=metadata_server_cfg, + server_start_timeout_secs=server_start_timeout) + logger.info(f"Coordinator serving on {public_host}:{coord_port} " + f"(uds={coord_uds}) (fleet on public port {public_port})") + asyncio.run( + CoordinatorServer(coordinator)(public_host, coord_port, uds=coord_uds)) + + +def _init_fleet_worker_process(): + """Per-process setup shared by every fleet worker (one OS process each). + + Restores logging/GC state for the fresh ``python -c`` interpreter and tags + every log line with this worker's PID so the shared-stdout fleet output stays + attributable.""" + # A worker is a plain HTTP process, never an MPI rank; drop WEB_CONCURRENCY so + # it is never itself re-forked into multiple uvicorn workers. + os.environ.pop("WEB_CONCURRENCY", None) + if os.getenv("TRTLLM_DISAGG_SERVER_DISABLE_GC", "1") == "1": + gc.disable() + + # This is a fresh Python process, so the TRT-LLM logger defaults to WARNING + # and would drop the workers' INFO logs (per-request [ttft_split] / + # [coord_api]). Restore the parent's level. + _worker_log_level = os.environ.get(DisaggWorkerEnvs.TLLM_DISAGG_LOG_LEVEL) + if _worker_log_level: + logger.set_level(_worker_log_level) + + # All N fleet workers share one stdout; tag every trtllm log line from this + # worker with its PID so interleaved fleet output is attributable. + import logging as _logging + for _h in _logging.getLogger("TRT-LLM").handlers: + _h.setFormatter( + _logging.Formatter( + fmt= + f"[%(asctime)s] [fleet-worker pid={os.getpid()}] %(message)s", + datefmt="%m/%d/%Y-%H:%M:%S")) + + +def _build_disagg_server_from_env() -> "OpenAIDisaggServer": + """Build one delegating disagg server from the env the launcher exported. + + Fully stateless -- reads config + coordinator URL from ``DisaggWorkerEnvs``; + the server holds a remote ``CoordinatorClient`` so routing/readiness are + delegated to the coordinator. The worker's process index (for the snowflake + disagg-id) is read from ``TLLM_DISAGG_WORKER_PROCESS_ID``, which the launcher + set explicitly per worker (see ``worker_local_process_id``).""" + config_file = os.environ[DisaggWorkerEnvs.TLLM_DISAGG_CONFIG_FILE] + coordinator_url = os.environ[DisaggWorkerEnvs.TLLM_DISAGG_COORDINATOR_URL] + metadata_config_file = os.environ.get( + DisaggWorkerEnvs.TLLM_DISAGG_METADATA_CONFIG_FILE) + request_timeout = int( + os.environ.get(DisaggWorkerEnvs.TLLM_DISAGG_REQUEST_TIMEOUT, "180")) + server_start_timeout = int( + os.environ.get(DisaggWorkerEnvs.TLLM_DISAGG_SERVER_START_TIMEOUT, + "180")) + + disagg_cfg = parse_disagg_config_file(config_file) + metadata_server_cfg = parse_metadata_server_config_file( + metadata_config_file) + + server = OpenAIDisaggServer(config=disagg_cfg, + req_timeout_secs=request_timeout, + server_start_timeout_secs=server_start_timeout, + metadata_server_cfg=metadata_server_cfg, + coordinator_url=coordinator_url) + logger.info(f"Disagg server built, coordinator={coordinator_url}") + return server + + +def create_disagg_server_app(): + """uvicorn import-string factory: build one disagg server's FastAPI app. + + Retained for the ``uvicorn --factory`` entry point; the SO_REUSEPORT fleet + launcher uses ``_run_fleet_worker`` instead.""" + _init_fleet_worker_process() + return _build_disagg_server_from_env().app + + +def _run_fleet_worker(): + """Entry point for a single fleet worker (one OS process per worker), launched + as its own ``Popen`` so it takes an explicit TLLM_DISAGG_WORKER_PROCESS_ID and + binds its own SO_REUSEPORT socket to the shared public port (the kernel then + 4-tuple-hashes connections across workers instead of the unfair shared-socket + accept() of ``uvicorn --workers N``).""" + _init_fleet_worker_process() + server = _build_disagg_server_from_env() + host, port = server._config.hostname, server._config.port + + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + # SO_REUSEPORT: every worker binds the same (host, port); the kernel spreads + # connections across the workers' accept queues by 4-tuple hash. + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + try: + s.bind((host, port)) + except OSError as e: + raise RuntimeError( + f"Fleet worker failed to SO_REUSEPORT-bind {host}:{port}: {e}") + pidx = os.environ.get(DisaggWorkerEnvs.TLLM_DISAGG_WORKER_PROCESS_ID, "0") + logger.info(f"Fleet worker process_id={pidx} bound {host}:{port} " + f"(SO_REUSEPORT)") + asyncio.run(server(host, port, sockets=[s])) + + def set_cuda_device(): if (os.getenv("OMPI_COMM_WORLD_RANK")): env_global_rank = int(os.environ["OMPI_COMM_WORLD_RANK"]) @@ -1779,6 +2057,28 @@ class DisaggLauncherEnvs(StrEnum): TLLM_DISAGG_ROLE = "TRTLLM_DISAGG_ROLE" +class DisaggWorkerEnvs(StrEnum): + # Passed from the `disaggregated` coordinator to the fleet of worker processes + # (one Popen per worker) via env, then read by _run_fleet_worker in each worker. + TLLM_DISAGG_COORDINATOR_URL = "TRTLLM_DISAGG_COORDINATOR_URL" + TLLM_DISAGG_CONFIG_FILE = "TRTLLM_DISAGG_CONFIG_FILE" + TLLM_DISAGG_METADATA_CONFIG_FILE = "TRTLLM_DISAGG_METADATA_CONFIG_FILE" + TLLM_DISAGG_REQUEST_TIMEOUT = "TRTLLM_DISAGG_REQUEST_TIMEOUT" + TLLM_DISAGG_SERVER_START_TIMEOUT = "TRTLLM_DISAGG_SERVER_START_TIMEOUT" + # Parent's logger level; the forked uvicorn fleet is a fresh process that + # otherwise defaults to WARNING and drops the workers' INFO logs. + TLLM_DISAGG_LOG_LEVEL = "TRTLLM_DISAGG_LOG_LEVEL" + # Coordinator transport (read in _serve_coordinator_and_fleet): PORT overrides + # the coordinator TCP port (default public_port-1); UDS="1" (default) routes the + # co-located fleet's hot /select,/finish over a Unix domain socket, "0" for TCP. + TLLM_DISAGG_COORDINATOR_PORT = "TRTLLM_DISAGG_COORDINATOR_PORT" + TLLM_DISAGG_COORDINATOR_UDS = "TRTLLM_DISAGG_COORDINATOR_UDS" + # Per-fleet-worker process index (0..N-1) → the snowflake disagg-id process_id + # field so co-located workers never collide; the launcher sets one distinct + # value per worker Popen (see worker_local_process_id). Defaults to 0 if unset. + TLLM_DISAGG_WORKER_PROCESS_ID = "TRTLLM_DISAGG_WORKER_PROCESS_ID" + + def _launch_disaggregated_server(disagg_config_file: str, llm_args: dict): # Launching the server instance_idx = os.environ.get(DisaggLauncherEnvs.TLLM_DISAGG_INSTANCE_IDX) diff --git a/tensorrt_llm/llmapi/disagg_utils.py b/tensorrt_llm/llmapi/disagg_utils.py index d1851364d7c5..55f953def917 100644 --- a/tensorrt_llm/llmapi/disagg_utils.py +++ b/tensorrt_llm/llmapi/disagg_utils.py @@ -1,4 +1,5 @@ import logging +import os import threading import time import uuid @@ -106,6 +107,14 @@ class DisaggServerConfig(): # the orchestrator relays a string instead of materializing the token-id list # on its event loop. Text-only, non-harmony deployments (see _get_ctx_request). gen_tokids_ctxbytes: bool = False + # Number of uvicorn disagg-server worker processes to fork on the public port. + # >1 means a fleet of delegating servers behind one coordinator. Replaces the + # WEB_CONCURRENCY env var (explicit config over implicit env). + num_workers: int = 1 + # URL of an already-running coordinator (e.g. "http://host:8332"). When set the + # fleet delegates to it; when absent, num_workers>1 starts an implicit in-process + # coordinator and num_workers==1 runs a single self-contained server. + disagg_coordinator_url: Optional[str] = None @dataclass @@ -192,6 +201,8 @@ def extract_disagg_cfg(hostname: str = 'localhost', allow_request_chat_template: bool = False, gen_strip_message_history: bool = False, gen_tokids_ctxbytes: bool = False, + num_workers: int = 1, + disagg_coordinator_url: Optional[str] = None, **kwargs: Any) -> DisaggServerConfig: context_servers = context_servers or {} generation_servers = generation_servers or {} @@ -243,6 +254,8 @@ def extract_disagg_cfg(hostname: str = 'localhost', allow_request_chat_template, "allow_request_chat_template") config.gen_strip_message_history = gen_strip_message_history config.gen_tokids_ctxbytes = gen_tokids_ctxbytes + config.num_workers = num_workers + config.disagg_coordinator_url = disagg_coordinator_url return config @@ -418,42 +431,59 @@ def parse_metadata_server_config_file( return MetadataServerConfig(**config) -MIN_GLOBAL_ID = 1 << 42 +# Snowflake global disagg request id, 64-bit / positive int64 (MSB reserved 0): +# [ 0 (1) | timestamp_ms (39) | node_id (8) | process_id (6) | counter (10) ] +# The (node_id, process_id) pair identifies a fleet worker process, so co-located +# workers never emit the same id in the same millisecond. See docs/source/ +# advanced/disaggregated-service.md for the full disagg-request-id design. +DISAGG_TIMESTAMP_BITS = 39 +DISAGG_NODE_ID_BITS = 8 +DISAGG_PROCESS_ID_BITS = 6 +DISAGG_COUNTER_BITS = 10 + +# Local ids [0, MIN_GLOBAL_ID) and global disagg ids [MIN_GLOBAL_ID, 2^63) are +# disjoint by construction so they never collide. Power of two (masked in +# get_local_request_id). +MIN_GLOBAL_ID = 1 << 40 # Consider GIL being removed in the future, use a lock to protect the counter _global_disagg_request_id_lock = threading.Lock() _global_disagg_request_id_counter = 0 -def get_global_disagg_request_id(machine_id: int) -> int: - """ - a snowflake global disagg request id that doesn't guarantee monotonicity - 0: positive integer - 1-41 41 bits: timestamp_ms - 42-51 10 bits: machine_id - 52-63 12 bits: counter +def get_global_disagg_request_id(node_id: int, process_id: int = 0) -> int: + """A snowflake global disagg request id (does not guarantee monotonicity). + + Layout: 0(1) | timestamp_ms(39) | node_id(8) | process_id(6) | counter(10). + node_id identifies the node, process_id the fleet worker process on it -- the + pair makes the id unique across co-located workers without any coordination. """ global _global_disagg_request_id_lock global _global_disagg_request_id_counter - COUNTER_BITS = 12 - MACHINE_ID_BITS = 10 - COUNTER_MASK = (1 << COUNTER_BITS) - 1 + NODE_ID_SPACE = 1 << DISAGG_NODE_ID_BITS + PROCESS_ID_SPACE = 1 << DISAGG_PROCESS_ID_BITS + COUNTER_MASK = (1 << DISAGG_COUNTER_BITS) - 1 + TIMESTAMP_MASK = (1 << DISAGG_TIMESTAMP_BITS) - 1 MAX_INT64 = (1 << 63) - 1 - if machine_id not in range(0, (1 << MACHINE_ID_BITS) - 1): - raise ValueError( - f"machine_id must be in range [0, {(1 << MACHINE_ID_BITS) - 1})") + if node_id not in range(0, NODE_ID_SPACE): + raise ValueError(f"node_id must be in range [0, {NODE_ID_SPACE})") + if process_id not in range(0, PROCESS_ID_SPACE): + raise ValueError(f"process_id must be in range [0, {PROCESS_ID_SPACE})") - timestamp_ms = int(time.monotonic() * 1000) + timestamp_ms = int(time.monotonic() * 1000) & TIMESTAMP_MASK with _global_disagg_request_id_lock: counter = _global_disagg_request_id_counter & COUNTER_MASK _global_disagg_request_id_counter += 1 - # Rotate in [MIN_GLOBAL_ID, MAX_INT64) - # [0, MIN_GLOBAL_ID) is reserved for local ids - global_id = (timestamp_ms << (MACHINE_ID_BITS + COUNTER_BITS)) | ( - machine_id << COUNTER_BITS) | counter + global_id = ( + (timestamp_ms << + (DISAGG_NODE_ID_BITS + DISAGG_PROCESS_ID_BITS + DISAGG_COUNTER_BITS)) + | (node_id << (DISAGG_PROCESS_ID_BITS + DISAGG_COUNTER_BITS)) + | (process_id << DISAGG_COUNTER_BITS) + | counter) + # Rotate into [MIN_GLOBAL_ID, MAX_INT64); [0, MIN_GLOBAL_ID) is local-id space. global_id_int64 = global_id % (MAX_INT64 - MIN_GLOBAL_ID) + MIN_GLOBAL_ID return global_id_int64 @@ -461,3 +491,16 @@ def get_global_disagg_request_id(machine_id: int) -> int: def get_local_request_id(last_id: int) -> int: """ increment the last_id by 1 and mod by MIN_GLOBAL_ID """ return (last_id + 1) & (MIN_GLOBAL_ID - 1) + + +def disagg_process_id_space() -> int: + """Number of distinct process_id slots in the snowflake id (2^bits).""" + return 1 << DISAGG_PROCESS_ID_BITS + + +def worker_local_process_id() -> int: + """Read this fleet worker's process index from TRTLLM_DISAGG_WORKER_PROCESS_ID + (set explicitly per worker by the fleet launcher, one distinct value 0..N-1 + per Popen). 0 if unset (e.g. single-process / non-fleet path).""" + return int(os.environ.get("TRTLLM_DISAGG_WORKER_PROCESS_ID", "0")) \ + % disagg_process_id_space() diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index cc6a66e53a08..98fb4f11b5be 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3523,6 +3523,22 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description="KV cache manager v2 block reuse policy. " "With SWA scratch reuse and 'all_reusable', only non-scratch " "blocks are saved for reuse.") + # Per-rank centralized KV cache routing + centralized_router_report_address: Optional[str] = Field( + default=None, + status="prototype", + description= + "ZMQ address to push KV cache events to for centralized routing " + "(e.g. 'tcp://orchestrator:5557'). When set with per_rank_routing=True, " + "each DP rank reports its own events independently.") + + per_rank_routing: bool = Field( + default=False, + status="prototype", + description= + "Enable per-rank KV cache event reporting. Each DP rank reports its " + "own cache blocks independently (no allgather). Requires " + "use_kv_cache_manager_v2=True and event_buffer_max_size > 0.") def _to_pybind(self): config = _KvCacheConfig( diff --git a/tensorrt_llm/serve/coordinator_server.py b/tensorrt_llm/serve/coordinator_server.py new file mode 100644 index 000000000000..d7b9a1dcb1b5 --- /dev/null +++ b/tensorrt_llm/serve/coordinator_server.py @@ -0,0 +1,163 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Coordinator HTTP server for disaggregated serving. + +One coordinator process owns all cluster state (routers, readiness, worker +events, and -- for the centralized router -- the single ZMQ event-ingest bind) +and answers the internal coordination API that the forked worker processes call: + + POST /select {"role", "routing_key", "req_id", "exclude_server"} + -> {"server": "host:port", "info": {...}, "req_id": } + POST /finish {"role", "req_id", "success"} -> {} + GET /cluster_info -> {...} + GET /health -> 200 when ready + GET /version + +The routing key is produced worker-side by ``Router.routing_key`` and consumed +here by ``Router.get_next_server_by_key`` (see ``serve/router.py``), so this +endpoint is generic across the stateful router types that use it (centralized -> +block hashes, conversation -> conversation_id). Single-process by design; it owns +the ZMQ ingest bind for centralized mode. +""" + +import asyncio +import time +from contextlib import asynccontextmanager +from typing import Optional + +import orjson +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import ORJSONResponse, Response + +from tensorrt_llm.logger import logger +from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService +from tensorrt_llm.version import __version__ as VERSION + +# Timing headers to decompose the coordinator IPC round-trip (co-located, so +# wall-clock time.time() is comparable): client stamps send time, coordinator +# stamps recv/send, from which the client derives req_wire/handler/resp_wire. +HDR_CLIENT_SEND = "x-client-send-time" +HDR_COORD_RECV = "x-coord-recv-time" +HDR_COORD_SEND = "x-coord-send-time" + +# The coordinator is a single event loop serving /select,/finish for the whole +# fleet; each /select body carries a routing_key (a flat list of block hashes, +# hundreds of int64 for long prompts). stdlib json parse+serialize on that hot +# path keeps the loop busy and lengthens the queue (client-observed /select was +# ~hundreds of ms while the handler itself is ~0.1ms). orjson is 2-5x faster at +# both ends, freeing the loop sooner. JSONResponse -> ORJSONResponse and +# request bodies parsed with orjson.loads (bytes, no text decode). + +TIMEOUT_KEEP_ALIVE = 10 # seconds + + +class CoordinatorServer: + """Serve a :class:`DisaggCoordinatorService`'s coordination API over HTTP.""" + + def __init__(self, coordinator: DisaggCoordinatorService) -> None: + self._coordinator = coordinator + + @asynccontextmanager + async def lifespan(app: FastAPI): + await self._coordinator.start() + yield + await self._coordinator.stop() + + self.app = FastAPI(lifespan=lifespan) + self.app.add_api_route("/select", self.select, methods=["POST"]) + self.app.add_api_route("/finish", self.finish, methods=["POST"]) + self.app.add_api_route("/cluster_info", self.cluster_info, methods=["GET"]) + self.app.add_api_route("/health", self.health, methods=["GET"]) + self.app.add_api_route("/version", self.version, methods=["GET"]) + + def _timing_headers(self, raw_req: Request, recv_t: float) -> dict: + """Echo the client send time + stamp coord recv/send times so the client + can decompose the IPC round-trip. recv_t is captured at handler entry; + send time is 'now' (just before responding). + """ + h = {HDR_COORD_RECV: repr(recv_t), HDR_COORD_SEND: repr(time.time())} + cs = raw_req.headers.get(HDR_CLIENT_SEND) + if cs is not None: + h[HDR_CLIENT_SEND] = cs + return h + + async def select(self, raw_req: Request) -> Response: + _recv = time.time() + try: + body = orjson.loads(await raw_req.body()) + except Exception as e: + return ORJSONResponse(status_code=400, content={"error": f"invalid JSON body: {e}"}) + if not isinstance(body, dict) or "role" not in body: + return ORJSONResponse( + status_code=400, content={"error": "body must include 'role' and 'routing_key'"} + ) + try: + server, info, req_id = await self._coordinator.select( + body["role"], + body.get("routing_key"), + body.get("req_id"), + body.get("exclude_server"), + ) + except ValueError as e: + return ORJSONResponse(status_code=503, content={"error": str(e)}) + except Exception as e: # noqa: BLE001 + logger.error(f"CoordinatorServer.select failed: {e}") + return ORJSONResponse(status_code=500, content={"error": str(e)}) + return ORJSONResponse( + content={"server": server, "info": info, "req_id": req_id}, + headers=self._timing_headers(raw_req, _recv), + ) + + async def finish(self, raw_req: Request) -> Response: + _recv = time.time() + try: + body = orjson.loads(await raw_req.body()) + except Exception as e: + return ORJSONResponse(status_code=400, content={"error": f"invalid JSON body: {e}"}) + await self._coordinator.finish( + body.get("role", "gen"), body.get("req_id"), body.get("success", True) + ) + return ORJSONResponse(content={}, headers=self._timing_headers(raw_req, _recv)) + + async def cluster_info(self) -> Response: + return ORJSONResponse(content=await self._coordinator.cluster_info()) + + async def health(self) -> Response: + return Response(status_code=200 if await self._coordinator.is_ready() else 503) + + async def version(self) -> Response: + return ORJSONResponse(content={"version": VERSION}) + + async def __call__(self, host: str, port: int, uds: Optional[str] = None) -> None: + # Single-process (owns routing state + the centralized ZMQ ingest bind); + # workers=1 forced so a leaked WEB_CONCURRENCY can't fork it. When ``uds`` + # is set the co-located fleet uses it for the hot /select,/finish path + # (avoids the TCP loopback overhead that dominated per-request latency). + kwargs = dict(workers=1, log_level="info", timeout_keep_alive=TIMEOUT_KEEP_ALIVE) + if uds: + # uvicorn.Config binds uds XOR host:port, so run two Servers: UDS for + # the fleet (hot path) and TCP for health/external clients. + import asyncio as _asyncio + + uds_cfg = uvicorn.Config(self.app, uds=uds, **kwargs) + tcp_cfg = uvicorn.Config(self.app, host=host, port=port, **kwargs) + await _asyncio.gather(uvicorn.Server(uds_cfg).serve(), uvicorn.Server(tcp_cfg).serve()) + else: + config = uvicorn.Config(self.app, host=host, port=port, **kwargs) + await uvicorn.Server(config).serve() + + +def serve_coordinator(host: str, port: int, coordinator: DisaggCoordinatorService) -> None: + asyncio.run(CoordinatorServer(coordinator)(host, port)) diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py new file mode 100644 index 000000000000..d45506d0331a --- /dev/null +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -0,0 +1,517 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Coordination for disaggregated serving. + +A :class:`DisaggCoordinator` owns everything that is *not* a completion: the +ctx/gen routers, readiness, cluster info, and worker/auto-scaling events. The +completions service holds one and reads ``ctx_router`` / ``gen_router`` off it, +then drives ``router.get_next_server`` / ``router.finish_request`` uniformly -- +so serving a completion is decoupled from managing the cluster and is identical +whether this process owns the routers or delegates to a remote coordinator. + +Two implementations for the coordinator/worker deployment: + +* :class:`DisaggCoordinatorService` -- runs in the coordinator (and in the + collapsed single-process path). Owns the real ctx/gen ``Router`` objects, + server preparation/monitoring, the auto-scaling ``DisaggClusterManager`` + + worker events, and readiness. Its :meth:`select` / :meth:`finish` are the + coordinator's ``/select`` / ``/finish`` handlers. +* :class:`CoordinatorClient` -- runs in each forked worker. Stateful routers + (conversation, centralized) are wrapped in a :class:`CoordinatorDelegatingRouter` + that posts the routing key to ``/select`` (finish -> ``/finish``); stateless + routers (round_robin, load_balancing) place locally in the worker. Readiness / + cluster_info proxy the coordinator over HTTP. +""" + +import asyncio +import time +from abc import ABC, abstractmethod +from typing import Any, Dict, Optional, Tuple + +import aiohttp + +from tensorrt_llm.llmapi.disagg_utils import ( + DisaggServerConfig, + MetadataServerConfig, + ServerRole, + get_ctx_gen_server_addrs, + get_global_disagg_request_id, + worker_local_process_id, # noqa: F401 +) +from tensorrt_llm.logger import logger +from tensorrt_llm.serve.cluster_storage import ( + ClusterStorage, + WatchEventType, + create_cluster_storage, +) +from tensorrt_llm.serve.disagg_auto_scaling import DisaggClusterManager, WorkerInfo +from tensorrt_llm.serve.metadata_server import create_metadata_server +from tensorrt_llm.serve.openai_client import OpenAIClient +from tensorrt_llm.serve.router import CoordinatorDelegatingRouter, Router, build_disagg_routers + +__all__ = [ + "DisaggCoordinator", + "DisaggCoordinatorService", + "CoordinatorClient", +] + + +class DisaggCoordinator(ABC): + """Abstract coordinator: ctx/gen routers + readiness + cluster info + lifecycle. + + Placement and finish are driven through ``ctx_router`` / ``gen_router`` + (``Router.get_next_server`` / ``Router.finish_request``), so this surface only + exposes the routers plus readiness/info/lifecycle. + """ + + @property + @abstractmethod + def ctx_router(self) -> Router: ... + + @property + @abstractmethod + def gen_router(self) -> Router: ... + + @abstractmethod + async def is_ready(self) -> bool: ... + + @abstractmethod + async def cluster_info(self) -> Dict[str, Any]: ... + + async def start(self) -> None: ... + + async def stop(self) -> None: ... + + @abstractmethod + async def get_disagg_request_id(self) -> int: ... + + +class DisaggCoordinatorService(DisaggCoordinator): + """In-process coordinator owning the ctx/gen routers and all cluster state. + + Used in the coordinator process and in the single-process (workers==1) path. + """ + + def __init__( + self, + config: DisaggServerConfig, + client_factory, + metadata_config: Optional[MetadataServerConfig] = None, + server_preparation_func=None, + server_start_timeout_secs: int = 180, + health_check_interval_secs: int = 3, + ): + self._config = config + self._client_factory = client_factory + self._metadata_config = metadata_config + # The coordinator owns routing state, so it builds the owner routers here + # (is_delegating_client=False): one shared namespace-aware core + a single + # ZMQ ingest server. Sole place owner routers are created (fleet workers + # hold only a CoordinatorClient's delegating surfaces). + self._metadata_server = create_metadata_server(metadata_config) + ctx_servers, gen_servers = get_ctx_gen_server_addrs(config.server_configs) + self._ctx_router, self._gen_router = build_disagg_routers( + config.ctx_router_config, + config.gen_router_config, + ctx_servers, + gen_servers, + metadata_config, + self._metadata_server, + server_preparation_func, + disagg_node_id=config.node_id, + is_delegating_client=False, + ) + # The coordinator owns the disagg cluster storage (auto-scaling backend): + # it drives the DisaggClusterManager below and, when the storage is an + # in-process HTTP server, its routes are mounted on the coordinator app. + self._cluster_storage: Optional[ClusterStorage] = ( + create_cluster_storage( + config.disagg_cluster_config.cluster_uri, config.disagg_cluster_config.cluster_name + ) + if config.disagg_cluster_config + else None + ) + self._server_start_timeout_secs = server_start_timeout_secs + self._health_check_interval_secs = health_check_interval_secs + + self._ctx_client: Optional[OpenAIClient] = None + self._gen_client: Optional[OpenAIClient] = None + self._disagg_cluster_manager: Optional[DisaggClusterManager] = None + + @property + def ctx_router(self) -> Router: + return self._ctx_router + + @property + def gen_router(self) -> Router: + return self._gen_router + + @property + def cluster_storage(self) -> Optional[ClusterStorage]: + return self._cluster_storage + + def set_clients(self, ctx_client: OpenAIClient, gen_client: OpenAIClient) -> None: + self._ctx_client = ctx_client + self._gen_client = gen_client + + # -- coordinator-path placement (workers call these via the HTTP server) -- + + def _api_lat(self, name: str): + """Lazily create a per-API PeriodicLatencyLogger. Measures the coordinator + owner's in-process handler time (excludes the fleet->coordinator HTTP hop, + which is captured client-side as [coord_api] client.*). + """ + cache = self.__dict__.setdefault("_coord_api_lat", {}) + if name not in cache: + from tensorrt_llm.serve.responses_utils import PeriodicLatencyLogger + + cache[name] = PeriodicLatencyLogger(f"owner.{name}") + return cache[name] + + async def select( + self, role: str, routing_key, req_id, exclude_server: Optional[str] + ) -> Tuple[str, dict, Optional[str]]: + _t0 = time.monotonic() + router = self._router_for_role(role) + if req_id is None: + # The coordinator owns IDs absent from generation requests. + req_id = get_global_disagg_request_id(self._config.node_id) + result = await router.get_next_server_by_key( + routing_key, req_id=req_id, exclude_server=exclude_server + ) + self._api_lat(f"select[{role}]").record(time.monotonic() - _t0) + return result + + async def get_disagg_request_id(self) -> int: + return get_global_disagg_request_id(self._config.node_id) + + async def finish(self, role: str, req_id, success: bool = True) -> None: + _t0 = time.monotonic() + await self._router_for_role(role).finish_request_by_id(req_id, success) + self._api_lat(f"finish[{role}]").record(time.monotonic() - _t0) + + def _router_for_role(self, role: str) -> Router: + return self._ctx_router if str(role).lower().startswith("c") else self._gen_router + + async def start(self) -> None: + await self._ctx_router.prepare_servers() + await self._gen_router.prepare_servers() + if self._ctx_client is None or self._gen_client is None: + self._ctx_client = self._client_factory( + self._ctx_router, ServerRole.CONTEXT, self._config.max_retries + ) + self._gen_client = self._client_factory( + self._gen_router, ServerRole.GENERATION, self._config.max_retries + ) + + if self._config.disagg_cluster_config and self._cluster_storage: + logger.info("Starting disagg cluster manager") + self._disagg_cluster_manager = DisaggClusterManager( + self._config.disagg_cluster_config, self._cluster_storage + ) + await self._disagg_cluster_manager.start() + await self._disagg_cluster_manager.watch_workers(on_event=self._on_worker_event) + logger.info("Disagg cluster manager started") + else: + if self._metadata_server and self._metadata_config: + logger.info("Starting server monitoring via metadata service") + await self._ctx_router.start_server_monitoring( + self._metadata_config.refresh_interval + ) + await self._gen_router.start_server_monitoring( + self._metadata_config.refresh_interval + ) + await self._wait_for_all_servers_ready() + + async def stop(self) -> None: + if self._disagg_cluster_manager: + await self._disagg_cluster_manager.stop() + if self._metadata_server: + await self._ctx_router.stop_server_monitoring() + await self._gen_router.stop_server_monitoring() + + async def is_ready(self) -> bool: + if self._disagg_cluster_manager: + return await self._disagg_cluster_manager.is_ready_with_router( + self._ctx_router.num_prepared_servers, + self._gen_router.num_prepared_servers, + ) + return True + + async def cluster_info(self) -> Dict[str, Any]: + info = {"is_ready": await self.is_ready()} + # Expose block-hash granularity so a delegating client (which doesn't + # monitor servers) hashes with the same tokens_per_block as the workers. + tpb = getattr(self._ctx_router, "_tokens_per_block", None) + if tpb is not None: + info["tokens_per_block"] = tpb + if self._disagg_cluster_manager: + info.update(await self._disagg_cluster_manager.cluster_info()) + return info + + async def _wait_for_all_servers_ready(self) -> None: + import os + + gen_only = os.getenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY") == "1" + + async def check_servers_ready(): + elapsed_time = 0 + interval = self._health_check_interval_secs + while elapsed_time < self._server_start_timeout_secs: + if gen_only: + unready_ctx_servers = [] + else: + _, unready_ctx_servers = await self._ctx_client.check_ready() + _, unready_gen_servers = await self._gen_client.check_ready() + if len(unready_ctx_servers) == 0 and len(unready_gen_servers) == 0: + logger.info( + "All servers are ready" + if not gen_only + else "Generation servers are ready (context skipped)" + ) + return + logger.info( + f"Waiting for servers, context: {unready_ctx_servers}, " + f"generation: {unready_gen_servers}" + ) + await asyncio.sleep(interval) + elapsed_time += interval + + try: + await asyncio.wait_for(check_servers_ready(), timeout=self._server_start_timeout_secs) + except asyncio.TimeoutError: + raise TimeoutError("Timeout waiting for context and generation servers to be ready") + + async def _on_worker_event(self, worker_info: WorkerInfo, event_type: WatchEventType): + router_map = { + ServerRole.CONTEXT: self._ctx_router, + ServerRole.GENERATION: self._gen_router, + } + worker_addr = f"{worker_info.host}:{worker_info.port}" + try: + router = router_map[worker_info.role] + if event_type == WatchEventType.SET: + await router.add_server(worker_addr) + elif event_type == WatchEventType.DELETE: + await router.remove_server(worker_addr) + logger.info(f"Worker {event_type.name} event: {worker_info.worker_id}, {worker_addr}") + except KeyError: + logger.error( + f"Unknown worker role: {worker_info.role}, Worker " + f"{worker_info.worker_id} event: {event_type.name}" + ) + + +COORDINATOR_UDS_SCHEME = "unix:" + + +def coordinator_base_url(remote_url: str) -> str: + """The URL prefix to build request URLs against. + + ``remote_url`` is either a TCP URL (``http://host:port``) or a Unix domain + socket (``unix:/abs/path.sock``). For UDS the socket path routes, so the HTTP + host is a dummy ``http://localhost`` (aiohttp still needs a valid http URL). + """ + if remote_url.startswith(COORDINATOR_UDS_SCHEME): + return "http://localhost" + return remote_url.rstrip("/") + + +def make_coordinator_session(remote_url: str) -> aiohttp.ClientSession: + """Aiohttp session for the coordinator endpoint. Uses a UnixConnector for a + ``unix:/path`` URL -- UDS avoids the TCP loopback stack, since the fleet is + co-located with the implicit coordinator. Created lazily by callers so it + binds to the running event loop. + + limit=0 (unlimited pool): each fleet worker has ~concurrency/num_workers + requests in flight (e.g. 320 at c1280/4), but aiohttp's default pool caps at + 100 -- so most /select,/finish calls were BLOCKING on a free pooled + connection, which was the real ~hundreds-of-ms client.select latency (not the + ~0.1ms handler, not the transport). Uncapping lets all concurrent calls hold + a connection. The coordinator is a trusted local socket, so no cap needed. + """ + if remote_url.startswith(COORDINATOR_UDS_SCHEME): + sock_path = remote_url[len(COORDINATOR_UDS_SCHEME) :] + return aiohttp.ClientSession(connector=aiohttp.UnixConnector(path=sock_path, limit=0)) + return aiohttp.ClientSession(connector=aiohttp.TCPConnector(limit=0)) + + +class CoordinatorClient(DisaggCoordinator): + """Worker-side coordinator: delegate stateful routing to the coordinator. + + A *stateful* router (conversation, centralized -- it exposes + ``get_next_server_by_key``) is wrapped in a :class:`CoordinatorDelegatingRouter` + so the worker computes the small routing key locally and the coordinator makes + the placement (placement -> ``/select``, finish -> ``/finish``). A *stateless* + router (round_robin, load_balancing) is used as-is and places locally in the + worker -- no coordinator round-trip. Readiness / cluster_info always proxy the + coordinator over HTTP. + + Args: + remote_url: Coordinator base URL (e.g. ``http://host:PORT``). + config: The disagg config; the client builds its own delegating routers + of the configured type (same config as the coordinator so the keys it + extracts line up). + """ + + def __init__( + self, + remote_url: str, + config: DisaggServerConfig, + metadata_config: Optional[MetadataServerConfig] = None, + request_timeout_s: float = 5.0, + startup_timeout_s: float = 180.0, + ): + # remote_url may be a TCP URL or unix:/path. Resolve to a request base + # URL now; the session (with the right connector) is created lazily so it + # binds to the running event loop. + self._remote_url_raw = remote_url + self._remote_url = coordinator_base_url(remote_url) + self._request_timeout_s = request_timeout_s + self._startup_timeout_s = startup_timeout_s + self._session: Optional[aiohttp.ClientSession] = None + # Local disagg-id generation (no HTTP hop): (node_id, per-worker process_id) + # keeps the snowflake unique across co-located workers. + self._node_id = config.node_id + self._process_id = worker_local_process_id() + # A delegating client builds coreless router surfaces (compute routing_key + # locally, delegate placement to the coordinator; no ingest port / core). + # Sole place delegating routers are created. + ctx_servers, gen_servers = get_ctx_gen_server_addrs(config.server_configs) + ctx_router, gen_router = build_disagg_routers( + config.ctx_router_config, + config.gen_router_config, + ctx_servers, + gen_servers, + metadata_config, + create_metadata_server(metadata_config), + disagg_node_id=config.node_id, + is_delegating_client=True, + ) + self._ctx_router = self._maybe_delegate(ctx_router, "context") + self._gen_router = self._maybe_delegate(gen_router, "generation") + + def _maybe_delegate(self, local_router: Router, role: str) -> Router: + # Stateful routers expose get_next_server_by_key -> delegate placement to + # the coordinator; stateless ones place locally (used unchanged). Pass the + # RAW url so the delegating router picks the UDS connector when applicable. + if hasattr(local_router, "get_next_server_by_key"): + return CoordinatorDelegatingRouter( + self._remote_url_raw, local_router, role, self._request_timeout_s + ) + return local_router + + @property + def ctx_router(self) -> Router: + return self._ctx_router + + @property + def gen_router(self) -> Router: + return self._gen_router + + @property + def session(self) -> aiohttp.ClientSession: + if self._session is None: + self._session = make_coordinator_session(self._remote_url_raw) + return self._session + + async def start(self) -> None: + # Fail fast: probe /cluster_info with bounded retry so a delegating server + # exits non-zero if its coordinator never comes up (vs 500-ing every req). + info = await self._await_coordinator() + # Adopt the coordinator's tokens_per_block so the client's block hashes + # line up with the workers (a delegating client can't learn it itself). + tpb = info.get("tokens_per_block") + if tpb is not None: + for router in (self._ctx_router, self._gen_router): + local = getattr(router, "_local", router) + if getattr(local, "_tokens_per_block", None) != tpb: + local._tokens_per_block = tpb + local._tpb_auto = False + logger.info( + f"CoordinatorClient: adopted coordinator " + f"tokens_per_block={tpb} for {getattr(local, '_namespace', '?')}" + ) + + async def _await_coordinator(self) -> Dict[str, Any]: + """Poll the coordinator's /cluster_info until reachable, or raise. + + Returns the cluster_info dict once the coordinator answers HTTP 200. A 200 + means the coordinator process is up (not that all workers are ready -- + that's is_ready's job). Raises RuntimeError if it stays unreachable past + startup_timeout_s so the delegating server fails fast instead of serving + against a missing coordinator. + """ + loop = asyncio.get_event_loop() + deadline = loop.time() + self._startup_timeout_s + attempt = 0 + while True: + try: + async with self.session.get( + f"{self._remote_url}/cluster_info", timeout=self._request_timeout_s + ) as resp: + if resp.status == 200: + logger.info( + f"CoordinatorClient: coordinator reachable at {self._remote_url}" + ) + return await resp.json() + last_err = f"HTTP {resp.status}" + except Exception as e: # noqa: BLE001 + last_err = str(e) + attempt += 1 + if loop.time() >= deadline: + raise RuntimeError( + f"Coordinator at {self._remote_url} not reachable after " + f"{self._startup_timeout_s}s ({attempt} attempts, last " + f"error: {last_err}); aborting delegating server startup" + ) + logger.info( + f"CoordinatorClient: waiting for coordinator at " + f"{self._remote_url} (attempt {attempt}, {last_err})" + ) + await asyncio.sleep(2.0) + + async def get_disagg_request_id(self) -> int: + # Generate locally: the snowflake id is self-contained (not shared state), + # so the per-worker (node_id, process_id) avoids collisions without any HTTP + # hop. See disagg_utils.get_global_disagg_request_id. + return get_global_disagg_request_id(self._node_id, self._process_id) + + async def is_ready(self) -> bool: + try: + async with self.session.get( + f"{self._remote_url}/health", timeout=self._request_timeout_s + ) as resp: + return resp.status == 200 + except Exception as e: # noqa: BLE001 + logger.warning(f"CoordinatorClient health check failed: {e}") + return False + + async def cluster_info(self) -> Dict[str, Any]: + try: + async with self.session.get( + f"{self._remote_url}/cluster_info", timeout=self._request_timeout_s + ) as resp: + if resp.status == 200: + return await resp.json() + except Exception as e: # noqa: BLE001 + logger.warning(f"CoordinatorClient cluster_info failed: {e}") + return {"is_ready": False} + + async def stop(self) -> None: + if self._session is not None: + await self._session.close() + self._session = None + await self._ctx_router.close() + await self._gen_router.close() diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 6acbefe103d7..2e3542a44a9a 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -17,7 +17,7 @@ import os import traceback from abc import ABC, abstractmethod -from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple, Type +from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, List, Optional, Tuple, Type import aiohttp @@ -120,7 +120,7 @@ def __init__( max_retries: int = 1, retry_interval_sec: int = 1, session: Optional[aiohttp.ClientSession] = None, - disagg_id_generator: Optional[Callable[[], int]] = None, + disagg_id_generator: Optional[Callable[[], Awaitable[int]]] = None, ): self._router = router self._role = role @@ -197,7 +197,7 @@ async def _post_with_retry( if attempt > 0 and self._disagg_id_generator is not None: dp = getattr(request, "disaggregated_params", None) if dp is not None and getattr(dp, "disagg_request_id", None) is not None: - dp.disagg_request_id = self._disagg_id_generator() + dp.disagg_request_id = await self._disagg_id_generator() # Serialize once on the orchestrator's single event-loop thread. if _MSGSPEC_ENABLED: # msgspec msgpack: encode the request dict to msgpack bytes. Keep diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 051bb4a1c8cf..4a081d4fc8dd 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -32,25 +32,25 @@ from tensorrt_llm.executor.executor import CppExecutorError from tensorrt_llm.llmapi import tracing from tensorrt_llm.llmapi.disagg_utils import (DisaggServerConfig, - MetadataServerConfig, ServerRole, - get_ctx_gen_server_addrs, - get_global_disagg_request_id) + MetadataServerConfig, ServerRole) from tensorrt_llm.logger import logger from tensorrt_llm.serve.cluster_storage import ( HttpClusterStorageServer, create_cluster_storage, validate_http_cluster_storage_scope) from tensorrt_llm.serve.conversation_id import resolve_request_conversation_id -from tensorrt_llm.serve.metadata_server import create_metadata_server +from tensorrt_llm.serve.disagg_coordinator import (CoordinatorClient, + DisaggCoordinatorService) from tensorrt_llm.serve.openai_client import OpenAIClient, OpenAIHttpClient from tensorrt_llm.serve.openai_disagg_service import ( OpenAIDisaggregatedService, ResponseHooks) from tensorrt_llm.serve.openai_protocol import ( - UCompletionRequest, UCompletionResponse, - ensure_request_chat_template_allowed) + ChatCompletionRequest, CompletionRequest, UCompletionRequest, + UCompletionResponse, ensure_request_chat_template_allowed) from tensorrt_llm.serve.perf_metrics import DisaggPerfMetricsCollector from tensorrt_llm.serve.responses_utils import (ServerArrivalTimeMiddleware, - get_steady_clock_now_in_seconds) -from tensorrt_llm.serve.router import Router, create_router + get_steady_clock_now_in_seconds, + ttft_split_logger) +from tensorrt_llm.serve.router import Router from tensorrt_llm.version import __version__ as VERSION # yapf: enale @@ -63,22 +63,45 @@ def __init__(self, raw_req: Request, perf_metrics_collector: DisaggPerfMetricsCo self.gen_server = "" self.request_arrival_time = raw_req.state.server_arrival_time self.server_first_token_time = 0 + # Orchestrator-side TTFT timeline stamps: arrival -> ctx_dispatch (pre-ctx + # wait) -> ctx_resp (ctx routing + prefill) -> first_token (KV transfer + + # gen). 0 until the corresponding hook fires. + self.ctx_dispatch_time = 0 + self.ctx_resp_time = 0 self.perf_metrics_collector = perf_metrics_collector def on_req_begin(self, request: UCompletionRequest): self.perf_metrics_collector.queue_latency_seconds.observe(get_steady_clock_now_in_seconds() - self.request_arrival_time) + def on_ctx_dispatch(self, request: UCompletionRequest): + self.ctx_dispatch_time = get_steady_clock_now_in_seconds() + def on_ctx_resp(self, ctx_server: str, response: UCompletionResponse): self.ctx_server = ctx_server + self.ctx_resp_time = get_steady_clock_now_in_seconds() def on_first_token(self, gen_server: str, request: UCompletionRequest, response: UCompletionResponse = None): self.gen_server = gen_server self.server_first_token_time = get_steady_clock_now_in_seconds() def on_resp_done(self, gen_server: str, request: UCompletionRequest, response: UCompletionResponse = None): + # Feed the TTFT breakdown to a periodic aggregator (p50/p95 every N reqs), + # not one log line per request -- per-request logging on the disagg event + # loop inflated TTFT ~3s (observer effect). + arr = self.request_arrival_time + disp = self.ctx_dispatch_time + cresp = self.ctx_resp_time + ft = self.server_first_token_time + if arr and ft: + ttft_split_logger().record({ + "pre_ctx_ms": (disp - arr) * 1000 if disp else -1.0, + "ctx_phase_ms": (cresp - disp) * 1000 if (disp and cresp) else -1.0, + "xfer_gen_ms": (ft - cresp) * 1000 if cresp else -1.0, + "total_ms": (ft - arr) * 1000, + }) if request.disaggregated_params: ctx_req_id = request.disaggregated_params.ctx_request_id - asyncio.create_task(self.perf_metrics_collector.add_per_request_metrics(self.ctx_server, gen_server, ctx_req_id, self.raw_req.state.server_arrival_time, self.server_first_token_time)) + asyncio.create_task(self.perf_metrics_collector.add_per_request_metrics(self.ctx_server, gen_server, ctx_req_id, self.raw_req.state.server_arrival_time, self.server_first_token_time, self.ctx_dispatch_time)) class OpenAIDisaggServer: @@ -87,7 +110,8 @@ def __init__(self, req_timeout_secs: int = 180, server_start_timeout_secs: int = 180, metadata_server_cfg: Optional[MetadataServerConfig] = None, - metrics_interval_secs: int = 0): + metrics_interval_secs: int = 0, + coordinator_url: Optional[str] = None): self._config = config self._req_timeout_secs = req_timeout_secs self._server_start_timeout_secs = server_start_timeout_secs @@ -95,11 +119,11 @@ def __init__(self, self._metrics_interval_secs = metrics_interval_secs self._allow_request_chat_template = getattr( config, "allow_request_chat_template", False) + # When set, this is a forked worker: routing/readiness are delegated to + # the coordinator at coordinator_url (CoordinatorClient). Otherwise this + # process owns the routers + cluster state (DisaggCoordinatorService). + self._coordinator_url = coordinator_url - self._ctx_servers, self._gen_servers = get_ctx_gen_server_addrs(config.server_configs) - self._ctx_router = create_router(config.ctx_router_config, self._ctx_servers, metadata_server_cfg, create_metadata_server(metadata_server_cfg), self._sync_server_clock, disagg_node_id=config.node_id) - self._gen_router = create_router(config.gen_router_config, self._gen_servers, metadata_server_cfg, create_metadata_server(metadata_server_cfg), self._sync_server_clock, disagg_node_id=config.node_id) - self._metadata_server = create_metadata_server(metadata_server_cfg) self._perf_metrics_collector = DisaggPerfMetricsCollector(config.perf_metrics_max_requests) self._disagg_cluster_storage = None @@ -109,15 +133,27 @@ def __init__(self, self._disagg_cluster_storage = create_cluster_storage( config.disagg_cluster_config.cluster_uri, config.disagg_cluster_config.cluster_name) + # The server doesn't build routers -- the coordinator object does: + # DisaggCoordinatorService (owner) or CoordinatorClient (delegating). The + # server just reads .ctx_router / .gen_router off whichever it holds. + if self._coordinator_url: + self._coordinator = CoordinatorClient( + self._coordinator_url, self._config, metadata_server_cfg, + request_timeout_s=self._req_timeout_secs, + startup_timeout_s=self._server_start_timeout_secs) + else: + self._coordinator = DisaggCoordinatorService( + self._config, self._create_client, + metadata_config=self._metadata_server_cfg, + server_preparation_func=self._sync_server_clock, + server_start_timeout_secs=self._server_start_timeout_secs) + self._ctx_router = self._coordinator.ctx_router + self._gen_router = self._coordinator.gen_router self._service = OpenAIDisaggregatedService( - self._config, self._ctx_router, self._gen_router, self._create_client, - metadata_server=self._metadata_server, - metadata_config=self._metadata_server_cfg, + self._config, self._coordinator, self._create_client, req_timeout_secs=self._req_timeout_secs, - server_start_timeout_secs=self._server_start_timeout_secs, - perf_metrics_collector=self._perf_metrics_collector, - disagg_cluster_storage=self._disagg_cluster_storage) + perf_metrics_collector=self._perf_metrics_collector) try: otlp_cfg = config.otlp_config @@ -132,9 +168,7 @@ def __init__(self, @asynccontextmanager async def lifespan(app) -> None: - # Prepare servers (sync server clock) when static ctx/gen server list is used - await self._ctx_router.prepare_servers() - await self._gen_router.prepare_servers() + # The cluster manager (via setup) owns server preparation + monitoring. await self._service.setup() yield await self._service.teardown() @@ -143,24 +177,44 @@ async def lifespan(app) -> None: self.app.add_middleware(ServerArrivalTimeMiddleware) + # Log request-body validation failures so a client/server schema mismatch + # shows up server-side. Throttled (first, then every 1000th) to avoid + # flooding the event loop when every request fails identically. + self._val_err_n = 0 @self.app.exception_handler(RequestValidationError) - async def validation_exception_handler(_, exc): + async def validation_exception_handler(request: Request, exc): self._perf_metrics_collector.validation_exceptions.inc() + self._val_err_n += 1 + if self._val_err_n == 1 or self._val_err_n % 1000 == 0: + try: + errs = exc.errors() + # Compact: [{loc, type, msg}] -- drops the (large) echoed input. + brief = [{"loc": e.get("loc"), "type": e.get("type"), + "msg": e.get("msg")} for e in errs][:8] + except Exception: # noqa: BLE001 + brief = str(exc)[:500] + logger.warning( + f"[validation] {request.method} {request.url.path} 400 " + f"(n={self._val_err_n}): {brief}") return JSONResponse(status_code=400, content={"error": str(exc)}) self.register_routes() def _create_client(self, router: Router, role: ServerRole, max_retries: int = 1) -> OpenAIClient: - node_id = self._config.node_id + async def disagg_id_generator(): + return await self._coordinator.get_disagg_request_id() client = OpenAIHttpClient( router, role, self._req_timeout_secs, max_retries, - disagg_id_generator=lambda: get_global_disagg_request_id(node_id)) + disagg_id_generator=disagg_id_generator) self._perf_metrics_collector.add_client(client) return client def register_routes(self): - self.app.add_api_route("/v1/completions", self._wrap_entry_point(self._service.openai_completion), methods=["POST"]) - self.app.add_api_route("/v1/chat/completions", self._wrap_entry_point(self._service.openai_chat_completion), methods=["POST"]) + # The disagg service owns only the request-serving endpoints (/v1/*) and + # perf metrics. Readiness / cluster topology are the coordinator's state, + # so /health and /cluster_info hook straight to self._coordinator. + self.app.add_api_route("/v1/completions", self._wrap_entry_point(self._service.openai_completion, CompletionRequest), methods=["POST"]) + self.app.add_api_route("/v1/chat/completions", self._wrap_entry_point(self._service.openai_chat_completion, ChatCompletionRequest), methods=["POST"]) self.app.add_api_route("/health", self.health, methods=["GET"]) self.app.add_api_route("/cluster_info", self.cluster_info, methods=["GET"]) self.app.add_api_route("/version", self.version, methods=["GET"]) @@ -168,8 +222,12 @@ def register_routes(self): # import prometheus_client lazily to break the `set_prometheus_multiproc_dir` from prometheus_client import make_asgi_app self.app.mount("/prometheus/metrics", make_asgi_app()) - if self._disagg_cluster_storage and isinstance(self._disagg_cluster_storage, HttpClusterStorageServer): - self._disagg_cluster_storage.add_routes(self.app) + # Single-process (local coordinator): mount the in-process HTTP cluster + # storage routes on this app. In worker mode the coordinator is remote and + # owns those routes (CoordinatorClient has no cluster_storage). + cluster_storage = getattr(self._coordinator, "cluster_storage", None) + if isinstance(cluster_storage, HttpClusterStorageServer): + cluster_storage.add_routes(self.app) @staticmethod def _extract_conversation_id(req: UCompletionRequest, raw_req: Request): @@ -180,8 +238,12 @@ def _extract_conversation_id(req: UCompletionRequest, raw_req: Request): """ resolve_request_conversation_id(req, raw_req.headers) - def _wrap_entry_point(self, entry_point: Callable) -> Callable: - async def wrapper(req: UCompletionRequest, raw_req: Request) -> Response: + def _wrap_entry_point(self, entry_point: Callable, request_type: type = UCompletionRequest) -> Callable: + # Bind the concrete request model per route so FastAPI validates against it. + # The bare Union UCompletionRequest (no discriminator) makes Pydantic try + # CompletionRequest first and 400 every chat body, so override the wrapper's + # annotation with request_type (as openai_server.py does). + async def wrapper(req: request_type, raw_req: Request) -> Response: try: self._perf_metrics_collector.total_requests.inc() if req.stream: @@ -220,12 +282,12 @@ def _handle_exception(self, exception): async def health(self) -> Response: - if not await self._service.is_ready(): + if not await self._coordinator.is_ready(): return Response(status_code=500) return Response(status_code=200) async def cluster_info(self) -> JSONResponse: - return JSONResponse(content=await self._service.cluster_info()) + return JSONResponse(content=await self._coordinator.cluster_info()) async def version(self) -> JSONResponse: return JSONResponse(content={"version": VERSION}) diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 286bf5661bd9..e771e1671f76 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -14,20 +14,12 @@ import asyncio import os -from typing import Any, Callable, Dict, Optional - -from tensorrt_llm.llmapi.disagg_utils import ( - ConditionalDisaggConfig, - DisaggClusterConfig, - DisaggServerConfig, - MetadataServerConfig, - ServerRole, - get_global_disagg_request_id, -) +import time as _time +from typing import Callable, Optional + +from tensorrt_llm.llmapi.disagg_utils import ConditionalDisaggConfig, DisaggServerConfig, ServerRole from tensorrt_llm.logger import logger -from tensorrt_llm.serve.cluster_storage import ClusterStorage, WatchEventType -from tensorrt_llm.serve.disagg_auto_scaling import DisaggClusterManager, WorkerInfo -from tensorrt_llm.serve.metadata_server import JsonDictionary +from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinator from tensorrt_llm.serve.openai_client import OpenAIClient from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, @@ -51,32 +43,45 @@ _GEN_PENDING_FINISH_REASONS = ("length", "not_finished") +def _prectx_logger(): + """Aggregated (p50/p95 every N) breakdown of the ORCHESTRATOR pre-ctx path, + isolating where the arrival->ctx_dispatch time goes on the fleet worker: + cond_ms : _check_conditional_disagg + _check_gen_only_disagg + idgen_ms : get_disagg_request_id (local now; was an HTTP hop) + ctxroute_ms: ctx_router.get_next_server (routing/select round trip) + No per-request logging (observer effect on the single fleet loop). + """ + global _PRECTX_LOGGER + try: + return _PRECTX_LOGGER + except NameError: + from tensorrt_llm.serve.responses_utils import PeriodicBreakdownLogger + + _PRECTX_LOGGER = PeriodicBreakdownLogger( + "orchestrator_prectx", ["cond_ms", "idgen_ms", "ctxroute_ms"] + ) + return _PRECTX_LOGGER + + class OpenAIDisaggregatedService(OpenAIService): def __init__( self, config: DisaggServerConfig, - ctx_router: Router, - gen_router: Router, + coordinator: "DisaggCoordinator", client_factory: Callable[[Router, ServerRole], OpenAIClient], - metadata_server: Optional[JsonDictionary] = None, - metadata_config: Optional[MetadataServerConfig] = None, req_timeout_secs: int = 180, - server_start_timeout_secs: int = 180, perf_metrics_collector: Optional[DisaggPerfMetricsCollector] = None, - disagg_cluster_storage: Optional[ClusterStorage] = None, - health_check_interval_secs: int = 3, ): self._config = config - self._ctx_router = ctx_router - self._gen_router = gen_router + # The service drives the coordinator's ctx/gen routers uniformly, so serving + # is identical whether the router is the real one (single-process) or a + # delegating one that forwards placement to a remote coordinator (worker). + self._coordinator = coordinator + self._ctx_router = coordinator.ctx_router + self._gen_router = coordinator.gen_router self._client_factory = client_factory - self._metadata_server = metadata_server - self._metadata_config = metadata_config self._req_timeout_secs = req_timeout_secs - self._server_start_timeout_secs = server_start_timeout_secs self._perf_metrics_collector = perf_metrics_collector - self._cluster_storage = disagg_cluster_storage - self._health_check_interval_secs = health_check_interval_secs # Opt-in body-shrink for generation_only requests; see _get_gen_request. self._strip_gen_message_history = config.gen_strip_message_history # Opt-in: ask context workers to return prompt_token_ids as base64 int32. @@ -84,7 +89,6 @@ def __init__( self._ctx_client = None self._gen_client = None - self._disagg_cluster_manager = None self._schedule_style = DisaggScheduleStyle.CONTEXT_FIRST match self._config.schedule_style: @@ -140,17 +144,33 @@ async def _send_disagg_request_ctx_first( # empty server means client decides which server to use ctx_server = None # reserve a gen_server if conditional disagg is needed + _t_cond = _time.monotonic() gen_server, need_ctx = await self._check_conditional_disagg(request) need_ctx = need_ctx and not await self._check_gen_only_disagg(request) + _t_idgen = _time.monotonic() ctx_response = None gen_req = request - disagg_request_id = get_global_disagg_request_id(self._config.node_id) + disagg_request_id = await self._coordinator.get_disagg_request_id() + _t_afterid = _time.monotonic() + # Fine-grained pre-ctx breakdown (aggregated). ctxroute is filled below. + _prectx = { + "cond_ms": (_t_idgen - _t_cond) * 1000, + "idgen_ms": (_t_afterid - _t_idgen) * 1000, + "ctxroute_ms": -1.0, + } if need_ctx: + # Mark ctx-dispatch start: arrival->here is the pre-ctx wait in the + # orchestrator/fleet (accept queue + event loop + pipeline). + if hooks: + hooks.on_ctx_dispatch(request) ctx_req = self._get_ctx_request(request, disagg_request_id) # ctx generator is empty + _t_route = _time.monotonic() ctx_server, _ = await self._ctx_router.get_next_server( ctx_req, exclude_server=gen_server ) + _prectx["ctxroute_ms"] = (_time.monotonic() - _t_route) * 1000 + _prectx_logger().record(_prectx) ctx_response = await self._ctx_client.send_request( ctx_req, server=ctx_server, hooks=hooks ) @@ -243,10 +263,9 @@ def _get_gen_request( else: request.prompt_token_ids = ctx_response.prompt_token_ids # Opt-in: drop conversation history so the gen worker doesn't - # re-parse the full conversation JSON (dominates its GIL at high - # concurrency). It uses prompt_token_ids and only reads the last - # message; tools are preserved. Config-gated because it's unsafe - # for harmony/multimodal workers (model type is fixed per deploy). + # re-parse the full conversation JSON (a GIL cost at high + # concurrency); it uses prompt_token_ids and only the last message, + # tools preserved. Config-gated (unsafe for harmony/multimodal). if ( self._strip_gen_message_history and request.messages @@ -315,110 +334,33 @@ async def _check_gen_only_disagg(self, request: UCompletionRequest) -> bool: return True return False - async def cluster_info(self) -> Dict[str, Any]: - cluster_info = {"is_ready": await self.is_ready()} - if self._disagg_cluster_manager: - cluster_info.update(await self._disagg_cluster_manager.cluster_info()) - return cluster_info - async def is_ready(self) -> bool: - if self._disagg_cluster_manager: - return await self._disagg_cluster_manager.is_ready_with_router( - self._ctx_router.num_prepared_servers, - self._gen_router.num_prepared_servers, - ) - return True - - @property - def disagg_cluster_config(self) -> Optional[DisaggClusterConfig]: - return self._config.disagg_cluster_config + # Per-request readiness gate for the /v1/ handlers (the server's /health + # and /cluster_info hook the coordinator directly). Cluster topology + # (cluster_info) is the coordinator's concern, not the request service's. + return await self._coordinator.is_ready() @property def conditional_disagg_config(self) -> Optional[ConditionalDisaggConfig]: return self._config.conditional_disagg_config async def setup(self) -> None: + # Build the request-sending clients from the coordinator's routers and share + # them with the coordinator service so its readiness checks use the same pool. self._ctx_client = self._client_factory( self._ctx_router, ServerRole.CONTEXT, self._config.max_retries ) self._gen_client = self._client_factory( self._gen_router, ServerRole.GENERATION, self._config.max_retries ) - - if self.disagg_cluster_config and self._cluster_storage: - logger.info("Starting disagg cluster manager") - self._disagg_cluster_manager = DisaggClusterManager( - self.disagg_cluster_config, self._cluster_storage - ) - await self._disagg_cluster_manager.start() - await self._disagg_cluster_manager.watch_workers(on_event=self._on_worker_event) - logger.info("Disagg cluster manager started") - else: - if self._metadata_server and self._metadata_config: - logger.info("Starting server monitoring via metadata service") - await self._ctx_router.start_server_monitoring( - self._metadata_config.refresh_interval - ) - await self._gen_router.start_server_monitoring( - self._metadata_config.refresh_interval - ) - await self._wait_for_all_servers_ready() + if hasattr(self._coordinator, "set_clients"): + self._coordinator.set_clients(self._ctx_client, self._gen_client) + await self._coordinator.start() async def teardown(self) -> None: await self._ctx_client.shutdown() await self._gen_client.shutdown() - - if self._disagg_cluster_manager: - await self._disagg_cluster_manager.stop() - - if self._metadata_server: - await self._ctx_router.stop_server_monitoring() - await self._gen_router.stop_server_monitoring() - - async def _wait_for_all_servers_ready(self) -> None: - # Skip context servers if TRTLLM_DISAGG_BENCHMARK_GEN_ONLY is set - gen_only = os.getenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY") == "1" - - async def check_servers_ready(): - elapsed_time = 0 - interval = self._health_check_interval_secs - while elapsed_time < self._server_start_timeout_secs: - if gen_only: - unready_ctx_servers = [] - else: - _, unready_ctx_servers = await self._ctx_client.check_ready() - _, unready_gen_servers = await self._gen_client.check_ready() - if len(unready_ctx_servers) == 0 and len(unready_gen_servers) == 0: - if gen_only: - logger.info("Generation servers are ready (context servers skipped)") - else: - logger.info("All servers are ready") - return - logger.info( - f"Waiting for servers, context: {unready_ctx_servers}, generation: {unready_gen_servers}" - ) - await asyncio.sleep(interval) - elapsed_time += interval - - try: - await asyncio.wait_for(check_servers_ready(), timeout=self._server_start_timeout_secs) - except asyncio.TimeoutError: - raise TimeoutError("Timeout waiting for context and generation servers to be ready") - - async def _on_worker_event(self, worker_info: WorkerInfo, event_type: WatchEventType): - router_map = {ServerRole.CONTEXT: self._ctx_router, ServerRole.GENERATION: self._gen_router} - worker_addr = f"{worker_info.host}:{worker_info.port}" - try: - router = router_map[worker_info.role] - if event_type == WatchEventType.SET: - await router.add_server(worker_addr) - elif event_type == WatchEventType.DELETE: - await router.remove_server(worker_addr) - logger.info(f"Worker {event_type.name} event: {worker_info.worker_id}, {worker_addr}") - except KeyError: - logger.error( - f"Unknown worker role: {worker_info.role}, Worker {worker_info.worker_id} event: {event_type.name}" - ) + await self._coordinator.stop() async def _verify_ctx_response(self, ctx_response: UCompletionResponse) -> None: if ctx_response: @@ -457,8 +399,13 @@ async def _send_disagg_request_gen_first( ctx_server, gen_server = None, None ctx_server_info = None ctx_req, gen_req = None, None - disagg_request_id = get_global_disagg_request_id(self._config.node_id) + # Single-issuer disagg id (see _send_disagg_request_ctx_first): fetch from + # the coordinator so fleet workers never mint colliding ids. + disagg_request_id = await self._coordinator.get_disagg_request_id() if need_ctx: + # arrival->here = pre-ctx wait in the orchestrator/fleet. + if hooks: + hooks.on_ctx_dispatch(request) ctx_server, ctx_server_info = await self._ctx_router.get_next_server(request) ctx_req = self._get_ctx_request(request, disagg_request_id) gen_req = self._get_gen_request( @@ -469,15 +416,10 @@ async def _send_disagg_request_gen_first( ) if request.stream and need_ctx: - # For streaming gen_first requests, the gen client returns a lazy - # async generator whose HTTP POST only fires when iterated. The ctx - # server blocks waiting for the gen server's rx session (gen_first - # protocol). Using asyncio.gather would deadlock: ctx waits for gen - # server, but gen POST is deferred until the generator is consumed, - # and the generator isn't consumed until gather returns. - # - # Fix: eagerly start consuming the gen generator in a background - # task so the HTTP POST fires, then pipe chunks through a queue. + # Streaming gen_first: the gen POST is deferred until its generator is + # iterated, but the ctx server blocks on the gen rx session, so gather + # would deadlock. Eagerly consume the gen generator in a background task + # (firing the POST) and pipe chunks through a queue. gen_response = await self._gen_client.send_request( gen_req, server=gen_server, hooks=hooks ) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 53ab14af4234..390651040042 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -337,6 +337,9 @@ def __init__( self.disagg_cluster_worker = None self.resource_governor = None + self._active_request_count = 0 + self._active_request_count_lock = __import__('threading').Lock() + # Skip loading AutoProcessor and model_config for VISUAL_GEN models # These are LLM-specific and can cause unnecessary memory usage if self._is_visual_gen: @@ -471,6 +474,30 @@ async def validation_exception_handler(_, exc): self.app.add_middleware(ServerArrivalTimeMiddleware) + # Track active /v1/ requests (in-flight request gauge). + server_ref = self + + @self.app.middleware("http") + async def _track_active_requests(request, call_next): + if request.url.path.startswith("/v1/"): + with server_ref._active_request_count_lock: + server_ref._active_request_count += 1 + try: + return await call_next(request) + finally: + with server_ref._active_request_count_lock: + server_ref._active_request_count -= 1 + return await call_next(request) + + def _get_iteration_stats_buffer_maxlen(self) -> Optional[int]: + if isinstance(self.generator, VisualGen): + return None + + max_iterations = self.generator.args.iter_stats_max_iterations + if max_iterations is None or max_iterations < 0: + return None + return max_iterations + def _init_visual_gen(self): self.processor = None self.model_config = None diff --git a/tensorrt_llm/serve/perf_metrics.py b/tensorrt_llm/serve/perf_metrics.py index e2a4ff7607bb..e2e44c04851d 100644 --- a/tensorrt_llm/serve/perf_metrics.py +++ b/tensorrt_llm/serve/perf_metrics.py @@ -170,6 +170,7 @@ async def add_per_request_metrics( ctx_request_id: int, server_arrival_time: float, server_first_token_time: float, + ctx_dispatch_time: float = 0, ): async with self._lock: self._request_meteics.append( @@ -179,6 +180,7 @@ async def add_per_request_metrics( ctx_request_id, server_arrival_time, server_first_token_time, + ctx_dispatch_time, ) ) @@ -212,6 +214,7 @@ async def get_perf_metrics(self) -> List[Dict[str, Any]]: ctx_request_id, server_arrival_time, server_first_token_time, + ctx_dispatch_time, ) in self._request_meteics: gen_perf_metrics = self._server_metrics[gen_server].pop(ctx_request_id, None) if gen_perf_metrics is None: @@ -223,6 +226,7 @@ async def get_perf_metrics(self) -> List[Dict[str, Any]]: ctx_request_id, server_arrival_time, server_first_token_time, + ctx_dispatch_time, ) ) continue @@ -233,6 +237,10 @@ async def get_perf_metrics(self) -> List[Dict[str, Any]]: "ctx_server": ctx_server, "gen_server": gen_server, "disagg_server_arrival_time": server_arrival_time, + # arrival->ctx_dispatch = pre-ctx wait in the + # orchestrator/fleet (accept queue + event loop + + # pipeline), the dominant TTFT term under fleet load. + "disagg_ctx_dispatch_time": ctx_dispatch_time, "disagg_server_first_token_time": server_first_token_time, "ctx_perf_metrics": ctx_perf_metrics, "gen_perf_metrics": gen_perf_metrics, diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index cfd363d71616..b71ba3940c04 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -1990,11 +1990,112 @@ async def __call__(self, scope, receive, send): scope["state"] = {} scope["state"][ "server_arrival_time"] = get_steady_clock_now_in_seconds() + # Measure client-send -> server-accept latency: the network + HTTP + # accept-queue backlog BEFORE the handler runs (invisible to per- + # request server metrics, which start at server_arrival_time). The + # client stamps x-client-send-time (wall-clock epoch s); compare to + # our wall clock (NTP-synced hosts). Only /v1/ request paths. + path = scope.get("path", "") + if path.startswith("/v1/"): + for k, v in scope.get("headers", ()): # raw ASGI header pairs + if k == b"x-client-send-time": + try: + sent = float(v.decode()) + _accept_latency_logger().record( + max(0.0, + time.time() - sent)) + except (ValueError, UnicodeDecodeError): + pass + break # Pass through the original receive/send - no wrapping! await self.app(scope, receive, send) +_ACCEPT_LAT = None +_TTFT_SPLIT = None + + +def _accept_latency_logger(): + """Lazy singleton for the client-send -> server-accept latency logger + (PeriodicLatencyLogger is defined below this class).""" + global _ACCEPT_LAT + if _ACCEPT_LAT is None: + _ACCEPT_LAT = PeriodicLatencyLogger("accept.client_send_to_arrival") + return _ACCEPT_LAT + + +def ttft_split_logger(): + """Lazy singleton aggregating the per-request TTFT breakdown (pre_ctx / + ctx_phase / xfer_gen / total), logged periodically instead of per request.""" + global _TTFT_SPLIT + if _TTFT_SPLIT is None: + _TTFT_SPLIT = PeriodicBreakdownLogger( + "orchestrator", + ["pre_ctx_ms", "ctx_phase_ms", "xfer_gen_ms", "total_ms"]) + return _TTFT_SPLIT + + +class PeriodicLatencyLogger: + """Accumulates per-call latencies for a named coordinator API and logs + percentiles every ``window`` calls. Cheap, lock-free (single asyncio loop), + and self-resetting. Used to profile the coordinator surface on both the + in-process owner (DisaggCoordinatorService) and the HTTP client + (CoordinatorDelegatingRouter / CoordinatorClient) without per-call log spam.""" + + def __init__(self, name: str, window: int = 500): + self._name = name + self._window = window + self._samples: List[float] = [] + self._n = 0 + + def record(self, dt_s: float) -> None: + self._samples.append(dt_s * 1000.0) # ms + self._n += 1 + if self._n % self._window == 0: + s = sorted(self._samples) + m = len(s) + p = lambda q: s[min(int(q * m), m - 1)] + logger.info( + f"[coord_api] {self._name} n={self._n} ms: " + f"mean={sum(s)/m:.2f} p50={p(0.5):.2f} p90={p(0.9):.2f} " + f"p99={p(0.99):.2f} max={s[-1]:.2f}") + self._samples = [] + + +class PeriodicBreakdownLogger: + """Like PeriodicLatencyLogger but for a multi-field breakdown: accumulates a + dict of named ms values per sample and logs p50/p95 of EACH field once every + ``window`` samples. Used for the per-request TTFT split so we get the + breakdown WITHOUT a synchronous log write per request (that per-request + logging inflated TTFT ~3s on the single-loop disagg server -- observer + effect). Negative values (stage not reached) are dropped per field.""" + + def __init__(self, name: str, fields, window: int = 1000): + self._name = name + self._fields = list(fields) + self._window = window + self._samples = {f: [] for f in self._fields} + self._n = 0 + + def record(self, values: dict) -> None: + self._n += 1 + for f in self._fields: + v = values.get(f) + if v is not None and v >= 0: + self._samples[f].append(v) + if self._n % self._window == 0: + parts = [] + for f in self._fields: + s = sorted(self._samples[f]) + if s: + p = lambda q: s[min(int(q * len(s)), len(s) - 1)] + parts.append(f"{f}(p50={p(0.5):.0f},p95={p(0.95):.0f})") + self._samples[f] = [] + logger.info(f"[ttft_split] {self._name} n={self._n} " + + " ".join(parts)) + + class ResponseHooks(ABC): """ Hooks for response processing and (disagg) service perf observability. @@ -2004,6 +2105,13 @@ class ResponseHooks(ABC): def on_req_begin(self, request: UCompletionRequest): pass + def on_ctx_dispatch(self, request: UCompletionRequest): + """Fired the instant the disagg service starts placing the request on a + ctx server (before routing / the ctx HTTP send). arrival->here measures + the pre-ctx wait in the orchestrator/fleet (accept queue + event loop + + pipeline), which dominates TTFT under high fleet concurrency. Default + no-op so non-instrumented hook implementations are unaffected.""" + @abstractmethod def on_ctx_resp(self, ctx_server: str, response: UCompletionResponse): pass diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 5941ed979e51..e6aa2335e150 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -13,43 +13,31 @@ # limitations under the License. import asyncio -import os import time from abc import ABC, abstractmethod from collections import OrderedDict -from typing import Awaitable, Callable, Dict, Iterable, List, Optional, Union +from typing import Awaitable, Callable, Dict, Iterable, List, Optional import aiohttp +import orjson + +# Content-Type for orjson-serialized bodies sent to the coordinator (we pass +# data=bytes, so aiohttp doesn't set the json content type itself). +_JSON_HEADERS = {"Content-Type": "application/json"} -from tensorrt_llm.bindings.internal.batch_manager import \ - BlockKey as _NativeBlockKey -from tensorrt_llm.bindings.internal.batch_manager import \ - BlockKeyHasher as _NativeBlockKeyHasher from tensorrt_llm.llmapi.disagg_utils import (MetadataServerConfig, RouterConfig, ServerRole) from tensorrt_llm.logger import logger -from tensorrt_llm.runtime import kv_cache_hash -from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import \ - Block as V2Block -from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import \ - ReuseScope -from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import \ - RootBlock as V2RootBlock from tensorrt_llm.serve.conversation_id import get_request_conversation_id from tensorrt_llm.serve.metadata_server import JsonDictionary -from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest, - CompletionRequest) - -KV_CACHE_HASH_ALGO_DEFAULT = kv_cache_hash.KV_CACHE_HASH_ALGO_DEFAULT -KV_CACHE_HASH_ALGO_V1 = kv_cache_hash.KV_CACHE_HASH_ALGO_V1 -KV_CACHE_HASH_ALGO_V2 = kv_cache_hash.KV_CACHE_HASH_ALGO_V2 -KV_CACHE_HASH_ALGO_V2_SHA256_64 = kv_cache_hash.KV_CACHE_HASH_ALGO_V2_SHA256_64 -get_cache_salt_id = kv_cache_hash.get_cache_salt_id -hash_v1_block_key = kv_cache_hash.hash_v1_block_key -truncate_sha256_hash_to_int64 = kv_cache_hash.truncate_sha256_hash_to_int64 - -OpenAIRequest = Union[CompletionRequest, ChatCompletionRequest] -BlockHash = Union[int, str] +from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest +# Shared tokenization / block-hashing utilities (single source of truth). +# Re-exported here for backward compat. +from tensorrt_llm.serve.router_utils import ( # noqa: F401 + KV_CACHE_HASH_ALGO_DEFAULT, KV_CACHE_HASH_ALGO_V1, KV_CACHE_HASH_ALGO_V2, + KV_CACHE_HASH_ALGO_V2_SHA256_64, BlockHash, BlockHashMixin, OpenAIRequest, + block_key_hasher, get_cache_salt_id, get_request_num_tokens, + hash_v1_block_key, truncate_sha256_hash_to_int64, v2_sha256_block_hasher) # Max number of conversations whose home-server pin is retained (LRU). ROUTE_AFFINITY_CACHE_SIZE = 50000 @@ -107,13 +95,17 @@ def _session(self) -> Optional[aiohttp.ClientSession]: return self._session_provider() if self._session_provider else None async def increment_load(self, request: OpenAIRequest): - num_tokens = get_request_num_tokens(request) if self._use_tokens else 0 + # request may be None on the coordinator-delegated path (no request + # object crosses the HTTP hop); token accounting is skipped then. + num_tokens = (get_request_num_tokens(request) + if self._use_tokens and request is not None else 0) async with self._lock: self._num_active_requests += 1 self._num_active_tokens += num_tokens async def decrement_load(self, request: OpenAIRequest): - num_tokens = get_request_num_tokens(request) if self._use_tokens else 0 + num_tokens = (get_request_num_tokens(request) + if self._use_tokens and request is not None else 0) async with self._lock: self._num_active_requests -= 1 self._num_active_tokens -= num_tokens @@ -141,6 +133,7 @@ def __init__( self._kv_cache_block_tables: dict[str, set[BlockHash]] = { KV_CACHE_HASH_ALGO_V1: self._kv_cache_block_table } + self._event_only_blocks: set[BlockHash] = set() self._kv_cache_hash_algo = KV_CACHE_HASH_ALGO_DEFAULT self._tokens_per_block = tokens_per_block self._poll_task: Optional[asyncio.Task] = None @@ -198,11 +191,14 @@ def update_with_events(self, events: Iterable[dict]): if event["type"] == "created": self.set_hash_algo(hash_algo) if event["type"] == "stored": - self.add_blocks( - (block["block_hash"] for block in event["blocks"]), - hash_algo=hash_algo) + block_hashes = [ + block["block_hash"] for block in event["blocks"] + ] + self.add_blocks(block_hashes, hash_algo=hash_algo) + self._event_only_blocks.update(block_hashes) elif event["type"] == "removed": self.remove_blocks(event["block_hashes"], hash_algo=hash_algo) + self._event_only_blocks.difference_update(event["block_hashes"]) async def poll_events(self, session: aiohttp.ClientSession): async with session.post( @@ -210,20 +206,33 @@ async def poll_events(self, session: aiohttp.ClientSession): events_raw = await response.json() return events_raw + _event_match_log_counter = 0 + async def matched_tokens( self, block_hashes: list[list[BlockHash]], hash_algo: str = KV_CACHE_HASH_ALGO_DEFAULT) -> int: match_count = 0 + event_match_count = 0 async with self._lock: block_table = self._block_table(hash_algo) for hash_list in block_hashes: for block_hash in hash_list: - # TODO: 1) parent hash verification, 2) partial matching if block_hash in block_table: match_count += self._tokens_per_block + if block_hash in self._event_only_blocks: + event_match_count += self._tokens_per_block else: break + KvCacheAwareServerState._event_match_log_counter += 1 + if KvCacheAwareServerState._event_match_log_counter <= 20 or KvCacheAwareServerState._event_match_log_counter % 100 == 0: + logger.info( + f"EVENT_MATCH_DIAG server={self._server} " + f"total_match={match_count} event_match={event_match_count} " + f"event_blocks={len(self._event_only_blocks)} " + f"total_blocks={len(block_table)} " + f"query_hashes={[h for hl in block_hashes for h in hl][:3]}" + ) return match_count async def decrement_load(self, request: OpenAIRequest): @@ -369,6 +378,31 @@ def __init__( self._health_check_timeout = metadata_server_cfg.health_check_timeout if metadata_server_cfg else None self._server_preparation_func = server_preparation_func self._prepared_ready_servers: set[str] = set() + # Routing-latency diagnostics (gated by TLLM_LOG_ROUTE_TIMING=1): record + # get_next_server wall time per request and log percentiles periodically. + import os + self._log_route_timing = (os.environ.get("TLLM_LOG_ROUTE_TIMING", + "0") == "1") + self._rt_samples: list = [] + self._rt_n = 0 + + def _record_route_timing(self, dt_s: float) -> None: + """Record one get_next_server latency sample; log percentiles every + 500 calls. No-op unless TLLM_LOG_ROUTE_TIMING=1.""" + if not self._log_route_timing: + return + self._rt_samples.append(dt_s * 1000.0) # ms + self._rt_n += 1 + if self._rt_n % 500 == 0: + import statistics + s = sorted(self._rt_samples) + n = len(s) + p = lambda q: s[min(int(q * n), n - 1)] + logger.info(f"[route_timing] {type(self).__name__} n={self._rt_n} " + f"get_next_server_ms: mean={statistics.mean(s):.2f} " + f"p50={p(0.5):.2f} p90={p(0.9):.2f} p99={p(0.99):.2f} " + f"max={s[-1]:.2f}") + self._rt_samples = [] # reset window async def close(self): """Close the shared HTTP session.""" @@ -1004,19 +1038,23 @@ def __init__(self, max_batch_size: int = 64, tokens_per_block: Optional[int] = None, custom_tokenizer: Optional[str] = None, + tokenizer_dir: Optional[str] = None, track_routed_blocks: bool = True, load_weight: float = 0.25, load_cap: float = float("inf"), **kwargs): super().__init__(server_role, servers, metadata_server_cfg, metadata_server, **kwargs) - self._init_block_hashing(tokens_per_block, custom_tokenizer) + self._init_block_hashing(tokens_per_block, custom_tokenizer, + tokenizer_dir) self._init_load_balancing(servers, use_tokens) # TODO: use max_num_tokens? per server? self._max_batch_size = max_batch_size self._load_weight = load_weight self._load_cap = load_cap self._track_routed_blocks = track_routed_blocks + # request key -> (flat block hashes, hash_algo). Key is id(request) on the + # standalone path, the disagg req_id on the coordinator path. self._pending_routed_blocks: dict[int, tuple[list[BlockHash], str]] = {} def _create_server_state(self, server: str) -> KvCacheAwareServerState: @@ -1029,19 +1067,18 @@ async def close(self): await state.cancel_poll_task() await super().close() - def _stash_routed_blocks_on_route(self, request: OpenAIRequest, + def _stash_routed_blocks_on_route(self, key: int, block_hashes: list[list[BlockHash]], hash_algo: str) -> None: if not self._track_routed_blocks: return flat = [h for hl in block_hashes for h in hl] - self._pending_routed_blocks[id(request)] = (flat, hash_algo) + self._pending_routed_blocks[key] = (flat, hash_algo) - def _apply_routed_blocks_on_finish(self, request: OpenAIRequest, - server: Optional[str], + def _apply_routed_blocks_on_finish(self, key: int, server: Optional[str], success: bool) -> None: # Pop unconditionally to avoid leaks; apply only when eligible. - entry = self._pending_routed_blocks.pop(id(request), None) + entry = self._pending_routed_blocks.pop(key, None) if not (self._track_routed_blocks and success): return if entry is None: @@ -1130,66 +1167,77 @@ async def get_next_server( self, request: OpenAIRequest, exclude_server: Optional[str] = None) -> tuple[str, dict]: - async with self._lock: - servers = list([ - server for server in self._server_state.keys() - if server != exclude_server - ]) - if not servers: - raise ValueError( - f"No available servers after excluding {exclude_server}") + # Standalone (in-process) entry point = routing_key(tokenize+hash) then + # the shared _route core -- the SAME core the coordinator path uses. + key = await asyncio.to_thread(self._routing_key_sync, request) + server, info, _handle = await self._route(key, + exclude_server=exclude_server, + request=request) + return server, info + + def _routing_key_sync(self, request: OpenAIRequest) -> dict: + """Tokenize + per-algo block-hash (CPU-bound; run in a thread). Returns a + plain dict so the coordinator path can send it over HTTP unchanged.""" cache_salt_id = self._get_request_cache_salt_id(request) - hash_algo_by_server = { - server: self._get_server_hash_algo(server) - for server in servers + # Hash for every algo any server might use (usually one). + algos = { + self._get_server_hash_algo(s) + for s in self._server_state.keys() + } or None + token_lists, block_hashes_by_algo = \ + self._tokenize_and_compute_block_hashes_by_algo( + request, algos, cache_salt_id) + return { + "token_lists": token_lists, + "block_hashes_by_algo": block_hashes_by_algo, + "conv_key": self._content_affinity_key(request) } - # Tokenize + block-hash is CPU-bound (~50 ms p50 for a 40 k-token - # chat request with a Rust-backed tokenizer). Running it directly - # inside the async handler blocks the orchestrator's event loop and - # serializes all concurrent requests through it; with HuggingFace - # tokenizers releasing the GIL, offloading to a thread lets multiple - # tokenize calls run in parallel and frees the event loop to - # dispatch HTTP traffic to the CTX/GEN workers meanwhile. - token_lists, block_hashes_by_algo = await asyncio.to_thread( - self._tokenize_and_compute_block_hashes_by_algo, request, - hash_algo_by_server.values(), cache_salt_id) - # select the server by (KV match - load), bounded by load_cap + + async def _route(self, key, exclude_server=None, request=None, req_id=None): + """THE single routing core, shared by the standalone and coordinator + paths. Scores each server by (matched_tokens/tokens_per_block - + load_weight*load), applies load_cap, conversation affinity and RR + tie-break, then registers load + remembers routed blocks. The ONLY + difference between the two callers is request identity: standalone passes + the request (load keyed by id(request), finished via finish_request); + the coordinator passes req_id (the disagg request id) and finishes via + finish_request_by_id(req_id). Returns (server, info, req_id-or-None).""" + import time as _time + _rt_t0 = _time.monotonic() + token_lists = (key or {}).get("token_lists") or [] + block_hashes_by_algo = (key or {}).get("block_hashes_by_algo") or {} + conv_key = (key or {}).get("conv_key") + async with self._lock: + servers = [ + s for s in self._server_state.keys() if s != exclude_server + ] + if not servers: + raise ValueError( + f"No available servers after excluding {exclude_server}") + + def _hashes(server): + algo = self._get_server_hash_algo(server) + return algo, block_hashes_by_algo.get(algo, []) + workloads = [ - self._server_state[server].num_active_requests() - for server in servers + self._server_state[s].num_active_requests() for s in servers ] load_fractions = [ workloads[i] / self._max_batch_size for i in range(len(servers)) ] - scores = [] - matches = [] - for i in range(len(servers)): - server = servers[i] - hash_algo = hash_algo_by_server[server] - block_hashes = block_hashes_by_algo[hash_algo] - # https://github.com/ai-dynamo/dynamo/blob/main/docs/kv_cache_routing.md#kv-cache-routing-and-load-balancing - matches.append(await self._server_state[server].matched_tokens( - block_hashes, hash_algo)) - score = matches[-1] / self._tokens_per_block - self._load_weight * \ - workloads[i] - scores.append(score) - # Optional hard cap: drop servers at/over load_cap; fall back to all if - # none remain. Disabled by default (load_cap=inf) to match the original - # score-only selection. + scores, matches = [], [] + for i, server in enumerate(servers): + algo, bh = _hashes(server) + matches.append(await + self._server_state[server].matched_tokens(bh, algo)) + scores.append(matches[-1] / self._tokens_per_block - + self._load_weight * workloads[i]) candidate_idx = [ i for i, lf in enumerate(load_fractions) if lf < self._load_cap - ] - if not candidate_idx: - candidate_idx = list(range(len(servers))) - # Conversation affinity: pin all turns of a conversation (keyed by a - # content-derived prefix hash, no conversation-id header) to the server - # it first landed on, so a worker eviction shrinking the match score - # cannot scatter the conversation off its warm home. New conversations - # (no pin yet) fall through to the score, which balances them by load. + ] or list(range(len(servers))) affinity = getattr(self, "_route_affinity", None) if affinity is None: affinity = self._route_affinity = OrderedDict() - conv_key = self._content_affinity_key(request) winner = None if conv_key is not None: pinned = affinity.get(conv_key) @@ -1208,33 +1256,87 @@ async def get_next_server( affinity.move_to_end(conv_key) while len(affinity) > ROUTE_AFFINITY_CACHE_SIZE: affinity.popitem(last=False) - hash_algo = hash_algo_by_server[server] - block_hashes = block_hashes_by_algo[hash_algo] + hash_algo, block_hashes = _hashes(server) + + # Same load/routing maps as the standalone path; only the key differs: + # id(request) standalone, disagg req_id under the coordinator (the id that + # crosses the /finish HTTP hop). + key = id(request) if req_id is None else req_id async with self._lock: - await self._register_request(server, request) - self._stash_routed_blocks_on_route(request, block_hashes, hash_algo) + await self._server_state[server].increment_load(request) + self._req_routing_table[key] = server + self._stash_routed_blocks_on_route(key, block_hashes, hash_algo) + self._record_route_timing(_time.monotonic() - _rt_t0) return server, { - "block_hashes": block_hashes, # list[list[int | str]] + "block_hashes": block_hashes, "hash_algo": hash_algo, - "token_lists": token_lists, # list[list[int]] - "matches": matches, # list[int] + "token_lists": token_lists, + "matches": matches, "server_info": self._server_info.get(server, {}), - } + }, req_id async def finish_request(self, request: OpenAIRequest, session: Optional[aiohttp.ClientSession] = None, success: bool = True): + # Standalone entry point: key by id(request); pass request so token-load + # accounting matches the increment_load(request) done at route time. + await self._finish(id(request), + success, + request=request, + session=session) + + async def _finish(self, key, success, request=None, session=None): + """THE single finish core, shared by finish_request (standalone, key = + id(request), request passed) and finish_request_by_id (coordinator, key = + disagg req_id, request=None -- symmetric with its increment_load(None)). + Pops the SAME maps _route filled, decrements load, commits routed blocks + on success, and refreshes the block table.""" async with self._lock: - server = self._req_routing_table.pop(id(request), None) + server = self._req_routing_table.pop(key, None) if server is not None and server in self._server_state: await self._server_state[server].decrement_load(request) - self._apply_routed_blocks_on_finish(request, server, success) + self._apply_routed_blocks_on_finish(key, server, success) + self._poll_server_on_finish(server, session) + + def _poll_server_on_finish(self, server, session=None): + """Refresh a server's KV-cache block table from /kv_cache_events after a + request finishes. Shared by finish_request (standalone) and + finish_request_by_id (coordinator) so the block table always stays warm -- + the delegated path is inert without this (matches score 0).""" if (server is not None and server in self._server_state and self._events_aligned(server)): # Fire-and-forget; poll runs in background and coalesces per server. self._server_state[server].schedule_poll_and_update(session) + # ---- coordinator delegation: thin wrappers over the shared _route core --- + # The fleet worker computes routing_key() locally; the coordinator (owns + # _server_state) runs get_next_server_by_key(). Both use the same _route core + # as the standalone get_next_server. + + def routing_key(self, request: OpenAIRequest): + """Worker-side: tokenize + block-hash. Same dict _route consumes, JSON- + serializable for the /select POST.""" + return self._routing_key_sync(request) + + async def get_next_server_by_key(self, + routing_key, + exclude_server=None, + req_id=None): + """Coordinator-side placement: the SAME _route core, keyed by the + caller's disagg request id (req_id) instead of id(request).""" + return await self._route(routing_key, + exclude_server=exclude_server, + request=None, + req_id=req_id) + + async def finish_request_by_id(self, req_id, success=True): + """Coordinator-side finish: the SAME _finish core, keyed by the disagg + request id instead of id(request).""" + if req_id is None: + return + await self._finish(req_id, success) + def _on_servers_updated(self, old_servers, new_servers): new_state = {} for server in new_servers: @@ -1383,6 +1485,9 @@ def __init__(self, } # id(request) -> (server, weight, monotonic_timestamp) self._req_content_entry: dict[int, tuple[str, int, float]] = {} + # Coordinator-delegated path only: disagg req_id -> server, between + # select and finish (id(request) can't cross the HTTP hop). + self._coord_pending: dict = {} # ── content-based load tracking ── @@ -1679,6 +1784,255 @@ async def finish_request(self, logger.debug(f"ConversationRouter: FINISH server={server}, " f"content_loads={loads}") + # -- coordinator-path: conversation_id-only sticky routing -- + # The coordinator has no request object, so per-request load is tracked by an + # opaque handle instead of id(request). Only explicit conversation_id sessions + # are supported over the coordinator (no implicit content match). + + def routing_key(self, request: OpenAIRequest): + """The conversation_id (or None); no tokenization on the worker.""" + return self._get_conversation_id(request) + + async def get_next_server_by_key(self, + routing_key, + exclude_server=None, + req_id=None): + conv_id = routing_key + self._validate_servers_available() + async with self._lock: + entry = self._session_table.get(conv_id) if conv_id else None + if (entry is not None and entry[0] in self._server_state + and entry[0] != exclude_server): + server = entry[0] + self._session_table.move_to_end(conv_id) + else: + server = self._select_least_loaded(exclude_server) + if server is None: + raise ValueError( + f"No available servers after excluding {exclude_server}" + ) + if conv_id: + self._update_session(conv_id, server, []) + # Request-count load (no request object at the coordinator). Keyed by + # the disagg req_id -- the sole id crossing the HTTP hop on /finish. + self._server_content_load[server] = ( + self._server_content_load.get(server, 0) + 1) + if req_id is not None: + self._coord_pending[req_id] = server + return server, { + "server_info": self._server_info.get(server, {}) + }, req_id + + async def finish_request_by_id(self, req_id, success=True): + del success + if req_id is None: + return + async with self._lock: + server = self._coord_pending.pop(req_id, None) + if server and server in self._server_content_load: + self._server_content_load[server] = max( + 0, self._server_content_load[server] - 1) + + +class CoordinatorDelegatingRouter(Router): + """Worker-side Router that delegates placement to the disagg coordinator. + + Used only for *stateful* routers (conversation, kv_cache_aware): the worker + must not keep its own copy of that state, so it wraps a local router of the same + type and, for each request, computes the small ``routing_key`` locally and + POSTs it to the coordinator's ``/select``. ``finish_request`` POSTs the + returned handle to ``/finish`` so the coordinator releases per-request state. + Server-pool / prepare / close operations delegate to the wrapped local router. + + Stateless routers (round_robin, load_balancing) are NOT wrapped -- the worker + holds the real router and places locally, so they never reach this class (see + ``CoordinatorClient``). ``OpenAIClient`` already drives + ``router.get_next_server`` / ``router.finish_request``, so the completions + service needs no worker-specific branching. + """ + + def __init__(self, + coordinator_url: str, + local_router: "Router", + role: str, + request_timeout_s: float = 5.0): + # Intentionally NOT calling Router.__init__: this is a thin proxy whose + # server-pool state lives on the wrapped local router (see __getattr__). + # coordinator_url may be a TCP URL or unix:/path (UDS avoids the TCP + # loopback stack for the hot /select,/finish calls). Resolve the request + # base URL now; the session (with UnixConnector when applicable) is made + # lazily so it binds to the running event loop. + from tensorrt_llm.serve.disagg_coordinator import coordinator_base_url + self._coordinator_url_raw = coordinator_url + self._coordinator_url = coordinator_base_url(coordinator_url) + self._local = local_router + self._role = role # "context" | "generation" + self._request_timeout_s = request_timeout_s + self._session: Optional[aiohttp.ClientSession] = None + # Coordinator HTTP-client API latency (includes network round-trip to the + # coordinator + its in-process handler). Compare against the owner-side + # [coord_api] to isolate the fleet /select|/finish HTTP overhead. + from tensorrt_llm.serve.responses_utils import (PeriodicBreakdownLogger, + PeriodicLatencyLogger) + self._select_lat = PeriodicLatencyLogger(f"client.select[{role}]") + self._finish_lat = PeriodicLatencyLogger(f"client.finish[{role}]") + # IPC latency breakdown (ms) for the coordinator client<->server round + # trip, decomposed via the X-Client-Send/X-Coord-Recv/X-Coord-Send headers + # (client + coordinator co-located => wall clock comparable): + # req_wire : client send -> coord recv (request transport + accept) + # handler : coord recv -> coord send (in-coordinator handler work) + # resp_wire : coord send -> client recv (response transport) + # rtt : client send -> client recv (total observed round trip) + # Aggregated p50/p95 per field every window samples (no per-call logging). + self._ipc_fields = [ + "req_wire_ms", "handler_ms", "resp_wire_ms", "rtt_ms" + ] + self._select_ipc = PeriodicBreakdownLogger(f"coord_ipc.select[{role}]", + self._ipc_fields) + self._finish_ipc = PeriodicBreakdownLogger(f"coord_ipc.finish[{role}]", + self._ipc_fields) + + def __getattr__(self, name): + # servers / prepare_servers / num_prepared_servers / start_server_monitoring + # / routing_key / ... all delegate to the local router. + return getattr(self._local, name) + + @property + def session(self) -> aiohttp.ClientSession: + if self._session is None: + from tensorrt_llm.serve.disagg_coordinator import \ + make_coordinator_session + self._session = make_coordinator_session(self._coordinator_url_raw) + return self._session + + def _on_servers_updated(self, old_servers, new_servers): + pass + + def _record_ipc(self, ipc_logger, send_t: float, recv_t: float, + resp) -> None: + """Decompose the coordinator round trip from the response timing headers. + send_t/recv_t are the client-side wall clocks (co-located => comparable). + Missing headers (older coordinator) => that field is dropped by the + breakdown logger (records -1). All values in ms.""" + from tensorrt_llm.serve.coordinator_server import (HDR_COORD_RECV, + HDR_COORD_SEND) + + def _f(name): + v = resp.headers.get(name) + try: + return float(v) if v is not None else None + except (TypeError, ValueError): + return None + + coord_recv = _f(HDR_COORD_RECV) + coord_send = _f(HDR_COORD_SEND) + ipc_logger.record({ + "req_wire_ms": (coord_recv - send_t) * 1000 if coord_recv else -1.0, + "handler_ms": (coord_send - coord_recv) * 1000 if + (coord_recv and coord_send) else -1.0, + "resp_wire_ms": + (recv_t - coord_send) * 1000 if coord_send else -1.0, + "rtt_ms": (recv_t - send_t) * 1000, + }) + + def _request_id(self, request: OpenAIRequest) -> int: + """The request's disagg id -- the sole cross-process key for select/finish. + Context requests carry disagg_request_id; generation requests inherit the + SAME id as ctx_request_id (set by OpenAIDisaggregatedService before + routing). This id is what the ctx worker registered its KV-transfer + TxSession under, so the gen request MUST keep it -- never re-issue a new + id here, or the gen transceiver waits on a key the ctx side never + registered (transfer never completes -> DISAGG_GENERATION_TRANS_IN_PROGRESS + fills the gen IndexMapper and throughput collapses).""" + dp = request.disaggregated_params + assert dp is not None, "delegated routing requires disaggregated_params" + rid = (dp.disagg_request_id + if self._role == "context" else dp.ctx_request_id) + assert rid is not None, ( + f"delegated {self._role} routing requires a disagg request id " + f"(disagg_request_id/ctx_request_id) on the request") + return rid + + async def get_next_server( + self, + request: OpenAIRequest, + exclude_server: Optional[str] = None) -> tuple[str, dict]: + # routing_key() tokenizes + block-hashes the prompt (CPU-bound for + # kv_cache_aware); run it in a thread so it doesn't block the fleet worker + # event loop driving the concurrent streams. + key = await asyncio.to_thread(self._local.routing_key, request) + # Send the request's existing disagg id as the cross-process key (the + # coordinator keys pending state by it for /finish); placement must not + # change it, since the ctx<->gen KV transfer is keyed by it. + payload = { + "role": self._role, + "routing_key": key, + "req_id": self._request_id(request), + "exclude_server": exclude_server + } + # orjson.dumps (bytes) instead of aiohttp json= (stdlib dumps): the + # routing_key is hundreds of int64, and this is the hot fleet->coordinator + # call. Parse the response with orjson too. X-Client-Send-Time (wall clock) + # lets the coordinator echo timing headers back for the IPC breakdown. + from tensorrt_llm.serve.coordinator_server import HDR_CLIENT_SEND + _t0 = time.monotonic() + _send_wall = time.time() + hdrs = {**_JSON_HEADERS, HDR_CLIENT_SEND: repr(_send_wall)} + async with self.session.post(f"{self._coordinator_url}/select", + data=orjson.dumps(payload), + headers=hdrs, + timeout=self._request_timeout_s) as resp: + if resp.status != 200: + raise ValueError(f"coordinator /select returned {resp.status}: " + f"{await resp.text()}") + body = orjson.loads(await resp.read()) + self._record_ipc(self._select_ipc, _send_wall, time.time(), resp) + self._select_lat.record(time.monotonic() - _t0) + info = body.get("info") or {} + return body["server"], info + + async def finish_request(self, + request: OpenAIRequest, + session: Optional[aiohttp.ClientSession] = None, + success: bool = True): + # Fire-and-forget: /finish only releases the coordinator's routed-load + # bookkeeping, so run it as a background task instead of awaiting the HTTP + # round trip on every request's critical path. Capture req_id first (the + # request may be freed after return); errors are logged, not propagated. + del session + req_id = self._request_id(request) + asyncio.create_task(self._finish_async(req_id, success)) + + async def _finish_async(self, req_id: int, success: bool): + from tensorrt_llm.serve.coordinator_server import HDR_CLIENT_SEND + _t0 = time.monotonic() + _send_wall = time.time() + hdrs = {**_JSON_HEADERS, HDR_CLIENT_SEND: repr(_send_wall)} + try: + async with self.session.post( + f"{self._coordinator_url}/finish", + data=orjson.dumps({ + "role": self._role, + "req_id": req_id, + "success": success + }), + headers=hdrs, + timeout=self._request_timeout_s) as resp: + if resp.status != 200: + logger.warning( + f"coordinator /finish returned {resp.status}") + self._record_ipc(self._finish_ipc, _send_wall, time.time(), + resp) + self._finish_lat.record(time.monotonic() - _t0) + except Exception as e: # noqa: BLE001 + logger.warning(f"CoordinatorDelegatingRouter finish failed: {e}") + + async def close(self): + if self._session is not None: + await self._session.close() + self._session = None + await self._local.close() + def create_router( router_config: Optional[RouterConfig], @@ -1725,3 +2079,35 @@ def create_router( metadata_server, server_preparation_func=server_preparation_func, **extra_args) + + +def build_disagg_routers( + ctx_router_config: Optional[RouterConfig], + gen_router_config: Optional[RouterConfig], + ctx_servers: Optional[List[str]], + gen_servers: Optional[List[str]], + metadata_server_cfg: Optional[MetadataServerConfig] = None, + metadata_server: Optional[JsonDictionary] = None, + server_preparation_func: Optional[Callable[[str], Awaitable[None]]] = None, + disagg_node_id: int = 0, + is_delegating_client: bool = False, +) -> tuple[Router, Router]: + """Build the ctx and gen routers for one disagg process. + + Each side is built independently via :func:`create_router`. Stateful router + types (conversation, kv_cache_aware) expose ``routing_key`` / + ``get_next_server_by_key``; when this process is a delegating client, the + caller (:class:`CoordinatorClient`) wraps those in a + :class:`CoordinatorDelegatingRouter` so placement is delegated to the + coordinator. Stateless types (round_robin, load_balancing) place locally. + ``is_delegating_client`` is accepted for call-site symmetry; router + construction itself does not depend on it. + """ + del is_delegating_client # no build-time behavior differs by this flag + ctx_router = create_router(ctx_router_config, ctx_servers, + metadata_server_cfg, metadata_server, + server_preparation_func, disagg_node_id) + gen_router = create_router(gen_router_config, gen_servers, + metadata_server_cfg, metadata_server, + server_preparation_func, disagg_node_id) + return ctx_router, gen_router diff --git a/tensorrt_llm/serve/router_utils.py b/tensorrt_llm/serve/router_utils.py new file mode 100644 index 000000000000..d691daccd8b2 --- /dev/null +++ b/tensorrt_llm/serve/router_utils.py @@ -0,0 +1,377 @@ +# Copyright (c) 2025-2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shared router utilities: request tokenization and KV-cache block hashing. + +Extracted from ``router.py`` so the surface routers there can share a single +implementation of block hashing without importing the whole router module. +""" + +import os +from collections import OrderedDict +from typing import Iterable, List, Optional, Union + +from transformers import AutoTokenizer + +from tensorrt_llm.bindings.internal.batch_manager import BlockKey as _NativeBlockKey +from tensorrt_llm.bindings.internal.batch_manager import BlockKeyHasher as _NativeBlockKeyHasher +from tensorrt_llm.logger import logger +from tensorrt_llm.runtime import kv_cache_hash +from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import Block as V2Block +from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ReuseScope +from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import RootBlock as V2RootBlock +from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest, CompletionRequest + +KV_CACHE_HASH_ALGO_DEFAULT = kv_cache_hash.KV_CACHE_HASH_ALGO_DEFAULT +KV_CACHE_HASH_ALGO_V1 = kv_cache_hash.KV_CACHE_HASH_ALGO_V1 +KV_CACHE_HASH_ALGO_V2 = kv_cache_hash.KV_CACHE_HASH_ALGO_V2 +KV_CACHE_HASH_ALGO_V2_SHA256_64 = kv_cache_hash.KV_CACHE_HASH_ALGO_V2_SHA256_64 +get_cache_salt_id = kv_cache_hash.get_cache_salt_id +hash_v1_block_key = kv_cache_hash.hash_v1_block_key +truncate_sha256_hash_to_int64 = kv_cache_hash.truncate_sha256_hash_to_int64 + +OpenAIRequest = Union[CompletionRequest, ChatCompletionRequest] +BlockHash = Union[int, str] + +__all__ = [ + "KV_CACHE_HASH_ALGO_DEFAULT", + "KV_CACHE_HASH_ALGO_V1", + "KV_CACHE_HASH_ALGO_V2", + "KV_CACHE_HASH_ALGO_V2_SHA256_64", + "get_cache_salt_id", + "hash_v1_block_key", + "truncate_sha256_hash_to_int64", + "OpenAIRequest", + "BlockHash", + "get_request_num_tokens", + "block_key_hasher", + "v2_sha256_block_hasher", + "BlockHashMixin", + "PrefixBlockSet", +] + + +class PrefixBlockSet: + """Single-owner block-hash index -- a flat ``set`` of held block hashes. + + A KV-cache block hash folds in its parent chain (the worker computes it with + ``BlockKeyHasher.hash(block_key, parent_hash)``), so every block-hash *value* + is globally unique to one position in one prefix path. A request's ordered + block-hash list is itself the prefix path, so longest-common-prefix against a + set of held blocks is just "walk the list until the first hash the owner + doesn't hold" -- no explicit tree needed. + + This is the exact structure the orchestrator ``KvCacheAwareServerState`` uses + (a ``set[block_hash]`` per server, walked until the first miss). It is the + right index whenever there is only ONE logical owner -- e.g. the centralized + router's per-instance ``combined_trie`` (owner = instance) and each rank's + trie (owner = that rank). Those are only ever queried for a single owner's + prefix depth (:meth:`match_one`), so a ``hash -> {owner}`` reverse map with + per-depth set intersections would be pure overhead here. + + The ``owner_id`` argument on :meth:`add` / :meth:`remove` / :meth:`match_one` + is accepted and ignored so this is a drop-in for the single-owner call sites. + """ + + __slots__ = ("_blocks",) + + def __init__(self) -> None: + self._blocks: set[int] = set() + + def add(self, owner_id: str, block_hashes: Iterable[int]) -> None: + self._blocks.update(block_hashes) + + def remove(self, owner_id: str, block_hashes: Iterable[int]) -> None: + self._blocks.difference_update(block_hashes) + + def remove_worker(self, owner_id: str) -> None: + self._blocks.clear() + + def match_one(self, owner_id: str, block_hashes: List[int]) -> int: + """Consecutive prefix-block count held by the (single) owner. + + Identical to ``KvCacheAwareServerState.matched_tokens``: walk the query + path, counting blocks present in the set, and stop at the first miss. + """ + blocks = self._blocks + depth = 0 + for h in block_hashes: + if h not in blocks: + break + depth += 1 + return depth + + def has_worker(self, owner_id: str) -> bool: + return bool(self._blocks) + + +def get_request_num_tokens(request: OpenAIRequest) -> int: + if ( + request.disaggregated_params is None + or request.disaggregated_params.request_type == "context_only" + ): + if isinstance(request, ChatCompletionRequest): + raise ValueError( + "LoadBalancing router with tokens doesn't support ChatCompletionRequest yet" + ) + + if isinstance(request.prompt, str) or ( + isinstance(request.prompt, list) and isinstance(request.prompt[0], int) + ): + prompts = [request.prompt] + else: + prompts = request.prompt + + num_tokens = sum(len(prompt) for prompt in prompts) + elif request.disaggregated_params.request_type == "generation_only": + raise ValueError( + "LoadBalancing router with tokens doesn't support generation_only requests" + ) + else: + raise ValueError(f"Unsupported request type: {request.disaggregated_params.request_type}") + + return num_tokens + + +def block_key_hasher( + token_ids: list[int], parent_hash: Optional[int] = None, cache_salt_id: Optional[int] = None +) -> int: + parent = 0 if parent_hash is None else parent_hash + # Fast path: the native C++ BlockKeyHasher is bit-exact with + # hash_v1_block_key and avoids the per-token Python loop. Its hash() binding + # takes no cache_salt_id, so fall back to Python only when a salt is set + # (rare opt-in; never in the unsalted agent/chat completion path). + if cache_salt_id is None: + return _NativeBlockKeyHasher.hash(_NativeBlockKey(token_ids), parent) + return hash_v1_block_key(token_ids, parent_hash=parent, cache_salt_id=cache_salt_id) + + +def v2_sha256_block_hasher( + token_ids: list[int], parent_hash: Optional[str] = None, cache_salt_id: Optional[int] = None +) -> str: + parent_key = ( + V2RootBlock.make_key(ReuseScope(salt=cache_salt_id)) + if parent_hash is None + else bytes.fromhex(parent_hash) + ) + return V2Block.make_key(parent_key, token_ids).hex() + + +class BlockHashMixin: + """Shared tokenization and block-hash computation. + + Used by routers that need KV-cache-aware prefix matching. + """ + + def _init_block_hashing( + self, + tokens_per_block: Optional[int] = None, + custom_tokenizer: Optional[str] = None, + tokenizer_dir: Optional[str] = None, + ): + env_tokens_per_block = os.environ.get("TRTLLM_KVCACHE_AWARE_ROUTER_HASH_TOKENS_PER_BLOCK") + if env_tokens_per_block is not None: + tokens_per_block = int(env_tokens_per_block) + self._tpb_auto = tokens_per_block is None + self._tokens_per_block = 32 if tokens_per_block is None else tokens_per_block + self._tokenizers: dict = {} + self._custom_tokenizer = custom_tokenizer + self._tokenizer_dir = tokenizer_dir + logger.info( + f"BlockHashMixin: tokens_per_block={self._tokens_per_block}" + f"{' (auto, adopts worker)' if self._tpb_auto else ''}" + f", custom_tokenizer={self._custom_tokenizer}" + ) + + def _get_tokenizer(self, model: str): + if model not in self._tokenizers: + model_path = self._tokenizer_dir or model + if self._custom_tokenizer: + from tensorrt_llm.tokenizer import load_custom_tokenizer + + self._tokenizers[model] = load_custom_tokenizer(self._custom_tokenizer, model_path) + else: + from tensorrt_llm.tokenizer import maybe_fix_byte_level_tokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + self._tokenizers[model] = maybe_fix_byte_level_tokenizer( + tokenizer, model_path, trust_remote_code=True + ) + return self._tokenizers[model] + + def _encode_with_prefix_cache(self, rendered: str, key: int, tokenizer) -> list[int]: + cache = getattr(self, "_tok_prefix_cache", None) + if cache is None: + cache = self._tok_prefix_cache = OrderedDict() + entry = cache.get(key) + if entry is not None and len(rendered) > len(entry[0]) and rendered.startswith(entry[0]): + ids = entry[1] + tokenizer.encode(rendered[len(entry[0]) :], add_special_tokens=False) + else: + ids = tokenizer.encode(rendered, add_special_tokens=False) + cache[key] = (rendered, ids) + cache.move_to_end(key) + while len(cache) > 1024: + cache.popitem(last=False) + return ids + + def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: + # Handle ChatCompletionRequest (has messages, not prompt) + if isinstance(request, ChatCompletionRequest): + if request.prompt_token_ids is not None: + return [request.prompt_token_ids] + tokenizer = self._get_tokenizer(request.model) + tool_dicts = ( + None + if getattr(request, "tools", None) is None + else [ + tool.model_dump() if hasattr(tool, "model_dump") else tool + for tool in request.tools + ] + ) + chat_template_kwargs = ( + request.chat_template_kwargs + if getattr(request, "chat_template_kwargs", None) + else {} + ) + rendered = tokenizer.apply_chat_template( + [msg if isinstance(msg, dict) else dict(msg) for msg in request.messages], + add_generation_prompt=request.add_generation_prompt, + tokenize=False, + return_dict=False, + tools=tool_dicts, + **chat_template_kwargs, + ) + if isinstance(rendered, str): + key = hash( + "".join( + str( + msg.get("content") + if isinstance(msg, dict) + else getattr(msg, "content", "") + ) + for msg in request.messages[:2] + ) + ) + result = self._encode_with_prefix_cache(rendered, key, tokenizer) + else: + result = list(rendered) + request.prompt_token_ids = result + return [result] + + # Handle CompletionRequest (has prompt) + prompts = request.prompt + if isinstance(prompts, list) and isinstance(prompts[0], list): + return prompts + elif isinstance(prompts, list) and isinstance(prompts[0], int): + return [prompts] + elif isinstance(prompts, str): + prompts = [prompts] + else: + assert isinstance(prompts, list) and isinstance(prompts[0], str) + + tokenizer = self._get_tokenizer(request.model) + token_lists = [tokenizer(prompt)["input_ids"] for prompt in prompts] + # Replace string prompts with token IDs so the worker server + # skips re-tokenization + request.prompt = token_lists if len(token_lists) > 1 else token_lists[0] + return token_lists + + def _compute_block_hashes( + self, + token_lists: list[list[int]], + hash_algo: str = KV_CACHE_HASH_ALGO_DEFAULT, + cache_salt_id: Optional[int] = None, + ) -> list[list[BlockHash]]: + if hash_algo == KV_CACHE_HASH_ALGO_V1: + block_hasher = block_key_hasher + elif hash_algo == KV_CACHE_HASH_ALGO_V2: + block_hasher = v2_sha256_block_hasher + elif hash_algo == KV_CACHE_HASH_ALGO_V2_SHA256_64: + reuse_scope = ReuseScope(salt=cache_salt_id) + block_hashes: list[list[BlockHash]] = [] + for token_list in token_lists: + hash_list = [] + parent_key = V2RootBlock.make_key(reuse_scope) + for t in range(0, len(token_list) - 1, self._tokens_per_block): + t_end = min(t + self._tokens_per_block, len(token_list) - 1) + parent_key = V2Block.make_key(parent_key, token_list[t:t_end]) + hash_list.append(truncate_sha256_hash_to_int64(parent_key)) + block_hashes.append(hash_list) + return block_hashes + else: + raise ValueError(f"Unsupported KV cache hash algorithm: {hash_algo}") + + block_hashes: list[list[BlockHash]] = [] + for token_list in token_lists: + hash_list = [] + # in KvCacheManager, the last token is not included in the block key + for t in range(0, len(token_list) - 1, self._tokens_per_block): + t_end = min(t + self._tokens_per_block, len(token_list) - 1) + hash_list.append( + block_hasher( + token_list[t:t_end], None if t == 0 else hash_list[-1], cache_salt_id + ) + ) + block_hashes.append(hash_list) + return block_hashes + + def _tokenize_and_compute_block_hashes( + self, request: OpenAIRequest + ) -> tuple[list[list[int]], list[list[int]]]: + """Synchronous tokenize + block-hash, combined for thread offload. + + Factored into one method so ``get_next_server`` can offload the whole + CPU-bound step via ``asyncio.to_thread`` in a single call, keeping + the orchestrator's asyncio event loop free to dispatch other + requests in parallel. + """ + token_lists = self._tokenize(request) + block_hashes = self._compute_block_hashes(token_lists) + return token_lists, block_hashes + + def _tokenize_and_compute_block_hashes_with_salt( + self, + request: OpenAIRequest, + cache_salt_id: Optional[int] = None, + ) -> tuple[list[list[int]], list[list[int]]]: + token_lists = self._tokenize(request) + block_hashes = self._compute_block_hashes(token_lists, cache_salt_id=cache_salt_id) + return token_lists, block_hashes + + def _tokenize_and_compute_block_hashes_by_algo( + self, + request: OpenAIRequest, + hash_algos: Iterable[str], + cache_salt_id: Optional[int] = None, + ) -> tuple[list[list[int]], dict[str, list[list[BlockHash]]]]: + """Synchronous tokenize + per-algorithm block hashes for thread offload.""" + token_lists = self._tokenize(request) + return token_lists, { + hash_algo: self._compute_block_hashes( + token_lists, hash_algo, cache_salt_id=cache_salt_id + ) + for hash_algo in set(hash_algos) + } + + @staticmethod + def _text_to_int_sequences(texts: list[str]) -> list[list[int]]: + """Convert text strings to lists of unicode code points. + + Usable as input to ``_compute_block_hashes``. + """ + return [[ord(c) for c in text] for text in texts] + + @staticmethod + def _get_request_cache_salt_id(request: OpenAIRequest) -> Optional[int]: + cache_salt = getattr(request, "cache_salt", None) + return None if cache_salt is None else get_cache_salt_id(cache_salt) diff --git a/tests/unittest/disaggregated/test_coordinator_e2e.py b/tests/unittest/disaggregated/test_coordinator_e2e.py new file mode 100644 index 000000000000..3fdae2cfedb8 --- /dev/null +++ b/tests/unittest/disaggregated/test_coordinator_e2e.py @@ -0,0 +1,430 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end coordinator/worker disagg serving with mocked ctx/gen workers. + +CPU-only, MPI-free, single-process (three uvicorn threads): + + * mocked ctx + gen HTTP workers serve ``/health`` + ``/v1/completions`` + (ctx returns a context_only response with disaggregated_params so the disagg + server proceeds to gen; gen returns the final completion text), + * a real ``CoordinatorServer`` (wrapping a ``DisaggCoordinatorService``) runs on + an internal port -- the gen router is a *stateful* conversation router, so gen + placement is delegated to it via ``/select``; the ctx router is round-robin + (placed locally in the disagg server), + * a real ``OpenAIDisaggServer`` in worker mode (``coordinator_url`` set, so it + holds a ``CoordinatorClient``) serves the public ``/v1/completions``. + +A real HTTP completion is sent to the disagg server and must round-trip +ctx -> (coordinator /select) -> gen, returning the gen worker's text. This +exercises the whole chain including the coordinator HTTP hop. +""" + +import asyncio +import os +import subprocess +import sys +import tempfile +import threading +import time + +import aiohttp +import pytest +import uvicorn +import yaml +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, Response + +from tensorrt_llm.llmapi.disagg_utils import ( + CtxGenServerConfig, + DisaggServerConfig, + RouterConfig, + ServerRole, +) +from tensorrt_llm.logger import logger +from tensorrt_llm.serve.coordinator_server import CoordinatorServer +from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService +from tensorrt_llm.serve.openai_client import OpenAIHttpClient +from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer + +GEN_TEXT = "HELLO_FROM_GEN" + +# The uvicorn worker threads / CLI-output pump thread are background threads that +# outlive a strict thread snapshot; exempt this module (same as the other e2e). +pytestmark = pytest.mark.threadleak(enabled=False) + + +@pytest.fixture(autouse=True) +def _reset_prometheus_registry(): + """Role-prefixed Prometheus counters are registered in the global default + registry; clear it between tests so a second server build in the same pytest + process does not hit duplicate-timeseries errors. + """ + from prometheus_client import REGISTRY + + yield + for collector in list(REGISTRY._collector_to_names): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +def _free_port(): + import socket + + s = socket.socket() + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +class _UvicornThread: + """Run a FastAPI app in a background uvicorn server thread.""" + + def __init__(self, app, port): + self.port = port + self._server = uvicorn.Server( + uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning") + ) + self._thread = threading.Thread(target=self._server.run, daemon=True) + + def __enter__(self): + self._thread.start() + for _ in range(100): + if self._server.started: + break + time.sleep(0.1) + return self + + def __exit__(self, *a): + self._server.should_exit = True + self._thread.join(timeout=10) + + +def _mock_worker_app(role: str) -> FastAPI: + """A ctx or gen worker: /health + /server_info + /v1/completions.""" + app = FastAPI() + + @app.get("/health") + async def health(): + return Response(status_code=200) + + @app.get("/server_info") + async def server_info(): + return JSONResponse({"kv_cache_hash_algo": "v1"}) + + @app.post("/v1/completions") + async def completions(raw: Request): + body = await raw.json() + dp = body.get("disaggregated_params") or {} + model = body.get("model", "m") + if dp.get("request_type") == "context_only": + # Context phase: return disagg params so the disagg server proceeds + # to the gen worker (finish_reason "length" => needs generation). + rid = dp.get("disagg_request_id") + return JSONResponse( + { + "id": "cmpl-ctx", + "object": "text_completion", + "created": 0, + "model": model, + "prompt_token_ids": [1, 2, 3], + "choices": [ + { + "index": 0, + "text": "", + "finish_reason": "length", + "disaggregated_params": { + "request_type": "context_only", + "ctx_request_id": rid, + "disagg_request_id": rid, + }, + } + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 0, "total_tokens": 3}, + } + ) + # Generation phase: final answer. + return JSONResponse( + { + "id": "cmpl-gen", + "object": "text_completion", + "created": 0, + "model": model, + "choices": [{"index": 0, "text": GEN_TEXT, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + } + ) + + return app + + +class _ReadinessClient: + """Minimal client the coordinator uses only for server readiness probing. + + In coordinator/worker mode the coordinator never sends completions (the disagg + servers do), so its readiness client needs no metrics -- reusing the real + ``check_ready_for_servers`` keeps the probe faithful while avoiding a second + set of role-prefixed Prometheus counters in this single-process test. + """ + + def __init__(self, router): + self._router = router + self._session = aiohttp.ClientSession() + + async def check_ready(self): + ready, unready = await OpenAIHttpClient.check_ready_for_servers( + self._session, self._router.servers + ) + if ready: + await self._router.prepare_servers(ready) + return ready, unready + + async def shutdown(self): + await self._session.close() + + +def _make_config(ctx_url, gen_url, public_port): + host_port = lambda u: (u.split(":")[0], int(u.split(":")[1])) + ctx_host, ctx_port = host_port(ctx_url) + gen_host, gen_port = host_port(gen_url) + return DisaggServerConfig( + server_configs=[ + CtxGenServerConfig(type="ctx", hostname=ctx_host, port=ctx_port), + CtxGenServerConfig(type="gen", hostname=gen_host, port=gen_port), + ], + hostname="127.0.0.1", + port=public_port, + # ctx: stateless (placed locally in the disagg server); + # gen: stateful conversation router (placement delegated to coordinator). + ctx_router_config=RouterConfig(type="round_robin", server_role=ServerRole.CONTEXT), + gen_router_config=RouterConfig(type="conversation", server_role=ServerRole.GENERATION), + ) + + +class _CoordinatorThread: + """Run a CoordinatorServer (DisaggCoordinatorService) in a uvicorn thread.""" + + def __init__(self, config): + self.port = _free_port() + self.url = f"http://127.0.0.1:{self.port}" + # The coordinator builds its own owner routers from config. + self._coordinator = DisaggCoordinatorService( + config, client_factory=lambda router, role, mr=1: _ReadinessClient(router) + ) + self._impl = _UvicornThread(CoordinatorServer(self._coordinator).app, self.port) + + def __enter__(self): + self._impl.__enter__() + return self + + def __exit__(self, *a): + self._impl.__exit__(*a) + + +async def _wait_healthy(url, timeout_s=30.0): + deadline = time.time() + timeout_s + async with aiohttp.ClientSession() as sess: + while time.time() < deadline: + try: + async with sess.get(f"{url}/health", timeout=1) as r: + if r.status == 200: + return True + except Exception: + pass + await asyncio.sleep(0.2) + return False + + +def test_disagg_completion_e2e_through_coordinator(): + with ( + _UvicornThread(_mock_worker_app("ctx"), _free_port()) as ctx, + _UvicornThread(_mock_worker_app("gen"), _free_port()) as gen, + ): + ctx_url = f"127.0.0.1:{ctx.port}" + gen_url = f"127.0.0.1:{gen.port}" + public_port = _free_port() + config = _make_config(ctx_url, gen_url, public_port) + + with _CoordinatorThread(config) as coord: + assert asyncio.run(_wait_healthy(coord.url)), "coordinator never became healthy" + + disagg = OpenAIDisaggServer(config=config, coordinator_url=coord.url) + with _UvicornThread(disagg.app, public_port) as server: + base = f"http://127.0.0.1:{server.port}" + assert asyncio.run(_wait_healthy(base)), "disagg server never became healthy" + + async def drive(): + async with aiohttp.ClientSession() as sess: + payload = {"model": "m", "prompt": "hello", "max_tokens": 8} + # X-Session-ID -> conversation_id, so the gen router + # (conversation) delegates placement to the coordinator. + headers = {"X-Session-ID": "conv-e2e"} + async with sess.post( + f"{base}/v1/completions", json=payload, headers=headers, timeout=30 + ) as r: + assert r.status == 200, await r.text() + return await r.json() + + body = asyncio.run(drive()) + + # The full ctx -> coordinator/select -> gen chain returned the gen text. + assert body["choices"][0]["text"] == GEN_TEXT, body + assert body["choices"][0]["finish_reason"] == "stop" + + +def _write_config(path, ctx_url, gen_url, public_port): + """A disagg config YAML: round-robin ctx, conversation gen (delegated).""" + cfg = { + "hostname": "127.0.0.1", + "port": public_port, + "context_servers": { + "num_instances": 1, + "urls": [ctx_url], + "router": {"type": "round_robin"}, + }, + "generation_servers": { + "num_instances": 1, + "urls": [gen_url], + "router": {"type": "conversation"}, + }, + } + with open(path, "w") as f: + yaml.safe_dump(cfg, f) + + +def test_disagg_completion_e2e_web_concurrency_4(): + """WEB_CONCURRENCY=4 through the real CLI: `trtllm-serve disaggregated` + forks a coordinator (port-1) + a uvicorn fleet of 4 disagg servers on the + public port. Mock ctx/gen workers run in this test process; a real HTTP + completion round-trips through one of the four workers -> coordinator -> gen. + """ + logger.set_level("info") # trtllm logger defaults to "error"; show progress + WORKERS = 4 + with ( + _UvicornThread(_mock_worker_app("ctx"), _free_port()) as ctx, + _UvicornThread(_mock_worker_app("gen"), _free_port()) as gen, + ): + ctx_url = f"127.0.0.1:{ctx.port}" + gen_url = f"127.0.0.1:{gen.port}" + # port-1 is the coordinator, so pick a public port with room below it. + public_port = _free_port() + coord_port = public_port - 1 + + with tempfile.TemporaryDirectory() as td: + cfg_path = os.path.join(td, "disagg.yaml") + _write_config(cfg_path, ctx_url, gen_url, public_port) + + env = dict(os.environ) + env["WEB_CONCURRENCY"] = str(WORKERS) + # Unbuffered so the child's launch logs stream out live (else stdout + # to a pipe is block-buffered and nothing shows until it exits). + env["PYTHONUNBUFFERED"] = "1" + # trtllm logger defaults to "error"; raise it so the coordinator/fleet + # launch logs are visible in the streamed [cli] output. + env["TLLM_LOG_LEVEL"] = "info" + # A plain HTTP fleet, never an MPI rank -- strip any launcher env so + # the CLI's own strip is not even relied upon. + for k in list(env): + if k.startswith( + ("SLURM_", "PMIX_", "PMI_", "OMPI_", "UCX_", "I_MPI_", "HYDRA_", "MPI_") + ): + env.pop(k) + + logger.info( + f"mock ctx={ctx_url} gen={gen_url}; launching " + f"`trtllm-serve disaggregated` WEB_CONCURRENCY={WORKERS}, " + f"public={public_port} coordinator={coord_port}" + ) + proc = subprocess.Popen( + [ + sys.executable, + "-m", + "tensorrt_llm.commands.serve", + "disaggregated", + "-c", + cfg_path, + ], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + start_new_session=True, + ) + + # Stream the CLI child's stdout live (prefixed) so coordinator/fleet + # startup is visible in real time instead of only at teardown. + def _pump(): + for line in proc.stdout: + logger.info(f"[cli] {line.rstrip()}") + + pump = threading.Thread(target=_pump, daemon=True) + pump.start() + try: + base = f"http://127.0.0.1:{public_port}" + coord = f"http://127.0.0.1:{coord_port}" + + async def _wait_all(): + # Both the coordinator (port-1) and the public fleet must be + # up before the fleet reports ready (fleet is_ready proxies + # the coordinator). + logger.info("waiting for coordinator health...") + assert await _wait_healthy(coord, 120.0), "coordinator never became healthy" + logger.info("coordinator healthy; waiting for fleet health...") + assert await _wait_healthy(base, 120.0), "disagg fleet never became healthy" + logger.info("fleet healthy") + + asyncio.run(_wait_all()) + + async def drive(): + # Fire several requests so the kernel spreads them across the + # 4 uvicorn workers. Every one must round-trip to GEN_TEXT. + async with aiohttp.ClientSession() as sess: + texts = [] + for i in range(8): + payload = {"model": "m", "prompt": f"hello-{i}", "max_tokens": 8} + headers = {"X-Session-ID": f"conv-{i}"} + async with sess.post( + f"{base}/v1/completions", json=payload, headers=headers, timeout=30 + ) as r: + assert r.status == 200, await r.text() + texts.append((await r.json())["choices"][0]["text"]) + logger.info(f"request {i} -> {texts[-1]!r}") + return texts + + texts = asyncio.run(drive()) + assert all(t == GEN_TEXT for t in texts), texts + logger.info(f"all {len(texts)} requests round-tripped to GEN_TEXT") + finally: + # Kill the whole process group: the CLI parent + coordinator + + # all uvicorn workers. Terminating only proc leaves the workers + # holding the stdout pipe open, so _pump never sees EOF. + logger.info("terminating CLI process group") + import signal + + pgid = os.getpgid(proc.pid) + os.killpg(pgid, signal.SIGTERM) + try: + proc.wait(timeout=15) + except subprocess.TimeoutExpired: + os.killpg(pgid, signal.SIGKILL) + proc.wait() + pump.join(timeout=10) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v", "-s"])) diff --git a/tests/unittest/disaggregated/test_coordinator_worker.py b/tests/unittest/disaggregated/test_coordinator_worker.py new file mode 100644 index 000000000000..37ba80c6049a --- /dev/null +++ b/tests/unittest/disaggregated/test_coordinator_worker.py @@ -0,0 +1,287 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Coordinator/worker disagg routing: cross-process placement contract. + +CPU-only, MPI-free. Wires the real coordinator surface to a real worker-side +coordinator: + + * fake ctx/gen HTTP workers answer ``/health`` (readiness only), + * a real ``CoordinatorServer`` (wrapping a ``DisaggCoordinatorService`` over the + configured routers) runs in a uvicorn thread on an internal port, + * a ``CoordinatorClient`` (what a worker holds) wraps only *stateful* routers + in a ``CoordinatorDelegatingRouter`` whose ``get_next_server`` computes the + routing key locally and POSTs it to the coordinator's ``/select``; + ``finish_request`` releases coordinator-side state via ``/finish`` and the + returned handle. *Stateless* routers (round_robin) are used as-is and place + locally in the worker. + +This proves the routing split: stateful routers (conversation, kv_cache_aware) +delegate to the coordinator via ``routing_key`` + ``get_next_server_by_key``, +while stateless routers never touch the coordinator. +""" + +import asyncio +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import aiohttp +import pytest +import uvicorn + +from tensorrt_llm.llmapi.disagg_utils import ( + CtxGenServerConfig, + DisaggServerConfig, + RouterConfig, + ServerRole, +) +from tensorrt_llm.serve.coordinator_server import CoordinatorServer +from tensorrt_llm.serve.disagg_coordinator import CoordinatorClient, DisaggCoordinatorService +from tensorrt_llm.serve.openai_protocol import CompletionRequest, DisaggregatedParams + + +@pytest.fixture(autouse=True) +def _reset_prometheus_registry(): + """Each coordinator builds role-prefixed Prometheus counters in the global + default registry (via its readiness OpenAIHttpClients). In production the + coordinator is a single process; here two coordinators share one pytest + process, so clear the registry between tests to avoid duplicate-timeseries + registration errors. + """ + from prometheus_client import REGISTRY + + yield + for collector in list(REGISTRY._collector_to_names): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +def _free_port(): + import socket + + s = socket.socket() + # SO_REUSEADDR so a port left in TIME_WAIT by a sibling server in the same + # suite can be rebound immediately (closes the alloc->bind race window). + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + s.close() + return port + + +class _FakeWorker: + """Minimal HTTP worker: /health -> 200 (used for readiness only).""" + + def __init__(self): + self.port = _free_port() + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *a): + pass + + def do_GET(self): + if self.path == "/health": + self.send_response(200) + else: + self.send_response(404) + self.end_headers() + + self._httpd = ThreadingHTTPServer(("127.0.0.1", self.port), Handler) + self._thread = threading.Thread(target=self._httpd.serve_forever, daemon=True) + + @property + def url(self): + return f"127.0.0.1:{self.port}" + + def __enter__(self): + self._thread.start() + return self + + def __exit__(self, *a): + self._httpd.shutdown() + + +def _make_config(ctx_urls, gen_urls, ctx_router_type, gen_router_type): + server_configs = [ + CtxGenServerConfig(type="ctx", hostname=u.split(":")[0], port=int(u.split(":")[1])) + for u in ctx_urls + ] + [ + CtxGenServerConfig(type="gen", hostname=u.split(":")[0], port=int(u.split(":")[1])) + for u in gen_urls + ] + return DisaggServerConfig( + server_configs=server_configs, + ctx_router_config=RouterConfig(type=ctx_router_type, server_role=ServerRole.CONTEXT), + gen_router_config=RouterConfig(type=gen_router_type, server_role=ServerRole.GENERATION), + ) + + +def _client_factory(router, role, max_retries=1): + from tensorrt_llm.serve.openai_client import OpenAIHttpClient + + return OpenAIHttpClient(router, role, 30, max_retries) + + +class _CoordinatorThread: + """Run a CoordinatorServer (DisaggCoordinatorService) in a background thread.""" + + def __init__(self, config): + self.port = _free_port() + self.url = f"http://127.0.0.1:{self.port}" + # The coordinator builds its own owner routers from config. + self._cluster = DisaggCoordinatorService(config, _client_factory) + self._server = uvicorn.Server( + uvicorn.Config( + CoordinatorServer(self._cluster).app, + host="127.0.0.1", + port=self.port, + log_level="warning", + ) + ) + self._thread = threading.Thread(target=self._server.run, daemon=True) + + def __enter__(self): + self._thread.start() + for _ in range(100): + if self._server.started: + break + time.sleep(0.1) + return self + + def __exit__(self, *a): + self._server.should_exit = True + self._thread.join(timeout=10) + + +async def _wait_coord_ready(url, timeout_s=30.0): + deadline = time.time() + timeout_s + async with aiohttp.ClientSession() as sess: + while time.time() < deadline: + try: + async with sess.get(f"{url}/health", timeout=1) as r: + if r.status == 200: + return True + except Exception: + pass + await asyncio.sleep(0.2) + return False + + +def test_stateless_router_places_locally_in_worker(): + """A round-robin (stateless) router is NOT wrapped: the worker places locally + with the real router and never calls the coordinator. + """ + from tensorrt_llm.serve.router import CoordinatorDelegatingRouter, RoundRobinRouter + + with _FakeWorker() as ctx0, _FakeWorker() as gen0, _FakeWorker() as gen1: + config = _make_config([ctx0.url], [gen0.url, gen1.url], "round_robin", "round_robin") + with _CoordinatorThread(config) as coord: + assert asyncio.run(_wait_coord_ready(coord.url)), "coordinator never became healthy" + + async def drive(): + remote = CoordinatorClient(coord.url, config) + # Stateless -> real local router, not a delegating proxy. + assert isinstance(remote.gen_router, RoundRobinRouter) + assert not isinstance(remote.gen_router, CoordinatorDelegatingRouter) + picks = [] + for _ in range(4): + req = CompletionRequest(model="m", prompt="hello") + server, _info = await remote.gen_router.get_next_server(req) + picks.append(server) + await remote.gen_router.finish_request(req) + await remote.stop() + return picks + + picks = asyncio.run(drive()) + assert set(picks) == {gen0.url, gen1.url}, ( + f"local round-robin should hit both gen workers, got {picks}" + ) + + +def test_conversation_coordinator_sticky_by_conv_id(): + """Same conversation_id sticks to one gen worker; a stateful (conversation) + router delegates placement to the coordinator via /select. + """ + from tensorrt_llm.serve.router import CoordinatorDelegatingRouter + + with _FakeWorker() as ctx0, _FakeWorker() as gen0, _FakeWorker() as gen1: + config = _make_config([ctx0.url], [gen0.url, gen1.url], "round_robin", "conversation") + with _CoordinatorThread(config) as coord: + assert asyncio.run(_wait_coord_ready(coord.url)) + + def _req(conv_id): + return CompletionRequest( + model="m", + prompt="hi", + disaggregated_params=DisaggregatedParams( + request_type="context_only", conversation_id=conv_id + ), + ) + + async def drive(): + remote = CoordinatorClient(coord.url, config) + # Stateful -> wrapped in a coordinator-delegating router. + assert isinstance(remote.gen_router, CoordinatorDelegatingRouter) + assert await remote.is_ready() is True + first, _ = await remote.gen_router.get_next_server(_req("conv-A")) + # Repeated conv-A requests must land on the same worker. + repeats = [] + for _ in range(3): + s, _ = await remote.gen_router.get_next_server(_req("conv-A")) + repeats.append(s) + await remote.stop() + return first, repeats + + first, repeats = asyncio.run(drive()) + assert all(s == first for s in repeats), ( + f"conv-A must be sticky, got first={first} repeats={repeats}" + ) + + +def test_coordinator_owns_generation_disagg_request_id(): + """Generation routing receives its request ID from the coordinator.""" + from tensorrt_llm.serve.router import CoordinatorDelegatingRouter + + with _FakeWorker() as ctx0, _FakeWorker() as gen0: + config = _make_config([ctx0.url], [gen0.url], "round_robin", "conversation") + with _CoordinatorThread(config) as coord: + assert asyncio.run(_wait_coord_ready(coord.url)) + + async def drive(): + remote = CoordinatorClient(coord.url, config) + assert isinstance(remote.gen_router, CoordinatorDelegatingRouter) + request = CompletionRequest( + model="m", + prompt="hello", + disaggregated_params=DisaggregatedParams( + request_type="generation_only", + ctx_request_id=123, + disagg_request_id=None, + conversation_id="conv-A", + ), + ) + await remote.gen_router.get_next_server(request) + assigned_id = request.disaggregated_params.disagg_request_id + await remote.gen_router.finish_request(request) + await remote.stop() + return assigned_id + + assigned_id = asyncio.run(drive()) + assert assigned_id is not None and assigned_id != 123 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v", "-s"])) diff --git a/tests/unittest/disaggregated/test_disagg_openai_client.py b/tests/unittest/disaggregated/test_disagg_openai_client.py index edc09809732b..cead306ae3ff 100644 --- a/tests/unittest/disaggregated/test_disagg_openai_client.py +++ b/tests/unittest/disaggregated/test_disagg_openai_client.py @@ -376,7 +376,11 @@ def _make_client(self, session, **kwargs): async def test_retry_regenerates_disagg_id(self): session = AsyncMock(spec=aiohttp.ClientSession) ids = iter(range(1000, 2000)) - client = self._make_client(session, disagg_id_generator=lambda: next(ids)) + + async def next_id(): + return next(ids) + + client = self._make_client(session, disagg_id_generator=next_id) session.post.side_effect = [ aiohttp.ClientError("transient"), diff --git a/tests/unittest/disaggregated/test_disagg_utils.py b/tests/unittest/disaggregated/test_disagg_utils.py index a086108187ad..35fe98f7676c 100644 --- a/tests/unittest/disaggregated/test_disagg_utils.py +++ b/tests/unittest/disaggregated/test_disagg_utils.py @@ -200,31 +200,48 @@ def test_get_server_configs_dict(): ids=["multithread", "singlethread"]) def test_get_global_disagg_request_id(multithread): iter = 10000 - node_ids = list(range(10)) - thread_num = len(node_ids) - - def get_ids(node_ids): - all_node_ids = [[] for _ in range(len(node_ids))] + # (node_id, process_id) pairs — the pair uniquely identifies a fleet worker. + # Mix of distinct nodes and distinct processes on the same node. + worker_ids = [(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (3, 5), (7, 2), + (10, 0), (10, 3), (255, 63)] + thread_num = len(worker_ids) + + def get_ids(worker_ids): + all_ids = [[] for _ in range(len(worker_ids))] for i in range(iter): if i % (4000 // thread_num) == 0: time.sleep(0.001) - for i, node_id in enumerate(node_ids): - all_node_ids[i].append(get_global_disagg_request_id(node_id)) - return all_node_ids + for j, (node_id, process_id) in enumerate(worker_ids): + all_ids[j].append( + get_global_disagg_request_id(node_id, process_id)) + return all_ids if multithread: - with ThreadPoolExecutor(max_workers=len(node_ids)) as executor: - all_node_ids = [ - ids[0] for ids in executor.map(get_ids, [[i] for i in node_ids]) + with ThreadPoolExecutor(max_workers=len(worker_ids)) as executor: + all_worker_ids = [ + ids[0] + for ids in executor.map(get_ids, [[w] for w in worker_ids]) ] else: - all_node_ids = get_ids(node_ids) + all_worker_ids = get_ids(worker_ids) - all_ids = set(i for ids in all_node_ids for i in ids) - assert len(all_ids) == iter * len(node_ids) + all_ids = set(i for ids in all_worker_ids for i in ids) + # Each (node_id, process_id) worker's ids must be globally unique across all. + assert len(all_ids) == iter * len(worker_ids) assert all(id >= MIN_GLOBAL_ID and id < ((1 << 63) - 1) for id in all_ids) +def test_get_global_disagg_request_id_range_validation(): + # node_id: 8 bits [0,256); process_id: 6 bits [0,64). + with pytest.raises(ValueError): + get_global_disagg_request_id(256, 0) + with pytest.raises(ValueError): + get_global_disagg_request_id(0, 64) + # valid extremes + assert get_global_disagg_request_id(255, 63) >= MIN_GLOBAL_ID + assert get_global_disagg_request_id(0, 0) >= MIN_GLOBAL_ID + + def test_get_local_request_id(): last_id = MIN_GLOBAL_ID - 100 ids = set() diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 7cc7fe4a649e..52d42e78fcae 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -29,6 +29,7 @@ ServerRole, ) from tensorrt_llm.serve.disagg_auto_scaling import DisaggClusterManager, WorkerInfo +from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService from tensorrt_llm.serve.openai_disagg_service import OpenAIDisaggregatedService from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, @@ -63,11 +64,18 @@ def _client_factory(*_args, **_kwargs): def _make_service(schedule_style: str) -> OpenAIDisaggregatedService: config = DisaggServerConfig(server_configs=[], schedule_style=schedule_style) + # The coordinator builds its own (empty) routers from config; override them + # with mocks so tests can stub placement / readiness directly. + cluster = DisaggCoordinatorService(config, client_factory=_client_factory) ctx_router = AsyncMock(spec=Router) gen_router = AsyncMock(spec=Router) - return OpenAIDisaggregatedService( - config, ctx_router, gen_router, client_factory=_client_factory - ) + cluster._ctx_router = ctx_router + cluster._gen_router = gen_router + service = OpenAIDisaggregatedService(config, cluster, client_factory=_client_factory) + # Convenience handles for tests that stub placement / readiness directly. + service._ctx_router = ctx_router + service._gen_router = gen_router + return service def _make_completion_response( @@ -213,16 +221,18 @@ async def test_is_ready_waits_for_router_preparation(): ), AsyncMock(), ) - service._disagg_cluster_manager = cluster_manager + # Readiness now lives on the DisaggCoordinatorService the service holds. + local = service._cluster + local._disagg_cluster_manager = cluster_manager cluster_manager._current_ctx_workers["ctx"] = WorkerInfo( worker_id="ctx", role=ServerRole.CONTEXT ) - service._ctx_router = SimpleNamespace(num_prepared_servers=0) - service._gen_router = SimpleNamespace(num_prepared_servers=1) + local._ctx_router = SimpleNamespace(num_prepared_servers=0) + local._gen_router = SimpleNamespace(num_prepared_servers=1) assert await service.is_ready() is False - service._ctx_router.num_prepared_servers = 1 + local._ctx_router.num_prepared_servers = 1 assert await service.is_ready() is False cluster_manager._current_gen_workers["gen"] = WorkerInfo( From ba2345b7bd9e508a1104e88251febeb1a1f01dbd Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 9 Jul 2026 02:02:13 +0000 Subject: [PATCH 02/18] [None][fix] address coordinator review feedback Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- docs/source/features/disagg-serving.md | 6 +- tensorrt_llm/commands/serve.py | 95 +++-- tensorrt_llm/llmapi/disagg_utils.py | 22 +- tensorrt_llm/llmapi/llm_args.py | 16 - tensorrt_llm/serve/coordinator_server.py | 55 ++- tensorrt_llm/serve/disagg_coordinator.py | 26 +- tensorrt_llm/serve/openai_disagg_server.py | 31 +- tensorrt_llm/serve/openai_disagg_service.py | 16 +- tensorrt_llm/serve/openai_server.py | 27 -- tensorrt_llm/serve/perf_metrics.py | 1 + tensorrt_llm/serve/responses_utils.py | 108 ++--- tensorrt_llm/serve/router.py | 397 ++++-------------- tensorrt_llm/serve/router_utils.py | 11 +- .../disaggregated/test_coordinator_e2e.py | 28 +- .../disaggregated/test_coordinator_worker.py | 77 +++- .../disaggregated/test_disagg_utils.py | 29 +- .../test_openai_disagg_service.py | 2 +- 17 files changed, 430 insertions(+), 517 deletions(-) diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index 9fb7926a0adb..48b35dac5ede 100644 --- a/docs/source/features/disagg-serving.md +++ b/docs/source/features/disagg-serving.md @@ -98,7 +98,7 @@ The context and generation phases of one request must share a single request ID: The disaggregated server generates this ID itself as a **snowflake** — a self-contained 64-bit positive integer that is unique without any cross-process coordination. The bit layout is: -``` +```text [ 0 (1 bit) | timestamp_ms (39 bits) | node_id (8 bits) | process_id (6 bits) | counter (10 bits) ] ``` @@ -185,7 +185,7 @@ Once the context and generation servers are launched, you can launch the disaggr server, which will accept requests from clients and do the orchestration between context and generation servers. The disaggregated server can be launched with: -``` +```bash trtllm-serve disaggregated -c disagg_config.yaml ``` where `disagg_config.yaml` contains information about the context and generation servers. For the current example, @@ -295,7 +295,7 @@ generation_servers: Launch it exactly as before — the coordinator and fleet are started for you: -``` +```bash trtllm-serve disaggregated -c disagg_config.yaml ``` diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index f96b0448e44a..e17e53e82217 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -120,8 +120,7 @@ def _signal_handler_cleanup_child(signum, frame): def is_non_default_or_required(param_name, value, backend, explicit_cli_keys): - """ - Check if a parameter should be explicitly included in llm_args. + """Check if a parameter should be explicitly included in llm_args. Returns True if parameter is either: 1. Always required (core params that must be present), OR @@ -436,8 +435,7 @@ def launch_grpc_server(host: str, port: int, llm_args: dict, served_model_name: Optional[str] = None): - """ - Launch a gRPC server for TensorRT-LLM. + """Launch a gRPC server for TensorRT-LLM. This provides a high-performance gRPC interface designed for external routers (e.g., sgl-router) using pre-tokenized input and raw token ID output. @@ -1624,7 +1622,6 @@ def disaggregated( schedule_style: str, ): """Running server in disaggregated mode""" - logger.set_level(log_level) if metrics_log_interval != 0: @@ -1713,12 +1710,14 @@ def disaggregated( def _launch_disagg_fleet(disagg_cfg, config_file, metadata_server_config_file, request_timeout, server_start_timeout, num_workers, coordinator_url): - """Launch ``num_workers`` delegating disagg servers (each its own ``Popen`` - running ``_run_fleet_worker``) on the shared public port, pointed at - ``coordinator_url``. Separate Popens (not ``uvicorn --workers N``) so each gets - an explicit per-worker TLLM_DISAGG_WORKER_PROCESS_ID and binds its own - SO_REUSEPORT socket for kernel-balanced connections. MPI/PMIX/SLURM env is - stripped (workers are plain HTTP processes). Returns the list of Popen handles. + """Launch delegating disaggregated-server workers. + + Each worker has its own ``Popen`` running ``_run_fleet_worker`` on the shared + public port and pointing at ``coordinator_url``. Separate processes allow an + explicit per-worker ``TLLM_DISAGG_WORKER_PROCESS_ID`` and independent + ``SO_REUSEPORT`` sockets. MPI/PMIX/SLURM environment variables are stripped. + + Returns the list of ``Popen`` handles. """ from tensorrt_llm.llmapi.disagg_utils import disagg_process_id_space public_host, public_port = disagg_cfg.hostname, disagg_cfg.port @@ -1755,18 +1754,19 @@ def _launch_disagg_fleet(disagg_cfg, config_file, metadata_server_config_file, f"{public_host}:{public_port}, coordinator={coordinator_url}") procid_space = disagg_process_id_space() + if not 1 <= num_workers <= procid_space: + raise ValueError( + f"num_workers must be between 1 and {procid_space}, got {num_workers}" + ) fleet = [] for i in range(num_workers): worker_env = dict(base_env) - # Explicit per-worker process index (no shared counter file). Wrap into - # the snowflake process_id space so ids stay in the 6-bit field. - worker_env[DisaggWorkerEnvs.TLLM_DISAGG_WORKER_PROCESS_ID] = str( - i % procid_space) + # Explicit per-worker process index (no shared counter file). + worker_env[DisaggWorkerEnvs.TLLM_DISAGG_WORKER_PROCESS_ID] = str(i) p = subprocess.Popen(cmd, env=worker_env, stdout=sys.stdout, - stderr=sys.stderr, - start_new_session=True) + stderr=sys.stderr) logger.info(f"Disagg fleet worker {i} launched (pid={p.pid})") fleet.append(p) @@ -1842,9 +1842,9 @@ def _serve_coordinator_and_fleet(disagg_cfg, config_file, # 1. Launch the delegating fleet pointed at the implicit coordinator we start # below (port-1 for TCP; UDS for the hot path). Workers hold # CoordinatorClients (no core), so they can't race the ZMQ ingest bind. - _launch_disagg_fleet(disagg_cfg, config_file, metadata_server_config_file, - request_timeout, server_start_timeout, num_workers, - coord_url) + fleet = _launch_disagg_fleet(disagg_cfg, config_file, + metadata_server_config_file, request_timeout, + server_start_timeout, num_workers, coord_url) # 2. Build + serve the coordinator in this process. It OWNS routing state and # builds the owner routers itself (single shared namespace-aware core + ONE @@ -1860,8 +1860,35 @@ def _client_factory(router, role, max_retries=1): server_start_timeout_secs=server_start_timeout) logger.info(f"Coordinator serving on {public_host}:{coord_port} " f"(uds={coord_uds}) (fleet on public port {public_port})") - asyncio.run( - CoordinatorServer(coordinator)(public_host, coord_port, uds=coord_uds)) + + async def _serve_and_monitor(): + server_task = asyncio.create_task( + CoordinatorServer(coordinator)(public_host, + coord_port, + uds=coord_uds)) + + async def _monitor_fleet(): + while True: + for i, process in enumerate(fleet): + return_code = process.poll() + if return_code is not None: + if return_code != 0: + raise RuntimeError( + f"Disagg fleet worker {i} (pid={process.pid}) " + f"exited with code {return_code}") + return + await asyncio.sleep(1) + + monitor_task = asyncio.create_task(_monitor_fleet()) + done, pending = await asyncio.wait((server_task, monitor_task), + return_when=asyncio.FIRST_COMPLETED) + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + for task in done: + task.result() + + asyncio.run(_serve_and_monitor()) def _init_fleet_worker_process(): @@ -1869,7 +1896,8 @@ def _init_fleet_worker_process(): Restores logging/GC state for the fresh ``python -c`` interpreter and tags every log line with this worker's PID so the shared-stdout fleet output stays - attributable.""" + attributable. + """ # A worker is a plain HTTP process, never an MPI rank; drop WEB_CONCURRENCY so # it is never itself re-forked into multiple uvicorn workers. os.environ.pop("WEB_CONCURRENCY", None) @@ -1901,7 +1929,8 @@ def _build_disagg_server_from_env() -> "OpenAIDisaggServer": the server holds a remote ``CoordinatorClient`` so routing/readiness are delegated to the coordinator. The worker's process index (for the snowflake disagg-id) is read from ``TLLM_DISAGG_WORKER_PROCESS_ID``, which the launcher - set explicitly per worker (see ``worker_local_process_id``).""" + set explicitly per worker (see ``worker_local_process_id``). + """ config_file = os.environ[DisaggWorkerEnvs.TLLM_DISAGG_CONFIG_FILE] coordinator_url = os.environ[DisaggWorkerEnvs.TLLM_DISAGG_COORDINATOR_URL] metadata_config_file = os.environ.get( @@ -1926,20 +1955,21 @@ def _build_disagg_server_from_env() -> "OpenAIDisaggServer": def create_disagg_server_app(): - """uvicorn import-string factory: build one disagg server's FastAPI app. + """Uvicorn import-string factory: build one disagg server's FastAPI app. Retained for the ``uvicorn --factory`` entry point; the SO_REUSEPORT fleet - launcher uses ``_run_fleet_worker`` instead.""" + launcher uses ``_run_fleet_worker`` instead. + """ _init_fleet_worker_process() return _build_disagg_server_from_env().app def _run_fleet_worker(): - """Entry point for a single fleet worker (one OS process per worker), launched - as its own ``Popen`` so it takes an explicit TLLM_DISAGG_WORKER_PROCESS_ID and - binds its own SO_REUSEPORT socket to the shared public port (the kernel then - 4-tuple-hashes connections across workers instead of the unfair shared-socket - accept() of ``uvicorn --workers N``).""" + """Run one fleet worker process. + + Each ``Popen`` receives an explicit ``TLLM_DISAGG_WORKER_PROCESS_ID`` and + binds its own ``SO_REUSEPORT`` socket on the shared public port. + """ _init_fleet_worker_process() server = _build_disagg_server_from_env() host, port = server._config.hostname, server._config.port @@ -1991,7 +2021,6 @@ def set_cuda_device(): status="beta") def disaggregated_mpi_worker(config_file: Optional[str], log_level: str): """Launching disaggregated MPI worker""" - from tensorrt_llm._utils import mpi_rank if os.environ.get(DisaggLauncherEnvs. TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT) != "1": @@ -2201,7 +2230,7 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str, f"Child process {_child_p_global.pid} failed to be killed even after 30s." ) assert _child_p_global.poll( - ) is not None, f"the subprocess should be terminated" + ) is not None, "the subprocess should be terminated" # Check if the process was launched and assert it's terminated if _child_p_global and hasattr(_child_p_global, diff --git a/tensorrt_llm/llmapi/disagg_utils.py b/tensorrt_llm/llmapi/disagg_utils.py index 55f953def917..dc9815faba70 100644 --- a/tensorrt_llm/llmapi/disagg_utils.py +++ b/tensorrt_llm/llmapi/disagg_utils.py @@ -352,8 +352,7 @@ def get_server_configs_dict( def extract_disagg_cluster_config( cluster_config_dict: Dict[str, Any], cluster_uri: Optional[str] = None) -> DisaggClusterConfig: - """ - Build the DisaggClusterConfig from the cluster_config_dict. + """Build the DisaggClusterConfig from the cluster_config_dict. Use the default value of DisaggClusterConfig and MinimalInstances if the corresponding fields are not provided. If cluster_uri is provided, it will override the cluster_uri in the cluster_config_dict. """ @@ -489,7 +488,7 @@ def get_global_disagg_request_id(node_id: int, process_id: int = 0) -> int: def get_local_request_id(last_id: int) -> int: - """ increment the last_id by 1 and mod by MIN_GLOBAL_ID """ + """Increment the last_id by 1 and mod by MIN_GLOBAL_ID""" return (last_id + 1) & (MIN_GLOBAL_ID - 1) @@ -499,8 +498,15 @@ def disagg_process_id_space() -> int: def worker_local_process_id() -> int: - """Read this fleet worker's process index from TRTLLM_DISAGG_WORKER_PROCESS_ID - (set explicitly per worker by the fleet launcher, one distinct value 0..N-1 - per Popen). 0 if unset (e.g. single-process / non-fleet path).""" - return int(os.environ.get("TRTLLM_DISAGG_WORKER_PROCESS_ID", "0")) \ - % disagg_process_id_space() + """Return this fleet worker's process index. + + The fleet launcher sets ``TRTLLM_DISAGG_WORKER_PROCESS_ID`` to a distinct + value per process. A standalone disaggregated server defaults to 0. + """ + process_id = int(os.environ.get("TRTLLM_DISAGG_WORKER_PROCESS_ID", "0")) + process_id_space = disagg_process_id_space() + if not 0 <= process_id < process_id_space: + raise ValueError( + "TRTLLM_DISAGG_WORKER_PROCESS_ID must be between 0 and " + f"{process_id_space - 1}, got {process_id}") + return process_id diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 98fb4f11b5be..cc6a66e53a08 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3523,22 +3523,6 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description="KV cache manager v2 block reuse policy. " "With SWA scratch reuse and 'all_reusable', only non-scratch " "blocks are saved for reuse.") - # Per-rank centralized KV cache routing - centralized_router_report_address: Optional[str] = Field( - default=None, - status="prototype", - description= - "ZMQ address to push KV cache events to for centralized routing " - "(e.g. 'tcp://orchestrator:5557'). When set with per_rank_routing=True, " - "each DP rank reports its own events independently.") - - per_rank_routing: bool = Field( - default=False, - status="prototype", - description= - "Enable per-rank KV cache event reporting. Each DP rank reports its " - "own cache blocks independently (no allgather). Requires " - "use_kv_cache_manager_v2=True and event_buffer_max_size > 0.") def _to_pybind(self): config = _KvCacheConfig( diff --git a/tensorrt_llm/serve/coordinator_server.py b/tensorrt_llm/serve/coordinator_server.py index d7b9a1dcb1b5..efd55f690507 100644 --- a/tensorrt_llm/serve/coordinator_server.py +++ b/tensorrt_llm/serve/coordinator_server.py @@ -83,9 +83,11 @@ async def lifespan(app: FastAPI): self.app.add_api_route("/version", self.version, methods=["GET"]) def _timing_headers(self, raw_req: Request, recv_t: float) -> dict: - """Echo the client send time + stamp coord recv/send times so the client - can decompose the IPC round-trip. recv_t is captured at handler entry; - send time is 'now' (just before responding). + """Return timestamps used to decompose the coordinator round trip. + + The client send time is echoed and coordinator receive/send times are + added. ``recv_t`` is captured at handler entry; send time is captured + immediately before responding. """ h = {HDR_COORD_RECV: repr(recv_t), HDR_COORD_SEND: repr(time.time())} cs = raw_req.headers.get(HDR_CLIENT_SEND) @@ -99,16 +101,24 @@ async def select(self, raw_req: Request) -> Response: body = orjson.loads(await raw_req.body()) except Exception as e: return ORJSONResponse(status_code=400, content={"error": f"invalid JSON body: {e}"}) - if not isinstance(body, dict) or "role" not in body: + if not isinstance(body, dict) or "role" not in body or "routing_key" not in body: return ORJSONResponse( status_code=400, content={"error": "body must include 'role' and 'routing_key'"} ) + role = body["role"] + if role not in ("context", "ctx", "generation", "gen"): + return ORJSONResponse(status_code=400, content={"error": f"invalid role: {role}"}) + req_id = body.get("req_id") + if req_id is not None and (not isinstance(req_id, int) or isinstance(req_id, bool)): + return ORJSONResponse(status_code=400, content={"error": "req_id must be an integer"}) + exclude_server = body.get("exclude_server") + if exclude_server is not None and not isinstance(exclude_server, str): + return ORJSONResponse( + status_code=400, content={"error": "exclude_server must be a string"} + ) try: server, info, req_id = await self._coordinator.select( - body["role"], - body.get("routing_key"), - body.get("req_id"), - body.get("exclude_server"), + role, body["routing_key"], req_id, exclude_server ) except ValueError as e: return ORJSONResponse(status_code=503, content={"error": str(e)}) @@ -126,9 +136,20 @@ async def finish(self, raw_req: Request) -> Response: body = orjson.loads(await raw_req.body()) except Exception as e: return ORJSONResponse(status_code=400, content={"error": f"invalid JSON body: {e}"}) - await self._coordinator.finish( - body.get("role", "gen"), body.get("req_id"), body.get("success", True) - ) + if not isinstance(body, dict) or "role" not in body or "req_id" not in body: + return ORJSONResponse( + status_code=400, content={"error": "body must include 'role' and 'req_id'"} + ) + role = body["role"] + if role not in ("context", "ctx", "generation", "gen"): + return ORJSONResponse(status_code=400, content={"error": f"invalid role: {role}"}) + req_id = body["req_id"] + if not isinstance(req_id, int) or isinstance(req_id, bool): + return ORJSONResponse(status_code=400, content={"error": "req_id must be an integer"}) + success = body.get("success", True) + if not isinstance(success, bool): + return ORJSONResponse(status_code=400, content={"error": "success must be a boolean"}) + await self._coordinator.finish(role, req_id, success) return ORJSONResponse(content={}, headers=self._timing_headers(raw_req, _recv)) async def cluster_info(self) -> Response: @@ -151,9 +172,15 @@ async def __call__(self, host: str, port: int, uds: Optional[str] = None) -> Non # the fleet (hot path) and TCP for health/external clients. import asyncio as _asyncio - uds_cfg = uvicorn.Config(self.app, uds=uds, **kwargs) - tcp_cfg = uvicorn.Config(self.app, host=host, port=port, **kwargs) - await _asyncio.gather(uvicorn.Server(uds_cfg).serve(), uvicorn.Server(tcp_cfg).serve()) + await self._coordinator.start() + try: + uds_cfg = uvicorn.Config(self.app, uds=uds, lifespan="off", **kwargs) + tcp_cfg = uvicorn.Config(self.app, host=host, port=port, lifespan="off", **kwargs) + await _asyncio.gather( + uvicorn.Server(uds_cfg).serve(), uvicorn.Server(tcp_cfg).serve() + ) + finally: + await self._coordinator.stop() else: config = uvicorn.Config(self.app, host=host, port=port, **kwargs) await uvicorn.Server(config).serve() diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py index d45506d0331a..06677aa8a00c 100644 --- a/tensorrt_llm/serve/disagg_coordinator.py +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -64,6 +64,8 @@ "DisaggCoordinator", "DisaggCoordinatorService", "CoordinatorClient", + "coordinator_base_url", + "make_coordinator_session", ] @@ -168,9 +170,11 @@ def set_clients(self, ctx_client: OpenAIClient, gen_client: OpenAIClient) -> Non # -- coordinator-path placement (workers call these via the HTTP server) -- def _api_lat(self, name: str): - """Lazily create a per-API PeriodicLatencyLogger. Measures the coordinator - owner's in-process handler time (excludes the fleet->coordinator HTTP hop, - which is captured client-side as [coord_api] client.*). + """Return a lazily created per-API latency logger. + + It measures the coordinator owner's in-process handler time, excluding + the fleet-to-coordinator HTTP hop captured client-side as + ``[coord_api] client.*``. """ cache = self.__dict__.setdefault("_coord_api_lat", {}) if name not in cache: @@ -202,7 +206,12 @@ async def finish(self, role: str, req_id, success: bool = True) -> None: self._api_lat(f"finish[{role}]").record(time.monotonic() - _t0) def _router_for_role(self, role: str) -> Router: - return self._ctx_router if str(role).lower().startswith("c") else self._gen_router + normalized_role = str(role).lower() + if normalized_role in ("context", "ctx"): + return self._ctx_router + if normalized_role in ("generation", "gen"): + return self._gen_router + raise ValueError(f"Unsupported coordinator role: {role}") async def start(self) -> None: await self._ctx_router.prepare_servers() @@ -329,10 +338,11 @@ def coordinator_base_url(remote_url: str) -> str: def make_coordinator_session(remote_url: str) -> aiohttp.ClientSession: - """Aiohttp session for the coordinator endpoint. Uses a UnixConnector for a - ``unix:/path`` URL -- UDS avoids the TCP loopback stack, since the fleet is - co-located with the implicit coordinator. Created lazily by callers so it - binds to the running event loop. + """Create an aiohttp session for the coordinator endpoint. + + A ``unix:/path`` URL uses a UnixConnector to avoid the TCP loopback stack + when the fleet is co-located with the implicit coordinator. Callers create + the session lazily so it binds to the running event loop. limit=0 (unlimited pool): each fleet worker has ~concurrency/num_workers requests in flight (e.g. 320 at c1280/4), but aiohttp's default pool caps at diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 4a081d4fc8dd..eb011b987dd1 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -55,6 +55,10 @@ # yapf: enale TIMEOUT_KEEP_ALIVE = 10 # seconds. +_LOG_CONTROL_CHARACTERS = { + code: f"\\x{code:02x}" + for code in (*range(32), 127) +} class RawRequestResponseHooks(ResponseHooks): def __init__(self, raw_req: Request, perf_metrics_collector: DisaggPerfMetricsCollector): @@ -101,7 +105,19 @@ def on_resp_done(self, gen_server: str, request: UCompletionRequest, response: U }) if request.disaggregated_params: ctx_req_id = request.disaggregated_params.ctx_request_id - asyncio.create_task(self.perf_metrics_collector.add_per_request_metrics(self.ctx_server, gen_server, ctx_req_id, self.raw_req.state.server_arrival_time, self.server_first_token_time, self.ctx_dispatch_time)) + task = asyncio.create_task( + self.perf_metrics_collector.add_per_request_metrics( + self.ctx_server, + gen_server, + ctx_req_id, + self.raw_req.state.server_arrival_time, + self.server_first_token_time, + self.ctx_dispatch_time, + ) + ) + background_tasks = self.perf_metrics_collector._background_tasks + background_tasks.add(task) + task.add_done_callback(background_tasks.discard) class OpenAIDisaggServer: @@ -172,6 +188,11 @@ async def lifespan(app) -> None: await self._service.setup() yield await self._service.teardown() + if self._perf_metrics_collector._background_tasks: + await asyncio.gather( + *self._perf_metrics_collector._background_tasks, + return_exceptions=True, + ) self.app = FastAPI(lifespan=lifespan) @@ -193,8 +214,10 @@ async def validation_exception_handler(request: Request, exc): "msg": e.get("msg")} for e in errs][:8] except Exception: # noqa: BLE001 brief = str(exc)[:500] + method = request.method.translate(_LOG_CONTROL_CHARACTERS) + path = request.url.path.translate(_LOG_CONTROL_CHARACTERS) logger.warning( - f"[validation] {request.method} {request.url.path} 400 " + f"[validation] {method} {path} 400 " f"(n={self._val_err_n}): {brief}") return JSONResponse(status_code=400, content={"error": str(exc)}) @@ -283,7 +306,7 @@ def _handle_exception(self, exception): async def health(self) -> Response: if not await self._coordinator.is_ready(): - return Response(status_code=500) + return Response(status_code=503) return Response(status_code=200) async def cluster_info(self) -> JSONResponse: @@ -301,7 +324,7 @@ async def __call__(self, host: str, port: int, sockets: list[socket.socket] | No await uvicorn.Server(config).serve(sockets=sockets) async def _sync_server_clock(self, server: str): - """ Sync the ctx/gen server's steady clock with the disagg-server's steady clock (in case NTP service is not running). """ + """Sync the ctx/gen server's steady clock with the disagg-server's steady clock (in case NTP service is not running).""" async def query_steady_clock_offset(session: aiohttp.ClientSession, server_url: str) -> tuple[Optional[float], Optional[float]]: try: originate_ts = get_steady_clock_now_in_seconds() diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index e771e1671f76..74a5417c12e6 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -43,24 +43,26 @@ _GEN_PENDING_FINISH_REASONS = ("length", "not_finished") +_PRECTX_LOGGER = None + + def _prectx_logger(): - """Aggregated (p50/p95 every N) breakdown of the ORCHESTRATOR pre-ctx path, - isolating where the arrival->ctx_dispatch time goes on the fleet worker: + """Return the orchestrator pre-context timing logger. + + The logger isolates where arrival-to-context-dispatch time is spent: cond_ms : _check_conditional_disagg + _check_gen_only_disagg idgen_ms : get_disagg_request_id (local now; was an HTTP hop) ctxroute_ms: ctx_router.get_next_server (routing/select round trip) No per-request logging (observer effect on the single fleet loop). """ global _PRECTX_LOGGER - try: - return _PRECTX_LOGGER - except NameError: + if _PRECTX_LOGGER is None: from tensorrt_llm.serve.responses_utils import PeriodicBreakdownLogger _PRECTX_LOGGER = PeriodicBreakdownLogger( "orchestrator_prectx", ["cond_ms", "idgen_ms", "ctxroute_ms"] ) - return _PRECTX_LOGGER + return _PRECTX_LOGGER class OpenAIDisaggregatedService(OpenAIService): @@ -170,7 +172,6 @@ async def _send_disagg_request_ctx_first( ctx_req, exclude_server=gen_server ) _prectx["ctxroute_ms"] = (_time.monotonic() - _t_route) * 1000 - _prectx_logger().record(_prectx) ctx_response = await self._ctx_client.send_request( ctx_req, server=ctx_server, hooks=hooks ) @@ -189,6 +190,7 @@ async def _send_disagg_request_ctx_first( and gen_req.disaggregated_params.request_type == "context_only" ): gen_req.disaggregated_params = None + _prectx_logger().record(_prectx) if ctx_response is None or self._need_gen(ctx_response): if not gen_server: gen_server, _ = await self._gen_router.get_next_server( diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 390651040042..53ab14af4234 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -337,9 +337,6 @@ def __init__( self.disagg_cluster_worker = None self.resource_governor = None - self._active_request_count = 0 - self._active_request_count_lock = __import__('threading').Lock() - # Skip loading AutoProcessor and model_config for VISUAL_GEN models # These are LLM-specific and can cause unnecessary memory usage if self._is_visual_gen: @@ -474,30 +471,6 @@ async def validation_exception_handler(_, exc): self.app.add_middleware(ServerArrivalTimeMiddleware) - # Track active /v1/ requests (in-flight request gauge). - server_ref = self - - @self.app.middleware("http") - async def _track_active_requests(request, call_next): - if request.url.path.startswith("/v1/"): - with server_ref._active_request_count_lock: - server_ref._active_request_count += 1 - try: - return await call_next(request) - finally: - with server_ref._active_request_count_lock: - server_ref._active_request_count -= 1 - return await call_next(request) - - def _get_iteration_stats_buffer_maxlen(self) -> Optional[int]: - if isinstance(self.generator, VisualGen): - return None - - max_iterations = self.generator.args.iter_stats_max_iterations - if max_iterations is None or max_iterations < 0: - return None - return max_iterations - def _init_visual_gen(self): self.processor = None self.model_config = None diff --git a/tensorrt_llm/serve/perf_metrics.py b/tensorrt_llm/serve/perf_metrics.py index e2e44c04851d..4eaed580ba8f 100644 --- a/tensorrt_llm/serve/perf_metrics.py +++ b/tensorrt_llm/serve/perf_metrics.py @@ -152,6 +152,7 @@ def __init__(self, max_requests: int): self._lock = asyncio.Lock() self._collect_lock = asyncio.Lock() self._clients = [] + self._background_tasks: set[asyncio.Task] = set() self._metrics = { definition.name: instance_metric(definition) for definition in SERVER_METRICS_DEFINITIONS diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index b71ba3940c04..a38e21b9aeae 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -124,7 +124,7 @@ def _parse_response_input( if not isinstance(input_msg, dict): input_msg = input_msg.model_dump() - _responses_debug_log(f"------- Parsing input -----------") + _responses_debug_log("------- Parsing input -----------") _responses_debug_log(input_msg) _responses_debug_log("") @@ -218,8 +218,7 @@ async def store_response(self, Union[list[Message], list[ChatCompletionMessageParam]]] = [], prev_resp_id: Optional[str] = None) -> None: - """ - Store the response and its messages(model output messages) in the conversation store. If the previous response id is provided, + """Store the response and its messages(model output messages) in the conversation store. If the previous response id is provided, the messages will be appended to the conversation. Otherwise, a new conversation will be created. Args: @@ -234,7 +233,7 @@ async def store_response(self, _responses_debug_log( f"ConversationHistoryStore storing resp: {resp_id}") if ENABLE_RESPONSES_DEBUG_MSG: - _responses_debug_log(f" -> resp_msgs:") + _responses_debug_log(" -> resp_msgs:") for msg in resp_msgs: _responses_debug_log(f" -> {msg}") @@ -278,8 +277,7 @@ async def store_messages(self, resp_id: str, msgs: Union[list[Message], list[ChatCompletionMessageParam]], prev_resp_id: Optional[str]) -> None: - """ - Store the messages in the conversation store. + """Store the messages in the conversation store. `msgs` should always contains the whole conversation messages, including the previous messages and the new messages. @@ -291,7 +289,7 @@ async def store_messages(self, resp_id: str, Returns: None """ - _responses_debug_log(f"ConversationHistoryStore storing msg:") + _responses_debug_log("ConversationHistoryStore storing msg:") if ENABLE_RESPONSES_DEBUG_MSG: for msg in msgs: _responses_debug_log(f" -> {msg}") @@ -317,7 +315,7 @@ async def store_messages(self, resp_id: str, async def get_conversation_history( self, resp_id: str ) -> Union[list[Message], list[ChatCompletionMessageParam]]: - _responses_debug_log(f"ConversationHistoryStore getting prev_msgs:") + _responses_debug_log("ConversationHistoryStore getting prev_msgs:") _responses_debug_log(f" -> prev_resp_id: {resp_id}") async with self.conversations_lock: if resp_id in self.response_to_conversation: @@ -330,8 +328,7 @@ async def get_conversation_history( return [] def _update_visited_conversation(self, conversation_id) -> None: - """ - Update the visited conversation to the front of the conversation store. + """Update the visited conversation to the front of the conversation store. This function is used to keep the conversation store sorted by the visited time. And also remove the least recently visited conversation if the number of conversations exceeds the limit. @@ -356,8 +353,7 @@ def _update_visited_conversation(self, conversation_id) -> None: self.conversation_to_response.pop(removed_id) def _pop_conversation(self, resp_id) -> None: - """ - Pop the oldest conversation messages from a conversation. + """Pop the oldest conversation messages from a conversation. The conversation is starting by a user message and ending by an assistant message. This function is used to keep the number of messages in a conversation within the limit. @@ -539,8 +535,7 @@ def _parse_output_tokens(tokens: list[int]) -> list[Message]: def _parse_output_message_harmony(message: Message) -> list[ResponseOutputItem]: - """ - Parse a Harmony message into a list of output response items. + """Parse a Harmony message into a list of output response items. """ if message.author.role != "assistant": # This is a message from a tool to the assistant (e.g., search result). @@ -721,8 +716,7 @@ async def _create_input_messages( def _create_output_messages( output_contents: dict[str, Any]) -> list[ChatCompletionMessageParam]: - """ - Convert output contents to chat completion messages for conversation store. + """Convert output contents to chat completion messages for conversation store. Reasoning content is not included in the output messages to reduce the token usage. @@ -804,8 +798,7 @@ async def _create_input_tokens( model_config: PretrainedConfig, processor: AutoProcessor, ) -> Tuple[list[int], Optional[dict[str, list[Any]]]]: - """ - Create input tokens for the model. Also return the mm data if the model is multimodal. + """Create input tokens for the model. Also return the mm data if the model is multimodal. Returns: Tuple[list[int], Optional[dict[str, list[Any]]]]: Input tokens and mm data. @@ -895,7 +888,7 @@ async def request_preprocess( prev_msgs = await conversation_store.get_conversation_history( prev_response_id) - _responses_debug_log(f"Prev msgs:") + _responses_debug_log("Prev msgs:") for msg in prev_msgs: _responses_debug_log(f" -> {msg}") @@ -1375,8 +1368,7 @@ def get_reasoning_text_delta_event( def _get_output_added_events( self, output_item: ResponseOutputMessage | ResponseReasoningItem ) -> list[StreamingResponsesResponse]: - """ - Get item added event and content part added event for a message item which is starting + """Get item added event and content part added event for a message item which is starting to be generated. Returns: @@ -1435,8 +1427,7 @@ def _should_send_done_events( streaming_events_helper: Optional[ResponsesStreamingEventsHelper] = None, finished_generation: bool = False, ) -> Tuple[bool, bool, Optional[str], Optional[str]]: - """ - Determine if done events should be sent for text or reasoning items. + """Determine if done events should be sent for text or reasoning items. Analyzes the complete output text to detect when reasoning or text sections have been completed and should receive done events. @@ -1964,8 +1955,7 @@ async def process_streaming_events( class ServerArrivalTimeMiddleware: - """ - Custom ASGI middleware to track server arrival time. + """Custom ASGI middleware to track server arrival time. We implement this as a pure ASGI middleware instead of using FastAPI's @app.middleware("http") decorator because the decorator internally uses @@ -2017,8 +2007,10 @@ async def __call__(self, scope, receive, send): def _accept_latency_logger(): - """Lazy singleton for the client-send -> server-accept latency logger - (PeriodicLatencyLogger is defined below this class).""" + """Return the client-send to server-accept latency logger. + + ``PeriodicLatencyLogger`` is defined below this function. + """ global _ACCEPT_LAT if _ACCEPT_LAT is None: _ACCEPT_LAT = PeriodicLatencyLogger("accept.client_send_to_arrival") @@ -2026,8 +2018,11 @@ def _accept_latency_logger(): def ttft_split_logger(): - """Lazy singleton aggregating the per-request TTFT breakdown (pre_ctx / - ctx_phase / xfer_gen / total), logged periodically instead of per request.""" + """Return the per-request TTFT breakdown logger. + + Pre-context, context, transfer/generation, and total times are aggregated + and logged periodically rather than once per request. + """ global _TTFT_SPLIT if _TTFT_SPLIT is None: _TTFT_SPLIT = PeriodicBreakdownLogger( @@ -2037,11 +2032,11 @@ def ttft_split_logger(): class PeriodicLatencyLogger: - """Accumulates per-call latencies for a named coordinator API and logs - percentiles every ``window`` calls. Cheap, lock-free (single asyncio loop), - and self-resetting. Used to profile the coordinator surface on both the - in-process owner (DisaggCoordinatorService) and the HTTP client - (CoordinatorDelegatingRouter / CoordinatorClient) without per-call log spam.""" + """Periodically log latency percentiles for a named coordinator API. + + This lock-free, self-resetting logger runs on one asyncio loop and profiles + the in-process owner and HTTP client without per-call log spam. + """ def __init__(self, name: str, window: int = 500): self._name = name @@ -2055,21 +2050,23 @@ def record(self, dt_s: float) -> None: if self._n % self._window == 0: s = sorted(self._samples) m = len(s) - p = lambda q: s[min(int(q * m), m - 1)] - logger.info( - f"[coord_api] {self._name} n={self._n} ms: " - f"mean={sum(s)/m:.2f} p50={p(0.5):.2f} p90={p(0.9):.2f} " - f"p99={p(0.99):.2f} max={s[-1]:.2f}") + + def percentile(q): + return s[min(int(q * m), m - 1)] + + logger.info(f"[coord_api] {self._name} n={self._n} ms: " + f"mean={sum(s)/m:.2f} p50={percentile(0.5):.2f} " + f"p90={percentile(0.9):.2f} " + f"p99={percentile(0.99):.2f} max={s[-1]:.2f}") self._samples = [] class PeriodicBreakdownLogger: - """Like PeriodicLatencyLogger but for a multi-field breakdown: accumulates a - dict of named ms values per sample and logs p50/p95 of EACH field once every - ``window`` samples. Used for the per-request TTFT split so we get the - breakdown WITHOUT a synchronous log write per request (that per-request - logging inflated TTFT ~3s on the single-loop disagg server -- observer - effect). Negative values (stage not reached) are dropped per field.""" + """Periodically log a multi-field latency breakdown. + + Named millisecond values are accumulated and each field's p50/p95 is logged + every ``window`` samples. Negative values for stages not reached are dropped. + """ def __init__(self, name: str, fields, window: int = 1000): self._name = name @@ -2089,16 +2086,19 @@ def record(self, values: dict) -> None: for f in self._fields: s = sorted(self._samples[f]) if s: - p = lambda q: s[min(int(q * len(s)), len(s) - 1)] - parts.append(f"{f}(p50={p(0.5):.0f},p95={p(0.95):.0f})") + + def percentile(q): + return s[min(int(q * len(s)), len(s) - 1)] + + parts.append(f"{f}(p50={percentile(0.5):.0f}," + f"p95={percentile(0.95):.0f})") self._samples[f] = [] logger.info(f"[ttft_split] {self._name} n={self._n} " + " ".join(parts)) class ResponseHooks(ABC): - """ - Hooks for response processing and (disagg) service perf observability. + """Hooks for response processing and (disagg) service perf observability. """ @abstractmethod @@ -2106,11 +2106,11 @@ def on_req_begin(self, request: UCompletionRequest): pass def on_ctx_dispatch(self, request: UCompletionRequest): - """Fired the instant the disagg service starts placing the request on a - ctx server (before routing / the ctx HTTP send). arrival->here measures - the pre-ctx wait in the orchestrator/fleet (accept queue + event loop + - pipeline), which dominates TTFT under high fleet concurrency. Default - no-op so non-instrumented hook implementations are unaffected.""" + """Record when the disaggregated service starts context placement. + + Arrival to this point measures the pre-context wait in the orchestrator + or fleet. The default is a no-op for non-instrumented implementations. + """ @abstractmethod def on_ctx_resp(self, ctx_server: str, response: UCompletionResponse): diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index e6aa2335e150..d4f356a8829a 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -13,6 +13,8 @@ # limitations under the License. import asyncio +import os +import statistics import time from abc import ABC, abstractmethod from collections import OrderedDict @@ -21,10 +23,6 @@ import aiohttp import orjson -# Content-Type for orjson-serialized bodies sent to the coordinator (we pass -# data=bytes, so aiohttp doesn't set the json content type itself). -_JSON_HEADERS = {"Content-Type": "application/json"} - from tensorrt_llm.llmapi.disagg_utils import (MetadataServerConfig, RouterConfig, ServerRole) from tensorrt_llm.logger import logger @@ -39,6 +37,10 @@ block_key_hasher, get_cache_salt_id, get_request_num_tokens, hash_v1_block_key, truncate_sha256_hash_to_int64, v2_sha256_block_hasher) +# Content-Type for orjson-serialized bodies sent to the coordinator (we pass +# data=bytes, so aiohttp doesn't set the json content type itself). +_JSON_HEADERS = {"Content-Type": "application/json"} + # Max number of conversations whose home-server pin is retained (LRU). ROUTE_AFFINITY_CACHE_SIZE = 50000 # Leading token-id count folded into the affinity key so pre-tokenized @@ -46,33 +48,6 @@ ROUTE_AFFINITY_TOKEN_PREFIX = 256 -def get_request_num_tokens(request: OpenAIRequest) -> int: - if request.disaggregated_params is None or request.disaggregated_params.request_type == "context_only": - if isinstance(request, ChatCompletionRequest): - raise ValueError( - "LoadBalancing router with tokens doesn't support ChatCompletionRequest yet" - ) - - if isinstance(request.prompt, str) or \ - (isinstance(request.prompt, list) and len(request.prompt) > 0 - and isinstance(request.prompt[0], int)): - prompts = [request.prompt] - else: - prompts = request.prompt - - num_tokens = sum(len(prompt) for prompt in prompts) - elif request.disaggregated_params.request_type == "generation_only": - raise ValueError( - "LoadBalancing router with tokens doesn't support generation_only requests" - ) - else: - raise ValueError( - f"Unsupported request type: {request.disaggregated_params.request_type}" - ) - - return num_tokens - - class ServerState: def __init__( @@ -380,27 +355,33 @@ def __init__( self._prepared_ready_servers: set[str] = set() # Routing-latency diagnostics (gated by TLLM_LOG_ROUTE_TIMING=1): record # get_next_server wall time per request and log percentiles periodically. - import os self._log_route_timing = (os.environ.get("TLLM_LOG_ROUTE_TIMING", "0") == "1") self._rt_samples: list = [] self._rt_n = 0 def _record_route_timing(self, dt_s: float) -> None: - """Record one get_next_server latency sample; log percentiles every - 500 calls. No-op unless TLLM_LOG_ROUTE_TIMING=1.""" + """Record one routing-latency sample. + + Percentiles are logged every 500 calls. This is a no-op unless + ``TLLM_LOG_ROUTE_TIMING=1``. + """ if not self._log_route_timing: return self._rt_samples.append(dt_s * 1000.0) # ms self._rt_n += 1 if self._rt_n % 500 == 0: - import statistics s = sorted(self._rt_samples) n = len(s) - p = lambda q: s[min(int(q * n), n - 1)] + + def percentile(q): + return s[min(int(q * n), n - 1)] + logger.info(f"[route_timing] {type(self).__name__} n={self._rt_n} " f"get_next_server_ms: mean={statistics.mean(s):.2f} " - f"p50={p(0.5):.2f} p90={p(0.9):.2f} p99={p(0.99):.2f} " + f"p50={percentile(0.5):.2f} " + f"p90={percentile(0.9):.2f} " + f"p99={percentile(0.99):.2f} " f"max={s[-1]:.2f}") self._rt_samples = [] # reset window @@ -506,7 +487,7 @@ async def get_next_server( self, request: OpenAIRequest, exclude_server: Optional[str] = None) -> tuple[str, dict]: - '''Select server by request and return some intermediate information, exclude_server is a server to exclude from the selection''' + """Select server by request and return some intermediate information, exclude_server is a server to exclude from the selection""" @abstractmethod async def finish_request(self, @@ -808,223 +789,6 @@ async def finish_request(self, await self._unregister_request(request) -def block_key_hasher(token_ids: list[int], - parent_hash: Optional[int] = None, - cache_salt_id: Optional[int] = None) -> int: - parent = 0 if parent_hash is None else parent_hash - # Fast path: the native C++ BlockKeyHasher is bit-exact with - # hash_v1_block_key and avoids the per-token Python loop. Its hash() binding - # takes no cache_salt_id, so fall back to Python only when a salt is set - # (rare opt-in; never in the unsalted agent/chat completion path). - if cache_salt_id is None: - return _NativeBlockKeyHasher.hash(_NativeBlockKey(token_ids), parent) - return hash_v1_block_key(token_ids, - parent_hash=parent, - cache_salt_id=cache_salt_id) - - -def v2_sha256_block_hasher(token_ids: list[int], - parent_hash: Optional[str] = None, - cache_salt_id: Optional[int] = None) -> str: - parent_key = (V2RootBlock.make_key(ReuseScope(salt=cache_salt_id)) - if parent_hash is None else bytes.fromhex(parent_hash)) - return V2Block.make_key(parent_key, token_ids).hex() - - -class BlockHashMixin: - """Shared tokenization and block-hash computation. - - Used by routers that need KV-cache-aware prefix matching. - """ - - def _init_block_hashing(self, - tokens_per_block: Optional[int] = None, - custom_tokenizer: Optional[str] = None): - env_tokens_per_block = os.environ.get( - "TRTLLM_KVCACHE_AWARE_ROUTER_HASH_TOKENS_PER_BLOCK") - if env_tokens_per_block is not None: - tokens_per_block = int(env_tokens_per_block) - self._tpb_auto = tokens_per_block is None - self._tokens_per_block = 32 if tokens_per_block is None \ - else tokens_per_block - self._tokenizers: dict = {} - self._custom_tokenizer = custom_tokenizer - logger.info(f"BlockHashMixin: tokens_per_block={self._tokens_per_block}" - f"{' (auto, adopts worker)' if self._tpb_auto else ''}" - f", custom_tokenizer={self._custom_tokenizer}") - - def _get_tokenizer(self, model: str): - if model not in self._tokenizers: - if self._custom_tokenizer: - from tensorrt_llm.tokenizer import load_custom_tokenizer - self._tokenizers[model] = load_custom_tokenizer( - self._custom_tokenizer, model) - else: - from tensorrt_llm.tokenizer import TransformersTokenizer - tokenizer = TransformersTokenizer.from_pretrained( - model, trust_remote_code=True) - self._tokenizers[model] = tokenizer.tokenizer - return self._tokenizers[model] - - def _encode_with_prefix_cache(self, rendered: str, key: int, - tokenizer) -> list[int]: - cache = getattr(self, "_tok_prefix_cache", None) - if cache is None: - cache = self._tok_prefix_cache = OrderedDict() - entry = cache.get(key) - if entry is not None and len(rendered) > len(entry[0]) and \ - rendered.startswith(entry[0]): - ids = entry[1] + tokenizer.encode(rendered[len(entry[0]):], - add_special_tokens=False) - else: - ids = tokenizer.encode(rendered, add_special_tokens=False) - cache[key] = (rendered, ids) - cache.move_to_end(key) - while len(cache) > 1024: - cache.popitem(last=False) - return ids - - def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: - # Handle ChatCompletionRequest (has messages, not prompt) - if isinstance(request, ChatCompletionRequest): - if request.prompt_token_ids is not None: - return [request.prompt_token_ids] - tokenizer = self._get_tokenizer(request.model) - tool_dicts = (None if getattr(request, "tools", None) is None else [ - tool.model_dump() if hasattr(tool, "model_dump") else tool - for tool in request.tools - ]) - chat_template_kwargs = (request.chat_template_kwargs if getattr( - request, "chat_template_kwargs", None) else {}) - rendered = tokenizer.apply_chat_template( - [ - msg if isinstance(msg, dict) else dict(msg) - for msg in request.messages - ], - add_generation_prompt=request.add_generation_prompt, - tokenize=False, - return_dict=False, - tools=tool_dicts, - **chat_template_kwargs, - ) - if isinstance(rendered, str): - key = hash("".join( - str( - msg.get("content") if isinstance(msg, dict) else - getattr(msg, "content", "")) - for msg in request.messages[:2])) - result = self._encode_with_prefix_cache(rendered, key, - tokenizer) - else: - result = list(rendered) - request.prompt_token_ids = result - return [result] - - # Handle CompletionRequest (has prompt) - prompts = request.prompt - if isinstance(prompts, list) and len(prompts) == 0: - return [] - if isinstance(prompts, list) and isinstance(prompts[0], list): - return prompts - elif isinstance(prompts, list) and isinstance(prompts[0], int): - return [prompts] - elif isinstance(prompts, str): - prompts = [prompts] - else: - assert isinstance(prompts, list) and isinstance(prompts[0], str) - - tokenizer = self._get_tokenizer(request.model) - token_lists = [tokenizer(prompt)["input_ids"] for prompt in prompts] - # Replace string prompts with token IDs so the worker server - # skips re-tokenization - request.prompt = (token_lists - if len(token_lists) > 1 else token_lists[0]) - return token_lists - - def _compute_block_hashes( - self, - token_lists: list[list[int]], - hash_algo: str = KV_CACHE_HASH_ALGO_DEFAULT, - cache_salt_id: Optional[int] = None, - ) -> list[list[BlockHash]]: - if hash_algo == KV_CACHE_HASH_ALGO_V1: - block_hasher = block_key_hasher - elif hash_algo == KV_CACHE_HASH_ALGO_V2: - block_hasher = v2_sha256_block_hasher - elif hash_algo == KV_CACHE_HASH_ALGO_V2_SHA256_64: - reuse_scope = ReuseScope(salt=cache_salt_id) - block_hashes: list[list[BlockHash]] = [] - for token_list in token_lists: - hash_list = [] - parent_key = V2RootBlock.make_key(reuse_scope) - for t in range(0, len(token_list) - 1, self._tokens_per_block): - t_end = min(t + self._tokens_per_block, len(token_list) - 1) - parent_key = V2Block.make_key(parent_key, - token_list[t:t_end]) - hash_list.append(truncate_sha256_hash_to_int64(parent_key)) - block_hashes.append(hash_list) - return block_hashes - else: - raise ValueError( - f"Unsupported KV cache hash algorithm: {hash_algo}") - - block_hashes: list[list[BlockHash]] = [] - for token_list in token_lists: - hash_list = [] - # in KvCacheManager, the last token is not included in the block key - for t in range(0, len(token_list) - 1, self._tokens_per_block): - t_end = min(t + self._tokens_per_block, len(token_list) - 1) - hash_list.append( - block_hasher(token_list[t:t_end], - None if t == 0 else hash_list[-1], - cache_salt_id)) - block_hashes.append(hash_list) - return block_hashes - - def _tokenize_and_compute_block_hashes( - self, - request: OpenAIRequest) -> tuple[list[list[int]], list[list[int]]]: - """Synchronous tokenize + block-hash, combined for thread offload. - - Factored into one method so ``get_next_server`` can offload the whole - CPU-bound step via ``asyncio.to_thread`` in a single call, keeping - the orchestrator's asyncio event loop free to dispatch other - requests in parallel. - """ - token_lists = self._tokenize(request) - block_hashes = self._compute_block_hashes(token_lists) - return token_lists, block_hashes - - def _tokenize_and_compute_block_hashes_by_algo( - self, - request: OpenAIRequest, - hash_algos: Iterable[str], - cache_salt_id: Optional[int] = None, - ) -> tuple[list[list[int]], dict[str, list[list[BlockHash]]]]: - """Synchronous tokenize + per-algorithm block hashes for thread offload.""" - token_lists = self._tokenize(request) - return token_lists, { - hash_algo: - self._compute_block_hashes(token_lists, - hash_algo, - cache_salt_id=cache_salt_id) - for hash_algo in set(hash_algos) - } - - @staticmethod - def _text_to_int_sequences(texts: list[str]) -> list[list[int]]: - """Convert text strings to lists of unicode code points. - - Usable as input to ``_compute_block_hashes``. - """ - return [[ord(c) for c in text] for text in texts] - - @staticmethod - def _get_request_cache_salt_id(request: OpenAIRequest) -> Optional[int]: - cache_salt = getattr(request, "cache_salt", None) - return None if cache_salt is None else get_cache_salt_id(cache_salt) - - class KvCacheAwareRouter(BlockHashMixin, LoadBalancingMixin, Router): _server_state_class = KvCacheAwareServerState @@ -1176,8 +940,11 @@ async def get_next_server( return server, info def _routing_key_sync(self, request: OpenAIRequest) -> dict: - """Tokenize + per-algo block-hash (CPU-bound; run in a thread). Returns a - plain dict so the coordinator path can send it over HTTP unchanged.""" + """Build the coordinator routing key. + + Tokenization and per-algorithm block hashing are CPU-bound and run in a + thread. The returned plain dict can be sent over HTTP unchanged. + """ cache_salt_id = self._get_request_cache_salt_id(request) # Hash for every algo any server might use (usually one). algos = { @@ -1194,33 +961,33 @@ def _routing_key_sync(self, request: OpenAIRequest) -> dict: } async def _route(self, key, exclude_server=None, request=None, req_id=None): - """THE single routing core, shared by the standalone and coordinator - paths. Scores each server by (matched_tokens/tokens_per_block - - load_weight*load), applies load_cap, conversation affinity and RR - tie-break, then registers load + remembers routed blocks. The ONLY - difference between the two callers is request identity: standalone passes - the request (load keyed by id(request), finished via finish_request); - the coordinator passes req_id (the disagg request id) and finishes via - finish_request_by_id(req_id). Returns (server, info, req_id-or-None).""" - import time as _time - _rt_t0 = _time.monotonic() + """Route through the core shared by standalone and coordinator paths. + + Servers are scored by cache match and load before applying the load cap, + conversation affinity, and round-robin tie-breaking. Standalone requests + are keyed by ``id(request)``; coordinator requests use ``req_id``. + """ + route_start = time.monotonic() token_lists = (key or {}).get("token_lists") or [] block_hashes_by_algo = (key or {}).get("block_hashes_by_algo") or {} conv_key = (key or {}).get("conv_key") async with self._lock: - servers = [ - s for s in self._server_state.keys() if s != exclude_server - ] + server_states = { + server: state + for server, state in self._server_state.items() + if server != exclude_server + } + servers = list(server_states) if not servers: raise ValueError( f"No available servers after excluding {exclude_server}") def _hashes(server): - algo = self._get_server_hash_algo(server) + algo = server_states[server].hash_algo return algo, block_hashes_by_algo.get(algo, []) workloads = [ - self._server_state[s].num_active_requests() for s in servers + server_states[server].num_active_requests() for server in servers ] load_fractions = [ workloads[i] / self._max_batch_size for i in range(len(servers)) @@ -1228,8 +995,7 @@ def _hashes(server): scores, matches = [], [] for i, server in enumerate(servers): algo, bh = _hashes(server) - matches.append(await - self._server_state[server].matched_tokens(bh, algo)) + matches.append(await server_states[server].matched_tokens(bh, algo)) scores.append(matches[-1] / self._tokens_per_block - self._load_weight * workloads[i]) candidate_idx = [ @@ -1263,10 +1029,13 @@ def _hashes(server): # crosses the /finish HTTP hop). key = id(request) if req_id is None else req_id async with self._lock: - await self._server_state[server].increment_load(request) + if self._server_state.get(server) is not server_states[server]: + raise ValueError( + f"Selected server {server} is no longer available") + await server_states[server].increment_load(request) self._req_routing_table[key] = server self._stash_routed_blocks_on_route(key, block_hashes, hash_algo) - self._record_route_timing(_time.monotonic() - _rt_t0) + self._record_route_timing(time.monotonic() - route_start) return server, { "block_hashes": block_hashes, "hash_algo": hash_algo, @@ -1287,11 +1056,11 @@ async def finish_request(self, session=session) async def _finish(self, key, success, request=None, session=None): - """THE single finish core, shared by finish_request (standalone, key = - id(request), request passed) and finish_request_by_id (coordinator, key = - disagg req_id, request=None -- symmetric with its increment_load(None)). - Pops the SAME maps _route filled, decrements load, commits routed blocks - on success, and refreshes the block table.""" + """Finish through the core shared by standalone and coordinator paths. + + This removes the maps populated by ``_route``, decrements load, commits + routed blocks on success, and refreshes the block table. + """ async with self._lock: server = self._req_routing_table.pop(key, None) if server is not None and server in self._server_state: @@ -1300,10 +1069,11 @@ async def _finish(self, key, success, request=None, session=None): self._poll_server_on_finish(server, session) def _poll_server_on_finish(self, server, session=None): - """Refresh a server's KV-cache block table from /kv_cache_events after a - request finishes. Shared by finish_request (standalone) and - finish_request_by_id (coordinator) so the block table always stays warm -- - the delegated path is inert without this (matches score 0).""" + """Refresh a server's KV-cache block table after a request finishes. + + This is shared by standalone and coordinator finish paths so the block + table remains warm for delegated routing. + """ if (server is not None and server in self._server_state and self._events_aligned(server)): # Fire-and-forget; poll runs in background and coalesces per server. @@ -1315,24 +1085,28 @@ def _poll_server_on_finish(self, server, session=None): # as the standalone get_next_server. def routing_key(self, request: OpenAIRequest): - """Worker-side: tokenize + block-hash. Same dict _route consumes, JSON- - serializable for the /select POST.""" + """Return the worker-side JSON-serializable routing key. + + The returned tokenization and block hashes are consumed by ``_route``. + """ return self._routing_key_sync(request) async def get_next_server_by_key(self, routing_key, exclude_server=None, req_id=None): - """Coordinator-side placement: the SAME _route core, keyed by the - caller's disagg request id (req_id) instead of id(request).""" + """Place a coordinator-side request. + + The shared routing core is keyed by the caller's disaggregated request + ID rather than ``id(request)``. + """ return await self._route(routing_key, exclude_server=exclude_server, request=None, req_id=req_id) async def finish_request_by_id(self, req_id, success=True): - """Coordinator-side finish: the SAME _finish core, keyed by the disagg - request id instead of id(request).""" + """Finish a coordinator-side request by its disaggregated request ID.""" if req_id is None: return await self._finish(req_id, success) @@ -1524,8 +1298,10 @@ def _get_content_load(self, server: str) -> int: return self._server_content_load.get(server, 0) def _get_server_load(self, server: str) -> int: - """Use content weight so ``_select_least_loaded`` balances by - estimated tokens rather than request count. + """Return content-weighted server load. + + This makes ``_select_least_loaded`` balance by estimated tokens rather + than request count. """ return self._get_content_load(server) @@ -1869,6 +1645,7 @@ def __init__(self, self._role = role # "context" | "generation" self._request_timeout_s = request_timeout_s self._session: Optional[aiohttp.ClientSession] = None + self._background_tasks: set[asyncio.Task] = set() # Coordinator HTTP-client API latency (includes network round-trip to the # coordinator + its in-process handler). Compare against the owner-side # [coord_api] to isolate the fleet /select|/finish HTTP overhead. @@ -1910,10 +1687,12 @@ def _on_servers_updated(self, old_servers, new_servers): def _record_ipc(self, ipc_logger, send_t: float, recv_t: float, resp) -> None: - """Decompose the coordinator round trip from the response timing headers. - send_t/recv_t are the client-side wall clocks (co-located => comparable). - Missing headers (older coordinator) => that field is dropped by the - breakdown logger (records -1). All values in ms.""" + """Record a coordinator round-trip breakdown. + + ``send_t`` and ``recv_t`` are client-side wall clocks. Missing headers + from an older coordinator are recorded as ``-1`` and dropped by the + breakdown logger. All values are milliseconds. + """ from tensorrt_llm.serve.coordinator_server import (HDR_COORD_RECV, HDR_COORD_SEND) @@ -1937,20 +1716,24 @@ def _f(name): def _request_id(self, request: OpenAIRequest) -> int: """The request's disagg id -- the sole cross-process key for select/finish. + Context requests carry disagg_request_id; generation requests inherit the SAME id as ctx_request_id (set by OpenAIDisaggregatedService before routing). This id is what the ctx worker registered its KV-transfer TxSession under, so the gen request MUST keep it -- never re-issue a new id here, or the gen transceiver waits on a key the ctx side never registered (transfer never completes -> DISAGG_GENERATION_TRANS_IN_PROGRESS - fills the gen IndexMapper and throughput collapses).""" + fills the gen IndexMapper and throughput collapses). + """ dp = request.disaggregated_params - assert dp is not None, "delegated routing requires disaggregated_params" + if dp is None: + raise ValueError("delegated routing requires disaggregated_params") rid = (dp.disagg_request_id if self._role == "context" else dp.ctx_request_id) - assert rid is not None, ( - f"delegated {self._role} routing requires a disagg request id " - f"(disagg_request_id/ctx_request_id) on the request") + if rid is None: + raise ValueError( + f"delegated {self._role} routing requires a disagg request id " + "(disagg_request_id/ctx_request_id) on the request") return rid async def get_next_server( @@ -2001,7 +1784,9 @@ async def finish_request(self, # request may be freed after return); errors are logged, not propagated. del session req_id = self._request_id(request) - asyncio.create_task(self._finish_async(req_id, success)) + task = asyncio.create_task(self._finish_async(req_id, success)) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) async def _finish_async(self, req_id: int, success: bool): from tensorrt_llm.serve.coordinator_server import HDR_CLIENT_SEND @@ -2028,6 +1813,9 @@ async def _finish_async(self, req_id: int, success: bool): logger.warning(f"CoordinatorDelegatingRouter finish failed: {e}") async def close(self): + if self._background_tasks: + await asyncio.gather(*self._background_tasks, + return_exceptions=True) if self._session is not None: await self._session.close() self._session = None @@ -2042,8 +1830,7 @@ def create_router( server_preparation_func: Optional[Callable[[str], Awaitable[None]]] = None, disagg_node_id: int = 0, ) -> Router: - """ - Factory function to create different types of router instances. + """Factory function to create different types of router instances. Args: router_type (str): Type of router to create. Supported values: diff --git a/tensorrt_llm/serve/router_utils.py b/tensorrt_llm/serve/router_utils.py index d691daccd8b2..a7e1d4b349ab 100644 --- a/tensorrt_llm/serve/router_utils.py +++ b/tensorrt_llm/serve/router_utils.py @@ -126,7 +126,8 @@ def get_request_num_tokens(request: OpenAIRequest) -> int: ) if isinstance(request.prompt, str) or ( - isinstance(request.prompt, list) and isinstance(request.prompt[0], int) + isinstance(request.prompt, list) + and (not request.prompt or isinstance(request.prompt[0], int)) ): prompts = [request.prompt] else: @@ -214,9 +215,11 @@ def _encode_with_prefix_cache(self, rendered: str, key: int, tokenizer) -> list[ if cache is None: cache = self._tok_prefix_cache = OrderedDict() entry = cache.get(key) - if entry is not None and len(rendered) > len(entry[0]) and rendered.startswith(entry[0]): - ids = entry[1] + tokenizer.encode(rendered[len(entry[0]) :], add_special_tokens=False) + if entry is not None and rendered == entry[0]: + ids = entry[1] else: + # Tokenizing a suffix independently is not generally composable: + # BPE/SentencePiece merges can cross the cached string boundary. ids = tokenizer.encode(rendered, add_special_tokens=False) cache[key] = (rendered, ids) cache.move_to_end(key) @@ -270,6 +273,8 @@ def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: # Handle CompletionRequest (has prompt) prompts = request.prompt + if isinstance(prompts, list) and not prompts: + return [prompts] if isinstance(prompts, list) and isinstance(prompts[0], list): return prompts elif isinstance(prompts, list) and isinstance(prompts[0], int): diff --git a/tests/unittest/disaggregated/test_coordinator_e2e.py b/tests/unittest/disaggregated/test_coordinator_e2e.py index 3fdae2cfedb8..a90a8b2cd86c 100644 --- a/tests/unittest/disaggregated/test_coordinator_e2e.py +++ b/tests/unittest/disaggregated/test_coordinator_e2e.py @@ -66,9 +66,10 @@ @pytest.fixture(autouse=True) def _reset_prometheus_registry(): - """Role-prefixed Prometheus counters are registered in the global default - registry; clear it between tests so a second server build in the same pytest - process does not hit duplicate-timeseries errors. + """Reset role-prefixed Prometheus counters. + + Counters use the global default registry, so clear it between tests to avoid + duplicate-timeseries errors when another server is built in this process. """ from prometheus_client import REGISTRY @@ -198,7 +199,9 @@ async def shutdown(self): def _make_config(ctx_url, gen_url, public_port): - host_port = lambda u: (u.split(":")[0], int(u.split(":")[1])) + def host_port(url): + return url.split(":")[0], int(url.split(":")[1]) + ctx_host, ctx_port = host_port(ctx_url) gen_host, gen_port = host_port(gen_url) return DisaggServerConfig( @@ -286,11 +289,12 @@ async def drive(): assert body["choices"][0]["finish_reason"] == "stop" -def _write_config(path, ctx_url, gen_url, public_port): +def _write_config(path, ctx_url, gen_url, public_port, num_workers=1): """A disagg config YAML: round-robin ctx, conversation gen (delegated).""" cfg = { "hostname": "127.0.0.1", "port": public_port, + "num_workers": num_workers, "context_servers": { "num_instances": 1, "urls": [ctx_url], @@ -307,10 +311,11 @@ def _write_config(path, ctx_url, gen_url, public_port): def test_disagg_completion_e2e_web_concurrency_4(): - """WEB_CONCURRENCY=4 through the real CLI: `trtllm-serve disaggregated` - forks a coordinator (port-1) + a uvicorn fleet of 4 disagg servers on the - public port. Mock ctx/gen workers run in this test process; a real HTTP - completion round-trips through one of the four workers -> coordinator -> gen. + """Exercise the real CLI disaggregated-fleet path. + + ``num_workers=4`` starts a coordinator and four disaggregated servers on the + public port. A completion traverses a fleet worker, coordinator, and mock + context/generation workers. """ logger.set_level("info") # trtllm logger defaults to "error"; show progress WORKERS = 4 @@ -326,10 +331,9 @@ def test_disagg_completion_e2e_web_concurrency_4(): with tempfile.TemporaryDirectory() as td: cfg_path = os.path.join(td, "disagg.yaml") - _write_config(cfg_path, ctx_url, gen_url, public_port) + _write_config(cfg_path, ctx_url, gen_url, public_port, WORKERS) env = dict(os.environ) - env["WEB_CONCURRENCY"] = str(WORKERS) # Unbuffered so the child's launch logs stream out live (else stdout # to a pipe is block-buffered and nothing shows until it exits). env["PYTHONUNBUFFERED"] = "1" @@ -346,7 +350,7 @@ def test_disagg_completion_e2e_web_concurrency_4(): logger.info( f"mock ctx={ctx_url} gen={gen_url}; launching " - f"`trtllm-serve disaggregated` WEB_CONCURRENCY={WORKERS}, " + f"`trtllm-serve disaggregated` num_workers={WORKERS}, " f"public={public_port} coordinator={coord_port}" ) proc = subprocess.Popen( diff --git a/tests/unittest/disaggregated/test_coordinator_worker.py b/tests/unittest/disaggregated/test_coordinator_worker.py index 37ba80c6049a..21646754b639 100644 --- a/tests/unittest/disaggregated/test_coordinator_worker.py +++ b/tests/unittest/disaggregated/test_coordinator_worker.py @@ -49,15 +49,15 @@ from tensorrt_llm.serve.coordinator_server import CoordinatorServer from tensorrt_llm.serve.disagg_coordinator import CoordinatorClient, DisaggCoordinatorService from tensorrt_llm.serve.openai_protocol import CompletionRequest, DisaggregatedParams +from tensorrt_llm.serve.router_utils import BlockHashMixin as SharedBlockHashMixin @pytest.fixture(autouse=True) def _reset_prometheus_registry(): - """Each coordinator builds role-prefixed Prometheus counters in the global - default registry (via its readiness OpenAIHttpClients). In production the - coordinator is a single process; here two coordinators share one pytest - process, so clear the registry between tests to avoid duplicate-timeseries - registration errors. + """Reset role-prefixed Prometheus counters. + + Tests create multiple coordinators in one process, so clear their shared + default registry between tests to avoid duplicate-timeseries errors. """ from prometheus_client import REGISTRY @@ -180,9 +180,38 @@ async def _wait_coord_ready(url, timeout_s=30.0): return False +def test_coordinator_rejects_unknown_role(): + config = _make_config([], [], "round_robin", "round_robin") + coordinator = DisaggCoordinatorService(config, _client_factory) + + with pytest.raises(ValueError, match="Unsupported coordinator role"): + coordinator._router_for_role("typo") + + +def test_prefix_token_cache_retokenizes_extended_text(): + class BoundarySensitiveTokenizer: + def __init__(self): + self.calls = [] + + def encode(self, text, add_special_tokens=False): + assert add_special_tokens is False + self.calls.append(text) + return {"ab": [1], "abc": [2], "c": [3]}[text] + + tokenizer = BoundarySensitiveTokenizer() + block_hashing = SharedBlockHashMixin() + block_hashing._init_block_hashing() + + assert block_hashing._encode_with_prefix_cache("ab", 1, tokenizer) == [1] + assert block_hashing._encode_with_prefix_cache("abc", 1, tokenizer) == [2] + assert block_hashing._encode_with_prefix_cache("abc", 1, tokenizer) == [2] + assert tokenizer.calls == ["ab", "abc"] + + def test_stateless_router_places_locally_in_worker(): - """A round-robin (stateless) router is NOT wrapped: the worker places locally - with the real router and never calls the coordinator. + """Verify stateless round-robin placement remains local. + + The worker uses the real router without calling the coordinator. """ from tensorrt_llm.serve.router import CoordinatorDelegatingRouter, RoundRobinRouter @@ -212,8 +241,10 @@ async def drive(): def test_conversation_coordinator_sticky_by_conv_id(): - """Same conversation_id sticks to one gen worker; a stateful (conversation) - router delegates placement to the coordinator via /select. + """Verify conversation IDs remain sticky through delegated routing. + + The stateful generation router delegates placement to coordinator + ``/select``. """ from tensorrt_llm.serve.router import CoordinatorDelegatingRouter @@ -222,12 +253,14 @@ def test_conversation_coordinator_sticky_by_conv_id(): with _CoordinatorThread(config) as coord: assert asyncio.run(_wait_coord_ready(coord.url)) - def _req(conv_id): + def _req(conv_id, request_id): return CompletionRequest( model="m", prompt="hi", disaggregated_params=DisaggregatedParams( - request_type="context_only", conversation_id=conv_id + request_type="generation_only", + ctx_request_id=request_id, + conversation_id=conv_id, ), ) @@ -236,12 +269,16 @@ async def drive(): # Stateful -> wrapped in a coordinator-delegating router. assert isinstance(remote.gen_router, CoordinatorDelegatingRouter) assert await remote.is_ready() is True - first, _ = await remote.gen_router.get_next_server(_req("conv-A")) + first_request = _req("conv-A", 1) + first, _ = await remote.gen_router.get_next_server(first_request) + await remote.gen_router.finish_request(first_request) # Repeated conv-A requests must land on the same worker. repeats = [] - for _ in range(3): - s, _ = await remote.gen_router.get_next_server(_req("conv-A")) + for request_id in range(2, 5): + request = _req("conv-A", request_id) + s, _ = await remote.gen_router.get_next_server(request) repeats.append(s) + await remote.gen_router.finish_request(request) await remote.stop() return first, repeats @@ -251,8 +288,8 @@ async def drive(): ) -def test_coordinator_owns_generation_disagg_request_id(): - """Generation routing receives its request ID from the coordinator.""" +def test_worker_generates_disagg_request_id_before_generation_routing(): + """Generation routing uses the ID generated by the coordinator client.""" from tensorrt_llm.serve.router import CoordinatorDelegatingRouter with _FakeWorker() as ctx0, _FakeWorker() as gen0: @@ -263,24 +300,26 @@ def test_coordinator_owns_generation_disagg_request_id(): async def drive(): remote = CoordinatorClient(coord.url, config) assert isinstance(remote.gen_router, CoordinatorDelegatingRouter) + assigned_id = await remote.get_disagg_request_id() request = CompletionRequest( model="m", prompt="hello", disaggregated_params=DisaggregatedParams( request_type="generation_only", - ctx_request_id=123, + ctx_request_id=assigned_id, disagg_request_id=None, conversation_id="conv-A", ), ) await remote.gen_router.get_next_server(request) - assigned_id = request.disaggregated_params.disagg_request_id + assert request.disaggregated_params.disagg_request_id is None + assert request.disaggregated_params.ctx_request_id == assigned_id await remote.gen_router.finish_request(request) await remote.stop() return assigned_id assigned_id = asyncio.run(drive()) - assert assigned_id is not None and assigned_id != 123 + assert assigned_id > 0 if __name__ == "__main__": diff --git a/tests/unittest/disaggregated/test_disagg_utils.py b/tests/unittest/disaggregated/test_disagg_utils.py index 35fe98f7676c..46e0718b6440 100644 --- a/tests/unittest/disaggregated/test_disagg_utils.py +++ b/tests/unittest/disaggregated/test_disagg_utils.py @@ -6,9 +6,10 @@ # isort: off from tensorrt_llm.llmapi.disagg_utils import ( - MIN_GLOBAL_ID, CtxGenServerConfig, DisaggServerConfig, extract_ctx_gen_cfgs, - extract_router_config, extract_disagg_cfg, get_global_disagg_request_id, - get_local_request_id, get_server_configs_dict, parse_disagg_config_file) + MIN_GLOBAL_ID, CtxGenServerConfig, DisaggServerConfig, + disagg_process_id_space, extract_ctx_gen_cfgs, extract_router_config, + extract_disagg_cfg, get_global_disagg_request_id, get_local_request_id, + get_server_configs_dict, parse_disagg_config_file, worker_local_process_id) # isort: on @@ -110,9 +111,20 @@ def test_parse_disagg_config_file(sample_yaml_file, sample_yaml_config): @pytest.mark.parametrize("sample_yaml_config", ["disagg_cluster", ""], indirect=True) def test_extract_disagg_cfg(sample_yaml_config): + sample_yaml_config.update({ + "gen_tokids_ctxbytes": + True, + "num_workers": + 4, + "disagg_coordinator_url": + "http://coordinator:7999", + }) config = extract_disagg_cfg(**sample_yaml_config) assert isinstance(config, DisaggServerConfig) verify_disagg_config(config, sample_yaml_config) + assert config.gen_tokids_ctxbytes is True + assert config.num_workers == 4 + assert config.disagg_coordinator_url == "http://coordinator:7999" def test_extract_ctx_gen_cfgs(): @@ -242,6 +254,17 @@ def test_get_global_disagg_request_id_range_validation(): assert get_global_disagg_request_id(0, 0) >= MIN_GLOBAL_ID +def test_worker_local_process_id_range_validation(monkeypatch): + process_id_space = disagg_process_id_space() + monkeypatch.setenv("TRTLLM_DISAGG_WORKER_PROCESS_ID", + str(process_id_space - 1)) + assert worker_local_process_id() == process_id_space - 1 + + monkeypatch.setenv("TRTLLM_DISAGG_WORKER_PROCESS_ID", str(process_id_space)) + with pytest.raises(ValueError): + worker_local_process_id() + + def test_get_local_request_id(): last_id = MIN_GLOBAL_ID - 100 ids = set() diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 52d42e78fcae..881d0037af0a 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -222,7 +222,7 @@ async def test_is_ready_waits_for_router_preparation(): AsyncMock(), ) # Readiness now lives on the DisaggCoordinatorService the service holds. - local = service._cluster + local = service._coordinator local._disagg_cluster_manager = cluster_manager cluster_manager._current_ctx_workers["ctx"] = WorkerInfo( From 634732bb1bc88f4ac48c6e94208bae32378fa5a0 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 9 Jul 2026 00:25:27 -0700 Subject: [PATCH 03/18] [None][fix] resolve disaggregated coordinator CI failures Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/llmapi/disagg_utils.py | 4 ++-- tensorrt_llm/serve/conversation_id.py | 7 ++++--- tensorrt_llm/serve/router.py | 11 +++++++++++ tensorrt_llm/serve/router_utils.py | 10 ++++------ tests/integration/test_lists/test-db/l0_a10.yml | 2 ++ tests/unittest/disaggregated/test_disagg_utils.py | 3 +-- tests/unittest/disaggregated/test_router.py | 7 +++---- 7 files changed, 27 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/llmapi/disagg_utils.py b/tensorrt_llm/llmapi/disagg_utils.py index dc9815faba70..4f9ef21e506a 100644 --- a/tensorrt_llm/llmapi/disagg_utils.py +++ b/tensorrt_llm/llmapi/disagg_utils.py @@ -94,8 +94,8 @@ class DisaggServerConfig(): perf_metrics_max_requests: int = 0 disagg_cluster_config: Optional[DisaggClusterConfig] = None node_id: int = uuid.getnode( - ) % 1021 # Assuming only one disagg-server is running on a machine, moding mac by the largest 10-bit prime - # If this causes collisions, users can set node_id manually within range [0, 1023] in config + ) % 256 # Assuming only one disagg-server is running on a machine, modulo 256. + # If this causes collisions, users can set node_id manually within range [0, 255] in config schedule_style: Literal['context_first', 'generation_first'] = 'context_first' allow_request_chat_template: bool = False diff --git a/tensorrt_llm/serve/conversation_id.py b/tensorrt_llm/serve/conversation_id.py index 7e0fed21640e..2acfcad8012e 100644 --- a/tensorrt_llm/serve/conversation_id.py +++ b/tensorrt_llm/serve/conversation_id.py @@ -36,9 +36,10 @@ class RequestWithConversationParams(Protocol): def get_request_conversation_id(request: RequestWithConversationParams) -> Optional[str]: conversation_params = request.conversation_params - if conversation_params is None: - return None - return conversation_params.conversation_id + if conversation_params is not None: + return conversation_params.conversation_id + disaggregated_params = getattr(request, "disaggregated_params", None) + return None if disaggregated_params is None else disaggregated_params.conversation_id def extract_conversation_id_from_headers(headers: Optional[Mapping[str, str]]) -> Optional[str]: diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index d4f356a8829a..305eb9423b28 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -793,6 +793,17 @@ class KvCacheAwareRouter(BlockHashMixin, LoadBalancingMixin, Router): _server_state_class = KvCacheAwareServerState + @staticmethod + def _get_request_cache_salt_id(request: OpenAIRequest) -> Optional[int]: + """Resolve request cache salt through this module's public shim. + + ``router.py`` historically exposed ``get_cache_salt_id``. Keeping the + lookup here preserves that compatibility while hashing helpers live in + ``router_utils.py``. + """ + cache_salt = getattr(request, "cache_salt", None) + return None if cache_salt is None else get_cache_salt_id(cache_salt) + def __init__(self, server_role: ServerRole = None, servers: list[str] = None, diff --git a/tensorrt_llm/serve/router_utils.py b/tensorrt_llm/serve/router_utils.py index a7e1d4b349ab..e27199a848d8 100644 --- a/tensorrt_llm/serve/router_utils.py +++ b/tensorrt_llm/serve/router_utils.py @@ -21,8 +21,6 @@ from collections import OrderedDict from typing import Iterable, List, Optional, Union -from transformers import AutoTokenizer - from tensorrt_llm.bindings.internal.batch_manager import BlockKey as _NativeBlockKey from tensorrt_llm.bindings.internal.batch_manager import BlockKeyHasher as _NativeBlockKeyHasher from tensorrt_llm.logger import logger @@ -202,12 +200,12 @@ def _get_tokenizer(self, model: str): self._tokenizers[model] = load_custom_tokenizer(self._custom_tokenizer, model_path) else: - from tensorrt_llm.tokenizer import maybe_fix_byte_level_tokenizer + from tensorrt_llm.tokenizer import TransformersTokenizer - tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) - self._tokenizers[model] = maybe_fix_byte_level_tokenizer( - tokenizer, model_path, trust_remote_code=True + tokenizer = TransformersTokenizer.from_pretrained( + model_path, trust_remote_code=True ) + self._tokenizers[model] = tokenizer.tokenizer return self._tokenizers[model] def _encode_with_prefix_cache(self, rendered: str, key: int, tokenizer) -> list[int]: diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index a6f112843fcf..a720fa3e9912 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -58,6 +58,8 @@ l0_a10: - unittest/others/test_tracing.py - unittest/disaggregated/test_disagg_openai_client.py - unittest/disaggregated/test_disagg_utils.py + - unittest/disaggregated/test_coordinator_e2e.py + - unittest/disaggregated/test_coordinator_worker.py - unittest/disaggregated/test_openai_disagg_server.py - unittest/disaggregated/test_openai_disagg_service.py - unittest/disaggregated/test_router.py diff --git a/tests/unittest/disaggregated/test_disagg_utils.py b/tests/unittest/disaggregated/test_disagg_utils.py index 46e0718b6440..3d2d286312c1 100644 --- a/tests/unittest/disaggregated/test_disagg_utils.py +++ b/tests/unittest/disaggregated/test_disagg_utils.py @@ -276,5 +276,4 @@ def test_get_local_request_id(): assert len(ids) == 1000 assert min(ids) == 0 assert max(ids) == MIN_GLOBAL_ID - 1 - assert max(ids) - min(ids) > ( - 1 << 40) # ensure there is enough space for local ids + assert max(ids) - min(ids) == MIN_GLOBAL_ID - 1 diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index b8d5980e087b..9009221e38aa 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -2051,7 +2051,7 @@ def test_prefix_cache_tokenize_matches_full_encode(servers): assert req.prompt_token_ids == reference -def test_prefix_cache_tokenize_encodes_only_delta(servers): +def test_prefix_cache_tokenize_uses_canonical_encoding(servers): router = KvCacheAwareRouter(server_role=None, servers=servers, tokens_per_block=4) @@ -2065,7 +2065,7 @@ def _recording_encode(text, add_special_tokens=False): tok.encode = _recording_encode with mock.patch.object(router, "_get_tokenizer", return_value=tok): - for index, convo in enumerate(_grow_conversation()): + for convo in _grow_conversation(): encoded_lengths.clear() req = ChatCompletionRequest(model="mock", messages=convo) rendered = tok.apply_chat_template( @@ -2074,8 +2074,7 @@ def _recording_encode(text, add_special_tokens=False): tokenize=False) router._tokenize(req) assert encoded_lengths - if index >= 1: - assert min(encoded_lengths) < len(rendered) + assert encoded_lengths == [len(rendered)] def test_prefix_cache_tokenize_falls_back_on_divergent_prefix(servers): From e90c74e67ff3a9e2731f7d048bbaa70a41a008f3 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 9 Jul 2026 00:25:27 -0700 Subject: [PATCH 04/18] [None][fix] resolve disaggregated coordinator CI failures Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/serve/router.py | 28 +++++---------------- tests/unittest/disaggregated/test_router.py | 26 ++++++++----------- 2 files changed, 17 insertions(+), 37 deletions(-) diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 305eb9423b28..6627eb8e0213 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -828,9 +828,6 @@ def __init__(self, self._load_weight = load_weight self._load_cap = load_cap self._track_routed_blocks = track_routed_blocks - # request key -> (flat block hashes, hash_algo). Key is id(request) on the - # standalone path, the disagg req_id on the coordinator path. - self._pending_routed_blocks: dict[int, tuple[list[BlockHash], str]] = {} def _create_server_state(self, server: str) -> KvCacheAwareServerState: return KvCacheAwareServerState(server, self._use_tokens, @@ -842,25 +839,12 @@ async def close(self): await state.cancel_poll_task() await super().close() - def _stash_routed_blocks_on_route(self, key: int, + def _apply_routed_blocks_on_route(self, server: str, block_hashes: list[list[BlockHash]], hash_algo: str) -> None: if not self._track_routed_blocks: return - flat = [h for hl in block_hashes for h in hl] - self._pending_routed_blocks[key] = (flat, hash_algo) - - def _apply_routed_blocks_on_finish(self, key: int, server: Optional[str], - success: bool) -> None: - # Pop unconditionally to avoid leaks; apply only when eligible. - entry = self._pending_routed_blocks.pop(key, None) - if not (self._track_routed_blocks and success): - return - if entry is None: - return - if server is None or server not in self._server_state: - return - flat_block_hashes, hash_algo = entry + flat_block_hashes = [h for hashes in block_hashes for h in hashes] self._server_state[server].add_blocks(flat_block_hashes, hash_algo=hash_algo) @@ -1045,7 +1029,7 @@ def _hashes(server): f"Selected server {server} is no longer available") await server_states[server].increment_load(request) self._req_routing_table[key] = server - self._stash_routed_blocks_on_route(key, block_hashes, hash_algo) + self._apply_routed_blocks_on_route(server, block_hashes, hash_algo) self._record_route_timing(time.monotonic() - route_start) return server, { "block_hashes": block_hashes, @@ -1069,14 +1053,14 @@ async def finish_request(self, async def _finish(self, key, success, request=None, session=None): """Finish through the core shared by standalone and coordinator paths. - This removes the maps populated by ``_route``, decrements load, commits - routed blocks on success, and refreshes the block table. + This removes the maps populated by ``_route``, decrements load, and + refreshes the block table. Routed blocks are applied at route time. """ + del success async with self._lock: server = self._req_routing_table.pop(key, None) if server is not None and server in self._server_state: await self._server_state[server].decrement_load(request) - self._apply_routed_blocks_on_finish(key, server, success) self._poll_server_on_finish(server, session) def _poll_server_on_finish(self, server, session=None): diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index 9009221e38aa..d2222d0d8982 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -786,7 +786,7 @@ def test_kv_cache_aware_server_state_remove_blocks_silent_on_missing(): @pytest.mark.asyncio -async def test_kv_cache_aware_router_tracks_routed_blocks_at_routing(servers): +async def test_kv_cache_aware_router_applies_routed_blocks_at_routing(servers): tokens_per_block = 4 token_lists = [[1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008]] router = KvCacheAwareRouter(server_role=None, @@ -800,19 +800,17 @@ async def test_kv_cache_aware_router_tracks_routed_blocks_at_routing(servers): prompt=copy.deepcopy(token_lists)) server, info = await router.get_next_server(request) - assert id(request) in router._pending_routed_blocks - stashed, stashed_algo = router._pending_routed_blocks[id(request)] expected_flat = [h for hl in info["block_hashes"] for h in hl] - assert stashed == expected_flat - assert stashed and not isinstance(stashed[0], list) - assert stashed_algo == info["hash_algo"] + block_table = router._server_state[server]._block_table(info["hash_algo"]) + assert block_table.issuperset(expected_flat) await router.finish_request(request) - assert id(request) not in router._pending_routed_blocks + assert block_table.issuperset(expected_flat) @pytest.mark.asyncio -async def test_kv_cache_aware_router_inserts_routed_blocks_on_finish(servers): +async def test_kv_cache_aware_router_does_not_wait_for_finish_to_apply_blocks( + servers): tokens_per_block = 4 token_lists = [[2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008]] router = KvCacheAwareRouter(server_role=None, @@ -829,7 +827,7 @@ async def test_kv_cache_aware_router_inserts_routed_blocks_on_finish(servers): hash_algo = info["hash_algo"] assert await router._server_state[server].matched_tokens( - info["block_hashes"], hash_algo=hash_algo) == 0 + info["block_hashes"], hash_algo=hash_algo) == 8 await router.finish_request(request) @@ -855,8 +853,6 @@ async def test_kv_cache_aware_router_routed_blocks_disabled_skips_pending( prompt=copy.deepcopy(token_lists)) server, info = await router.get_next_server(request) - assert router._pending_routed_blocks == {} - await router.finish_request(request) assert await router._server_state[server].matched_tokens( @@ -864,7 +860,7 @@ async def test_kv_cache_aware_router_routed_blocks_disabled_skips_pending( @pytest.mark.asyncio -async def test_kv_cache_aware_router_drops_routed_blocks_on_failure(servers): +async def test_kv_cache_aware_router_keeps_routed_blocks_on_failure(servers): tokens_per_block = 4 token_lists = [[4000, 4001, 4002, 4003, 4004, 4005, 4006, 4007, 4008]] router = KvCacheAwareRouter(server_role=None, @@ -877,13 +873,13 @@ async def test_kv_cache_aware_router_drops_routed_blocks_on_failure(servers): request = CompletionRequest(model="TinyLlama", prompt=copy.deepcopy(token_lists)) server, info = await router.get_next_server(request) - assert id(request) in router._pending_routed_blocks + assert await router._server_state[server].matched_tokens( + info["block_hashes"], hash_algo=info["hash_algo"]) == 8 await router.finish_request(request, success=False) - assert id(request) not in router._pending_routed_blocks assert await router._server_state[server].matched_tokens( - info["block_hashes"], hash_algo=info["hash_algo"]) == 0 + info["block_hashes"], hash_algo=info["hash_algo"]) == 8 @pytest.mark.asyncio From 8370c960215e54ab8911729aa61faf55d23fd7a0 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 9 Jul 2026 22:10:18 -0700 Subject: [PATCH 05/18] [None][fix] reduce coordinator diff and refresh telemetry manifest Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- docs/source/features/disagg-serving.md | 6 +-- tensorrt_llm/commands/serve.py | 20 +++------ tensorrt_llm/llmapi/disagg_utils.py | 5 ++- tensorrt_llm/serve/openai_disagg_server.py | 2 +- tensorrt_llm/serve/openai_disagg_service.py | 20 ++++++--- tensorrt_llm/serve/responses_utils.py | 49 +++++++++++++-------- tensorrt_llm/serve/router.py | 19 ++------ tests/unittest/disaggregated/test_router.py | 25 +---------- 8 files changed, 62 insertions(+), 84 deletions(-) diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index 48b35dac5ede..9ed621158a26 100644 --- a/docs/source/features/disagg-serving.md +++ b/docs/source/features/disagg-serving.md @@ -98,7 +98,7 @@ The context and generation phases of one request must share a single request ID: The disaggregated server generates this ID itself as a **snowflake** — a self-contained 64-bit positive integer that is unique without any cross-process coordination. The bit layout is: -```text +``` [ 0 (1 bit) | timestamp_ms (39 bits) | node_id (8 bits) | process_id (6 bits) | counter (10 bits) ] ``` @@ -185,7 +185,7 @@ Once the context and generation servers are launched, you can launch the disaggr server, which will accept requests from clients and do the orchestration between context and generation servers. The disaggregated server can be launched with: -```bash +``` trtllm-serve disaggregated -c disagg_config.yaml ``` where `disagg_config.yaml` contains information about the context and generation servers. For the current example, @@ -209,7 +209,7 @@ When routing requests to the context servers, the disaggregated server will mark when routing requests to the generation servers, the disaggregated server will mark the requests as "generation-only" to skip the context phase. Clients can then send requests to the disaggregated server at `localhost:8000`, which is an OpenAI-compatible endpoint. For example, you can send requests to the disaggregated server using curl: -```bash +``` curl http://localhost:8000/v1/completions \ -H "Content-Type: application/json" \ -d '{ diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index e17e53e82217..02f7d8ee0e2b 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -120,7 +120,8 @@ def _signal_handler_cleanup_child(signum, frame): def is_non_default_or_required(param_name, value, backend, explicit_cli_keys): - """Check if a parameter should be explicitly included in llm_args. + """ + Check if a parameter should be explicitly included in llm_args. Returns True if parameter is either: 1. Always required (core params that must be present), OR @@ -435,7 +436,8 @@ def launch_grpc_server(host: str, port: int, llm_args: dict, served_model_name: Optional[str] = None): - """Launch a gRPC server for TensorRT-LLM. + """ + Launch a gRPC server for TensorRT-LLM. This provides a high-performance gRPC interface designed for external routers (e.g., sgl-router) using pre-tokenized input and raw token ID output. @@ -1622,6 +1624,7 @@ def disaggregated( schedule_style: str, ): """Running server in disaggregated mode""" + logger.set_level(log_level) if metrics_log_interval != 0: @@ -1954,16 +1957,6 @@ def _build_disagg_server_from_env() -> "OpenAIDisaggServer": return server -def create_disagg_server_app(): - """Uvicorn import-string factory: build one disagg server's FastAPI app. - - Retained for the ``uvicorn --factory`` entry point; the SO_REUSEPORT fleet - launcher uses ``_run_fleet_worker`` instead. - """ - _init_fleet_worker_process() - return _build_disagg_server_from_env().app - - def _run_fleet_worker(): """Run one fleet worker process. @@ -2021,6 +2014,7 @@ def set_cuda_device(): status="beta") def disaggregated_mpi_worker(config_file: Optional[str], log_level: str): """Launching disaggregated MPI worker""" + from tensorrt_llm._utils import mpi_rank if os.environ.get(DisaggLauncherEnvs. TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT) != "1": @@ -2230,7 +2224,7 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str, f"Child process {_child_p_global.pid} failed to be killed even after 30s." ) assert _child_p_global.poll( - ) is not None, "the subprocess should be terminated" + ) is not None, f"the subprocess should be terminated" # Check if the process was launched and assert it's terminated if _child_p_global and hasattr(_child_p_global, diff --git a/tensorrt_llm/llmapi/disagg_utils.py b/tensorrt_llm/llmapi/disagg_utils.py index 4f9ef21e506a..f55504b309f1 100644 --- a/tensorrt_llm/llmapi/disagg_utils.py +++ b/tensorrt_llm/llmapi/disagg_utils.py @@ -352,7 +352,8 @@ def get_server_configs_dict( def extract_disagg_cluster_config( cluster_config_dict: Dict[str, Any], cluster_uri: Optional[str] = None) -> DisaggClusterConfig: - """Build the DisaggClusterConfig from the cluster_config_dict. + """ + Build the DisaggClusterConfig from the cluster_config_dict. Use the default value of DisaggClusterConfig and MinimalInstances if the corresponding fields are not provided. If cluster_uri is provided, it will override the cluster_uri in the cluster_config_dict. """ @@ -488,7 +489,7 @@ def get_global_disagg_request_id(node_id: int, process_id: int = 0) -> int: def get_local_request_id(last_id: int) -> int: - """Increment the last_id by 1 and mod by MIN_GLOBAL_ID""" + """ increment the last_id by 1 and mod by MIN_GLOBAL_ID """ return (last_id + 1) & (MIN_GLOBAL_ID - 1) diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index eb011b987dd1..7ed8b56c5dbb 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -324,7 +324,7 @@ async def __call__(self, host: str, port: int, sockets: list[socket.socket] | No await uvicorn.Server(config).serve(sockets=sockets) async def _sync_server_clock(self, server: str): - """Sync the ctx/gen server's steady clock with the disagg-server's steady clock (in case NTP service is not running).""" + """ Sync the ctx/gen server's steady clock with the disagg-server's steady clock (in case NTP service is not running). """ async def query_steady_clock_offset(session: aiohttp.ClientSession, server_url: str) -> tuple[Optional[float], Optional[float]]: try: originate_ts = get_steady_clock_now_in_seconds() diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 74a5417c12e6..978850aa4ed0 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -265,9 +265,10 @@ def _get_gen_request( else: request.prompt_token_ids = ctx_response.prompt_token_ids # Opt-in: drop conversation history so the gen worker doesn't - # re-parse the full conversation JSON (a GIL cost at high - # concurrency); it uses prompt_token_ids and only the last message, - # tools preserved. Config-gated (unsafe for harmony/multimodal). + # re-parse the full conversation JSON (dominates its GIL at high + # concurrency). It uses prompt_token_ids and only reads the last + # message; tools are preserved. Config-gated because it's unsafe + # for harmony/multimodal workers (model type is fixed per deploy). if ( self._strip_gen_message_history and request.messages @@ -418,10 +419,15 @@ async def _send_disagg_request_gen_first( ) if request.stream and need_ctx: - # Streaming gen_first: the gen POST is deferred until its generator is - # iterated, but the ctx server blocks on the gen rx session, so gather - # would deadlock. Eagerly consume the gen generator in a background task - # (firing the POST) and pipe chunks through a queue. + # For streaming gen_first requests, the gen client returns a lazy + # async generator whose HTTP POST only fires when iterated. The ctx + # server blocks waiting for the gen server's rx session (gen_first + # protocol). Using asyncio.gather would deadlock: ctx waits for gen + # server, but gen POST is deferred until the generator is consumed, + # and the generator isn't consumed until gather returns. + # + # Fix: eagerly start consuming the gen generator in a background + # task so the HTTP POST fires, then pipe chunks through a queue. gen_response = await self._gen_client.send_request( gen_req, server=gen_server, hooks=hooks ) diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index a38e21b9aeae..10e468f5a270 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -124,7 +124,7 @@ def _parse_response_input( if not isinstance(input_msg, dict): input_msg = input_msg.model_dump() - _responses_debug_log("------- Parsing input -----------") + _responses_debug_log(f"------- Parsing input -----------") _responses_debug_log(input_msg) _responses_debug_log("") @@ -218,8 +218,10 @@ async def store_response(self, Union[list[Message], list[ChatCompletionMessageParam]]] = [], prev_resp_id: Optional[str] = None) -> None: - """Store the response and its messages(model output messages) in the conversation store. If the previous response id is provided, - the messages will be appended to the conversation. Otherwise, a new conversation will be created. + """Store a response and its model-output messages. + + If the previous response ID is provided, the messages are appended to + that conversation. Otherwise, a new conversation is created. Args: resp: ResponsesResponse @@ -233,7 +235,7 @@ async def store_response(self, _responses_debug_log( f"ConversationHistoryStore storing resp: {resp_id}") if ENABLE_RESPONSES_DEBUG_MSG: - _responses_debug_log(" -> resp_msgs:") + _responses_debug_log(f" -> resp_msgs:") for msg in resp_msgs: _responses_debug_log(f" -> {msg}") @@ -277,7 +279,8 @@ async def store_messages(self, resp_id: str, msgs: Union[list[Message], list[ChatCompletionMessageParam]], prev_resp_id: Optional[str]) -> None: - """Store the messages in the conversation store. + """ + Store the messages in the conversation store. `msgs` should always contains the whole conversation messages, including the previous messages and the new messages. @@ -289,7 +292,7 @@ async def store_messages(self, resp_id: str, Returns: None """ - _responses_debug_log("ConversationHistoryStore storing msg:") + _responses_debug_log(f"ConversationHistoryStore storing msg:") if ENABLE_RESPONSES_DEBUG_MSG: for msg in msgs: _responses_debug_log(f" -> {msg}") @@ -315,7 +318,7 @@ async def store_messages(self, resp_id: str, async def get_conversation_history( self, resp_id: str ) -> Union[list[Message], list[ChatCompletionMessageParam]]: - _responses_debug_log("ConversationHistoryStore getting prev_msgs:") + _responses_debug_log(f"ConversationHistoryStore getting prev_msgs:") _responses_debug_log(f" -> prev_resp_id: {resp_id}") async with self.conversations_lock: if resp_id in self.response_to_conversation: @@ -328,7 +331,8 @@ async def get_conversation_history( return [] def _update_visited_conversation(self, conversation_id) -> None: - """Update the visited conversation to the front of the conversation store. + """Move the visited conversation to the front of the store. + This function is used to keep the conversation store sorted by the visited time. And also remove the least recently visited conversation if the number of conversations exceeds the limit. @@ -353,7 +357,8 @@ def _update_visited_conversation(self, conversation_id) -> None: self.conversation_to_response.pop(removed_id) def _pop_conversation(self, resp_id) -> None: - """Pop the oldest conversation messages from a conversation. + """Pop the oldest messages from a conversation. + The conversation is starting by a user message and ending by an assistant message. This function is used to keep the number of messages in a conversation within the limit. @@ -535,7 +540,8 @@ def _parse_output_tokens(tokens: list[int]) -> list[Message]: def _parse_output_message_harmony(message: Message) -> list[ResponseOutputItem]: - """Parse a Harmony message into a list of output response items. + """ + Parse a Harmony message into a list of output response items. """ if message.author.role != "assistant": # This is a message from a tool to the assistant (e.g., search result). @@ -716,7 +722,8 @@ async def _create_input_messages( def _create_output_messages( output_contents: dict[str, Any]) -> list[ChatCompletionMessageParam]: - """Convert output contents to chat completion messages for conversation store. + """ + Convert output contents to chat completion messages for conversation store. Reasoning content is not included in the output messages to reduce the token usage. @@ -798,7 +805,8 @@ async def _create_input_tokens( model_config: PretrainedConfig, processor: AutoProcessor, ) -> Tuple[list[int], Optional[dict[str, list[Any]]]]: - """Create input tokens for the model. Also return the mm data if the model is multimodal. + """ + Create input tokens for the model. Also return the mm data if the model is multimodal. Returns: Tuple[list[int], Optional[dict[str, list[Any]]]]: Input tokens and mm data. @@ -888,7 +896,7 @@ async def request_preprocess( prev_msgs = await conversation_store.get_conversation_history( prev_response_id) - _responses_debug_log("Prev msgs:") + _responses_debug_log(f"Prev msgs:") for msg in prev_msgs: _responses_debug_log(f" -> {msg}") @@ -1368,8 +1376,10 @@ def get_reasoning_text_delta_event( def _get_output_added_events( self, output_item: ResponseOutputMessage | ResponseReasoningItem ) -> list[StreamingResponsesResponse]: - """Get item added event and content part added event for a message item which is starting - to be generated. + """Get the added events for a message item. + + Returns the item-added and content-part-added events when generation + starts. Returns: list[StreamingResponsesResponse]: A list of streaming responses responses @@ -1427,7 +1437,8 @@ def _should_send_done_events( streaming_events_helper: Optional[ResponsesStreamingEventsHelper] = None, finished_generation: bool = False, ) -> Tuple[bool, bool, Optional[str], Optional[str]]: - """Determine if done events should be sent for text or reasoning items. + """ + Determine if done events should be sent for text or reasoning items. Analyzes the complete output text to detect when reasoning or text sections have been completed and should receive done events. @@ -1955,7 +1966,8 @@ async def process_streaming_events( class ServerArrivalTimeMiddleware: - """Custom ASGI middleware to track server arrival time. + """ + Custom ASGI middleware to track server arrival time. We implement this as a pure ASGI middleware instead of using FastAPI's @app.middleware("http") decorator because the decorator internally uses @@ -2098,7 +2110,8 @@ def percentile(q): class ResponseHooks(ABC): - """Hooks for response processing and (disagg) service perf observability. + """ + Hooks for response processing and (disagg) service perf observability. """ @abstractmethod diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 6627eb8e0213..ec318ced058b 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -487,7 +487,7 @@ async def get_next_server( self, request: OpenAIRequest, exclude_server: Optional[str] = None) -> tuple[str, dict]: - """Select server by request and return some intermediate information, exclude_server is a server to exclude from the selection""" + """Select server by request and return some intermediate information""" @abstractmethod async def finish_request(self, @@ -793,17 +793,6 @@ class KvCacheAwareRouter(BlockHashMixin, LoadBalancingMixin, Router): _server_state_class = KvCacheAwareServerState - @staticmethod - def _get_request_cache_salt_id(request: OpenAIRequest) -> Optional[int]: - """Resolve request cache salt through this module's public shim. - - ``router.py`` historically exposed ``get_cache_salt_id``. Keeping the - lookup here preserves that compatibility while hashing helpers live in - ``router_utils.py``. - """ - cache_salt = getattr(request, "cache_salt", None) - return None if cache_salt is None else get_cache_salt_id(cache_salt) - def __init__(self, server_role: ServerRole = None, servers: list[str] = None, @@ -1293,10 +1282,8 @@ def _get_content_load(self, server: str) -> int: return self._server_content_load.get(server, 0) def _get_server_load(self, server: str) -> int: - """Return content-weighted server load. - - This makes ``_select_least_loaded`` balance by estimated tokens rather - than request count. + """Use content weight so ``_select_least_loaded`` balances by + estimated tokens rather than request count. """ return self._get_content_load(server) diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index d2222d0d8982..99184cb3296a 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -517,7 +517,7 @@ def __init__(self, prompt): use_tokens=False, max_batch_size=32, tokens_per_block=tokens_per_block) - monkeypatch.setattr("tensorrt_llm.serve.router.get_cache_salt_id", + monkeypatch.setattr("tensorrt_llm.serve.router_utils.get_cache_salt_id", lambda cache_salt: cache_salt_id) for server in servers: router._server_info[server] = { @@ -785,29 +785,6 @@ def test_kv_cache_aware_server_state_remove_blocks_silent_on_missing(): assert state._block_table(KV_CACHE_HASH_ALGO_V1) == {30} -@pytest.mark.asyncio -async def test_kv_cache_aware_router_applies_routed_blocks_at_routing(servers): - tokens_per_block = 4 - token_lists = [[1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008]] - router = KvCacheAwareRouter(server_role=None, - servers=servers, - use_tokens=False, - max_batch_size=32, - tokens_per_block=tokens_per_block, - track_routed_blocks=True) - - request = CompletionRequest(model="TinyLlama", - prompt=copy.deepcopy(token_lists)) - server, info = await router.get_next_server(request) - - expected_flat = [h for hl in info["block_hashes"] for h in hl] - block_table = router._server_state[server]._block_table(info["hash_algo"]) - assert block_table.issuperset(expected_flat) - - await router.finish_request(request) - assert block_table.issuperset(expected_flat) - - @pytest.mark.asyncio async def test_kv_cache_aware_router_does_not_wait_for_finish_to_apply_blocks( servers): From 192ec5442037e57484158ce87626b3f458cc032e Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Fri, 10 Jul 2026 00:00:58 -0700 Subject: [PATCH 06/18] remove orjson and debug logs Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/serve/coordinator_server.py | 93 ++++------- tensorrt_llm/serve/disagg_coordinator.py | 28 +++- tensorrt_llm/serve/openai_disagg_server.py | 22 +-- tensorrt_llm/serve/openai_disagg_service.py | 35 ---- tensorrt_llm/serve/responses_utils.py | 82 ---------- tensorrt_llm/serve/router.py | 153 +++--------------- tensorrt_llm/serve/router_utils.py | 5 +- .../disagg_config_multi_orchestrator.yaml | 24 +++ .../defs/disaggregated/test_disaggregated.py | 34 +++- .../test_lists/test-db/l0_dgx_h100.yml | 1 + 10 files changed, 143 insertions(+), 334 deletions(-) create mode 100644 tests/integration/defs/disaggregated/test_configs/disagg_config_multi_orchestrator.yaml diff --git a/tensorrt_llm/serve/coordinator_server.py b/tensorrt_llm/serve/coordinator_server.py index efd55f690507..3217720e1553 100644 --- a/tensorrt_llm/serve/coordinator_server.py +++ b/tensorrt_llm/serve/coordinator_server.py @@ -32,35 +32,21 @@ """ import asyncio -import time from contextlib import asynccontextmanager from typing import Optional -import orjson +import msgpack import uvicorn from fastapi import FastAPI, Request -from fastapi.responses import ORJSONResponse, Response +from fastapi.responses import Response from tensorrt_llm.logger import logger +from tensorrt_llm.serve.cluster_storage import HttpClusterStorageServer from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService from tensorrt_llm.version import __version__ as VERSION -# Timing headers to decompose the coordinator IPC round-trip (co-located, so -# wall-clock time.time() is comparable): client stamps send time, coordinator -# stamps recv/send, from which the client derives req_wire/handler/resp_wire. -HDR_CLIENT_SEND = "x-client-send-time" -HDR_COORD_RECV = "x-coord-recv-time" -HDR_COORD_SEND = "x-coord-send-time" - -# The coordinator is a single event loop serving /select,/finish for the whole -# fleet; each /select body carries a routing_key (a flat list of block hashes, -# hundreds of int64 for long prompts). stdlib json parse+serialize on that hot -# path keeps the loop busy and lengthens the queue (client-observed /select was -# ~hundreds of ms while the handler itself is ~0.1ms). orjson is 2-5x faster at -# both ends, freeing the loop sooner. JSONResponse -> ORJSONResponse and -# request bodies parsed with orjson.loads (bytes, no text decode). - TIMEOUT_KEEP_ALIVE = 10 # seconds +MSGPACK_MEDIA_TYPE = "application/msgpack" class CoordinatorServer: @@ -81,85 +67,76 @@ async def lifespan(app: FastAPI): self.app.add_api_route("/cluster_info", self.cluster_info, methods=["GET"]) self.app.add_api_route("/health", self.health, methods=["GET"]) self.app.add_api_route("/version", self.version, methods=["GET"]) - - def _timing_headers(self, raw_req: Request, recv_t: float) -> dict: - """Return timestamps used to decompose the coordinator round trip. - - The client send time is echoed and coordinator receive/send times are - added. ``recv_t`` is captured at handler entry; send time is captured - immediately before responding. - """ - h = {HDR_COORD_RECV: repr(recv_t), HDR_COORD_SEND: repr(time.time())} - cs = raw_req.headers.get(HDR_CLIENT_SEND) - if cs is not None: - h[HDR_CLIENT_SEND] = cs - return h + cluster_storage = self._coordinator.cluster_storage + if isinstance(cluster_storage, HttpClusterStorageServer): + cluster_storage.add_routes(self.app) + + @staticmethod + def _response(content: object, status_code: int = 200) -> Response: + return Response( + content=msgpack.packb(content, use_bin_type=True), + status_code=status_code, + media_type=MSGPACK_MEDIA_TYPE, + ) async def select(self, raw_req: Request) -> Response: - _recv = time.time() try: - body = orjson.loads(await raw_req.body()) + body = msgpack.unpackb(await raw_req.body(), raw=False) except Exception as e: - return ORJSONResponse(status_code=400, content={"error": f"invalid JSON body: {e}"}) + return self._response({"error": f"invalid MessagePack body: {e}"}, status_code=400) if not isinstance(body, dict) or "role" not in body or "routing_key" not in body: - return ORJSONResponse( - status_code=400, content={"error": "body must include 'role' and 'routing_key'"} + return self._response( + {"error": "body must include 'role' and 'routing_key'"}, status_code=400 ) role = body["role"] if role not in ("context", "ctx", "generation", "gen"): - return ORJSONResponse(status_code=400, content={"error": f"invalid role: {role}"}) + return self._response({"error": f"invalid role: {role}"}, status_code=400) req_id = body.get("req_id") if req_id is not None and (not isinstance(req_id, int) or isinstance(req_id, bool)): - return ORJSONResponse(status_code=400, content={"error": "req_id must be an integer"}) + return self._response({"error": "req_id must be an integer"}, status_code=400) exclude_server = body.get("exclude_server") if exclude_server is not None and not isinstance(exclude_server, str): - return ORJSONResponse( - status_code=400, content={"error": "exclude_server must be a string"} - ) + return self._response({"error": "exclude_server must be a string"}, status_code=400) try: server, info, req_id = await self._coordinator.select( role, body["routing_key"], req_id, exclude_server ) except ValueError as e: - return ORJSONResponse(status_code=503, content={"error": str(e)}) + return self._response({"error": str(e)}, status_code=503) except Exception as e: # noqa: BLE001 logger.error(f"CoordinatorServer.select failed: {e}") - return ORJSONResponse(status_code=500, content={"error": str(e)}) - return ORJSONResponse( - content={"server": server, "info": info, "req_id": req_id}, - headers=self._timing_headers(raw_req, _recv), - ) + return self._response({"error": str(e)}, status_code=500) + return self._response({"server": server, "info": info, "req_id": req_id}) async def finish(self, raw_req: Request) -> Response: - _recv = time.time() try: - body = orjson.loads(await raw_req.body()) + body = msgpack.unpackb(await raw_req.body(), raw=False) except Exception as e: - return ORJSONResponse(status_code=400, content={"error": f"invalid JSON body: {e}"}) + return self._response({"error": f"invalid MessagePack body: {e}"}, status_code=400) if not isinstance(body, dict) or "role" not in body or "req_id" not in body: - return ORJSONResponse( - status_code=400, content={"error": "body must include 'role' and 'req_id'"} + return self._response( + {"error": "body must include 'role' and 'req_id'"}, status_code=400 ) role = body["role"] if role not in ("context", "ctx", "generation", "gen"): - return ORJSONResponse(status_code=400, content={"error": f"invalid role: {role}"}) + return self._response({"error": f"invalid role: {role}"}, status_code=400) req_id = body["req_id"] if not isinstance(req_id, int) or isinstance(req_id, bool): - return ORJSONResponse(status_code=400, content={"error": "req_id must be an integer"}) + return self._response({"error": "req_id must be an integer"}, status_code=400) success = body.get("success", True) if not isinstance(success, bool): - return ORJSONResponse(status_code=400, content={"error": "success must be a boolean"}) + return self._response({"error": "success must be a boolean"}, status_code=400) await self._coordinator.finish(role, req_id, success) - return ORJSONResponse(content={}, headers=self._timing_headers(raw_req, _recv)) + return self._response({}) async def cluster_info(self) -> Response: - return ORJSONResponse(content=await self._coordinator.cluster_info()) + return self._response(await self._coordinator.cluster_info()) async def health(self) -> Response: return Response(status_code=200 if await self._coordinator.is_ready() else 503) async def version(self) -> Response: - return ORJSONResponse(content={"version": VERSION}) + return self._response({"version": VERSION}) async def __call__(self, host: str, port: int, uds: Optional[str] = None) -> None: # Single-process (owns routing state + the centralized ZMQ ingest bind); diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py index 06677aa8a00c..16f51ab5d723 100644 --- a/tensorrt_llm/serve/disagg_coordinator.py +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -40,6 +40,7 @@ from typing import Any, Dict, Optional, Tuple import aiohttp +import msgpack from tensorrt_llm.llmapi.disagg_utils import ( DisaggServerConfig, @@ -475,7 +476,7 @@ async def _await_coordinator(self) -> Dict[str, Any]: logger.info( f"CoordinatorClient: coordinator reachable at {self._remote_url}" ) - return await resp.json() + return msgpack.unpackb(await resp.read(), raw=False) last_err = f"HTTP {resp.status}" except Exception as e: # noqa: BLE001 last_err = str(e) @@ -503,18 +504,39 @@ async def is_ready(self) -> bool: async with self.session.get( f"{self._remote_url}/health", timeout=self._request_timeout_s ) as resp: - return resp.status == 200 + if resp.status != 200: + return False + info = await self.cluster_info() + await self._sync_stateless_routers(info) + return info.get("is_ready", False) except Exception as e: # noqa: BLE001 logger.warning(f"CoordinatorClient health check failed: {e}") return False + async def _sync_stateless_routers(self, info: Dict[str, Any]) -> None: + current_workers = info.get("current_workers") + if current_workers is None: + return + router_workers = ( + (self._ctx_router, current_workers.get("context_servers", [])), + (self._gen_router, current_workers.get("generation_servers", [])), + ) + for router, workers in router_workers: + if isinstance(router, CoordinatorDelegatingRouter): + continue + desired = {f"{worker['host']}:{worker['port']}" for worker in workers} + for server in set(router.servers) - desired: + await router.remove_server(server) + for server in desired - set(router.servers): + await router.add_server(server) + async def cluster_info(self) -> Dict[str, Any]: try: async with self.session.get( f"{self._remote_url}/cluster_info", timeout=self._request_timeout_s ) as resp: if resp.status == 200: - return await resp.json() + return msgpack.unpackb(await resp.read(), raw=False) except Exception as e: # noqa: BLE001 logger.warning(f"CoordinatorClient cluster_info failed: {e}") return {"is_ready": False} diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 7ed8b56c5dbb..0c07961d3248 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -48,8 +48,7 @@ UCompletionResponse, ensure_request_chat_template_allowed) from tensorrt_llm.serve.perf_metrics import DisaggPerfMetricsCollector from tensorrt_llm.serve.responses_utils import (ServerArrivalTimeMiddleware, - get_steady_clock_now_in_seconds, - ttft_split_logger) + get_steady_clock_now_in_seconds) from tensorrt_llm.serve.router import Router from tensorrt_llm.version import __version__ as VERSION @@ -67,11 +66,7 @@ def __init__(self, raw_req: Request, perf_metrics_collector: DisaggPerfMetricsCo self.gen_server = "" self.request_arrival_time = raw_req.state.server_arrival_time self.server_first_token_time = 0 - # Orchestrator-side TTFT timeline stamps: arrival -> ctx_dispatch (pre-ctx - # wait) -> ctx_resp (ctx routing + prefill) -> first_token (KV transfer + - # gen). 0 until the corresponding hook fires. self.ctx_dispatch_time = 0 - self.ctx_resp_time = 0 self.perf_metrics_collector = perf_metrics_collector def on_req_begin(self, request: UCompletionRequest): @@ -82,27 +77,12 @@ def on_ctx_dispatch(self, request: UCompletionRequest): def on_ctx_resp(self, ctx_server: str, response: UCompletionResponse): self.ctx_server = ctx_server - self.ctx_resp_time = get_steady_clock_now_in_seconds() def on_first_token(self, gen_server: str, request: UCompletionRequest, response: UCompletionResponse = None): self.gen_server = gen_server self.server_first_token_time = get_steady_clock_now_in_seconds() def on_resp_done(self, gen_server: str, request: UCompletionRequest, response: UCompletionResponse = None): - # Feed the TTFT breakdown to a periodic aggregator (p50/p95 every N reqs), - # not one log line per request -- per-request logging on the disagg event - # loop inflated TTFT ~3s (observer effect). - arr = self.request_arrival_time - disp = self.ctx_dispatch_time - cresp = self.ctx_resp_time - ft = self.server_first_token_time - if arr and ft: - ttft_split_logger().record({ - "pre_ctx_ms": (disp - arr) * 1000 if disp else -1.0, - "ctx_phase_ms": (cresp - disp) * 1000 if (disp and cresp) else -1.0, - "xfer_gen_ms": (ft - cresp) * 1000 if cresp else -1.0, - "total_ms": (ft - arr) * 1000, - }) if request.disaggregated_params: ctx_req_id = request.disaggregated_params.ctx_request_id task = asyncio.create_task( diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 978850aa4ed0..30c1e56a8cb1 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -14,7 +14,6 @@ import asyncio import os -import time as _time from typing import Callable, Optional from tensorrt_llm.llmapi.disagg_utils import ConditionalDisaggConfig, DisaggServerConfig, ServerRole @@ -43,28 +42,6 @@ _GEN_PENDING_FINISH_REASONS = ("length", "not_finished") -_PRECTX_LOGGER = None - - -def _prectx_logger(): - """Return the orchestrator pre-context timing logger. - - The logger isolates where arrival-to-context-dispatch time is spent: - cond_ms : _check_conditional_disagg + _check_gen_only_disagg - idgen_ms : get_disagg_request_id (local now; was an HTTP hop) - ctxroute_ms: ctx_router.get_next_server (routing/select round trip) - No per-request logging (observer effect on the single fleet loop). - """ - global _PRECTX_LOGGER - if _PRECTX_LOGGER is None: - from tensorrt_llm.serve.responses_utils import PeriodicBreakdownLogger - - _PRECTX_LOGGER = PeriodicBreakdownLogger( - "orchestrator_prectx", ["cond_ms", "idgen_ms", "ctxroute_ms"] - ) - return _PRECTX_LOGGER - - class OpenAIDisaggregatedService(OpenAIService): def __init__( self, @@ -146,20 +123,11 @@ async def _send_disagg_request_ctx_first( # empty server means client decides which server to use ctx_server = None # reserve a gen_server if conditional disagg is needed - _t_cond = _time.monotonic() gen_server, need_ctx = await self._check_conditional_disagg(request) need_ctx = need_ctx and not await self._check_gen_only_disagg(request) - _t_idgen = _time.monotonic() ctx_response = None gen_req = request disagg_request_id = await self._coordinator.get_disagg_request_id() - _t_afterid = _time.monotonic() - # Fine-grained pre-ctx breakdown (aggregated). ctxroute is filled below. - _prectx = { - "cond_ms": (_t_idgen - _t_cond) * 1000, - "idgen_ms": (_t_afterid - _t_idgen) * 1000, - "ctxroute_ms": -1.0, - } if need_ctx: # Mark ctx-dispatch start: arrival->here is the pre-ctx wait in the # orchestrator/fleet (accept queue + event loop + pipeline). @@ -167,11 +135,9 @@ async def _send_disagg_request_ctx_first( hooks.on_ctx_dispatch(request) ctx_req = self._get_ctx_request(request, disagg_request_id) # ctx generator is empty - _t_route = _time.monotonic() ctx_server, _ = await self._ctx_router.get_next_server( ctx_req, exclude_server=gen_server ) - _prectx["ctxroute_ms"] = (_time.monotonic() - _t_route) * 1000 ctx_response = await self._ctx_client.send_request( ctx_req, server=ctx_server, hooks=hooks ) @@ -190,7 +156,6 @@ async def _send_disagg_request_ctx_first( and gen_req.disaggregated_params.request_type == "context_only" ): gen_req.disaggregated_params = None - _prectx_logger().record(_prectx) if ctx_response is None or self._need_gen(ctx_response): if not gen_server: gen_server, _ = await self._gen_router.get_next_server( diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index 10e468f5a270..b3b1c935f77d 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -1992,57 +1992,11 @@ async def __call__(self, scope, receive, send): scope["state"] = {} scope["state"][ "server_arrival_time"] = get_steady_clock_now_in_seconds() - # Measure client-send -> server-accept latency: the network + HTTP - # accept-queue backlog BEFORE the handler runs (invisible to per- - # request server metrics, which start at server_arrival_time). The - # client stamps x-client-send-time (wall-clock epoch s); compare to - # our wall clock (NTP-synced hosts). Only /v1/ request paths. - path = scope.get("path", "") - if path.startswith("/v1/"): - for k, v in scope.get("headers", ()): # raw ASGI header pairs - if k == b"x-client-send-time": - try: - sent = float(v.decode()) - _accept_latency_logger().record( - max(0.0, - time.time() - sent)) - except (ValueError, UnicodeDecodeError): - pass - break # Pass through the original receive/send - no wrapping! await self.app(scope, receive, send) -_ACCEPT_LAT = None -_TTFT_SPLIT = None - - -def _accept_latency_logger(): - """Return the client-send to server-accept latency logger. - - ``PeriodicLatencyLogger`` is defined below this function. - """ - global _ACCEPT_LAT - if _ACCEPT_LAT is None: - _ACCEPT_LAT = PeriodicLatencyLogger("accept.client_send_to_arrival") - return _ACCEPT_LAT - - -def ttft_split_logger(): - """Return the per-request TTFT breakdown logger. - - Pre-context, context, transfer/generation, and total times are aggregated - and logged periodically rather than once per request. - """ - global _TTFT_SPLIT - if _TTFT_SPLIT is None: - _TTFT_SPLIT = PeriodicBreakdownLogger( - "orchestrator", - ["pre_ctx_ms", "ctx_phase_ms", "xfer_gen_ms", "total_ms"]) - return _TTFT_SPLIT - - class PeriodicLatencyLogger: """Periodically log latency percentiles for a named coordinator API. @@ -2073,42 +2027,6 @@ def percentile(q): self._samples = [] -class PeriodicBreakdownLogger: - """Periodically log a multi-field latency breakdown. - - Named millisecond values are accumulated and each field's p50/p95 is logged - every ``window`` samples. Negative values for stages not reached are dropped. - """ - - def __init__(self, name: str, fields, window: int = 1000): - self._name = name - self._fields = list(fields) - self._window = window - self._samples = {f: [] for f in self._fields} - self._n = 0 - - def record(self, values: dict) -> None: - self._n += 1 - for f in self._fields: - v = values.get(f) - if v is not None and v >= 0: - self._samples[f].append(v) - if self._n % self._window == 0: - parts = [] - for f in self._fields: - s = sorted(self._samples[f]) - if s: - - def percentile(q): - return s[min(int(q * len(s)), len(s) - 1)] - - parts.append(f"{f}(p50={percentile(0.5):.0f}," - f"p95={percentile(0.95):.0f})") - self._samples[f] = [] - logger.info(f"[ttft_split] {self._name} n={self._n} " + - " ".join(parts)) - - class ResponseHooks(ABC): """ Hooks for response processing and (disagg) service perf observability. diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index ec318ced058b..9a7c97b196a1 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -13,15 +13,13 @@ # limitations under the License. import asyncio -import os -import statistics import time from abc import ABC, abstractmethod from collections import OrderedDict from typing import Awaitable, Callable, Dict, Iterable, List, Optional import aiohttp -import orjson +import msgpack from tensorrt_llm.llmapi.disagg_utils import (MetadataServerConfig, RouterConfig, ServerRole) @@ -37,9 +35,7 @@ block_key_hasher, get_cache_salt_id, get_request_num_tokens, hash_v1_block_key, truncate_sha256_hash_to_int64, v2_sha256_block_hasher) -# Content-Type for orjson-serialized bodies sent to the coordinator (we pass -# data=bytes, so aiohttp doesn't set the json content type itself). -_JSON_HEADERS = {"Content-Type": "application/json"} +_MSGPACK_HEADERS = {"Content-Type": "application/msgpack"} # Max number of conversations whose home-server pin is retained (LRU). ROUTE_AFFINITY_CACHE_SIZE = 50000 @@ -70,17 +66,13 @@ def _session(self) -> Optional[aiohttp.ClientSession]: return self._session_provider() if self._session_provider else None async def increment_load(self, request: OpenAIRequest): - # request may be None on the coordinator-delegated path (no request - # object crosses the HTTP hop); token accounting is skipped then. - num_tokens = (get_request_num_tokens(request) - if self._use_tokens and request is not None else 0) + num_tokens = get_request_num_tokens(request) if self._use_tokens else 0 async with self._lock: self._num_active_requests += 1 self._num_active_tokens += num_tokens async def decrement_load(self, request: OpenAIRequest): - num_tokens = (get_request_num_tokens(request) - if self._use_tokens and request is not None else 0) + num_tokens = get_request_num_tokens(request) if self._use_tokens else 0 async with self._lock: self._num_active_requests -= 1 self._num_active_tokens -= num_tokens @@ -108,7 +100,6 @@ def __init__( self._kv_cache_block_tables: dict[str, set[BlockHash]] = { KV_CACHE_HASH_ALGO_V1: self._kv_cache_block_table } - self._event_only_blocks: set[BlockHash] = set() self._kv_cache_hash_algo = KV_CACHE_HASH_ALGO_DEFAULT self._tokens_per_block = tokens_per_block self._poll_task: Optional[asyncio.Task] = None @@ -170,10 +161,8 @@ def update_with_events(self, events: Iterable[dict]): block["block_hash"] for block in event["blocks"] ] self.add_blocks(block_hashes, hash_algo=hash_algo) - self._event_only_blocks.update(block_hashes) elif event["type"] == "removed": self.remove_blocks(event["block_hashes"], hash_algo=hash_algo) - self._event_only_blocks.difference_update(event["block_hashes"]) async def poll_events(self, session: aiohttp.ClientSession): async with session.post( @@ -181,33 +170,19 @@ async def poll_events(self, session: aiohttp.ClientSession): events_raw = await response.json() return events_raw - _event_match_log_counter = 0 - async def matched_tokens( self, block_hashes: list[list[BlockHash]], hash_algo: str = KV_CACHE_HASH_ALGO_DEFAULT) -> int: match_count = 0 - event_match_count = 0 async with self._lock: block_table = self._block_table(hash_algo) for hash_list in block_hashes: for block_hash in hash_list: if block_hash in block_table: match_count += self._tokens_per_block - if block_hash in self._event_only_blocks: - event_match_count += self._tokens_per_block else: break - KvCacheAwareServerState._event_match_log_counter += 1 - if KvCacheAwareServerState._event_match_log_counter <= 20 or KvCacheAwareServerState._event_match_log_counter % 100 == 0: - logger.info( - f"EVENT_MATCH_DIAG server={self._server} " - f"total_match={match_count} event_match={event_match_count} " - f"event_blocks={len(self._event_only_blocks)} " - f"total_blocks={len(block_table)} " - f"query_hashes={[h for hl in block_hashes for h in hl][:3]}" - ) return match_count async def decrement_load(self, request: OpenAIRequest): @@ -353,37 +328,6 @@ def __init__( self._health_check_timeout = metadata_server_cfg.health_check_timeout if metadata_server_cfg else None self._server_preparation_func = server_preparation_func self._prepared_ready_servers: set[str] = set() - # Routing-latency diagnostics (gated by TLLM_LOG_ROUTE_TIMING=1): record - # get_next_server wall time per request and log percentiles periodically. - self._log_route_timing = (os.environ.get("TLLM_LOG_ROUTE_TIMING", - "0") == "1") - self._rt_samples: list = [] - self._rt_n = 0 - - def _record_route_timing(self, dt_s: float) -> None: - """Record one routing-latency sample. - - Percentiles are logged every 500 calls. This is a no-op unless - ``TLLM_LOG_ROUTE_TIMING=1``. - """ - if not self._log_route_timing: - return - self._rt_samples.append(dt_s * 1000.0) # ms - self._rt_n += 1 - if self._rt_n % 500 == 0: - s = sorted(self._rt_samples) - n = len(s) - - def percentile(q): - return s[min(int(q * n), n - 1)] - - logger.info(f"[route_timing] {type(self).__name__} n={self._rt_n} " - f"get_next_server_ms: mean={statistics.mean(s):.2f} " - f"p50={percentile(0.5):.2f} " - f"p90={percentile(0.9):.2f} " - f"p99={percentile(0.99):.2f} " - f"max={s[-1]:.2f}") - self._rt_samples = [] # reset window async def close(self): """Close the shared HTTP session.""" @@ -951,7 +895,6 @@ async def _route(self, key, exclude_server=None, request=None, req_id=None): conversation affinity, and round-robin tie-breaking. Standalone requests are keyed by ``id(request)``; coordinator requests use ``req_id``. """ - route_start = time.monotonic() token_lists = (key or {}).get("token_lists") or [] block_hashes_by_algo = (key or {}).get("block_hashes_by_algo") or {} conv_key = (key or {}).get("conv_key") @@ -1019,7 +962,6 @@ def _hashes(server): await server_states[server].increment_load(request) self._req_routing_table[key] = server self._apply_routed_blocks_on_route(server, block_hashes, hash_algo) - self._record_route_timing(time.monotonic() - route_start) return server, { "block_hashes": block_hashes, "hash_algo": hash_algo, @@ -1631,25 +1573,9 @@ def __init__(self, # Coordinator HTTP-client API latency (includes network round-trip to the # coordinator + its in-process handler). Compare against the owner-side # [coord_api] to isolate the fleet /select|/finish HTTP overhead. - from tensorrt_llm.serve.responses_utils import (PeriodicBreakdownLogger, - PeriodicLatencyLogger) + from tensorrt_llm.serve.responses_utils import PeriodicLatencyLogger self._select_lat = PeriodicLatencyLogger(f"client.select[{role}]") self._finish_lat = PeriodicLatencyLogger(f"client.finish[{role}]") - # IPC latency breakdown (ms) for the coordinator client<->server round - # trip, decomposed via the X-Client-Send/X-Coord-Recv/X-Coord-Send headers - # (client + coordinator co-located => wall clock comparable): - # req_wire : client send -> coord recv (request transport + accept) - # handler : coord recv -> coord send (in-coordinator handler work) - # resp_wire : coord send -> client recv (response transport) - # rtt : client send -> client recv (total observed round trip) - # Aggregated p50/p95 per field every window samples (no per-call logging). - self._ipc_fields = [ - "req_wire_ms", "handler_ms", "resp_wire_ms", "rtt_ms" - ] - self._select_ipc = PeriodicBreakdownLogger(f"coord_ipc.select[{role}]", - self._ipc_fields) - self._finish_ipc = PeriodicBreakdownLogger(f"coord_ipc.finish[{role}]", - self._ipc_fields) def __getattr__(self, name): # servers / prepare_servers / num_prepared_servers / start_server_monitoring @@ -1667,35 +1593,6 @@ def session(self) -> aiohttp.ClientSession: def _on_servers_updated(self, old_servers, new_servers): pass - def _record_ipc(self, ipc_logger, send_t: float, recv_t: float, - resp) -> None: - """Record a coordinator round-trip breakdown. - - ``send_t`` and ``recv_t`` are client-side wall clocks. Missing headers - from an older coordinator are recorded as ``-1`` and dropped by the - breakdown logger. All values are milliseconds. - """ - from tensorrt_llm.serve.coordinator_server import (HDR_COORD_RECV, - HDR_COORD_SEND) - - def _f(name): - v = resp.headers.get(name) - try: - return float(v) if v is not None else None - except (TypeError, ValueError): - return None - - coord_recv = _f(HDR_COORD_RECV) - coord_send = _f(HDR_COORD_SEND) - ipc_logger.record({ - "req_wire_ms": (coord_recv - send_t) * 1000 if coord_recv else -1.0, - "handler_ms": (coord_send - coord_recv) * 1000 if - (coord_recv and coord_send) else -1.0, - "resp_wire_ms": - (recv_t - coord_send) * 1000 if coord_send else -1.0, - "rtt_ms": (recv_t - send_t) * 1000, - }) - def _request_id(self, request: OpenAIRequest) -> int: """The request's disagg id -- the sole cross-process key for select/finish. @@ -1735,23 +1632,16 @@ async def get_next_server( "req_id": self._request_id(request), "exclude_server": exclude_server } - # orjson.dumps (bytes) instead of aiohttp json= (stdlib dumps): the - # routing_key is hundreds of int64, and this is the hot fleet->coordinator - # call. Parse the response with orjson too. X-Client-Send-Time (wall clock) - # lets the coordinator echo timing headers back for the IPC breakdown. - from tensorrt_llm.serve.coordinator_server import HDR_CLIENT_SEND _t0 = time.monotonic() - _send_wall = time.time() - hdrs = {**_JSON_HEADERS, HDR_CLIENT_SEND: repr(_send_wall)} async with self.session.post(f"{self._coordinator_url}/select", - data=orjson.dumps(payload), - headers=hdrs, + data=msgpack.packb(payload, + use_bin_type=True), + headers=_MSGPACK_HEADERS, timeout=self._request_timeout_s) as resp: + body = msgpack.unpackb(await resp.read(), raw=False) if resp.status != 200: raise ValueError(f"coordinator /select returned {resp.status}: " - f"{await resp.text()}") - body = orjson.loads(await resp.read()) - self._record_ipc(self._select_ipc, _send_wall, time.time(), resp) + f"{body.get('error', body)}") self._select_lat.record(time.monotonic() - _t0) info = body.get("info") or {} return body["server"], info @@ -1771,25 +1661,24 @@ async def finish_request(self, task.add_done_callback(self._background_tasks.discard) async def _finish_async(self, req_id: int, success: bool): - from tensorrt_llm.serve.coordinator_server import HDR_CLIENT_SEND _t0 = time.monotonic() - _send_wall = time.time() - hdrs = {**_JSON_HEADERS, HDR_CLIENT_SEND: repr(_send_wall)} try: async with self.session.post( f"{self._coordinator_url}/finish", - data=orjson.dumps({ - "role": self._role, - "req_id": req_id, - "success": success - }), - headers=hdrs, + data=msgpack.packb( + { + "role": self._role, + "req_id": req_id, + "success": success + }, + use_bin_type=True), + headers=_MSGPACK_HEADERS, timeout=self._request_timeout_s) as resp: if resp.status != 200: + body = msgpack.unpackb(await resp.read(), raw=False) logger.warning( - f"coordinator /finish returned {resp.status}") - self._record_ipc(self._finish_ipc, _send_wall, time.time(), - resp) + f"coordinator /finish returned {resp.status}: " + f"{body.get('error', body)}") self._finish_lat.record(time.monotonic() - _t0) except Exception as e: # noqa: BLE001 logger.warning(f"CoordinatorDelegatingRouter finish failed: {e}") diff --git a/tensorrt_llm/serve/router_utils.py b/tensorrt_llm/serve/router_utils.py index e27199a848d8..b9d8c9265aa8 100644 --- a/tensorrt_llm/serve/router_utils.py +++ b/tensorrt_llm/serve/router_utils.py @@ -113,7 +113,10 @@ def has_worker(self, owner_id: str) -> bool: return bool(self._blocks) -def get_request_num_tokens(request: OpenAIRequest) -> int: +def get_request_num_tokens(request: Optional[OpenAIRequest]) -> int: + if request is None: + return 0 + if ( request.disaggregated_params is None or request.disaggregated_params.request_type == "context_only" diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_multi_orchestrator.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_multi_orchestrator.yaml new file mode 100644 index 000000000000..970c2e276647 --- /dev/null +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_multi_orchestrator.yaml @@ -0,0 +1,24 @@ +hostname: localhost +model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 +num_workers: 4 +free_gpu_memory_fraction: 0.25 +backend: pytorch +disable_overlap_scheduler: true +context_servers: + num_instances: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + kv_cache_config: + free_gpu_memory_fraction: 0.2 + cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON +generation_servers: + num_instances: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON + router: + type: conversation diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index fa6321f4207c..62e318d3c1ad 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -149,6 +149,7 @@ def build_worker_config(base_config: dict[str, Any], EXCLUDE_FROM_WORKER = { 'hostname', 'port', + 'num_workers', 'num_instances', 'urls', 'router', @@ -206,6 +207,8 @@ def get_test_config(test_desc, example_dir, test_root): f"{test_configs_root}/disagg_config_load_balancing.yaml", "conversation": f"{test_configs_root}/disagg_config_conversation.yaml", + "multi_orchestrator": + f"{test_configs_root}/disagg_config_multi_orchestrator.yaml", "4_ranks": f"{test_configs_root}/disagg_config_ctxtp2_gentp1.yaml", "cuda_graph": @@ -644,7 +647,16 @@ def setup_disagg_cluster( server_host = config.get("hostname", "localhost") server_port = get_free_port() work_dir = tempfile.mkdtemp() - disagg_cluster["cluster_uri"] = f"http://{server_host}:{server_port}" + server_env = env + coordinator_url = f"http://{server_host}:{server_port}" + if config.get("num_workers", 1) > 1: + coordinator_port = get_free_port() + while coordinator_port == server_port: + coordinator_port = get_free_port() + coordinator_url = f"http://{server_host}:{coordinator_port}" + server_env = (env or os.environ).copy() + server_env["TRTLLM_DISAGG_COORDINATOR_PORT"] = str(coordinator_port) + disagg_cluster["cluster_uri"] = coordinator_url # Auto-deduce minimal_instances from num_instances ctx_servers = config.get("context_servers", {}) @@ -721,6 +733,8 @@ def setup_disagg_cluster( server_host, "port": server_port, + "num_workers": + config.get("num_workers", 1), "disagg_cluster": disagg_cluster, "context_servers": { @@ -740,7 +754,7 @@ def setup_disagg_cluster( work_dir, server_port, save_log=save_log, - env=env, + env=server_env, cwd=cwd) async def _wait_with_ticker(): @@ -869,6 +883,22 @@ def test_disaggregated_single_gpu(disaggregated_test_root, cwd=llm_venv.get_working_directory()) +@pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], + indirect=True) +def test_disaggregated_tinyllama_multi_orchestrator(disaggregated_test_root, + disaggregated_example_root, + llm_venv, llama_model_root): + setup_model_symlink(llm_venv, llama_model_root, + "TinyLlama/TinyLlama-1.1B-Chat-v1.0") + + run_disaggregated_test(disaggregated_example_root, + "multi_orchestrator", + num_iters=1, + env=llm_venv._new_env, + model_path=llama_model_root, + cwd=llm_venv.get_working_directory()) + + @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], indirect=True) def test_disaggregated_benchmark_gen_only(disaggregated_test_root, diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index e940d149d02f..f35937f21f98 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -120,6 +120,7 @@ l0_dgx_h100: - unittest/disaggregated/test_py_cache_transceiver_mp.py -k "gen_first1" - unittest/disaggregated/test_py_cache_transceiver_mp.py -k "gen_first2" - disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu[TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated.py::test_disaggregated_tinyllama_multi_orchestrator[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_ctxtp2_genpp2[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_gentp2[TinyLlama-1.1B-Chat-v1.0] From 01ef20eeb70e1b110d5e475ae80f2e16a36a85b3 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Sun, 12 Jul 2026 19:32:03 -0700 Subject: [PATCH 07/18] [None][fix] address coordinator review comments Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/commands/serve.py | 44 ++++++++++----- tensorrt_llm/serve/coordinator_server.py | 14 +++-- tensorrt_llm/serve/disagg_coordinator.py | 3 - tensorrt_llm/serve/openai_disagg_service.py | 61 ++++++++++++++------- tensorrt_llm/serve/router.py | 8 ++- 5 files changed, 89 insertions(+), 41 deletions(-) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 02f7d8ee0e2b..50c6b6506ddb 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1777,13 +1777,21 @@ def _cleanup(): for p in fleet: if p.poll() is None: p.terminate() + deadline = time.monotonic() + 10 for p in fleet: try: - p.wait(timeout=10) + p.wait(timeout=max(0, deadline - time.monotonic())) except Exception: p.kill() + p.wait() + + def _handle_signal(signum, _frame): + _cleanup() + raise SystemExit(128 + signum) atexit.register(_cleanup) + signal.signal(signal.SIGTERM, _handle_signal) + signal.signal(signal.SIGINT, _handle_signal) return fleet @@ -1800,17 +1808,22 @@ def _serve_disagg_fleet(disagg_cfg, config_file, metadata_server_config_file, server_start_timeout, num_workers, coordinator_url) # Block until any worker exits; a nonzero exit from any worker is a failure. - while True: - for i, p in enumerate(fleet): - rc = p.poll() - if rc is not None: - if rc != 0: - raise RuntimeError( - f"Disagg fleet worker {i} (pid={p.pid}) exited with " - f"code {rc}") - # A clean exit of one worker ends the fleet. - return - time.sleep(1) + try: + while True: + for i, p in enumerate(fleet): + rc = p.poll() + if rc is not None: + if rc != 0: + raise RuntimeError( + f"Disagg fleet worker {i} (pid={p.pid}) exited with " + f"code {rc}") + # A clean exit of one worker ends the fleet. + return + time.sleep(1) + finally: + for process in fleet: + if process.poll() is None: + process.terminate() def _serve_coordinator_and_fleet(disagg_cfg, config_file, @@ -1891,7 +1904,12 @@ async def _monitor_fleet(): for task in done: task.result() - asyncio.run(_serve_and_monitor()) + try: + asyncio.run(_serve_and_monitor()) + finally: + for process in fleet: + if process.poll() is None: + process.terminate() def _init_fleet_worker_process(): diff --git a/tensorrt_llm/serve/coordinator_server.py b/tensorrt_llm/serve/coordinator_server.py index 3217720e1553..8eb081fd5e34 100644 --- a/tensorrt_llm/serve/coordinator_server.py +++ b/tensorrt_llm/serve/coordinator_server.py @@ -84,15 +84,21 @@ async def select(self, raw_req: Request) -> Response: body = msgpack.unpackb(await raw_req.body(), raw=False) except Exception as e: return self._response({"error": f"invalid MessagePack body: {e}"}, status_code=400) - if not isinstance(body, dict) or "role" not in body or "routing_key" not in body: + if ( + not isinstance(body, dict) + or "role" not in body + or "routing_key" not in body + or "req_id" not in body + ): return self._response( - {"error": "body must include 'role' and 'routing_key'"}, status_code=400 + {"error": "body must include 'role', 'routing_key', and 'req_id'"}, + status_code=400, ) role = body["role"] if role not in ("context", "ctx", "generation", "gen"): return self._response({"error": f"invalid role: {role}"}, status_code=400) - req_id = body.get("req_id") - if req_id is not None and (not isinstance(req_id, int) or isinstance(req_id, bool)): + req_id = body["req_id"] + if not isinstance(req_id, int) or isinstance(req_id, bool): return self._response({"error": "req_id must be an integer"}, status_code=400) exclude_server = body.get("exclude_server") if exclude_server is not None and not isinstance(exclude_server, str): diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py index 16f51ab5d723..abbb00fb59bc 100644 --- a/tensorrt_llm/serve/disagg_coordinator.py +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -189,9 +189,6 @@ async def select( ) -> Tuple[str, dict, Optional[str]]: _t0 = time.monotonic() router = self._router_for_role(role) - if req_id is None: - # The coordinator owns IDs absent from generation requests. - req_id = get_global_disagg_request_id(self._config.node_id) result = await router.get_next_server_by_key( routing_key, req_id=req_id, exclude_server=exclude_server ) diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 30c1e56a8cb1..7c9c53cfaaa1 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -35,7 +35,7 @@ UCompletionResponseOrGenerator, done_generator, ) -from tensorrt_llm.serve.router import KvCacheAwareRouter, Router +from tensorrt_llm.serve.router import CoordinatorDelegatingRouter, KvCacheAwareRouter, Router # Finish reasons for which a GEN handoff is still pending; any other reason means # the CTX request already completed and the disagg KV-cache handoff was never set up. @@ -122,30 +122,41 @@ async def _send_disagg_request_ctx_first( hooks.on_req_begin(request) # empty server means client decides which server to use ctx_server = None + disagg_request_id = await self._coordinator.get_disagg_request_id() + if self.conditional_disagg_config: + if request.disaggregated_params is None: + self._get_gen_request(request, None, disagg_request_id) + else: + request.disaggregated_params.disagg_request_id = disagg_request_id + request.disaggregated_params.ctx_request_id = disagg_request_id # reserve a gen_server if conditional disagg is needed gen_server, need_ctx = await self._check_conditional_disagg(request) need_ctx = need_ctx and not await self._check_gen_only_disagg(request) ctx_response = None gen_req = request - disagg_request_id = await self._coordinator.get_disagg_request_id() if need_ctx: - # Mark ctx-dispatch start: arrival->here is the pre-ctx wait in the - # orchestrator/fleet (accept queue + event loop + pipeline). - if hooks: - hooks.on_ctx_dispatch(request) - ctx_req = self._get_ctx_request(request, disagg_request_id) - # ctx generator is empty - ctx_server, _ = await self._ctx_router.get_next_server( - ctx_req, exclude_server=gen_server - ) - ctx_response = await self._ctx_client.send_request( - ctx_req, server=ctx_server, hooks=hooks - ) - await self._verify_ctx_response(ctx_response) - ctx_response_disagg_params = ctx_response.choices[0].disaggregated_params - if ctx_response_disagg_params.disagg_request_id is not None: - disagg_request_id = ctx_response_disagg_params.disagg_request_id - gen_req = self._get_gen_request(request, ctx_response, disagg_request_id) + try: + # Mark ctx-dispatch start: arrival->here is the pre-ctx wait in the + # orchestrator/fleet (accept queue + event loop + pipeline). + if hooks: + hooks.on_ctx_dispatch(request) + ctx_req = self._get_ctx_request(request, disagg_request_id) + # ctx generator is empty + ctx_server, _ = await self._ctx_router.get_next_server( + ctx_req, exclude_server=gen_server + ) + ctx_response = await self._ctx_client.send_request( + ctx_req, server=ctx_server, hooks=hooks + ) + await self._verify_ctx_response(ctx_response) + ctx_response_disagg_params = ctx_response.choices[0].disaggregated_params + if ctx_response_disagg_params.disagg_request_id is not None: + disagg_request_id = ctx_response_disagg_params.disagg_request_id + gen_req = self._get_gen_request(request, ctx_response, disagg_request_id) + except Exception: + if gen_server: + await self._gen_router.finish_request(request, success=False) + raise else: # When need_ctx=False the gen server handles full generation and # must not see a stale request_type="context_only". @@ -166,6 +177,8 @@ async def _send_disagg_request_ctx_first( ) return gen_response else: + if gen_server: + await self._gen_router.finish_request(request) if request.stream: # ctx client will never return a generator when streaming is requested # make up for this by returning a done generator @@ -266,7 +279,15 @@ def _get_gen_request( async def _check_conditional_disagg(self, request: UCompletionRequest) -> bool: if self.conditional_disagg_config: - assert isinstance(self._gen_router, KvCacheAwareRouter) + local_gen_router = ( + self._gen_router._local + if isinstance(self._gen_router, CoordinatorDelegatingRouter) + else self._gen_router + ) + if not isinstance(local_gen_router, KvCacheAwareRouter): + raise TypeError( + "conditional disaggregation requires a KV-cache-aware generation router" + ) # Query kv cache status and select a best gen_server. # The server is reserved for generation request gen_server, info = await self._gen_router.get_next_server(request) diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 9a7c97b196a1..6221230c90cc 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -961,7 +961,13 @@ def _hashes(server): f"Selected server {server} is no longer available") await server_states[server].increment_load(request) self._req_routing_table[key] = server - self._apply_routed_blocks_on_route(server, block_hashes, hash_algo) + try: + self._apply_routed_blocks_on_route(server, block_hashes, + hash_algo) + except Exception: + self._req_routing_table.pop(key, None) + await server_states[server].decrement_load(request) + raise return server, { "block_hashes": block_hashes, "hash_algo": hash_algo, From 18d1fb9c598025ce8f07b3a310a6621bf5cecdc7 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Mon, 13 Jul 2026 18:25:40 -0700 Subject: [PATCH 08/18] [None][fix] preserve conditional disaggregation routing Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/serve/openai_client.py | 42 +++++++--- tensorrt_llm/serve/openai_disagg_service.py | 53 ++++++++----- tensorrt_llm/serve/router.py | 87 ++++++++++++--------- 3 files changed, 115 insertions(+), 67 deletions(-) diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 2e3542a44a9a..8834eca15dad 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -64,14 +64,20 @@ async def send_request( request: UCompletionRequest, server: Optional[str] = None, hooks: Optional[ResponseHooks] = None, + req_id: Optional[int] = None, ) -> UCompletionResponseOrGenerator: if isinstance(request, CompletionRequest): return await self._send_request( - "v1/completions", request, CompletionResponse, server, hooks + "v1/completions", request, CompletionResponse, server, hooks, req_id ) elif isinstance(request, ChatCompletionRequest): return await self._send_request( - "v1/chat/completions", request, ChatCompletionResponse, server, hooks + "v1/chat/completions", + request, + ChatCompletionResponse, + server, + hooks, + req_id, ) else: raise ValueError(f"Invalid request type: {type(request)}") @@ -84,6 +90,7 @@ async def _send_request( response_type: Type[UCompletionResponse], server: Optional[str] = None, hooks: Optional[ResponseHooks] = None, + req_id: Optional[int] = None, ) -> UCompletionResponseOrGenerator: """Send a request to the server and return the response and the body generator. @@ -102,7 +109,12 @@ async def check_ready(self) -> Tuple[List[str], List[str]]: async def shutdown(self) -> None: ... @abstractmethod - async def _finish_request(self, request: UCompletionRequest, success: bool = True) -> None: + async def _finish_request( + self, + request: UCompletionRequest, + success: bool = True, + req_id: Optional[int] = None, + ) -> None: """Finish the request in the router. ``success`` lets the router distinguish completed vs failed requests @@ -146,9 +158,10 @@ async def _send_request( response_type: Type[UCompletionResponse], server: Optional[str] = None, hooks: Optional[ResponseHooks] = None, + req_id: Optional[int] = None, ) -> UCompletionResponseOrGenerator: if server is None: - server, _ = await self._router.get_next_server(request) + server, _ = await self._router.get_next_server(request, req_id=req_id) url = f"http://{server}/{endpoint}" # disaggregated_params is None when conditional_disagg bypasses ctx. _dp = request.disaggregated_params @@ -156,7 +169,7 @@ async def _send_request( logger.debug(f"Sending {self._role} request {_ctx_rid} to {url}") try: self._metrics_collector.total_requests.inc() - resp_generator = self._post_with_retry(server, url, request, hooks) + resp_generator = self._post_with_retry(server, url, request, hooks, req_id) if request.stream: # return the response generator, the request is not done yet return resp_generator @@ -175,7 +188,7 @@ async def _send_request( except Exception: self._metrics_collector.error_requests.inc() # finish the request upon error - await self._finish_request(request, success=False) + await self._finish_request(request, success=False, req_id=req_id) raise async def _post_with_retry( @@ -184,6 +197,7 @@ async def _post_with_retry( url: str, request: UCompletionRequest, hooks: Optional[ResponseHooks] = None, + req_id: Optional[int] = None, ) -> AsyncGenerator[Any, None]: is_stream = request.stream # Loop range must cover the transient-TCP extended budget (up to 5) @@ -227,7 +241,7 @@ async def _post_with_retry( # do NOT return generator directly here or the response will go # out of scope and get destroyed async for line in self._response_generator( - request, http_response, start_time, server, hooks + request, http_response, start_time, server, hooks, req_id ): lines_yielded += 1 yield line @@ -246,7 +260,7 @@ async def _post_with_retry( # yield here since python forbids return statements in async generators yield response_dict # finish the request after the successful response - await self._finish_request(request) + await self._finish_request(request, req_id=req_id) self._metrics_collector.complete_latency_seconds.observe( get_steady_clock_now_in_seconds() - start_time ) @@ -296,6 +310,7 @@ async def _response_generator( start_time: float, server: str, hooks: Optional[ResponseHooks] = None, + req_id: Optional[int] = None, ) -> AsyncGenerator[Any, None]: assert request.stream, "Request is not streaming" assert "text/event-stream" in http_response.headers.get("Content-Type", ""), ( @@ -340,11 +355,16 @@ async def _response_generator( raise finally: # finish the request after streaming response is done or error is raised - await self._finish_request(request, success=success) + await self._finish_request(request, success=success, req_id=req_id) - async def _finish_request(self, request: UCompletionRequest, success: bool = True) -> None: + async def _finish_request( + self, + request: UCompletionRequest, + success: bool = True, + req_id: Optional[int] = None, + ) -> None: self._metrics_collector.completed_requests.inc() - await self._router.finish_request(request, self._session, success=success) + await self._router.finish_request(request, self._session, success=success, req_id=req_id) async def collect_metrics(self) -> Dict[str, Any]: metrics = {} diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 7c9c53cfaaa1..34cfef562681 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -123,14 +123,8 @@ async def _send_disagg_request_ctx_first( # empty server means client decides which server to use ctx_server = None disagg_request_id = await self._coordinator.get_disagg_request_id() - if self.conditional_disagg_config: - if request.disaggregated_params is None: - self._get_gen_request(request, None, disagg_request_id) - else: - request.disaggregated_params.disagg_request_id = disagg_request_id - request.disaggregated_params.ctx_request_id = disagg_request_id # reserve a gen_server if conditional disagg is needed - gen_server, need_ctx = await self._check_conditional_disagg(request) + gen_server, need_ctx = await self._check_conditional_disagg(request, disagg_request_id) need_ctx = need_ctx and not await self._check_gen_only_disagg(request) ctx_response = None gen_req = request @@ -143,10 +137,10 @@ async def _send_disagg_request_ctx_first( ctx_req = self._get_ctx_request(request, disagg_request_id) # ctx generator is empty ctx_server, _ = await self._ctx_router.get_next_server( - ctx_req, exclude_server=gen_server + ctx_req, exclude_server=gen_server, req_id=disagg_request_id ) ctx_response = await self._ctx_client.send_request( - ctx_req, server=ctx_server, hooks=hooks + ctx_req, server=ctx_server, hooks=hooks, req_id=disagg_request_id ) await self._verify_ctx_response(ctx_response) ctx_response_disagg_params = ctx_response.choices[0].disaggregated_params @@ -155,7 +149,9 @@ async def _send_disagg_request_ctx_first( gen_req = self._get_gen_request(request, ctx_response, disagg_request_id) except Exception: if gen_server: - await self._gen_router.finish_request(request, success=False) + await self._gen_router.finish_request( + request, success=False, req_id=disagg_request_id + ) raise else: # When need_ctx=False the gen server handles full generation and @@ -170,15 +166,15 @@ async def _send_disagg_request_ctx_first( if ctx_response is None or self._need_gen(ctx_response): if not gen_server: gen_server, _ = await self._gen_router.get_next_server( - gen_req, exclude_server=ctx_server + gen_req, exclude_server=ctx_server, req_id=disagg_request_id ) gen_response = await self._gen_client.send_request( - gen_req, server=gen_server, hooks=hooks + gen_req, server=gen_server, hooks=hooks, req_id=disagg_request_id ) return gen_response else: if gen_server: - await self._gen_router.finish_request(request) + await self._gen_router.finish_request(request, req_id=disagg_request_id) if request.stream: # ctx client will never return a generator when streaming is requested # make up for this by returning a done generator @@ -277,7 +273,7 @@ def _get_gen_request( request.disaggregated_params.disagg_request_id = disagg_request_id return request - async def _check_conditional_disagg(self, request: UCompletionRequest) -> bool: + async def _check_conditional_disagg(self, request: UCompletionRequest, req_id: int) -> bool: if self.conditional_disagg_config: local_gen_router = ( self._gen_router._local @@ -290,7 +286,7 @@ async def _check_conditional_disagg(self, request: UCompletionRequest) -> bool: ) # Query kv cache status and select a best gen_server. # The server is reserved for generation request - gen_server, info = await self._gen_router.get_next_server(request) + gen_server, info = await self._gen_router.get_next_server(request, req_id=req_id) match_length = sum(info["matches"]) total_length = sum(len(token_list) for token_list in info["token_lists"]) need_ctx_decision = ( @@ -395,7 +391,9 @@ async def _send_disagg_request_gen_first( # arrival->here = pre-ctx wait in the orchestrator/fleet. if hooks: hooks.on_ctx_dispatch(request) - ctx_server, ctx_server_info = await self._ctx_router.get_next_server(request) + ctx_server, ctx_server_info = await self._ctx_router.get_next_server( + request, req_id=disagg_request_id + ) ctx_req = self._get_ctx_request(request, disagg_request_id) gen_req = self._get_gen_request( request, @@ -415,7 +413,7 @@ async def _send_disagg_request_gen_first( # Fix: eagerly start consuming the gen generator in a background # task so the HTTP POST fires, then pipe chunks through a queue. gen_response = await self._gen_client.send_request( - gen_req, server=gen_server, hooks=hooks + gen_req, server=gen_server, hooks=hooks, req_id=disagg_request_id ) queue: asyncio.Queue = asyncio.Queue() @@ -432,7 +430,12 @@ async def _consume_gen(): # Now send ctx request — gen server has received its request try: - await self._ctx_client.send_request(ctx_req, server=ctx_server, hooks=hooks) + await self._ctx_client.send_request( + ctx_req, + server=ctx_server, + hooks=hooks, + req_id=disagg_request_id, + ) except Exception: consume_task.cancel() try: @@ -466,12 +469,22 @@ async def _yield_from_queue(): if need_ctx: tasks.append( asyncio.create_task( - self._ctx_client.send_request(ctx_req, server=ctx_server, hooks=hooks) + self._ctx_client.send_request( + ctx_req, + server=ctx_server, + hooks=hooks, + req_id=disagg_request_id, + ) ) ) tasks.append( asyncio.create_task( - self._gen_client.send_request(gen_req, server=gen_server, hooks=hooks) + self._gen_client.send_request( + gen_req, + server=gen_server, + hooks=hooks, + req_id=disagg_request_id, + ) ) ) responses = await asyncio.gather(*tasks) diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 6221230c90cc..32177b47c080 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -427,17 +427,18 @@ async def remove_server(self, server: str): f"Removed server {server}, current server list: {self._servers}") @abstractmethod - async def get_next_server( - self, - request: OpenAIRequest, - exclude_server: Optional[str] = None) -> tuple[str, dict]: + async def get_next_server(self, + request: OpenAIRequest, + exclude_server: Optional[str] = None, + req_id: Optional[int] = None) -> tuple[str, dict]: """Select server by request and return some intermediate information""" @abstractmethod async def finish_request(self, request: OpenAIRequest, session: Optional[aiohttp.ClientSession] = None, - success: bool = True): + success: bool = True, + req_id: Optional[int] = None): pass @property @@ -660,10 +661,11 @@ def _get_next_server(self) -> str: self._server_idx += 1 return server - async def get_next_server( - self, - request: OpenAIRequest, - exclude_server: Optional[str] = None) -> tuple[str, dict]: + async def get_next_server(self, + request: OpenAIRequest, + exclude_server: Optional[str] = None, + req_id: Optional[int] = None) -> tuple[str, dict]: + del req_id if not self._servers: if self._metadata_server: raise ValueError( @@ -685,8 +687,9 @@ async def get_next_server( async def finish_request(self, request: OpenAIRequest, session: Optional[aiohttp.ClientSession] = None, - success: bool = True): - del request, session, success + success: bool = True, + req_id: Optional[int] = None): + del request, session, success, req_id class LoadBalancingRouter(LoadBalancingMixin, Router): @@ -709,10 +712,11 @@ def _on_servers_updated(self, old_servers, new_servers): or self._create_server_state(server)) self._server_state = new_state - async def get_next_server( - self, - request: OpenAIRequest, - exclude_server: Optional[str] = None) -> tuple[str, dict]: + async def get_next_server(self, + request: OpenAIRequest, + exclude_server: Optional[str] = None, + req_id: Optional[int] = None) -> tuple[str, dict]: + del req_id self._validate_servers_available() async with self._lock: @@ -727,8 +731,9 @@ async def get_next_server( async def finish_request(self, request: OpenAIRequest, session: Optional[aiohttp.ClientSession] = None, - success: bool = True): - del session, success + success: bool = True, + req_id: Optional[int] = None): + del session, success, req_id async with self._lock: await self._unregister_request(request) @@ -855,10 +860,11 @@ def _content_affinity_key(request: OpenAIRequest) -> Optional[int]: parts.append(str(list(token_ids[:ROUTE_AFFINITY_TOKEN_PREFIX]))) return hash("".join(parts)) - async def get_next_server( - self, - request: OpenAIRequest, - exclude_server: Optional[str] = None) -> tuple[str, dict]: + async def get_next_server(self, + request: OpenAIRequest, + exclude_server: Optional[str] = None, + req_id: Optional[int] = None) -> tuple[str, dict]: + del req_id # Standalone (in-process) entry point = routing_key(tokenize+hash) then # the shared _route core -- the SAME core the coordinator path uses. key = await asyncio.to_thread(self._routing_key_sync, request) @@ -979,7 +985,9 @@ def _hashes(server): async def finish_request(self, request: OpenAIRequest, session: Optional[aiohttp.ClientSession] = None, - success: bool = True): + success: bool = True, + req_id: Optional[int] = None): + del req_id # Standalone entry point: key by id(request); pass request so token-load # accounting matches the increment_load(request) done at route time. await self._finish(id(request), @@ -1406,10 +1414,11 @@ def _evict_oldest_session(self): # ── public interface ── - async def get_next_server( - self, - request: OpenAIRequest, - exclude_server: Optional[str] = None) -> tuple[str, dict]: + async def get_next_server(self, + request: OpenAIRequest, + exclude_server: Optional[str] = None, + req_id: Optional[int] = None) -> tuple[str, dict]: + del req_id self._validate_servers_available() conv_id = self._get_conversation_id(request) @@ -1481,8 +1490,9 @@ async def get_next_server( async def finish_request(self, request: OpenAIRequest, session: Optional[aiohttp.ClientSession] = None, - success: bool = True): - del session, success + success: bool = True, + req_id: Optional[int] = None): + del session, success, req_id async with self._lock: server = await self._unregister_request(request) self._remove_content_load(server, request) @@ -1599,7 +1609,9 @@ def session(self) -> aiohttp.ClientSession: def _on_servers_updated(self, old_servers, new_servers): pass - def _request_id(self, request: OpenAIRequest) -> int: + def _request_id(self, + request: OpenAIRequest, + req_id: Optional[int] = None) -> int: """The request's disagg id -- the sole cross-process key for select/finish. Context requests carry disagg_request_id; generation requests inherit the @@ -1610,6 +1622,8 @@ def _request_id(self, request: OpenAIRequest) -> int: registered (transfer never completes -> DISAGG_GENERATION_TRANS_IN_PROGRESS fills the gen IndexMapper and throughput collapses). """ + if req_id is not None: + return req_id dp = request.disaggregated_params if dp is None: raise ValueError("delegated routing requires disaggregated_params") @@ -1621,10 +1635,10 @@ def _request_id(self, request: OpenAIRequest) -> int: "(disagg_request_id/ctx_request_id) on the request") return rid - async def get_next_server( - self, - request: OpenAIRequest, - exclude_server: Optional[str] = None) -> tuple[str, dict]: + async def get_next_server(self, + request: OpenAIRequest, + exclude_server: Optional[str] = None, + req_id: Optional[int] = None) -> tuple[str, dict]: # routing_key() tokenizes + block-hashes the prompt (CPU-bound for # kv_cache_aware); run it in a thread so it doesn't block the fleet worker # event loop driving the concurrent streams. @@ -1635,7 +1649,7 @@ async def get_next_server( payload = { "role": self._role, "routing_key": key, - "req_id": self._request_id(request), + "req_id": self._request_id(request, req_id), "exclude_server": exclude_server } _t0 = time.monotonic() @@ -1655,13 +1669,14 @@ async def get_next_server( async def finish_request(self, request: OpenAIRequest, session: Optional[aiohttp.ClientSession] = None, - success: bool = True): + success: bool = True, + req_id: Optional[int] = None): # Fire-and-forget: /finish only releases the coordinator's routed-load # bookkeeping, so run it as a background task instead of awaiting the HTTP # round trip on every request's critical path. Capture req_id first (the # request may be freed after return); errors are logged, not propagated. del session - req_id = self._request_id(request) + req_id = self._request_id(request, req_id) task = asyncio.create_task(self._finish_async(req_id, success)) self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) From 7ca2102dd7cebf13d0bfc3ef44b54376ef4581ca Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 14 Jul 2026 04:23:03 -0700 Subject: [PATCH 09/18] fix: preserve router calls without request id Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/serve/openai_client.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 8834eca15dad..26646a54344c 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -161,7 +161,10 @@ async def _send_request( req_id: Optional[int] = None, ) -> UCompletionResponseOrGenerator: if server is None: - server, _ = await self._router.get_next_server(request, req_id=req_id) + if req_id is None: + server, _ = await self._router.get_next_server(request) + else: + server, _ = await self._router.get_next_server(request, req_id=req_id) url = f"http://{server}/{endpoint}" # disaggregated_params is None when conditional_disagg bypasses ctx. _dp = request.disaggregated_params @@ -364,7 +367,12 @@ async def _finish_request( req_id: Optional[int] = None, ) -> None: self._metrics_collector.completed_requests.inc() - await self._router.finish_request(request, self._session, success=success, req_id=req_id) + if req_id is None: + await self._router.finish_request(request, self._session, success=success) + else: + await self._router.finish_request( + request, self._session, success=success, req_id=req_id + ) async def collect_metrics(self) -> Dict[str, Any]: metrics = {} From a43906c5041706cfb64ad3d76026737950b9912d Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 14 Jul 2026 04:39:58 -0700 Subject: [PATCH 10/18] test: run multi-orchestrator case on one GPU Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tests/integration/defs/disaggregated/test_disaggregated.py | 4 +++- tests/integration/test_lists/test-db/l0_dgx_h100.yml | 1 - tests/integration/test_lists/test-db/l0_h100.yml | 1 + 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 62e318d3c1ad..f998b87a373e 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -891,10 +891,12 @@ def test_disaggregated_tinyllama_multi_orchestrator(disaggregated_test_root, setup_model_symlink(llm_venv, llama_model_root, "TinyLlama/TinyLlama-1.1B-Chat-v1.0") + env = llm_venv._new_env.copy() + env["CUDA_VISIBLE_DEVICES"] = "0" run_disaggregated_test(disaggregated_example_root, "multi_orchestrator", num_iters=1, - env=llm_venv._new_env, + env=env, model_path=llama_model_root, cwd=llm_venv.get_working_directory()) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index f35937f21f98..e940d149d02f 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -120,7 +120,6 @@ l0_dgx_h100: - unittest/disaggregated/test_py_cache_transceiver_mp.py -k "gen_first1" - unittest/disaggregated/test_py_cache_transceiver_mp.py -k "gen_first2" - disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu[TinyLlama-1.1B-Chat-v1.0] - - disaggregated/test_disaggregated.py::test_disaggregated_tinyllama_multi_orchestrator[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_ctxtp2_genpp2[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_gentp2[TinyLlama-1.1B-Chat-v1.0] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 3fc1f846b163..d9e67fd5274e 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -174,6 +174,7 @@ l0_h100: - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_tp1_two_mtp[DeepSeek-V3-Lite-fp8] - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_ucx_tp1_single_gpu[DeepSeek-V3-Lite-fp8] - disaggregated/test_disaggregated.py::test_disaggregated_load_balance[TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated.py::test_disaggregated_tinyllama_multi_orchestrator[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_deepseek[False-False-DeepSeek-V3-Lite-fp8/fp8] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_deepseek[False-True-DeepSeek-V3-Lite-fp8/fp8] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_deepseek[True-False-DeepSeek-V3-Lite-fp8/fp8] From 1dc68338718fe6b600561329c9d1614d44e8860a Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 15 Jul 2026 01:16:47 +0000 Subject: [PATCH 11/18] [None][fix] address coordinator retry feedback Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/serve/disagg_coordinator.py | 20 +++++++ tensorrt_llm/serve/openai_disagg_server.py | 26 +++++++++ tensorrt_llm/serve/openai_disagg_service.py | 11 +++- tensorrt_llm/serve/router.py | 57 ++++++++++++------- .../test_openai_disagg_server.py | 35 ++++++++++++ .../test_openai_disagg_service.py | 31 ++++++++++ tests/unittest/disaggregated/test_router.py | 34 +++++++++++ 7 files changed, 192 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py index abbb00fb59bc..1407296af509 100644 --- a/tensorrt_llm/serve/disagg_coordinator.py +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -538,6 +538,26 @@ async def cluster_info(self) -> Dict[str, Any]: logger.warning(f"CoordinatorClient cluster_info failed: {e}") return {"is_ready": False} + async def proxy_cluster_storage_request( + self, + method: str, + path: str, + query: list[tuple[str, str]], + body: bytes, + content_type: Optional[str], + ) -> Tuple[bytes, int, Optional[str]]: + """Forward a public HTTP-storage request to the coordinator.""" + headers = {"Content-Type": content_type} if content_type else None + async with self.session.request( + method, + f"{self._remote_url}{path}", + params=query, + data=body, + headers=headers, + timeout=self._request_timeout_s, + ) as resp: + return await resp.read(), resp.status, resp.headers.get("Content-Type") + async def stop(self) -> None: if self._session is not None: await self._session.close() diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 0c07961d3248..6df7651fb1dd 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -231,6 +231,32 @@ def register_routes(self): cluster_storage = getattr(self._coordinator, "cluster_storage", None) if isinstance(cluster_storage, HttpClusterStorageServer): cluster_storage.add_routes(self.app) + elif (isinstance(self._coordinator, CoordinatorClient) + and isinstance(self._disagg_cluster_storage, + HttpClusterStorageServer)): + # Keep the configured public cluster_uri valid in fleet mode while + # the coordinator remains the sole owner of the HTTP storage state. + for path, method in (("/set", "POST"), ("/get", "GET"), + ("/delete", "DELETE"), ("/expire", "GET"), + ("/get_prefix", "GET")): + self.app.add_api_route(path, + self._proxy_cluster_storage_request, + methods=[method]) + + async def _proxy_cluster_storage_request(self, + raw_req: Request) -> Response: + try: + body, status, content_type = ( + await self._coordinator.proxy_cluster_storage_request( + raw_req.method, raw_req.url.path, + list(raw_req.query_params.multi_items()), + await raw_req.body(), raw_req.headers.get("Content-Type"))) + except (aiohttp.ClientError, asyncio.TimeoutError, OSError) as e: + logger.warning(f"Failed to proxy cluster storage request: {e}") + return JSONResponse(status_code=502, + content={"error": "coordinator unavailable"}) + headers = {"Content-Type": content_type} if content_type else None + return Response(content=body, status_code=status, headers=headers) @staticmethod def _extract_conversation_id(req: UCompletionRequest, raw_req: Request): diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 34cfef562681..da48405cf556 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -125,6 +125,10 @@ async def _send_disagg_request_ctx_first( disagg_request_id = await self._coordinator.get_disagg_request_id() # reserve a gen_server if conditional disagg is needed gen_server, need_ctx = await self._check_conditional_disagg(request, disagg_request_id) + # Context retries may replace disagg_request_id for the KV-transfer + # handshake. Keep the ID used to reserve the generation server separate + # so its coordinator-side load is released under the original key. + gen_reservation_id = disagg_request_id if gen_server else None need_ctx = need_ctx and not await self._check_gen_only_disagg(request) ctx_response = None gen_req = request @@ -150,7 +154,7 @@ async def _send_disagg_request_ctx_first( except Exception: if gen_server: await self._gen_router.finish_request( - request, success=False, req_id=disagg_request_id + request, success=False, req_id=gen_reservation_id ) raise else: @@ -168,13 +172,14 @@ async def _send_disagg_request_ctx_first( gen_server, _ = await self._gen_router.get_next_server( gen_req, exclude_server=ctx_server, req_id=disagg_request_id ) + gen_reservation_id = disagg_request_id gen_response = await self._gen_client.send_request( - gen_req, server=gen_server, hooks=hooks, req_id=disagg_request_id + gen_req, server=gen_server, hooks=hooks, req_id=gen_reservation_id ) return gen_response else: if gen_server: - await self._gen_router.finish_request(request, req_id=disagg_request_id) + await self._gen_router.finish_request(request, req_id=gen_reservation_id) if request.stream: # ctx client will never return a generator when streaming is requested # make up for this by returning a done generator diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 32177b47c080..1fb29583070b 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -36,6 +36,8 @@ hash_v1_block_key, truncate_sha256_hash_to_int64, v2_sha256_block_hasher) _MSGPACK_HEADERS = {"Content-Type": "application/msgpack"} +COORDINATOR_FINISH_MAX_ATTEMPTS = 3 +COORDINATOR_FINISH_RETRY_DELAY_S = 0.1 # Max number of conversations whose home-server pin is retained (LRU). ROUTE_AFFINITY_CACHE_SIZE = 50000 @@ -1683,26 +1685,43 @@ async def finish_request(self, async def _finish_async(self, req_id: int, success: bool): _t0 = time.monotonic() - try: - async with self.session.post( - f"{self._coordinator_url}/finish", - data=msgpack.packb( - { - "role": self._role, - "req_id": req_id, - "success": success - }, - use_bin_type=True), - headers=_MSGPACK_HEADERS, - timeout=self._request_timeout_s) as resp: - if resp.status != 200: - body = msgpack.unpackb(await resp.read(), raw=False) + for attempt in range(1, COORDINATOR_FINISH_MAX_ATTEMPTS + 1): + try: + async with self.session.post( + f"{self._coordinator_url}/finish", + data=msgpack.packb( + { + "role": self._role, + "req_id": req_id, + "success": success + }, + use_bin_type=True), + headers=_MSGPACK_HEADERS, + timeout=self._request_timeout_s) as resp: + if resp.status != 200: + raw_body = await resp.read() + try: + body = msgpack.unpackb(raw_body, raw=False) + except Exception: # noqa: BLE001 + body = raw_body.decode(errors="replace") + error = body.get("error", body) if isinstance( + body, dict) else body + raise RuntimeError( + f"coordinator /finish returned {resp.status}: {error}" + ) + self._finish_lat.record(time.monotonic() - _t0) + return + except Exception as e: # noqa: BLE001 + if attempt == COORDINATOR_FINISH_MAX_ATTEMPTS: logger.warning( - f"coordinator /finish returned {resp.status}: " - f"{body.get('error', body)}") - self._finish_lat.record(time.monotonic() - _t0) - except Exception as e: # noqa: BLE001 - logger.warning(f"CoordinatorDelegatingRouter finish failed: {e}") + f"CoordinatorDelegatingRouter finish failed after " + f"{attempt} attempts: {e}") + return + logger.warning( + f"CoordinatorDelegatingRouter finish attempt {attempt} " + f"failed: {e}; retrying") + await asyncio.sleep(COORDINATOR_FINISH_RETRY_DELAY_S * + (2**(attempt - 1))) async def close(self): if self._background_tasks: diff --git a/tests/unittest/disaggregated/test_openai_disagg_server.py b/tests/unittest/disaggregated/test_openai_disagg_server.py index beff804f9ed8..5b865027c2bf 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_server.py +++ b/tests/unittest/disaggregated/test_openai_disagg_server.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest +from fastapi import Request from starlette.datastructures import Headers from tensorrt_llm.llmapi.disagg_utils import extract_disagg_cfg @@ -29,6 +31,39 @@ def _raw_request(headers: dict[str, str]): return SimpleNamespace(headers=Headers(headers=headers)) +@pytest.mark.asyncio +async def test_http_cluster_storage_request_is_proxied_to_coordinator(): + payload = b'{"key":"worker","value":"ready"}' + + async def receive(): + return {"type": "http.request", "body": payload, "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/set", + "query_string": b"source=worker", + "headers": [(b"content-type", b"application/json")], + }, + receive, + ) + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._coordinator = SimpleNamespace( + proxy_cluster_storage_request=AsyncMock( + return_value=(b'{"result":true}', 200, "application/json") + ) + ) + + response = await server._proxy_cluster_storage_request(request) + + server._coordinator.proxy_cluster_storage_request.assert_awaited_once_with( + "POST", "/set", [("source", "worker")], payload, "application/json" + ) + assert response.status_code == 200 + assert response.body == b'{"result":true}' + + def test_extract_conversation_id_from_headers(): cases = [ ({"X-Session-ID": "session-id"}, "session-id"), diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 881d0037af0a..b9cec8475a94 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -408,6 +408,37 @@ async def _gen_response(*_args, **_kwargs): ) +@pytest.mark.asyncio +async def test_context_retry_preserves_generation_reservation_id(): + service = _make_service("context_first") + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._coordinator.get_disagg_request_id = AsyncMock(return_value=101) + service._check_conditional_disagg = AsyncMock(return_value=("gen:9001", True)) + service._check_gen_only_disagg = AsyncMock(return_value=False) + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) + + async def _ctx_response(request, *_args, **_kwargs): + # OpenAIHttpClient regenerates this field before a successful retry. + request.disaggregated_params.disagg_request_id = 202 + return _make_completion_response("", finish_reason="length", disagg_request_id=202) + + service._ctx_client.send_request = AsyncMock(side_effect=_ctx_response) + service._gen_client.send_request = AsyncMock( + return_value=_make_completion_response( + "done", finish_reason="stop", disagg_request_id=202, context_only=False + ) + ) + + request = CompletionRequest(model="test-model", prompt="hello") + await service._send_disagg_request(request) + + service._gen_router.get_next_server.assert_not_awaited() + gen_call = service._gen_client.send_request.call_args + assert gen_call.kwargs["req_id"] == 101 + assert gen_call.args[0].disaggregated_params.ctx_request_id == 202 + + def test_generation_postprocessor_rewrites_usage_from_disaggregated_params(): ctx_usage = UsageInfo( prompt_tokens=128, diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index 99184cb3296a..d165fbe5a040 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -5,6 +5,7 @@ from unittest import mock import aiohttp +import msgpack import pytest from tensorrt_llm.llmapi.disagg_utils import RouterConfig @@ -24,6 +25,7 @@ KV_CACHE_HASH_ALGO_V2, KV_CACHE_HASH_ALGO_V2_SHA256_64, BlockHashMixin, ConversationRouter, + CoordinatorDelegatingRouter, KvCacheAwareRouter, KvCacheAwareServerState, LoadBalancingRouter, RoundRobinRouter, @@ -74,6 +76,38 @@ def _make_mock_aiohttp_session(return_value=None): return mock_session +@pytest.mark.asyncio +async def test_coordinator_finish_retries_failed_release(): + local_router = RoundRobinRouter(server_role=None, servers=["server1"]) + router = CoordinatorDelegatingRouter("http://coordinator", local_router, + "generation") + + def _response(status, body): + response = mock.AsyncMock() + response.status = status + response.read = mock.AsyncMock( + return_value=msgpack.packb(body, use_bin_type=True)) + context = mock.AsyncMock() + context.__aenter__ = mock.AsyncMock(return_value=response) + context.__aexit__ = mock.AsyncMock(return_value=False) + return context + + session = mock.MagicMock() + session.post = mock.MagicMock(side_effect=[ + _response(503, {"error": "temporarily unavailable"}), + _response(200, {}), + ]) + session.close = mock.AsyncMock() + router._session = session + + with mock.patch("tensorrt_llm.serve.router.asyncio.sleep", + new_callable=mock.AsyncMock) as sleep: + await router._finish_async(123, True) + + assert session.post.call_count == 2 + sleep.assert_awaited_once() + + @pytest.fixture(autouse=True) def mock_aiohttp_session(request): """Auto-mock aiohttp.ClientSession so poll_events doesn't make real HTTP calls.""" From 1278df7f5c4bf21b2792ca9d80d5cc9378da16b2 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 15 Jul 2026 01:46:39 -0700 Subject: [PATCH 12/18] [None][fix] synchronize delegated router hashing Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/serve/disagg_coordinator.py | 43 +++++----- tensorrt_llm/serve/router.py | 56 +++++++++++-- tensorrt_llm/serve/router_utils.py | 8 ++ .../disaggregated/test_coordinator_worker.py | 81 ++++++++++++++++++- 4 files changed, 160 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py index 1407296af509..0e01411b597a 100644 --- a/tensorrt_llm/serve/disagg_coordinator.py +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -59,7 +59,12 @@ from tensorrt_llm.serve.disagg_auto_scaling import DisaggClusterManager, WorkerInfo from tensorrt_llm.serve.metadata_server import create_metadata_server from tensorrt_llm.serve.openai_client import OpenAIClient -from tensorrt_llm.serve.router import CoordinatorDelegatingRouter, Router, build_disagg_routers +from tensorrt_llm.serve.router import ( + CoordinatorDelegatingRouter, + KvCacheAwareRouter, + Router, + build_disagg_routers, +) __all__ = [ "DisaggCoordinator", @@ -258,11 +263,14 @@ async def is_ready(self) -> bool: async def cluster_info(self) -> Dict[str, Any]: info = {"is_ready": await self.is_ready()} - # Expose block-hash granularity so a delegating client (which doesn't - # monitor servers) hashes with the same tokens_per_block as the workers. - tpb = getattr(self._ctx_router, "_tokens_per_block", None) - if tpb is not None: - info["tokens_per_block"] = tpb + routing_key_configs = {} + for role, router in (("context", self._ctx_router), ("generation", self._gen_router)): + if isinstance(router, KvCacheAwareRouter): + config = router.routing_key_config() + if config is not None: + routing_key_configs[role] = config + if routing_key_configs: + info["routing_key_configs"] = routing_key_configs if self._disagg_cluster_manager: info.update(await self._disagg_cluster_manager.cluster_info()) return info @@ -438,19 +446,7 @@ async def start(self) -> None: # Fail fast: probe /cluster_info with bounded retry so a delegating server # exits non-zero if its coordinator never comes up (vs 500-ing every req). info = await self._await_coordinator() - # Adopt the coordinator's tokens_per_block so the client's block hashes - # line up with the workers (a delegating client can't learn it itself). - tpb = info.get("tokens_per_block") - if tpb is not None: - for router in (self._ctx_router, self._gen_router): - local = getattr(router, "_local", router) - if getattr(local, "_tokens_per_block", None) != tpb: - local._tokens_per_block = tpb - local._tpb_auto = False - logger.info( - f"CoordinatorClient: adopted coordinator " - f"tokens_per_block={tpb} for {getattr(local, '_namespace', '?')}" - ) + self._sync_delegating_router_configs(info) async def _await_coordinator(self) -> Dict[str, Any]: """Poll the coordinator's /cluster_info until reachable, or raise. @@ -504,12 +500,21 @@ async def is_ready(self) -> bool: if resp.status != 200: return False info = await self.cluster_info() + self._sync_delegating_router_configs(info) await self._sync_stateless_routers(info) return info.get("is_ready", False) except Exception as e: # noqa: BLE001 logger.warning(f"CoordinatorClient health check failed: {e}") return False + def _sync_delegating_router_configs(self, info: Dict[str, Any]) -> None: + configs = info.get("routing_key_configs", {}) + for role, router in (("context", self._ctx_router), ("generation", self._gen_router)): + config = configs.get(role) + local = getattr(router, "_local", None) + if config is not None and isinstance(local, KvCacheAwareRouter): + local.set_routing_key_config(config) + async def _sync_stateless_routers(self, info: Dict[str, Any]) -> None: current_workers = info.get("current_workers") if current_workers is None: diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 1fb29583070b..52802cba3e36 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -20,6 +20,7 @@ import aiohttp import msgpack +from blake3 import blake3 from tensorrt_llm.llmapi.disagg_utils import (MetadataServerConfig, RouterConfig, ServerRole) @@ -31,9 +32,10 @@ # Re-exported here for backward compat. from tensorrt_llm.serve.router_utils import ( # noqa: F401 KV_CACHE_HASH_ALGO_DEFAULT, KV_CACHE_HASH_ALGO_V1, KV_CACHE_HASH_ALGO_V2, - KV_CACHE_HASH_ALGO_V2_SHA256_64, BlockHash, BlockHashMixin, OpenAIRequest, - block_key_hasher, get_cache_salt_id, get_request_num_tokens, - hash_v1_block_key, truncate_sha256_hash_to_int64, v2_sha256_block_hasher) + KV_CACHE_HASH_ALGO_V2_SHA256_64, KV_CACHE_HASH_ALGOS, BlockHash, + BlockHashMixin, OpenAIRequest, block_key_hasher, get_cache_salt_id, + get_request_num_tokens, hash_v1_block_key, truncate_sha256_hash_to_int64, + v2_sha256_block_hasher) _MSGPACK_HEADERS = {"Content-Type": "application/msgpack"} COORDINATOR_FINISH_MAX_ATTEMPTS = 3 @@ -41,6 +43,7 @@ # Max number of conversations whose home-server pin is retained (LRU). ROUTE_AFFINITY_CACHE_SIZE = 50000 +ROUTE_AFFINITY_HASH_SEED = b"TensorRT-LLM-route-affinity-v1!!" # Leading token-id count folded into the affinity key so pre-tokenized # requests (placeholder message content) still key per conversation. ROUTE_AFFINITY_TOKEN_PREFIX = 256 @@ -768,6 +771,10 @@ def __init__(self, self._load_weight = load_weight self._load_cap = load_cap self._track_routed_blocks = track_routed_blocks + # A coordinator client has no local server pool. The coordinator injects + # the role's effective hash algorithm so this router can act solely as a + # routing-key encoder. + self._routing_hash_algo: Optional[str] = None def _create_server_state(self, server: str) -> KvCacheAwareServerState: return KvCacheAwareServerState(server, self._use_tokens, @@ -793,6 +800,32 @@ def _get_server_hash_algo(self, server: str) -> str: # by update_with_events. return self._server_state[server].hash_algo + def routing_key_config(self) -> Optional[dict[str, int | str]]: + """Return the encoding configuration shared by prepared servers.""" + hash_algos = { + self._get_server_hash_algo(server) + for server in self._prepared_ready_servers + } + if len(hash_algos) > 1: + raise RuntimeError( + f"KV-cache-aware routing requires one hash algorithm per role; " + f"found {sorted(hash_algos)} for {self._server_role}") + if not hash_algos: + return None + return { + "tokens_per_block": self._tokens_per_block, + "kv_cache_hash_algo": hash_algos.pop(), + } + + def set_routing_key_config(self, config: dict[str, int | str]) -> None: + hash_algo = str(config["kv_cache_hash_algo"]) + if hash_algo not in KV_CACHE_HASH_ALGOS: + raise ValueError(f"Unknown KV cache hash algorithm {hash_algo!r}; " + f"expected one of {sorted(KV_CACHE_HASH_ALGOS)}") + self._tokens_per_block = int(config["tokens_per_block"]) + self._tpb_auto = False + self._routing_hash_algo = hash_algo + def _events_aligned(self, server: str) -> bool: worker_tpb = self._server_info.get(server, {}).get("tokens_per_block") return worker_tpb is None or worker_tpb == self._tokens_per_block @@ -860,7 +893,9 @@ def _content_affinity_key(request: OpenAIRequest) -> Optional[int]: token_ids = getattr(request, "prompt_token_ids", None) if token_ids: parts.append(str(list(token_ids[:ROUTE_AFFINITY_TOKEN_PREFIX]))) - return hash("".join(parts)) + digest = blake3("".join(parts).encode("utf-8"), + key=ROUTE_AFFINITY_HASH_SEED).digest(length=8) + return int.from_bytes(digest, "little", signed=False) async def get_next_server(self, request: OpenAIRequest, @@ -883,10 +918,15 @@ def _routing_key_sync(self, request: OpenAIRequest) -> dict: """ cache_salt_id = self._get_request_cache_salt_id(request) # Hash for every algo any server might use (usually one). - algos = { - self._get_server_hash_algo(s) - for s in self._server_state.keys() - } or None + algos = ({self._routing_hash_algo} + if self._routing_hash_algo is not None else { + self._get_server_hash_algo(s) + for s in self._server_state.keys() + }) + if not algos: + raise RuntimeError( + "KV-cache routing-key encoder has no hash algorithm; " + "synchronize it with the coordinator before routing") token_lists, block_hashes_by_algo = \ self._tokenize_and_compute_block_hashes_by_algo( request, algos, cache_salt_id) diff --git a/tensorrt_llm/serve/router_utils.py b/tensorrt_llm/serve/router_utils.py index b9d8c9265aa8..a96d087acf06 100644 --- a/tensorrt_llm/serve/router_utils.py +++ b/tensorrt_llm/serve/router_utils.py @@ -34,6 +34,13 @@ KV_CACHE_HASH_ALGO_V1 = kv_cache_hash.KV_CACHE_HASH_ALGO_V1 KV_CACHE_HASH_ALGO_V2 = kv_cache_hash.KV_CACHE_HASH_ALGO_V2 KV_CACHE_HASH_ALGO_V2_SHA256_64 = kv_cache_hash.KV_CACHE_HASH_ALGO_V2_SHA256_64 +KV_CACHE_HASH_ALGOS = frozenset( + { + KV_CACHE_HASH_ALGO_V1, + KV_CACHE_HASH_ALGO_V2, + KV_CACHE_HASH_ALGO_V2_SHA256_64, + } +) get_cache_salt_id = kv_cache_hash.get_cache_salt_id hash_v1_block_key = kv_cache_hash.hash_v1_block_key truncate_sha256_hash_to_int64 = kv_cache_hash.truncate_sha256_hash_to_int64 @@ -46,6 +53,7 @@ "KV_CACHE_HASH_ALGO_V1", "KV_CACHE_HASH_ALGO_V2", "KV_CACHE_HASH_ALGO_V2_SHA256_64", + "KV_CACHE_HASH_ALGOS", "get_cache_salt_id", "hash_v1_block_key", "truncate_sha256_hash_to_int64", diff --git a/tests/unittest/disaggregated/test_coordinator_worker.py b/tests/unittest/disaggregated/test_coordinator_worker.py index 21646754b639..c32f9013ff2f 100644 --- a/tests/unittest/disaggregated/test_coordinator_worker.py +++ b/tests/unittest/disaggregated/test_coordinator_worker.py @@ -48,7 +48,17 @@ ) from tensorrt_llm.serve.coordinator_server import CoordinatorServer from tensorrt_llm.serve.disagg_coordinator import CoordinatorClient, DisaggCoordinatorService -from tensorrt_llm.serve.openai_protocol import CompletionRequest, DisaggregatedParams +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, + CompletionRequest, + DisaggregatedParams, +) +from tensorrt_llm.serve.router import ( + KV_CACHE_HASH_ALGO_V1, + KV_CACHE_HASH_ALGO_V2, + CoordinatorDelegatingRouter, + KvCacheAwareRouter, +) from tensorrt_llm.serve.router_utils import BlockHashMixin as SharedBlockHashMixin @@ -188,6 +198,75 @@ def test_coordinator_rejects_unknown_role(): coordinator._router_for_role("typo") +def test_kv_router_rejects_mixed_prepared_hash_algorithms(): + router = KvCacheAwareRouter(server_role=ServerRole.CONTEXT, servers=["server-a", "server-b"]) + router._prepared_ready_servers.update(router.servers) + router._server_state["server-a"].set_hash_algo(KV_CACHE_HASH_ALGO_V1) + router._server_state["server-b"].set_hash_algo(KV_CACHE_HASH_ALGO_V2) + + with pytest.raises(RuntimeError, match="one hash algorithm per role"): + router.routing_key_config() + + +@pytest.mark.asyncio +async def test_coordinator_exposes_role_hash_algorithm(): + ctx_server = "127.0.0.1:1234" + gen_server = "127.0.0.1:1235" + config = _make_config([ctx_server], [gen_server], "kv_cache_aware", "kv_cache_aware") + coordinator = DisaggCoordinatorService(config, _client_factory) + ctx_router = coordinator.ctx_router + gen_router = coordinator.gen_router + assert isinstance(ctx_router, KvCacheAwareRouter) + assert isinstance(gen_router, KvCacheAwareRouter) + ctx_router._prepared_ready_servers.add(ctx_server) + gen_router._prepared_ready_servers.add(gen_server) + ctx_router._server_state[ctx_server].set_hash_algo(KV_CACHE_HASH_ALGO_V2) + gen_router._server_state[gen_server].set_hash_algo(KV_CACHE_HASH_ALGO_V1) + + info = await coordinator.cluster_info() + + assert info["routing_key_configs"]["context"] == { + "tokens_per_block": 32, + "kv_cache_hash_algo": KV_CACHE_HASH_ALGO_V2, + } + assert info["routing_key_configs"]["generation"] == { + "tokens_per_block": 32, + "kv_cache_hash_algo": KV_CACHE_HASH_ALGO_V1, + } + + +def test_coordinator_client_configures_empty_delegating_kv_router(): + config = _make_config([], [], "kv_cache_aware", "round_robin") + client = CoordinatorClient("http://coordinator", config) + assert isinstance(client.ctx_router, CoordinatorDelegatingRouter) + local = client.ctx_router._local + assert isinstance(local, KvCacheAwareRouter) + assert local.servers == [] + + client._sync_delegating_router_configs( + { + "routing_key_configs": { + "context": { + "tokens_per_block": 64, + "kv_cache_hash_algo": KV_CACHE_HASH_ALGO_V2, + } + } + } + ) + + request = CompletionRequest(model="m", prompt=[1, 2, 3]) + routing_key = local.routing_key(request) + assert local.servers == [] + assert local._tokens_per_block == 64 + assert set(routing_key["block_hashes_by_algo"]) == {KV_CACHE_HASH_ALGO_V2} + + +def test_content_affinity_key_uses_fixed_seed(): + request = ChatCompletionRequest(model="m", messages=[{"role": "user", "content": "hello"}]) + + assert KvCacheAwareRouter._content_affinity_key(request) == 7306401829117098140 + + def test_prefix_token_cache_retokenizes_extended_text(): class BoundarySensitiveTokenizer: def __init__(self): From a780efdc78b533919aa274eca1956b7c0c3d3239 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 15 Jul 2026 04:58:00 -0700 Subject: [PATCH 13/18] [None][fix] address coordinator reservation feedback Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/llmapi/disagg_utils.py | 4 ++++ tensorrt_llm/serve/router.py | 18 ++++++++---------- .../disaggregated/test_disagg_utils.py | 6 ++++++ tests/unittest/disaggregated/test_router.py | 6 ++++-- 4 files changed, 22 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/llmapi/disagg_utils.py b/tensorrt_llm/llmapi/disagg_utils.py index f55504b309f1..a8272b494645 100644 --- a/tensorrt_llm/llmapi/disagg_utils.py +++ b/tensorrt_llm/llmapi/disagg_utils.py @@ -247,6 +247,10 @@ def extract_disagg_cfg(hostname: str = 'localhost', max_retries, perf_metrics_max_requests, disagg_cluster_config) if node_id is not None: + node_id_space = 1 << DISAGG_NODE_ID_BITS + if not 0 <= node_id < node_id_space: + raise ValueError( + f"node_id must be in range [0, {node_id_space}), got {node_id}") config.node_id = node_id if schedule_style: config.schedule_style = schedule_style diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 52802cba3e36..289ae0af2cf1 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -38,8 +38,8 @@ v2_sha256_block_hasher) _MSGPACK_HEADERS = {"Content-Type": "application/msgpack"} -COORDINATOR_FINISH_MAX_ATTEMPTS = 3 COORDINATOR_FINISH_RETRY_DELAY_S = 0.1 +COORDINATOR_FINISH_MAX_RETRY_DELAY_S = 5.0 # Max number of conversations whose home-server pin is retained (LRU). ROUTE_AFFINITY_CACHE_SIZE = 50000 @@ -1725,7 +1725,10 @@ async def finish_request(self, async def _finish_async(self, req_id: int, success: bool): _t0 = time.monotonic() - for attempt in range(1, COORDINATOR_FINISH_MAX_ATTEMPTS + 1): + attempt = 0 + delay = COORDINATOR_FINISH_RETRY_DELAY_S + while True: + attempt += 1 try: async with self.session.post( f"{self._coordinator_url}/finish", @@ -1752,16 +1755,11 @@ async def _finish_async(self, req_id: int, success: bool): self._finish_lat.record(time.monotonic() - _t0) return except Exception as e: # noqa: BLE001 - if attempt == COORDINATOR_FINISH_MAX_ATTEMPTS: - logger.warning( - f"CoordinatorDelegatingRouter finish failed after " - f"{attempt} attempts: {e}") - return logger.warning( f"CoordinatorDelegatingRouter finish attempt {attempt} " - f"failed: {e}; retrying") - await asyncio.sleep(COORDINATOR_FINISH_RETRY_DELAY_S * - (2**(attempt - 1))) + f"failed: {e}; retrying in {delay:.1f}s") + await asyncio.sleep(delay) + delay = min(delay * 2, COORDINATOR_FINISH_MAX_RETRY_DELAY_S) async def close(self): if self._background_tasks: diff --git a/tests/unittest/disaggregated/test_disagg_utils.py b/tests/unittest/disaggregated/test_disagg_utils.py index 3d2d286312c1..f3556bae1fe5 100644 --- a/tests/unittest/disaggregated/test_disagg_utils.py +++ b/tests/unittest/disaggregated/test_disagg_utils.py @@ -127,6 +127,12 @@ def test_extract_disagg_cfg(sample_yaml_config): assert config.disagg_coordinator_url == "http://coordinator:7999" +@pytest.mark.parametrize("node_id", [-1, 256]) +def test_extract_disagg_cfg_rejects_out_of_range_node_id(node_id): + with pytest.raises(ValueError, match="node_id must be in range"): + extract_disagg_cfg(node_id=node_id) + + def test_extract_ctx_gen_cfgs(): configs = extract_ctx_gen_cfgs( type="ctx", diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index d165fbe5a040..710a4fc33acc 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -94,6 +94,8 @@ def _response(status, body): session = mock.MagicMock() session.post = mock.MagicMock(side_effect=[ + _response(503, {"error": "temporarily unavailable"}), + _response(503, {"error": "temporarily unavailable"}), _response(503, {"error": "temporarily unavailable"}), _response(200, {}), ]) @@ -104,8 +106,8 @@ def _response(status, body): new_callable=mock.AsyncMock) as sleep: await router._finish_async(123, True) - assert session.post.call_count == 2 - sleep.assert_awaited_once() + assert session.post.call_count == 4 + assert sleep.await_count == 3 @pytest.fixture(autouse=True) From 0186c597ba033d71f6734e7d233f2f6220f59fb9 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 15 Jul 2026 20:02:49 -0700 Subject: [PATCH 14/18] [None][fix] address fleet coordinator review feedback Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/commands/serve.py | 6 ++ tensorrt_llm/serve/disagg_coordinator.py | 86 ++++++++++++++--- tensorrt_llm/serve/openai_disagg_service.py | 2 +- tensorrt_llm/serve/router.py | 26 +++--- .../disaggregated/test_coordinator_worker.py | 93 ++++++++++++++++++- tests/unittest/disaggregated/test_router.py | 7 +- .../llmapi/test_disagg_telemetry_launcher.py | 90 +++++++++++++++++- 7 files changed, 276 insertions(+), 34 deletions(-) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 50c6b6506ddb..d753689c65f4 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1743,6 +1743,8 @@ def _launch_disagg_fleet(disagg_cfg, config_file, metadata_server_config_file, request_timeout) base_env[DisaggWorkerEnvs.TLLM_DISAGG_SERVER_START_TIMEOUT] = str( server_start_timeout) + base_env[DisaggWorkerEnvs.TLLM_DISAGG_SCHEDULE_STYLE] = \ + disagg_cfg.schedule_style # Propagate the parent's log level so fleet workers' INFO logs (e.g. the # per-request [ttft_split] / [coord_api] breakdowns) are not dropped. base_env[DisaggWorkerEnvs.TLLM_DISAGG_LOG_LEVEL] = logger.level @@ -1963,6 +1965,9 @@ def _build_disagg_server_from_env() -> "OpenAIDisaggServer": "180")) disagg_cfg = parse_disagg_config_file(config_file) + schedule_style = os.environ.get(DisaggWorkerEnvs.TLLM_DISAGG_SCHEDULE_STYLE) + if schedule_style: + disagg_cfg.schedule_style = schedule_style metadata_server_cfg = parse_metadata_server_config_file( metadata_config_file) @@ -2106,6 +2111,7 @@ class DisaggWorkerEnvs(StrEnum): TLLM_DISAGG_METADATA_CONFIG_FILE = "TRTLLM_DISAGG_METADATA_CONFIG_FILE" TLLM_DISAGG_REQUEST_TIMEOUT = "TRTLLM_DISAGG_REQUEST_TIMEOUT" TLLM_DISAGG_SERVER_START_TIMEOUT = "TRTLLM_DISAGG_SERVER_START_TIMEOUT" + TLLM_DISAGG_SCHEDULE_STYLE = "TRTLLM_DISAGG_SCHEDULE_STYLE" # Parent's logger level; the forked uvicorn fleet is a fresh process that # otherwise defaults to WARNING and drops the workers' INFO logs. TLLM_DISAGG_LOG_LEVEL = "TRTLLM_DISAGG_LOG_LEVEL" diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py index 0e01411b597a..9123767bfa98 100644 --- a/tensorrt_llm/serve/disagg_coordinator.py +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -35,6 +35,7 @@ """ import asyncio +import os import time from abc import ABC, abstractmethod from typing import Any, Dict, Optional, Tuple @@ -74,6 +75,18 @@ "make_coordinator_session", ] +COORDINATOR_RESERVATION_TIMEOUT_ENV = "TRTLLM_DISAGG_COORDINATOR_RESERVATION_TIMEOUT" +COORDINATOR_RESERVATION_TIMEOUT_DEFAULT_S = 180.0 + + +def coordinator_reservation_timeout() -> float: + return float( + os.environ.get( + COORDINATOR_RESERVATION_TIMEOUT_ENV, + COORDINATOR_RESERVATION_TIMEOUT_DEFAULT_S, + ) + ) + class DisaggCoordinator(ABC): """Abstract coordinator: ctx/gen routers + readiness + cluster info + lifecycle. @@ -119,6 +132,7 @@ def __init__( server_preparation_func=None, server_start_timeout_secs: int = 180, health_check_interval_secs: int = 3, + reservation_timeout_secs: Optional[float] = None, ): self._config = config self._client_factory = client_factory @@ -152,6 +166,12 @@ def __init__( ) self._server_start_timeout_secs = server_start_timeout_secs self._health_check_interval_secs = health_check_interval_secs + self._reservation_timeout_secs = ( + coordinator_reservation_timeout() + if reservation_timeout_secs is None + else reservation_timeout_secs + ) + self._reservation_tasks: dict[tuple[str, int], asyncio.Task] = {} self._ctx_client: Optional[OpenAIClient] = None self._gen_client: Optional[OpenAIClient] = None @@ -194,9 +214,17 @@ async def select( ) -> Tuple[str, dict, Optional[str]]: _t0 = time.monotonic() router = self._router_for_role(role) + reservation_key = (self._normalize_role(role), req_id) + previous = self._reservation_tasks.pop(reservation_key, None) + if previous is not None: + previous.cancel() + await router.finish_request_by_id(req_id, False) result = await router.get_next_server_by_key( routing_key, req_id=req_id, exclude_server=exclude_server ) + self._reservation_tasks[reservation_key] = asyncio.create_task( + self._expire_reservation(reservation_key, router) + ) self._api_lat(f"select[{role}]").record(time.monotonic() - _t0) return result @@ -205,17 +233,39 @@ async def get_disagg_request_id(self) -> int: async def finish(self, role: str, req_id, success: bool = True) -> None: _t0 = time.monotonic() + reservation = self._reservation_tasks.pop((self._normalize_role(role), req_id), None) + if reservation is not None: + reservation.cancel() await self._router_for_role(role).finish_request_by_id(req_id, success) self._api_lat(f"finish[{role}]").record(time.monotonic() - _t0) - def _router_for_role(self, role: str) -> Router: + async def _expire_reservation(self, key: tuple[str, int], router: Router) -> None: + try: + await asyncio.sleep(self._reservation_timeout_secs) + logger.warning( + f"Releasing stale coordinator reservation for role={key[0]}, " + f"req_id={key[1]} after {self._reservation_timeout_secs}s" + ) + await router.finish_request_by_id(key[1], False) + finally: + if self._reservation_tasks.get(key) is asyncio.current_task(): + self._reservation_tasks.pop(key, None) + + @staticmethod + def _normalize_role(role: str) -> str: normalized_role = str(role).lower() if normalized_role in ("context", "ctx"): - return self._ctx_router + return "context" if normalized_role in ("generation", "gen"): - return self._gen_router + return "generation" raise ValueError(f"Unsupported coordinator role: {role}") + def _router_for_role(self, role: str) -> Router: + normalized_role = self._normalize_role(role) + if normalized_role == "context": + return self._ctx_router + return self._gen_router + async def start(self) -> None: await self._ctx_router.prepare_servers() await self._gen_router.prepare_servers() @@ -247,6 +297,12 @@ async def start(self) -> None: await self._wait_for_all_servers_ready() async def stop(self) -> None: + reservations = list(self._reservation_tasks.values()) + self._reservation_tasks.clear() + for reservation in reservations: + reservation.cancel() + if reservations: + await asyncio.gather(*reservations, return_exceptions=True) if self._disagg_cluster_manager: await self._disagg_cluster_manager.stop() if self._metadata_server: @@ -262,7 +318,13 @@ async def is_ready(self) -> bool: return True async def cluster_info(self) -> Dict[str, Any]: - info = {"is_ready": await self.is_ready()} + info = { + "is_ready": await self.is_ready(), + "server_lists": { + "context": list(self._ctx_router.servers), + "generation": list(self._gen_router.servers), + }, + } routing_key_configs = {} for role, router in (("context", self._ctx_router), ("generation", self._gen_router)): if isinstance(router, KvCacheAwareRouter): @@ -447,6 +509,7 @@ async def start(self) -> None: # exits non-zero if its coordinator never comes up (vs 500-ing every req). info = await self._await_coordinator() self._sync_delegating_router_configs(info) + await self._sync_stateless_routers(info) async def _await_coordinator(self) -> Dict[str, Any]: """Poll the coordinator's /cluster_info until reachable, or raise. @@ -516,21 +579,22 @@ def _sync_delegating_router_configs(self, info: Dict[str, Any]) -> None: local.set_routing_key_config(config) async def _sync_stateless_routers(self, info: Dict[str, Any]) -> None: - current_workers = info.get("current_workers") - if current_workers is None: + server_lists = info.get("server_lists") + if server_lists is None: return - router_workers = ( - (self._ctx_router, current_workers.get("context_servers", [])), - (self._gen_router, current_workers.get("generation_servers", [])), + router_servers = ( + (self._ctx_router, server_lists.get("context", [])), + (self._gen_router, server_lists.get("generation", [])), ) - for router, workers in router_workers: + for router, servers in router_servers: if isinstance(router, CoordinatorDelegatingRouter): continue - desired = {f"{worker['host']}:{worker['port']}" for worker in workers} + desired = set(servers) for server in set(router.servers) - desired: await router.remove_server(server) for server in desired - set(router.servers): await router.add_server(server) + await router.prepare_servers() async def cluster_info(self) -> Dict[str, Any]: try: diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index da48405cf556..1dc5cef0d1ff 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -293,7 +293,7 @@ async def _check_conditional_disagg(self, request: UCompletionRequest, req_id: i # The server is reserved for generation request gen_server, info = await self._gen_router.get_next_server(request, req_id=req_id) match_length = sum(info["matches"]) - total_length = sum(len(token_list) for token_list in info["token_lists"]) + total_length = info["num_tokens"] need_ctx_decision = ( match_length == 0 or total_length - match_length diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 289ae0af2cf1..79bc14bb40e4 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -38,8 +38,8 @@ v2_sha256_block_hasher) _MSGPACK_HEADERS = {"Content-Type": "application/msgpack"} +COORDINATOR_FINISH_MAX_ATTEMPTS = 3 COORDINATOR_FINISH_RETRY_DELAY_S = 0.1 -COORDINATOR_FINISH_MAX_RETRY_DELAY_S = 5.0 # Max number of conversations whose home-server pin is retained (LRU). ROUTE_AFFINITY_CACHE_SIZE = 50000 @@ -931,9 +931,9 @@ def _routing_key_sync(self, request: OpenAIRequest) -> dict: self._tokenize_and_compute_block_hashes_by_algo( request, algos, cache_salt_id) return { - "token_lists": token_lists, "block_hashes_by_algo": block_hashes_by_algo, - "conv_key": self._content_affinity_key(request) + "conv_key": self._content_affinity_key(request), + "num_tokens": sum(len(token_list) for token_list in token_lists), } async def _route(self, key, exclude_server=None, request=None, req_id=None): @@ -943,9 +943,9 @@ async def _route(self, key, exclude_server=None, request=None, req_id=None): conversation affinity, and round-robin tie-breaking. Standalone requests are keyed by ``id(request)``; coordinator requests use ``req_id``. """ - token_lists = (key or {}).get("token_lists") or [] block_hashes_by_algo = (key or {}).get("block_hashes_by_algo") or {} conv_key = (key or {}).get("conv_key") + num_tokens = (key or {}).get("num_tokens", 0) async with self._lock: server_states = { server: state @@ -1019,8 +1019,8 @@ def _hashes(server): return server, { "block_hashes": block_hashes, "hash_algo": hash_algo, - "token_lists": token_lists, "matches": matches, + "num_tokens": num_tokens, "server_info": self._server_info.get(server, {}), }, req_id @@ -1725,10 +1725,7 @@ async def finish_request(self, async def _finish_async(self, req_id: int, success: bool): _t0 = time.monotonic() - attempt = 0 - delay = COORDINATOR_FINISH_RETRY_DELAY_S - while True: - attempt += 1 + for attempt in range(1, COORDINATOR_FINISH_MAX_ATTEMPTS + 1): try: async with self.session.post( f"{self._coordinator_url}/finish", @@ -1755,11 +1752,16 @@ async def _finish_async(self, req_id: int, success: bool): self._finish_lat.record(time.monotonic() - _t0) return except Exception as e: # noqa: BLE001 + if attempt == COORDINATOR_FINISH_MAX_ATTEMPTS: + logger.warning( + f"CoordinatorDelegatingRouter finish failed after " + f"{attempt} attempts: {e}") + return logger.warning( f"CoordinatorDelegatingRouter finish attempt {attempt} " - f"failed: {e}; retrying in {delay:.1f}s") - await asyncio.sleep(delay) - delay = min(delay * 2, COORDINATOR_FINISH_MAX_RETRY_DELAY_S) + f"failed: {e}; retrying") + await asyncio.sleep(COORDINATOR_FINISH_RETRY_DELAY_S * + (2**(attempt - 1))) async def close(self): if self._background_tasks: diff --git a/tests/unittest/disaggregated/test_coordinator_worker.py b/tests/unittest/disaggregated/test_coordinator_worker.py index c32f9013ff2f..615d671530c4 100644 --- a/tests/unittest/disaggregated/test_coordinator_worker.py +++ b/tests/unittest/disaggregated/test_coordinator_worker.py @@ -32,9 +32,11 @@ """ import asyncio +import json import threading import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from unittest.mock import AsyncMock import aiohttp import pytest @@ -47,7 +49,12 @@ ServerRole, ) from tensorrt_llm.serve.coordinator_server import CoordinatorServer -from tensorrt_llm.serve.disagg_coordinator import CoordinatorClient, DisaggCoordinatorService +from tensorrt_llm.serve.disagg_coordinator import ( + COORDINATOR_RESERVATION_TIMEOUT_ENV, + CoordinatorClient, + DisaggCoordinatorService, + coordinator_reservation_timeout, +) from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, CompletionRequest, @@ -93,10 +100,11 @@ def _free_port(): class _FakeWorker: - """Minimal HTTP worker: /health -> 200 (used for readiness only).""" + """Minimal HTTP worker exposing health and routing metadata.""" - def __init__(self): + def __init__(self, server_info=None): self.port = _free_port() + server_info_body = json.dumps(server_info or {}).encode() class Handler(BaseHTTPRequestHandler): def log_message(self, *a): @@ -105,9 +113,15 @@ def log_message(self, *a): def do_GET(self): if self.path == "/health": self.send_response(200) + elif self.path == "/server_info": + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(server_info_body))) else: self.send_response(404) self.end_headers() + if self.path == "/server_info": + self.wfile.write(server_info_body) self._httpd = ThreadingHTTPServer(("127.0.0.1", self.port), Handler) self._thread = threading.Thread(target=self._httpd.serve_forever, daemon=True) @@ -235,6 +249,32 @@ async def test_coordinator_exposes_role_hash_algorithm(): } +@pytest.mark.asyncio +async def test_coordinator_expires_stale_reservation(): + config = _make_config([], ["gen:8000"], "round_robin", "conversation") + coordinator = DisaggCoordinatorService( + config, + _client_factory, + reservation_timeout_secs=0.01, + ) + + await coordinator.select("generation", "conversation", 123, None) + assert coordinator.gen_router._server_content_load["gen:8000"] == 1 + + await asyncio.sleep(0.02) + + assert coordinator.gen_router._server_content_load["gen:8000"] == 0 + assert coordinator._reservation_tasks == {} + + +def test_coordinator_reservation_timeout_env(monkeypatch): + monkeypatch.delenv(COORDINATOR_RESERVATION_TIMEOUT_ENV, raising=False) + assert coordinator_reservation_timeout() == 180 + + monkeypatch.setenv(COORDINATOR_RESERVATION_TIMEOUT_ENV, "60") + assert coordinator_reservation_timeout() == 60 + + def test_coordinator_client_configures_empty_delegating_kv_router(): config = _make_config([], [], "kv_cache_aware", "round_robin") client = CoordinatorClient("http://coordinator", config) @@ -259,6 +299,8 @@ def test_coordinator_client_configures_empty_delegating_kv_router(): assert local.servers == [] assert local._tokens_per_block == 64 assert set(routing_key["block_hashes_by_algo"]) == {KV_CACHE_HASH_ALGO_V2} + assert routing_key["num_tokens"] == 3 + assert "token_lists" not in routing_key def test_content_affinity_key_uses_fixed_seed(): @@ -319,6 +361,51 @@ async def drive(): ) +def test_static_stateless_router_prepares_generation_first_server_info(): + """Fleet startup prepares static local routers for generation-first.""" + from tensorrt_llm.serve.router import RoundRobinRouter + + ctx_info_endpoint = "tcp://127.0.0.1:12345" + ctx_server_info = {"disaggregated_params": {"ctx_info_endpoint": ctx_info_endpoint}} + with _FakeWorker(ctx_server_info) as ctx0, _FakeWorker() as gen0: + config = _make_config([ctx0.url], [gen0.url], "round_robin", "round_robin") + config.schedule_style = "generation_first" + with _CoordinatorThread(config) as coord: + assert asyncio.run(_wait_coord_ready(coord.url)) + + async def drive(): + remote = CoordinatorClient(coord.url, config) + await remote.start() + assert isinstance(remote.ctx_router, RoundRobinRouter) + request = CompletionRequest(model="m", prompt="hello") + server, info = await remote.ctx_router.get_next_server(request) + await remote.stop() + return server, info + + server, info = asyncio.run(drive()) + assert server == ctx0.url + assert ( + info["server_info"]["disaggregated_params"]["ctx_info_endpoint"] + == ctx_info_endpoint + ) + + +@pytest.mark.asyncio +async def test_stateless_router_syncs_coordinator_server_add_remove(): + """Coordinator server lists propagate metadata topology changes.""" + config = _make_config(["ctx-old:8000"], [], "round_robin", "round_robin") + client = CoordinatorClient("http://coordinator", config) + client.ctx_router._fetch_server_info = AsyncMock(return_value={}) + + await client._sync_stateless_routers( + {"server_lists": {"context": ["ctx-new:8001"], "generation": []}} + ) + + assert client.ctx_router.servers == ["ctx-new:8001"] + client.ctx_router._fetch_server_info.assert_awaited_once_with("ctx-new:8001", None) + await client.stop() + + def test_conversation_coordinator_sticky_by_conv_id(): """Verify conversation IDs remain sticky through delegated routing. diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index 710a4fc33acc..f8d3c0ef903b 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -77,7 +77,7 @@ def _make_mock_aiohttp_session(return_value=None): @pytest.mark.asyncio -async def test_coordinator_finish_retries_failed_release(): +async def test_coordinator_finish_retry_is_bounded(): local_router = RoundRobinRouter(server_role=None, servers=["server1"]) router = CoordinatorDelegatingRouter("http://coordinator", local_router, "generation") @@ -97,7 +97,6 @@ def _response(status, body): _response(503, {"error": "temporarily unavailable"}), _response(503, {"error": "temporarily unavailable"}), _response(503, {"error": "temporarily unavailable"}), - _response(200, {}), ]) session.close = mock.AsyncMock() router._session = session @@ -106,8 +105,8 @@ def _response(status, body): new_callable=mock.AsyncMock) as sleep: await router._finish_async(123, True) - assert session.post.call_count == 4 - assert sleep.await_count == 3 + assert session.post.call_count == 3 + assert sleep.await_count == 2 @pytest.fixture(autouse=True) diff --git a/tests/unittest/llmapi/test_disagg_telemetry_launcher.py b/tests/unittest/llmapi/test_disagg_telemetry_launcher.py index d67a3f84faae..5d8530c53a1e 100644 --- a/tests/unittest/llmapi/test_disagg_telemetry_launcher.py +++ b/tests/unittest/llmapi/test_disagg_telemetry_launcher.py @@ -49,7 +49,13 @@ def test_disaggregated_command_sets_shared_deployment_id(monkeypatch) -> None: raising=False, ) - disagg_config = SimpleNamespace(hostname="127.0.0.1", port=0, schedule_style=None) + disagg_config = SimpleNamespace( + hostname="127.0.0.1", + port=0, + schedule_style=None, + num_workers=1, + disagg_coordinator_url=None, + ) fake_socket = mock.MagicMock() fake_socket.__enter__.return_value = fake_socket deployment_id = SimpleNamespace(hex="deploy123") @@ -60,7 +66,7 @@ def test_disaggregated_command_sets_shared_deployment_id(monkeypatch) -> None: mock.patch.object(serve.socket, "socket", return_value=fake_socket), mock.patch.object(serve, "parse_metadata_server_config_file", return_value=None), mock.patch.object(serve, "OpenAIDisaggServer"), - mock.patch.object(serve.asyncio, "run"), + mock.patch.object(serve.uvloop, "run"), ): serve.disaggregated.callback( config_file="disagg.yaml", @@ -76,6 +82,80 @@ def test_disaggregated_command_sets_shared_deployment_id(monkeypatch) -> None: fake_socket.bind.assert_called_once_with(("127.0.0.1", 0)) +@pytest.mark.parametrize("schedule_style", ["context_first", "generation_first"]) +def test_launch_disagg_fleet_propagates_resolved_schedule_style( + monkeypatch, + schedule_style, +) -> None: + """Fleet children receive the resolved CLI-or-config schedule style.""" + observed_envs = [] + + class _FakePopen: + pid = 12345 + + def __init__(self, _command, **kwargs): + observed_envs.append(kwargs["env"]) + + def poll(self): + return None + + disagg_config = SimpleNamespace( + hostname="127.0.0.1", + port=8000, + schedule_style=schedule_style, + ) + monkeypatch.setattr(serve.subprocess, "Popen", _FakePopen) + monkeypatch.setattr(serve.atexit, "register", lambda *_: None) + monkeypatch.setattr(serve.signal, "signal", lambda *_: None) + + serve._launch_disagg_fleet( + disagg_config, + "disagg.yaml", + None, + 180, + 180, + 1, + "http://coordinator:8001", + ) + + assert observed_envs[0][serve.DisaggWorkerEnvs.TLLM_DISAGG_SCHEDULE_STYLE] == schedule_style + + +@pytest.mark.parametrize( + ("config_style", "resolved_style"), + [ + ("context_first", "generation_first"), + ("generation_first", "context_first"), + ], +) +def test_build_fleet_worker_applies_resolved_schedule_style( + monkeypatch, + config_style, + resolved_style, +) -> None: + """A fleet worker overrides its reparsed YAML with the parent resolution.""" + disagg_config = SimpleNamespace(schedule_style=config_style) + monkeypatch.setenv(serve.DisaggWorkerEnvs.TLLM_DISAGG_CONFIG_FILE, "disagg.yaml") + monkeypatch.setenv( + serve.DisaggWorkerEnvs.TLLM_DISAGG_COORDINATOR_URL, + "http://coordinator:8001", + ) + monkeypatch.setenv( + serve.DisaggWorkerEnvs.TLLM_DISAGG_SCHEDULE_STYLE, + resolved_style, + ) + + with ( + mock.patch.object(serve, "parse_disagg_config_file", return_value=disagg_config), + mock.patch.object(serve, "parse_metadata_server_config_file", return_value=None), + mock.patch.object(serve, "OpenAIDisaggServer") as mock_server, + ): + serve._build_disagg_server_from_env() + + assert disagg_config.schedule_style == resolved_style + assert mock_server.call_args.kwargs["config"] is disagg_config + + def test_launch_disaggregated_leader_propagates_deployment_id(monkeypatch) -> None: """Leader subprocess env keeps the shared telemetry deployment id.""" observed = {} @@ -153,7 +233,10 @@ def test_launch_disaggregated_server_sets_worker_role( llm_args = {"model": "dummy/model"} server_config = SimpleNamespace(type=server_type, hostname="127.0.0.1", port=8000) - disagg_config = SimpleNamespace(server_configs=[server_config]) + disagg_config = SimpleNamespace( + server_configs=[server_config], + allow_request_chat_template=False, + ) with ( mock.patch.object(serve, "parse_disagg_config_file", return_value=disagg_config), @@ -167,4 +250,5 @@ def test_launch_disaggregated_server_sets_worker_role( host="127.0.0.1", port=8000, llm_args=llm_args, + allow_request_chat_template=False, ) From a882b1738e26ef4642bc18717007758abc9ef9b9 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 15 Jul 2026 21:41:08 -0700 Subject: [PATCH 15/18] [None][fix] move fleet state sync off request path Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/serve/disagg_coordinator.py | 53 +++++++++----- .../disaggregated/test_coordinator_worker.py | 70 +++++++++++++++++++ 2 files changed, 106 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py index 9123767bfa98..74be46d48ece 100644 --- a/tensorrt_llm/serve/disagg_coordinator.py +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -77,6 +77,7 @@ COORDINATOR_RESERVATION_TIMEOUT_ENV = "TRTLLM_DISAGG_COORDINATOR_RESERVATION_TIMEOUT" COORDINATOR_RESERVATION_TIMEOUT_DEFAULT_S = 180.0 +COORDINATOR_STATE_SYNC_INTERVAL_S = 3.0 def coordinator_reservation_timeout() -> float: @@ -433,8 +434,8 @@ class CoordinatorClient(DisaggCoordinator): so the worker computes the small routing key locally and the coordinator makes the placement (placement -> ``/select``, finish -> ``/finish``). A *stateless* router (round_robin, load_balancing) is used as-is and places locally in the - worker -- no coordinator round-trip. Readiness / cluster_info always proxy the - coordinator over HTTP. + worker -- no coordinator round-trip. A background ``/cluster_info`` poll keeps + readiness and stateless-router server lists synchronized with the coordinator. Args: remote_url: Coordinator base URL (e.g. ``http://host:PORT``). @@ -459,6 +460,14 @@ def __init__( self._request_timeout_s = request_timeout_s self._startup_timeout_s = startup_timeout_s self._session: Optional[aiohttp.ClientSession] = None + self._sync_task: Optional[asyncio.Task] = None + self._is_ready = False + if config.disagg_cluster_config is not None: + self._state_sync_interval_s = config.disagg_cluster_config.heartbeat_interval_sec + elif metadata_config is not None: + self._state_sync_interval_s = metadata_config.refresh_interval + else: + self._state_sync_interval_s = COORDINATOR_STATE_SYNC_INTERVAL_S # Local disagg-id generation (no HTTP hop): (node_id, per-worker process_id) # keeps the snowflake unique across co-located workers. self._node_id = config.node_id @@ -508,8 +517,21 @@ async def start(self) -> None: # Fail fast: probe /cluster_info with bounded retry so a delegating server # exits non-zero if its coordinator never comes up (vs 500-ing every req). info = await self._await_coordinator() - self._sync_delegating_router_configs(info) - await self._sync_stateless_routers(info) + await self._apply_cluster_info(info) + self._sync_task = asyncio.create_task( + self._sync_coordinator_state(self._state_sync_interval_s) + ) + + async def _sync_coordinator_state(self, interval_s: float) -> None: + while True: + try: + await asyncio.sleep(interval_s) + await self._apply_cluster_info(await self.cluster_info()) + except asyncio.CancelledError: + raise + except Exception as error: # noqa: BLE001 + self._is_ready = False + logger.warning(f"CoordinatorClient state sync failed: {error}") async def _await_coordinator(self) -> Dict[str, Any]: """Poll the coordinator's /cluster_info until reachable, or raise. @@ -556,19 +578,12 @@ async def get_disagg_request_id(self) -> int: return get_global_disagg_request_id(self._node_id, self._process_id) async def is_ready(self) -> bool: - try: - async with self.session.get( - f"{self._remote_url}/health", timeout=self._request_timeout_s - ) as resp: - if resp.status != 200: - return False - info = await self.cluster_info() - self._sync_delegating_router_configs(info) - await self._sync_stateless_routers(info) - return info.get("is_ready", False) - except Exception as e: # noqa: BLE001 - logger.warning(f"CoordinatorClient health check failed: {e}") - return False + return self._is_ready + + async def _apply_cluster_info(self, info: Dict[str, Any]) -> None: + self._is_ready = info.get("is_ready", False) + self._sync_delegating_router_configs(info) + await self._sync_stateless_routers(info) def _sync_delegating_router_configs(self, info: Dict[str, Any]) -> None: configs = info.get("routing_key_configs", {}) @@ -628,6 +643,10 @@ async def proxy_cluster_storage_request( return await resp.read(), resp.status, resp.headers.get("Content-Type") async def stop(self) -> None: + if self._sync_task is not None: + self._sync_task.cancel() + await asyncio.gather(self._sync_task, return_exceptions=True) + self._sync_task = None if self._session is not None: await self._session.close() self._session = None diff --git a/tests/unittest/disaggregated/test_coordinator_worker.py b/tests/unittest/disaggregated/test_coordinator_worker.py index 615d671530c4..894d08ff0682 100644 --- a/tests/unittest/disaggregated/test_coordinator_worker.py +++ b/tests/unittest/disaggregated/test_coordinator_worker.py @@ -44,6 +44,7 @@ from tensorrt_llm.llmapi.disagg_utils import ( CtxGenServerConfig, + DisaggClusterConfig, DisaggServerConfig, RouterConfig, ServerRole, @@ -303,6 +304,75 @@ def test_coordinator_client_configures_empty_delegating_kv_router(): assert "token_lists" not in routing_key +@pytest.mark.asyncio +async def test_coordinator_client_readiness_is_cached(): + config = _make_config([], [], "round_robin", "round_robin") + client = CoordinatorClient("http://coordinator", config) + client._is_ready = True + + assert await client.is_ready() is True + assert client._session is None + await client.stop() + + +@pytest.mark.asyncio +async def test_coordinator_state_sync_starts_in_background(): + config = _make_config([], [], "round_robin", "round_robin") + client = CoordinatorClient("http://coordinator", config) + client._state_sync_interval_s = 10 + client._await_coordinator = AsyncMock( + return_value={ + "is_ready": True, + "server_lists": {"context": [], "generation": []}, + } + ) + client._sync_coordinator_state = AsyncMock() + + await client.start() + + assert client._is_ready is True + client._sync_coordinator_state.assert_called_once_with(10) + await client.stop() + + +def test_service_discovery_sets_coordinator_state_sync_interval(): + config = _make_config([], [], "round_robin", "round_robin") + config.disagg_cluster_config = DisaggClusterConfig( + cluster_uri="http://cluster-storage", + heartbeat_interval_sec=7, + ) + + client = CoordinatorClient("http://coordinator", config) + + assert client._state_sync_interval_s == 7 + + +@pytest.mark.asyncio +async def test_cluster_info_updates_readiness_and_stateless_servers(): + config = _make_config(["ctx-old:8001"], [], "round_robin", "round_robin") + client = CoordinatorClient("http://coordinator", config) + client.ctx_router.remove_server = AsyncMock(wraps=client.ctx_router.remove_server) + client.ctx_router.add_server = AsyncMock(wraps=client.ctx_router.add_server) + client.ctx_router.prepare_servers = AsyncMock() + + await client._apply_cluster_info( + { + "is_ready": True, + "server_lists": { + "context": ["ctx-new:8001"], + "generation": [], + }, + } + ) + + assert await client.is_ready() is True + assert client.ctx_router.servers == ["ctx-new:8001"] + client.ctx_router.remove_server.assert_awaited_once_with("ctx-old:8001") + client.ctx_router.add_server.assert_awaited_once_with("ctx-new:8001") + client.ctx_router.prepare_servers.assert_awaited_once() + await client.stop() + + def test_content_affinity_key_uses_fixed_seed(): request = ChatCompletionRequest(model="m", messages=[{"role": "user", "content": "hello"}]) From 9b2a8cd8a01d2adc52f7ba40fe6f7f1d74f0bded Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 15 Jul 2026 21:49:09 -0700 Subject: [PATCH 16/18] [None][test] start coordinator client before readiness check Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tests/unittest/disaggregated/test_coordinator_worker.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/disaggregated/test_coordinator_worker.py b/tests/unittest/disaggregated/test_coordinator_worker.py index 894d08ff0682..cfe435b22a97 100644 --- a/tests/unittest/disaggregated/test_coordinator_worker.py +++ b/tests/unittest/disaggregated/test_coordinator_worker.py @@ -502,6 +502,7 @@ def _req(conv_id, request_id): async def drive(): remote = CoordinatorClient(coord.url, config) + await remote.start() # Stateful -> wrapped in a coordinator-delegating router. assert isinstance(remote.gen_router, CoordinatorDelegatingRouter) assert await remote.is_ready() is True From 647370cf1d7b4dc85083f9257ab443bfcb14c1f5 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 16 Jul 2026 00:11:15 -0700 Subject: [PATCH 17/18] [None][fix] harden coordinator routing hot path Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/serve/disagg_coordinator.py | 27 +++- tensorrt_llm/serve/openai_disagg_service.py | 2 +- tensorrt_llm/serve/router.py | 141 ++++++++++++++---- .../disaggregated/test_coordinator_worker.py | 41 +++++ .../test_openai_disagg_service.py | 26 +++- tests/unittest/disaggregated/test_router.py | 28 +++- 6 files changed, 223 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py index 74be46d48ece..86475af498fa 100644 --- a/tensorrt_llm/serve/disagg_coordinator.py +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -220,14 +220,22 @@ async def select( if previous is not None: previous.cancel() await router.finish_request_by_id(req_id, False) - result = await router.get_next_server_by_key( + server, info, request_id = await router.get_next_server_by_key( routing_key, req_id=req_id, exclude_server=exclude_server ) self._reservation_tasks[reservation_key] = asyncio.create_task( self._expire_reservation(reservation_key, router) ) self._api_lat(f"select[{role}]").record(time.monotonic() - _t0) - return result + return server, self._compact_route_info(info), request_id + + @staticmethod + def _compact_route_info(info: dict) -> dict: + compact = {key: info[key] for key in ("match_length", "num_tokens") if key in info} + disaggregated_params = info.get("server_info", {}).get("disaggregated_params") + if disaggregated_params is not None: + compact["server_info"] = {"disaggregated_params": disaggregated_params} + return compact async def get_disagg_request_id(self) -> int: return get_global_disagg_request_id(self._config.node_id) @@ -605,11 +613,20 @@ async def _sync_stateless_routers(self, info: Dict[str, Any]) -> None: if isinstance(router, CoordinatorDelegatingRouter): continue desired = set(servers) + for server in desired - set(router.servers): + if not await router.add_server(server): + raise RuntimeError( + f"Failed to prepare {server} before adding it to the " + f"{router.server_role.name} router" + ) + await router.prepare_servers(list(desired)) + unprepared = desired - router.prepared_servers + if unprepared: + raise RuntimeError( + f"Servers are not prepared for {router.server_role.name}: {sorted(unprepared)}" + ) for server in set(router.servers) - desired: await router.remove_server(server) - for server in desired - set(router.servers): - await router.add_server(server) - await router.prepare_servers() async def cluster_info(self) -> Dict[str, Any]: try: diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index 1dc5cef0d1ff..e2ee00257e98 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -292,7 +292,7 @@ async def _check_conditional_disagg(self, request: UCompletionRequest, req_id: i # Query kv cache status and select a best gen_server. # The server is reserved for generation request gen_server, info = await self._gen_router.get_next_server(request, req_id=req_id) - match_length = sum(info["matches"]) + match_length = info["match_length"] total_length = info["num_tokens"] need_ctx_decision = ( match_length == 0 diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 79bc14bb40e4..f0588a7e82f6 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -40,6 +40,10 @@ _MSGPACK_HEADERS = {"Content-Type": "application/msgpack"} COORDINATOR_FINISH_MAX_ATTEMPTS = 3 COORDINATOR_FINISH_RETRY_DELAY_S = 0.1 +COORDINATOR_FINISH_TIMEOUT_S = 5.0 +COORDINATOR_FINISH_WORKERS = 16 +COORDINATOR_FINISH_QUEUE_SIZE = 4096 +COORDINATOR_FINISH_DRAIN_TIMEOUT_S = 5.0 # Max number of conversations whose home-server pin is retained (LRU). ROUTE_AFFINITY_CACHE_SIZE = 50000 @@ -269,6 +273,13 @@ def _create_server_state(self, server: str) -> ServerState: return self._server_state_class(server, self._use_tokens, lambda: self.session) + def _stage_server(self, server: str) -> None: + self._server_state.setdefault(server, self._create_server_state(server)) + + def _unstage_server(self, server: str) -> None: + if server not in self._servers: + self._server_state.pop(server, None) + def _get_server_load(self, server: str) -> int: state = self._server_state[server] return state._num_active_tokens if self._use_tokens \ @@ -299,7 +310,7 @@ def _select_least_loaded(self, exclude_server: Optional[str] = None ) -> Optional[str]: """Pick the server with the lowest load. Round-robin breaks ties.""" - candidates = [s for s in self._server_state if s != exclude_server] + candidates = [s for s in self._servers if s != exclude_server] if not candidates: return None loads = {s: self._get_server_load(s) for s in candidates} @@ -365,6 +376,14 @@ def servers(self) -> List[str]: def num_prepared_servers(self) -> int: return len(self._prepared_ready_servers) + @property + def prepared_servers(self) -> set[str]: + return set(self._prepared_ready_servers) + + @property + def server_role(self) -> ServerRole: + return self._server_role + @staticmethod def _ensure_url(server: str) -> str: return server if server.startswith("http") else f"http://{server}" @@ -398,23 +417,45 @@ async def _prepare_server(self, server: str): logger.warning(f"Error preparing server {server}: {e}") async def prepare_servers(self, servers: Optional[List[str]] = None): - for server in servers or self._servers: + targets = self._servers if servers is None else servers + for server in targets: if server not in self._servers: continue await self._prepare_server(server) - async def add_server(self, server: str): + def _stage_server(self, server: str) -> None: + pass + + def _unstage_server(self, server: str) -> None: + pass + + async def add_server(self, server: str) -> bool: if server in self._servers: logger.warning(f"Server {server} already exists") - return + return True async with self._lock: - old_servers = self._servers.copy() - self._servers = [*old_servers, server] - self._on_servers_updated(old_servers, self._servers) - await self._prepare_server(server) + self._stage_server(server) + try: + await self._prepare_server(server) + except Exception: + async with self._lock: + self._unstage_server(server) + raise + if server not in self._prepared_ready_servers: + async with self._lock: + self._unstage_server(server) + logger.warning( + f"Server {server} was not added because preparation failed") + return False + async with self._lock: + if server not in self._servers: + old_servers = self._servers.copy() + self._servers = [*old_servers, server] + self._on_servers_updated(old_servers, self._servers) logger.debug( f"Added server {server}, {self._server_role.name} current server list: {self._servers}" ) + return True async def remove_server(self, server: str): if server not in self._servers: @@ -771,6 +812,7 @@ def __init__(self, self._load_weight = load_weight self._load_cap = load_cap self._track_routed_blocks = track_routed_blocks + self._pending_routed_blocks: dict[int, tuple[list[BlockHash], str]] = {} # A coordinator client has no local server pool. The coordinator injects # the role's effective hash algorithm so this router can act solely as a # routing-key encoder. @@ -786,14 +828,13 @@ async def close(self): await state.cancel_poll_task() await super().close() - def _apply_routed_blocks_on_route(self, server: str, - block_hashes: list[list[BlockHash]], - hash_algo: str) -> None: + def _stash_routed_blocks(self, key: int, + block_hashes: list[list[BlockHash]], + hash_algo: str) -> None: if not self._track_routed_blocks: return flat_block_hashes = [h for hashes in block_hashes for h in hashes] - self._server_state[server].add_blocks(flat_block_hashes, - hash_algo=hash_algo) + self._pending_routed_blocks[key] = (flat_block_hashes, hash_algo) def _get_server_hash_algo(self, server: str) -> str: # Lock-free attribute read; state is seeded at handshake and refreshed @@ -950,7 +991,7 @@ async def _route(self, key, exclude_server=None, request=None, req_id=None): server_states = { server: state for server, state in self._server_state.items() - if server != exclude_server + if server in self._servers and server != exclude_server } servers = list(server_states) if not servers: @@ -1010,8 +1051,7 @@ def _hashes(server): await server_states[server].increment_load(request) self._req_routing_table[key] = server try: - self._apply_routed_blocks_on_route(server, block_hashes, - hash_algo) + self._stash_routed_blocks(key, block_hashes, hash_algo) except Exception: self._req_routing_table.pop(key, None) await server_states[server].decrement_load(request) @@ -1020,6 +1060,7 @@ def _hashes(server): "block_hashes": block_hashes, "hash_algo": hash_algo, "matches": matches, + "match_length": matches[winner], "num_tokens": num_tokens, "server_info": self._server_info.get(server, {}), }, req_id @@ -1041,13 +1082,18 @@ async def _finish(self, key, success, request=None, session=None): """Finish through the core shared by standalone and coordinator paths. This removes the maps populated by ``_route``, decrements load, and - refreshes the block table. Routed blocks are applied at route time. + applies routed blocks only after successful completion. """ - del success async with self._lock: server = self._req_routing_table.pop(key, None) + pending = self._pending_routed_blocks.pop(key, None) if server is not None and server in self._server_state: await self._server_state[server].decrement_load(request) + if (success and pending is not None and server is not None + and server in self._server_state): + block_hashes, hash_algo = pending + self._server_state[server].add_blocks(block_hashes, + hash_algo=hash_algo) self._poll_server_on_finish(server, session) def _poll_server_on_finish(self, server, session=None): @@ -1627,7 +1673,10 @@ def __init__(self, self._role = role # "context" | "generation" self._request_timeout_s = request_timeout_s self._session: Optional[aiohttp.ClientSession] = None - self._background_tasks: set[asyncio.Task] = set() + self._finish_queue: asyncio.Queue[tuple[int, bool]] = asyncio.Queue( + maxsize=COORDINATOR_FINISH_QUEUE_SIZE) + self._finish_workers: set[asyncio.Task] = set() + self._dropped_finishes = 0 # Coordinator HTTP-client API latency (includes network round-trip to the # coordinator + its in-process handler). Compare against the owner-side # [coord_api] to isolate the fleet /select|/finish HTTP overhead. @@ -1713,15 +1762,36 @@ async def finish_request(self, session: Optional[aiohttp.ClientSession] = None, success: bool = True, req_id: Optional[int] = None): - # Fire-and-forget: /finish only releases the coordinator's routed-load - # bookkeeping, so run it as a background task instead of awaiting the HTTP - # round trip on every request's critical path. Capture req_id first (the - # request may be freed after return); errors are logged, not propagated. + # /finish only releases coordinator bookkeeping, so enqueue it off the + # request path. A fixed worker pool bounds tasks and connections during a + # coordinator outage; overflow relies on coordinator-side expiration. del session req_id = self._request_id(request, req_id) - task = asyncio.create_task(self._finish_async(req_id, success)) - self._background_tasks.add(task) - task.add_done_callback(self._background_tasks.discard) + self._ensure_finish_workers() + try: + self._finish_queue.put_nowait((req_id, success)) + except asyncio.QueueFull: + self._dropped_finishes += 1 + if self._dropped_finishes == 1 or self._dropped_finishes % 1000 == 0: + logger.warning( + f"CoordinatorDelegatingRouter finish queue full; " + f"coordinator expiration will release dropped requests " + f"(dropped={self._dropped_finishes})") + + def _ensure_finish_workers(self) -> None: + if self._finish_workers: + return + for _ in range(COORDINATOR_FINISH_WORKERS): + task = asyncio.create_task(self._finish_worker()) + self._finish_workers.add(task) + + async def _finish_worker(self) -> None: + while True: + req_id, success = await self._finish_queue.get() + try: + await self._finish_async(req_id, success) + finally: + self._finish_queue.task_done() async def _finish_async(self, req_id: int, success: bool): _t0 = time.monotonic() @@ -1737,7 +1807,8 @@ async def _finish_async(self, req_id: int, success: bool): }, use_bin_type=True), headers=_MSGPACK_HEADERS, - timeout=self._request_timeout_s) as resp: + timeout=min(self._request_timeout_s, + COORDINATOR_FINISH_TIMEOUT_S)) as resp: if resp.status != 200: raw_body = await resp.read() try: @@ -1764,9 +1835,19 @@ async def _finish_async(self, req_id: int, success: bool): (2**(attempt - 1))) async def close(self): - if self._background_tasks: - await asyncio.gather(*self._background_tasks, - return_exceptions=True) + if self._finish_workers: + try: + await asyncio.wait_for( + self._finish_queue.join(), + timeout=COORDINATOR_FINISH_DRAIN_TIMEOUT_S) + except asyncio.TimeoutError: + logger.warning( + "Timed out draining coordinator finish queue; " + "coordinator expiration will release remaining requests") + for task in self._finish_workers: + task.cancel() + await asyncio.gather(*self._finish_workers, return_exceptions=True) + self._finish_workers.clear() if self._session is not None: await self._session.close() self._session = None diff --git a/tests/unittest/disaggregated/test_coordinator_worker.py b/tests/unittest/disaggregated/test_coordinator_worker.py index cfe435b22a97..39ae11bc33a4 100644 --- a/tests/unittest/disaggregated/test_coordinator_worker.py +++ b/tests/unittest/disaggregated/test_coordinator_worker.py @@ -268,6 +268,28 @@ async def test_coordinator_expires_stale_reservation(): assert coordinator._reservation_tasks == {} +def test_coordinator_compacts_route_info(): + compact = DisaggCoordinatorService._compact_route_info( + { + "block_hashes": [["large-hash"]], + "hash_algo": KV_CACHE_HASH_ALGO_V2, + "matches": [64, 32], + "match_length": 64, + "num_tokens": 128, + "server_info": { + "tokens_per_block": 32, + "disaggregated_params": {"ctx_info_endpoint": "tcp://ctx"}, + }, + } + ) + + assert compact == { + "match_length": 64, + "num_tokens": 128, + "server_info": {"disaggregated_params": {"ctx_info_endpoint": "tcp://ctx"}}, + } + + def test_coordinator_reservation_timeout_env(monkeypatch): monkeypatch.delenv(COORDINATOR_RESERVATION_TIMEOUT_ENV, raising=False) assert coordinator_reservation_timeout() == 180 @@ -354,6 +376,7 @@ async def test_cluster_info_updates_readiness_and_stateless_servers(): client.ctx_router.remove_server = AsyncMock(wraps=client.ctx_router.remove_server) client.ctx_router.add_server = AsyncMock(wraps=client.ctx_router.add_server) client.ctx_router.prepare_servers = AsyncMock() + client.ctx_router._fetch_server_info = AsyncMock(return_value={}) await client._apply_cluster_info( { @@ -476,6 +499,24 @@ async def test_stateless_router_syncs_coordinator_server_add_remove(): await client.stop() +@pytest.mark.asyncio +async def test_stateless_router_keeps_old_server_when_replacement_is_unprepared(): + config = _make_config(["ctx-old:8000"], [], "round_robin", "round_robin") + client = CoordinatorClient("http://coordinator", config) + client.ctx_router._prepared_ready_servers.add("ctx-old:8000") + client.ctx_router._fetch_server_info = AsyncMock( + side_effect=RuntimeError("server info unavailable") + ) + + with pytest.raises(RuntimeError, match="Failed to prepare ctx-new:8001"): + await client._sync_stateless_routers( + {"server_lists": {"context": ["ctx-new:8001"], "generation": []}} + ) + + assert client.ctx_router.servers == ["ctx-old:8000"] + await client.stop() + + def test_conversation_coordinator_sticky_by_conv_id(): """Verify conversation IDs remain sticky through delegated routing. diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index b9cec8475a94..04defdd1cd37 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -23,6 +23,7 @@ from tensorrt_llm.disaggregated_params import DisaggregatedParams as LlmDisaggregatedParams from tensorrt_llm.executor.result import Logprob from tensorrt_llm.llmapi.disagg_utils import ( + ConditionalDisaggConfig, DisaggClusterConfig, DisaggServerConfig, MinimalInstances, @@ -55,7 +56,7 @@ chat_response_post_processor, completion_response_post_processor, ) -from tensorrt_llm.serve.router import Router +from tensorrt_llm.serve.router import KvCacheAwareRouter, Router def _client_factory(*_args, **_kwargs): @@ -78,6 +79,29 @@ def _make_service(schedule_style: str) -> OpenAIDisaggregatedService: return service +@pytest.mark.asyncio +async def test_conditional_disagg_uses_selected_server_match_length(): + service = _make_service("context_first") + service._config.conditional_disagg_config = ConditionalDisaggConfig(max_local_prefill_length=32) + router = KvCacheAwareRouter(server_role=ServerRole.GENERATION, servers=[]) + router.get_next_server = AsyncMock( + return_value=( + "gen:8000", + { + "match_length": 64, + "num_tokens": 96, + }, + ) + ) + service._gen_router = router + request = CompletionRequest(model="model", prompt=[1] * 96) + + server, need_context = await service._check_conditional_disagg(request, 123) + + assert server == "gen:8000" + assert need_context is False + + def _make_completion_response( text: str, finish_reason: str, diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index f8d3c0ef903b..ddaf12a0c8e2 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -107,6 +107,24 @@ def _response(status, body): assert session.post.call_count == 3 assert sleep.await_count == 2 + assert all(call.kwargs["timeout"] == 5 + for call in session.post.call_args_list) + + +@pytest.mark.asyncio +async def test_coordinator_finish_queue_is_bounded(): + local_router = RoundRobinRouter(server_role=None, servers=["server1"]) + router = CoordinatorDelegatingRouter("http://coordinator", local_router, + "generation") + router._finish_queue = asyncio.Queue(maxsize=1) + router._ensure_finish_workers = mock.Mock() + + request = mock.Mock() + await router.finish_request(request, req_id=1) + await router.finish_request(request, req_id=2) + + assert router._finish_queue.qsize() == 1 + assert router._dropped_finishes == 1 @pytest.fixture(autouse=True) @@ -821,7 +839,7 @@ def test_kv_cache_aware_server_state_remove_blocks_silent_on_missing(): @pytest.mark.asyncio -async def test_kv_cache_aware_router_does_not_wait_for_finish_to_apply_blocks( +async def test_kv_cache_aware_router_applies_blocks_after_successful_finish( servers): tokens_per_block = 4 token_lists = [[2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008]] @@ -839,7 +857,7 @@ async def test_kv_cache_aware_router_does_not_wait_for_finish_to_apply_blocks( hash_algo = info["hash_algo"] assert await router._server_state[server].matched_tokens( - info["block_hashes"], hash_algo=hash_algo) == 8 + info["block_hashes"], hash_algo=hash_algo) == 0 await router.finish_request(request) @@ -872,7 +890,7 @@ async def test_kv_cache_aware_router_routed_blocks_disabled_skips_pending( @pytest.mark.asyncio -async def test_kv_cache_aware_router_keeps_routed_blocks_on_failure(servers): +async def test_kv_cache_aware_router_discards_routed_blocks_on_failure(servers): tokens_per_block = 4 token_lists = [[4000, 4001, 4002, 4003, 4004, 4005, 4006, 4007, 4008]] router = KvCacheAwareRouter(server_role=None, @@ -886,12 +904,12 @@ async def test_kv_cache_aware_router_keeps_routed_blocks_on_failure(servers): prompt=copy.deepcopy(token_lists)) server, info = await router.get_next_server(request) assert await router._server_state[server].matched_tokens( - info["block_hashes"], hash_algo=info["hash_algo"]) == 8 + info["block_hashes"], hash_algo=info["hash_algo"]) == 0 await router.finish_request(request, success=False) assert await router._server_state[server].matched_tokens( - info["block_hashes"], hash_algo=info["hash_algo"]) == 8 + info["block_hashes"], hash_algo=info["hash_algo"]) == 0 @pytest.mark.asyncio From 2922df335e37835289c49e47c2c5f15df4be6343 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Fri, 17 Jul 2026 00:31:10 -0700 Subject: [PATCH 18/18] [None][fix] update KV router integration test metadata usage Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../defs/disaggregated/test_workers.py | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/tests/integration/defs/disaggregated/test_workers.py b/tests/integration/defs/disaggregated/test_workers.py index 35149391bf32..75259f782ddb 100644 --- a/tests/integration/defs/disaggregated/test_workers.py +++ b/tests/integration/defs/disaggregated/test_workers.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import asyncio import contextlib import copy @@ -409,10 +423,10 @@ async def multi_round_request(self, prompt=request["prompt"], disaggregated_params=DisaggregatedParams( request_type="context_only")) - ctx_server, ctx_info = await self.ctx_router.get_next_server( - openai_request) + ctx_server, _ = await self.ctx_router.get_next_server(openai_request + ) prompt_str = request["prompt"] - request["prompt"] = ctx_info["token_lists"][0] + request["prompt"] = openai_request.prompt openai_request.disaggregated_params.request_type = "generation_only" gen_server, _ = await self.gen_router.get_next_server(openai_request )