diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index f38e3c818f57..9ed621158a26 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 @@ -201,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 '{ @@ -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: + +```bash +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..d753689c65f4 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,302 @@ 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 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 + 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) + 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 + + 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() + 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). + worker_env[DisaggWorkerEnvs.TLLM_DISAGG_WORKER_PROCESS_ID] = str(i) + p = subprocess.Popen(cmd, + env=worker_env, + stdout=sys.stdout, + stderr=sys.stderr) + 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() + deadline = time.monotonic() + 10 + for p in fleet: + try: + 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 + + +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. + 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, + 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. + 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 + # 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})") + + 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() + + try: + asyncio.run(_serve_and_monitor()) + finally: + for process in fleet: + if process.poll() is None: + process.terminate() + + +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) + 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) + + 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 _run_fleet_worker(): + """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 + + 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 +2103,29 @@ 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" + 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" + # 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..a8272b494645 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 @@ -93,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 @@ -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 {} @@ -236,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 @@ -243,6 +258,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 +435,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 +495,23 @@ 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: + """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/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/coordinator_server.py b/tensorrt_llm/serve/coordinator_server.py new file mode 100644 index 000000000000..8eb081fd5e34 --- /dev/null +++ b/tensorrt_llm/serve/coordinator_server.py @@ -0,0 +1,173 @@ +# 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 +from contextlib import asynccontextmanager +from typing import Optional + +import msgpack +import uvicorn +from fastapi import FastAPI, Request +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 + +TIMEOUT_KEEP_ALIVE = 10 # seconds +MSGPACK_MEDIA_TYPE = "application/msgpack" + + +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"]) + 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: + try: + 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 + or "req_id" not in body + ): + return self._response( + {"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["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): + 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 self._response({"error": str(e)}, status_code=503) + except Exception as e: # noqa: BLE001 + logger.error(f"CoordinatorServer.select failed: {e}") + 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: + try: + 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 "req_id" not in body: + 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 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 self._response({"error": "req_id must be an integer"}, status_code=400) + success = body.get("success", True) + if not isinstance(success, bool): + return self._response({"error": "success must be a boolean"}, status_code=400) + await self._coordinator.finish(role, req_id, success) + return self._response({}) + + async def cluster_info(self) -> Response: + 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 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); + # 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 + + 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() + + +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..86475af498fa --- /dev/null +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -0,0 +1,671 @@ +# 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 os +import time +from abc import ABC, abstractmethod +from typing import Any, Dict, Optional, Tuple + +import aiohttp +import msgpack + +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, + KvCacheAwareRouter, + Router, + build_disagg_routers, +) + +__all__ = [ + "DisaggCoordinator", + "DisaggCoordinatorService", + "CoordinatorClient", + "coordinator_base_url", + "make_coordinator_session", +] + +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: + 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. + + 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, + reservation_timeout_secs: Optional[float] = None, + ): + 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._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 + 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): + """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: + 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) + 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) + 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 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) + + 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) + + 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 "context" + if normalized_role in ("generation", "gen"): + 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() + 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: + 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: + 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(), + "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): + 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 + + 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: + """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 + 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. 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``). + 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 + 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 + 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() + 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. + + 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 msgpack.unpackb(await resp.read(), raw=False) + 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: + 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", {}) + 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: + server_lists = info.get("server_lists") + if server_lists is None: + return + router_servers = ( + (self._ctx_router, server_lists.get("context", [])), + (self._gen_router, server_lists.get("generation", [])), + ) + for router, servers in router_servers: + 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) + + 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 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} + + 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._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 + 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..26646a54344c 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 @@ -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 @@ -120,7 +132,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 @@ -146,9 +158,13 @@ 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) + 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 @@ -156,7 +172,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 +191,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 +200,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) @@ -197,7 +214,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 @@ -227,7 +244,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 +263,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 +313,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 +358,21 @@ 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) + 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 = {} diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 051bb4a1c8cf..6df7651fb1dd 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -32,29 +32,32 @@ 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 +from tensorrt_llm.serve.router import Router from tensorrt_llm.version import __version__ as VERSION # 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): @@ -63,11 +66,15 @@ 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 + self.ctx_dispatch_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 @@ -78,7 +85,19 @@ def on_first_token(self, gen_server: str, request: UCompletionRequest, response: def on_resp_done(self, gen_server: str, request: UCompletionRequest, response: UCompletionResponse = None): 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)) + 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: @@ -87,7 +106,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 +115,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 +129,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,35 +164,60 @@ 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() + if self._perf_metrics_collector._background_tasks: + await asyncio.gather( + *self._perf_metrics_collector._background_tasks, + return_exceptions=True, + ) self.app = FastAPI(lifespan=lifespan) 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] + method = request.method.translate(_LOG_CONTROL_CHARACTERS) + path = request.url.path.translate(_LOG_CONTROL_CHARACTERS) + logger.warning( + f"[validation] {method} {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 +225,38 @@ 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) + 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): @@ -180,8 +267,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 +311,12 @@ def _handle_exception(self, exception): async def health(self) -> Response: - if not await self._service.is_ready(): - return Response(status_code=500) + if not await self._coordinator.is_ready(): + return Response(status_code=503) 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..e2ee00257e98 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -14,20 +14,11 @@ 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, -) +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, @@ -44,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. @@ -55,28 +46,21 @@ 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 +68,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: @@ -139,26 +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() # 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) + # 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 - disagg_request_id = get_global_disagg_request_id(self._config.node_id) if need_ctx: - 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, req_id=disagg_request_id + ) + ctx_response = await self._ctx_client.send_request( + 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 + 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, req_id=gen_reservation_id + ) + raise else: # When need_ctx=False the gen server handles full generation and # must not see a stale request_type="context_only". @@ -172,13 +170,16 @@ 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_reservation_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=gen_reservation_id ) return gen_response else: + if gen_server: + 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 @@ -277,14 +278,22 @@ 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: - 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) - match_length = sum(info["matches"]) - total_length = sum(len(token_list) for token_list in info["token_lists"]) + gen_server, info = await self._gen_router.get_next_server(request, req_id=req_id) + match_length = info["match_length"] + total_length = info["num_tokens"] need_ctx_decision = ( match_length == 0 or total_length - match_length @@ -315,110 +324,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,9 +389,16 @@ 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: - ctx_server, ctx_server_info = await self._ctx_router.get_next_server(request) + # 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, req_id=disagg_request_id + ) ctx_req = self._get_ctx_request(request, disagg_request_id) gen_req = self._get_gen_request( request, @@ -479,7 +418,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() @@ -496,7 +435,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: @@ -530,12 +474,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/perf_metrics.py b/tensorrt_llm/serve/perf_metrics.py index e2a4ff7607bb..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 @@ -170,6 +171,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 +181,7 @@ async def add_per_request_metrics( ctx_request_id, server_arrival_time, server_first_token_time, + ctx_dispatch_time, ) ) @@ -212,6 +215,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 +227,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 +238,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..b3b1c935f77d 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -218,9 +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 @@ -330,8 +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. @@ -356,8 +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. @@ -1375,9 +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 @@ -1995,6 +1997,36 @@ async def __call__(self, scope, receive, send): await self.app(scope, receive, send) +class PeriodicLatencyLogger: + """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 + 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) + + 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 ResponseHooks(ABC): """ Hooks for response processing and (disagg) service perf observability. @@ -2004,6 +2036,13 @@ class ResponseHooks(ABC): def on_req_begin(self, request: UCompletionRequest): pass + def on_ctx_dispatch(self, request: UCompletionRequest): + """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): pass diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 5941ed979e51..f0588a7e82f6 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -13,78 +13,46 @@ # 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 msgpack +from blake3 import blake3 -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, 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 +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 +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 -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__( @@ -198,9 +166,10 @@ 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) elif event["type"] == "removed": self.remove_blocks(event["block_hashes"], hash_algo=hash_algo) @@ -219,7 +188,6 @@ async def matched_tokens( 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 else: @@ -305,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 \ @@ -335,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} @@ -401,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}" @@ -434,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: @@ -468,17 +473,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]: - '''Select server by request and return some intermediate information, exclude_server is a server to exclude from the selection''' + 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 @@ -701,10 +707,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( @@ -726,8 +733,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): @@ -750,10 +758,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: @@ -768,229 +777,13 @@ 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) -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 @@ -1004,13 +797,15 @@ 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 @@ -1018,6 +813,10 @@ def __init__(self, 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. + self._routing_hash_algo: Optional[str] = None def _create_server_state(self, server: str) -> KvCacheAwareServerState: return KvCacheAwareServerState(server, self._use_tokens, @@ -1029,34 +828,45 @@ async def close(self): await state.cancel_poll_task() await super().close() - def _stash_routed_blocks_on_route(self, request: OpenAIRequest, - 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 = [h for hl in block_hashes for h in hl] - self._pending_routed_blocks[id(request)] = (flat, hash_algo) - - def _apply_routed_blocks_on_finish(self, request: OpenAIRequest, - server: Optional[str], - success: bool) -> None: - # Pop unconditionally to avoid leaks; apply only when eligible. - entry = self._pending_routed_blocks.pop(id(request), 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 - self._server_state[server].add_blocks(flat_block_hashes, - hash_algo=hash_algo) + flat_block_hashes = [h for hashes in block_hashes for h in hashes] + 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 # 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 @@ -1124,72 +934,92 @@ 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)) - - 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}") + 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, + 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) + server, info, _handle = await self._route(key, + exclude_server=exclude_server, + request=request) + return server, info + + def _routing_key_sync(self, request: OpenAIRequest) -> dict: + """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_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._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) + return { + "block_hashes_by_algo": block_hashes_by_algo, + "conv_key": self._content_affinity_key(request), + "num_tokens": sum(len(token_list) for token_list in token_lists), } - # 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): + """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``. + """ + 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 + for server, state in self._server_state.items() + if server in self._servers and server != exclude_server + } + servers = list(server_states) + if not servers: + raise ValueError( + f"No available servers after excluding {exclude_server}") + + def _hashes(server): + algo = server_states[server].hash_algo + return algo, block_hashes_by_algo.get(algo, []) + workloads = [ - self._server_state[server].num_active_requests() - for server 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)) ] - 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 server_states[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 +1038,107 @@ 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) + 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 + try: + 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) + raise 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] + "matches": matches, + "match_length": matches[winner], + "num_tokens": num_tokens, "server_info": self._server_info.get(server, {}), - } + }, req_id 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), + success, + request=request, + session=session) + + 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 + applies routed blocks only after successful completion. + """ async with self._lock: - server = self._req_routing_table.pop(id(request), None) + 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) - self._apply_routed_blocks_on_finish(request, server, success) + 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): + """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. 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): + """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): + """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): + """Finish a coordinator-side request by its disaggregated request ID.""" + 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 +1287,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 ── @@ -1595,10 +1502,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) @@ -1670,8 +1578,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) @@ -1679,6 +1588,271 @@ 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 + 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. + from tensorrt_llm.serve.responses_utils import PeriodicLatencyLogger + self._select_lat = PeriodicLatencyLogger(f"client.select[{role}]") + self._finish_lat = PeriodicLatencyLogger(f"client.finish[{role}]") + + 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 _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 + 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). + """ + if req_id is not None: + return req_id + dp = request.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) + 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(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. + 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, req_id), + "exclude_server": exclude_server + } + _t0 = time.monotonic() + async with self.session.post(f"{self._coordinator_url}/select", + 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"{body.get('error', body)}") + 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, + req_id: Optional[int] = None): + # /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) + 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() + 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=min(self._request_timeout_s, + COORDINATOR_FINISH_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"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._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 + await self._local.close() + def create_router( router_config: Optional[RouterConfig], @@ -1688,8 +1862,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: @@ -1725,3 +1898,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..a96d087acf06 --- /dev/null +++ b/tensorrt_llm/serve/router_utils.py @@ -0,0 +1,391 @@ +# 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 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 +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 + +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", + "KV_CACHE_HASH_ALGOS", + "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: Optional[OpenAIRequest]) -> int: + if request is None: + return 0 + + 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 (not request.prompt or 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 TransformersTokenizer + + 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]: + 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 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) + 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 not prompts: + return [prompts] + 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/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 6b78b3bfa788..d22bf863a5a4 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": @@ -653,7 +656,16 @@ def setup_disagg_cluster( else: work_dir = tempfile.mkdtemp() logger.info(f"Disagg cluster work_dir (worker logs): {work_dir}") - 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", {}) @@ -734,6 +746,8 @@ def setup_disagg_cluster( server_host, "port": server_port, + "num_workers": + config.get("num_workers", 1), "disagg_cluster": disagg_cluster, "context_servers": { @@ -753,7 +767,7 @@ def setup_disagg_cluster( work_dir, server_port, save_log=save_log, - env=env, + env=server_env, cwd=cwd) all_workers = ctx_workers + gen_workers @@ -975,6 +989,24 @@ 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") + + env = llm_venv._new_env.copy() + env["CUDA_VISIBLE_DEVICES"] = "0" + run_disaggregated_test(disaggregated_example_root, + "multi_orchestrator", + num_iters=1, + env=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/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 ) 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/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 67fb75c29aaf..5602f66498ea 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -176,6 +176,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] diff --git a/tests/unittest/disaggregated/test_coordinator_e2e.py b/tests/unittest/disaggregated/test_coordinator_e2e.py new file mode 100644 index 000000000000..a90a8b2cd86c --- /dev/null +++ b/tests/unittest/disaggregated/test_coordinator_e2e.py @@ -0,0 +1,434 @@ +# 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(): + """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 + + 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): + 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( + 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, 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], + "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(): + """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 + 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, WORKERS) + + env = dict(os.environ) + # 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` num_workers={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..39ae11bc33a4 --- /dev/null +++ b/tests/unittest/disaggregated/test_coordinator_worker.py @@ -0,0 +1,604 @@ +# 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 json +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from unittest.mock import AsyncMock + +import aiohttp +import pytest +import uvicorn + +from tensorrt_llm.llmapi.disagg_utils import ( + CtxGenServerConfig, + DisaggClusterConfig, + DisaggServerConfig, + RouterConfig, + ServerRole, +) +from tensorrt_llm.serve.coordinator_server import CoordinatorServer +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, + 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 + + +@pytest.fixture(autouse=True) +def _reset_prometheus_registry(): + """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 + + 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 exposing health and routing metadata.""" + + 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): + pass + + 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) + + @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_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_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, + } + + +@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_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 + + 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) + 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} + assert routing_key["num_tokens"] == 3 + 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() + client.ctx_router._fetch_server_info = AsyncMock(return_value={}) + + 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"}]) + + assert KvCacheAwareRouter._content_affinity_key(request) == 7306401829117098140 + + +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(): + """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 + + 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_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() + + +@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. + + The stateful generation router delegates placement to coordinator + ``/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, request_id): + return CompletionRequest( + model="m", + prompt="hi", + disaggregated_params=DisaggregatedParams( + request_type="generation_only", + ctx_request_id=request_id, + conversation_id=conv_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 + 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 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 + + 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_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: + 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) + assigned_id = await remote.get_disagg_request_id() + request = CompletionRequest( + model="m", + prompt="hello", + disaggregated_params=DisaggregatedParams( + request_type="generation_only", + ctx_request_id=assigned_id, + disagg_request_id=None, + conversation_id="conv-A", + ), + ) + await remote.gen_router.get_next_server(request) + 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 > 0 + + +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..f3556bae1fe5 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,26 @@ 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" + + +@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(): @@ -200,31 +218,59 @@ 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_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() @@ -236,5 +282,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_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 7cc7fe4a649e..04defdd1cd37 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -23,12 +23,14 @@ 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, 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, @@ -54,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): @@ -63,11 +65,41 @@ 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 + + +@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( @@ -213,16 +245,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._coordinator + 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( @@ -398,6 +432,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 b8d5980e087b..ddaf12a0c8e2 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,57 @@ def _make_mock_aiohttp_session(return_value=None): return mock_session +@pytest.mark.asyncio +async def test_coordinator_finish_retry_is_bounded(): + 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(503, {"error": "temporarily unavailable"}), + _response(503, {"error": "temporarily unavailable"}), + ]) + 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 == 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) def mock_aiohttp_session(request): """Auto-mock aiohttp.ClientSession so poll_events doesn't make real HTTP calls.""" @@ -517,7 +570,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] = { @@ -786,33 +839,8 @@ 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): - 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) - - 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"] - - await router.finish_request(request) - assert id(request) not in router._pending_routed_blocks - - -@pytest.mark.asyncio -async def test_kv_cache_aware_router_inserts_routed_blocks_on_finish(servers): +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]] router = KvCacheAwareRouter(server_role=None, @@ -855,8 +883,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 +890,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_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, @@ -877,11 +903,11 @@ 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"]) == 0 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 @@ -2051,7 +2077,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 +2091,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 +2100,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): 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, )