diff --git a/docs/mooncake-kv-store.md b/docs/mooncake-kv-store.md index 8f5bbf29e..fd185e62d 100644 --- a/docs/mooncake-kv-store.md +++ b/docs/mooncake-kv-store.md @@ -23,13 +23,14 @@ This mapping is indexed by **physical GPU ID on each node**, not by the CUDA ord This is **process-local**, not nested-rank-local: all TP ranks spawned inside one process still inherit its subset. A one-HCA-per-nested-rank policy requires separate connector support. This option does not patch vLLM, set private vLLM environment variables, or change store capacity. Use homogeneous mappings across nodes, and verify device locality on the target cluster. Restricting a process to its assigned HCAs can avoid unnecessary RDMA registration fanout; performance equivalence to other HCA policies must be measured, not assumed. -First-class support for [Mooncake](https://github.com/kvcache-ai/Mooncake) as the KV transfer backend for prefill-decode disaggregation. A `mooncake-master` entry under `services:` in an SGLang or vLLM recipe makes srtslurm launch and configure the mooncake master automatically and wire up worker env vars so peer-to-peer transfers work across multiple nodes. +First-class support for [Mooncake](https://github.com/kvcache-ai/Mooncake) as the KV transfer backend for prefill-decode disaggregation. A `mooncake-master` entry under `services:` in an SGLang, vLLM, or TensorRT-LLM recipe makes srtslurm launch and configure the mooncake master automatically and wire up worker env vars so peer-to-peer transfers work across multiple nodes. ## Table of Contents - [Overview](#overview) - [Quick Start (SGLang)](#quick-start-sglang) - [Quick Start (vLLM)](#quick-start-vllm) +- [Quick Start (TensorRT-LLM)](#quick-start-tensorrt-llm) - [What srtslurm Owns vs What You Set](#what-srtslurm-owns-vs-what-you-set) - [Configuration Reference](#configuration-reference) - [Standalone Store Services](#standalone-store-services) @@ -82,7 +83,7 @@ services: type: mooncake-master container: nvcr.io/nvidia/mooncake:latest # optional; default: job container args: [] # appended to mooncake_master - # options.store_config: ... # vLLM only; rendered into MOONCAKE_CONFIG_PATH + # options.store_config: ... # vLLM, TensorRT-LLM; rendered into MOONCAKE_CONFIG_PATH # placement.node: dedicated # share the reserved infra node with etcd/nats # external: 10.0.0.5:8700 # a master that already runs; nothing launches ``` @@ -168,6 +169,59 @@ srtslurm stamps `MOONCAKE_MASTER`, `MOONCAKE_TE_META_DATA_SERVER`, `MOONCAKE_LOC The role `env:` maps are injected on the vLLM workers (not on the standalone `mooncake_master` daemon; the master gets the service entry's own `env`). Use them for in-process Mooncake C++ knobs like `MC_ENABLE_DEST_DEVICE_AFFINITY`, `MC_STORE_CLIENT_METRIC`, `MC_TE_METRIC`. +## Quick Start (TensorRT-LLM) + +TensorRT-LLM's `mooncake-store` KV connector shares one Mooncake Store pool between context and generation servers. It needs a TensorRT-LLM build with the refined connector API ([NVIDIA/TensorRT-LLM#19706](https://github.com/NVIDIA/TensorRT-LLM/pull/19706)). Every rank reads its Mooncake client config from the JSON file `MOONCAKE_CONFIG_PATH` names before it looks at `kv_connector_config.mooncake_store` ([`MooncakeStoreConnectorConfig.resolve`](https://github.com/NVIDIA/TensorRT-LLM/blob/1ab6bbc990a30de41bd0e30f1af38761ec342b8c/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py#L282-L302)). So srtslurm runs the same `mooncake_master` it runs for SGLang and vLLM, and gives each role its own client config: + +```yaml +engine: trtllm +roles: + prefill: + args: + kv_cache_config: + use_kv_cache_manager_v2: true # the connector needs KV cache manager v2; "auto" may not pick it + cache_transceiver_config: # v2 in disaggregated serving needs the Python NIXL transceiver + backend: NIXL + transceiver_runtime: PYTHON + kv_connector_config: + connector: mooncake-store # no mooncake_store block + decode: + args: + kv_cache_config: + use_kv_cache_manager_v2: true + cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON + kv_connector_config: + connector: mooncake-store + mooncake_store_config: # this role's client config keys, over store_config + global_segment_size: 32GiB +services: + - name: mooncake-master + type: mooncake-master + options: + store_config: # TensorRT-LLM Mooncake client config keys, every role + protocol: rdma + device_name: "mlx5_0,mlx5_1" # optional; empty lets Mooncake discover RDMA devices + global_segment_size: 16GiB # host memory each rank lends; byte counts or binary units + stage_through_host: true # copy through host memory; no GPUDirect RDMA needed +``` + +For every role whose `args` select the connector, srtslurm writes `/logs/mooncake_store_config_.json` before workers start and sets `MOONCAKE_CONFIG_PATH` to it on that role's srun task, so every rank `trtllm-llmapi-launch` starts inherits it. The file layers, later wins: a per-role `role` default and `model_key` (the served model name), the service's `store_config`, then `roles..mooncake_store_config`; `master_server_address` is always `:8700`: + +| Role | Default `role` | What the server does with the pool | +| --- | --- | --- | +| `prefill`, `agg` | `both` | Looks up prefixes before computing them and stores what it computed. | +| `decode` | `capacity` | Lends its ranks' segments and never transfers: prompt KV reaches it over the cache transceiver. | + +`model_key` names the checkpoint in the pool's keys. TensorRT-LLM gives it no default, so that two checkpoints sharing a pool cannot read each other's KV; an srtslurm pool belongs to one job, so the served model name (`engine.served_model_name`, else the model directory name) is safe, and a `store_config.model_key` overrides it. Keys every server of the pool must agree on (`model_key`, `namespace`, `metadata_server`, `protocol`, `master_server_address`) cannot be set per role. The rest of the keys pass through to TensorRT-LLM's reader, which supplies its own defaults (`metadata_server: P2PHANDSHAKE`, `protocol: rdma`, `namespace: trtllm`, ...). Notes: + +- The same recipe works with `frontend.type: trtllm_serve` and `frontend.type: dynamo`. trtllm-serve skips its own pool provisioning when `MOONCAKE_CONFIG_PATH` is already set ([`provision_pool`](https://github.com/NVIDIA/TensorRT-LLM/blob/1ab6bbc990a30de41bd0e30f1af38761ec342b8c/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py#L810-L833)), and `dynamo.trtllm` provisions nothing, so no wrapper is needed around either worker. +- No `trtllm-serve mooncake_master` runs: it starts this same `mooncake_master` binary ([`_launch_master`](https://github.com/NVIDIA/TensorRT-LLM/blob/1ab6bbc990a30de41bd0e30f1af38761ec342b8c/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py#L469-L503)). There are no donor processes either; a `capacity` role's ranks lend the memory. +- `global_segment_size` is per rank: a node needs its ranks times the segment size of host memory. TensorRT-LLM refuses `GB`/`MB`, because vLLM's reader of the same file scales them by 1024 and TensorRT-LLM's by 1000; write `GiB` or a byte count. +- Leave `local_hostname` unset so each rank uses its own address; a value in `store_config` would apply to every rank of the role. TensorRT-LLM does not read `MOONCAKE_LOCAL_HOSTNAME`, `MOONCAKE_MASTER`, or `MOONCAKE_TE_META_DATA_SERVER`; srtslurm still stamps them, and they are inert for these workers. +- Engine-side limits are TensorRT-LLM's and depend on the build, so srtslurm does not repeat them. On the `feat/m3_with_msa` branch, a disaggregated server on KV cache manager v2 needs `cache_transceiver_config: {backend: NIXL, transceiver_runtime: PYTHON}`; the C++ transceiver is refused at startup ([`kv_cache_transceiver.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/1ab6bbc990a30de41bd0e30f1af38761ec342b8c/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py#L188-L196)); `main` requires it only for hybrid Mamba models. On a cluster without GPUDirect RDMA, set `stage_through_host: true`. On `main` ([NVIDIA/TensorRT-LLM#19171](https://github.com/NVIDIA/TensorRT-LLM/pull/19171)) a server with any KV connector also rejects attention DP, and speculative decoding on KV cache manager v2. + ## What srtslurm Owns vs What You Set | Concern | Owner | Notes | @@ -176,12 +230,15 @@ The role `env:` maps are injected on the vLLM workers (not on the standalone `mo | `MOONCAKE_MASTER` env var on workers | srtslurm | Always computed as `:8700`. User values in role `env` are overridden. | | `MOONCAKE_TE_META_DATA_SERVER` env var | srtslurm | Always computed as `http://:8701/metadata`. | | `MOONCAKE_LOCAL_HOSTNAME` env var | srtslurm | Auto-resolved per-worker via `runtime.network_interface`. User can override in role `env` for custom NICs. | -| `MOONCAKE_CONFIG_PATH` (vLLM only) | srtslurm | Always points to the JSON file srtslurm renders from `options.store_config:`. Mounted under `/logs` in every worker. | -| `master_server_address` in `store_config` (vLLM)| srtslurm | Always overridden with `:8700`. User values are ignored. | +| `MOONCAKE_CONFIG_PATH` (vLLM, TensorRT-LLM) | srtslurm | Always points to the JSON file srtslurm renders from `options.store_config:` (TensorRT-LLM: one file per role, `/logs/mooncake_store_config_.json`). Mounted under `/logs` in every worker. | +| `master_server_address` in `store_config` | srtslurm | (vLLM, TensorRT-LLM) Always overridden with `:8700`. User values are ignored. | +| `role` in the TensorRT-LLM client config | srtslurm | Defaults to `both` for prefill and agg and `capacity` for decode; set it in `roles..mooncake_store_config` to override. | +| `model_key` in the TensorRT-LLM client config | srtslurm | Defaults to the served model name; set it in `options.store_config` to override. | | `MOONCAKE_PROTOCOL`, `MOONCAKE_DEVICE`, etc. | User | Set in `roles.prefill.env` and `roles.decode.env`. | | `disaggregation-transfer-backend: mooncake` | User | (SGLang only) Set in `roles.prefill.args` and `roles.decode.args`. srtslurm validates this is present. | | `disaggregation-ib-device` | User | (SGLang only) Set in `roles.prefill.args` and `roles.decode.args`. Format: `"mlx5_0,mlx5_1"` or JSON map. | | `kv-transfer-config` | User | (vLLM only) Set in `roles.prefill.args` and `roles.decode.args` to wire vLLM's `MooncakeStoreConnector`. | +| `kv_connector_config: {connector: mooncake-store}` | User | (TensorRT-LLM only) Set in the `args` of each role that joins the pool, without a `mooncake_store` block. | ## Configuration Reference @@ -253,10 +310,12 @@ On the `mooncake-master` service entry (the generic service fields such as `plac - **`container`** (`str`, optional): Container image used for the `mooncake_master` srun. Defaults to the job container if unset. Useful when mooncake needs a different runtime than your worker container. - **`args`** (`list[str]`, optional): Extra arguments appended to the standalone `mooncake_master` command. Use this for flags supported only by the Mooncake version in the selected container. Do not use it to override the RPC, HTTP metadata, or metrics ports because srtslurm configures worker endpoints and readiness checks from its own port values. -- **`options.store_config`** (vLLM only, `dict[str, Any]`): Pass-through dict rendered as JSON into the file pointed to by `MOONCAKE_CONFIG_PATH`. Keys map 1:1 to vLLM's `MooncakeStoreConfig` dataclass, a mix of `str` (e.g. `protocol`), `int` (e.g. `port`), and human-readable size strings (e.g. `"4GB"`). srtslurm does not default these fields; values like `global_segment_size`, `protocol`, and `device_name` are hardware-specific and silently using a srtslurm-picked default is worse than failing loudly, so set them explicitly. `master_server_address` is auto-filled and any user value is ignored. +- **`options.store_config`** (vLLM and TensorRT-LLM, `dict[str, Any]`): Pass-through dict rendered as JSON into the file pointed to by `MOONCAKE_CONFIG_PATH`. For TensorRT-LLM the keys are those of its Mooncake client config (`protocol`, `device_name`, `global_segment_size`, `local_buffer_size`, `namespace`, `metadata_server`, `transfer_batch_size`, `stage_through_host`, and `model_key`, which defaults to the served model name); sizes are byte counts or binary units, `role` defaults per role, and `mooncake_store`-block names (`segment_size`, `pool`, `run_dir`, `master_timeout`) are refused because the client config would ignore them. +- **`options.store_config`, vLLM**: Keys map 1:1 to vLLM's `MooncakeStoreConfig` dataclass, a mix of `str` (e.g. `protocol`), `int` (e.g. `port`), and human-readable size strings (e.g. `"4GB"`). srtslurm does not default these fields; values like `global_segment_size`, `protocol`, and `device_name` are hardware-specific and silently using a srtslurm-picked default is worse than failing loudly, so set them explicitly. `master_server_address` is auto-filled and any user value is ignored. On the roles: +- **`roles..mooncake_store_config`** (TensorRT-LLM only, `dict[str, Any]`, optional): Client config keys for that role's workers, layered over `options.store_config`, for example `role` or a larger `global_segment_size`. The role must select the `mooncake-store` connector, and keys every server of the pool must agree on (`model_key`, `namespace`, `metadata_server`, `protocol`, `master_server_address`) stay in `options.store_config`. SGLang takes per-role Mooncake settings in `roles..env` instead. - **`roles..env`** (`dict[str, str]`, optional): Env vars injected on that role's workers. - For **SGLang**, keys map directly to mooncake's environment variable names. See the [SGLang server_args.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/environ.py) and [mooncake_store.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py) for the full list. - For **vLLM**, this is for in-process Mooncake C++ knobs (`MC_*`) read by the transfer engine / store client. vLLM's connector itself reads configuration from `MOONCAKE_CONFIG_PATH` (the JSON rendered from `options.store_config:`), not from these env vars. @@ -323,6 +382,15 @@ use the mooncake master srtslurm launches for you. Both dash and underscore forms (`disaggregation-transfer-backend`, `disaggregation_transfer_backend`) are accepted. The vLLM check looks for a Mooncake connector in each role's `kv-transfer-config`. +For TensorRT-LLM, srtslurm checks at load what the workers would otherwise reject in every rank after the model has loaded: + +- A `mooncake-master` service needs at least one role with `kv_connector_config.connector: mooncake-store`, and those roles must not also set a `mooncake_store` block: TensorRT-LLM would read the srtslurm config instead of it. +- Each rendered config needs a non-empty `model_key` and a known `role`; sizes must be byte counts or binary units (`local_buffer_size` at least 1, `global_segment_size` at least 0), `transfer_batch_size` must be positive, and `stage_through_host` must not be a string (TensorRT-LLM reads `"false"` as true). +- `store_config` and `roles..mooncake_store_config` must not carry `mooncake_store`-block keys (`segment_size`, `pool`, `run_dir`, `master_timeout`), which the client config would ignore. +- `roles..mooncake_store_config` needs a `mooncake-master` service and a role that selects the connector, must not set keys the whole pool shares, and is refused on SGLang and vLLM roles. +- A role that selects the connector without a `mooncake-master` service must name its pool itself: `MOONCAKE_CONFIG_PATH` in its `env`, or, when trtllm-serve runs the engine (`frontend.type: trtllm_serve`, or a Dynamo sidecar), a `mooncake_store` block with its own `run_dir` and one worker. Every TensorRT-LLM rank runs under `trtllm-llmapi-launch`, so it reads a config trtllm-serve rendered only back from `run_dir`, which one server owns; `dynamo.trtllm` renders none. +- `external` on the `mooncake-master` service is refused for TensorRT-LLM: the rendered `master_server_address` always names the infra node. + ## Common Configurations ### RDMA / InfiniBand diff --git a/docs/schema-reference.md b/docs/schema-reference.md index 45781d6ce..0c2dddd27 100644 --- a/docs/schema-reference.md +++ b/docs/schema-reference.md @@ -111,6 +111,7 @@ One worker role of the recipe: `roles.prefill`, `roles.decode`, or `roles.agg`. | `engine` | str \| mapping | `None` | Engine type or mapping with engine options. Set on every role when no top-level `engine` is declared; the two forms cannot be mixed, and role engines do not inherit options from each other. | | `container` | str \| None | `None` | Optional role image; accepts cluster container aliases. Defaults to `model.container`. | | `kv_events` | bool \| dict[str, Any] \| None | `None` | `true` for the default ZMQ publisher, or a mapping with `publisher` / `topic`. | +| `mooncake_store_config` | dict[str, Any] | `{}` | TRT-LLM only: Mooncake client config keys for this role's workers (`role`, `global_segment_size`, ...), layered over the `mooncake-master` service's `options.store_config`. Pool-wide keys stay there. | | `sidecar` | bool \| None | `None` | Run the native engine with a Dynamo sidecar (turns on `dynamo.sidecar`); every role must agree. | | `critical` | bool | `True` | A worker of this role exiting fails the run. `false` keeps the run alive for probes that kill workers. | | `restart` | [RestartPolicy](#restartpolicy) | `RestartPolicy()` | Relaunch exited workers in place: `never`, `on-failure`, `always`, or a mapping with `policy`, `max_restarts`, `backoff_seconds`, and `max_backoff_seconds`. | @@ -345,7 +346,7 @@ One entry of the top-level ``services:`` list. | `build_timeout_seconds` | int | `1800` | Kill ``build_command`` after this many seconds. | | `enabled` | bool | `True` | ``false`` drops the service, including an implicit one (``etcd`` / ``nats`` under the Dynamo frontend, the default exporters) declared here by name. | | `external` | str \| None | `None` | For discovery-plane kinds (``etcd``, ``nats``, ``mooncake-master``): use this already-running endpoint and launch nothing; the URL is what the job's processes are pointed at. | -| `options` | dict[str, Any] | `{}` | Kind-specific settings (``nats``: ``max_payload_mb``; ``mooncake-master``: ``store_config`` and ``device_names_by_gpu`` for vLLM). Unknown keys are rejected by the kind. | +| `options` | dict[str, Any] | `{}` | Kind-specific settings (``nats``: ``max_payload_mb``; ``mooncake-master``: ``store_config`` for vLLM and TRT-LLM, ``device_names_by_gpu`` for vLLM). Unknown keys are rejected by the kind. | | `metrics` | list[[ServiceMetricsConfig](#servicemetricsconfig)] | `[]` | Prometheus endpoints this service serves: one mapping or a list of ``{port, path, nodes, name}`` (``path`` defaults to ``/metrics``, ``nodes`` to ``all``). Tachometer scrapes each on every node the service runs on, or on its first node with ``nodes: first``, as endpoint ``_`` where ``name`` defaults to the service name. The exporter kinds declare theirs; write it for a generic service that publishes metrics, or on a ``ray`` service whose head serves a trainer's collector and router. | ### PostEvalConfig diff --git a/docs/schema/recipe.schema.json b/docs/schema/recipe.schema.json index e12ed1756..bf09c7f80 100644 --- a/docs/schema/recipe.schema.json +++ b/docs/schema/recipe.schema.json @@ -2538,6 +2538,11 @@ "description": "`true` for the default ZMQ publisher, or a mapping with `publisher` / `topic`.", "default": null }, + "mooncake_store_config": { + "type": "object", + "description": "TRT-LLM only: Mooncake client config keys for this role's workers (`role`, `global_segment_size`, ...), layered over the `mooncake-master` service's `options.store_config`. Pool-wide keys stay there.", + "default": {} + }, "sidecar": { "anyOf": [ { @@ -2924,7 +2929,7 @@ }, "options": { "type": "object", - "description": "Kind-specific settings (``nats``: ``max_payload_mb``; ``mooncake-master``: ``store_config`` and ``device_names_by_gpu`` for vLLM). Unknown keys are rejected by the kind.", + "description": "Kind-specific settings (``nats``: ``max_payload_mb``; ``mooncake-master``: ``store_config`` for vLLM and TRT-LLM, ``device_names_by_gpu`` for vLLM). Unknown keys are rejected by the kind.", "default": {} }, "metrics": { diff --git a/docs/services.md b/docs/services.md index 0105ba292..c424fd8b0 100644 --- a/docs/services.md +++ b/docs/services.md @@ -95,7 +95,7 @@ services: | `type` | `generic` | Selects a [service type](#service-types) that supplies defaults, and for the typed kinds the command. | | `enabled` | `true` | `false` drops the service. Declaring an implied name with `enabled: false` switches it off. | | `external` | none | `etcd`, `nats`, `mooncake-master` only: an address of an already-running instance. Nothing launches; the address is injected where the job's own would have been. | -| `options` | `{}` | Kind-specific knobs; unknown keys are rejected. `nats`: `max_payload_mb`. `dcgm-exporter`, `node-exporter`: `port`, `collect_interval_ms`. `mooncake-master`: `store_config` (vLLM). | +| `options` | `{}` | Kind-specific knobs; unknown keys are rejected. `nats`: `max_payload_mb`. `dcgm-exporter`, `node-exporter`: `port`, `collect_interval_ms`. `mooncake-master`: `store_config` (vLLM, TRT-LLM), `device_names_by_gpu` (vLLM). | | `command` | type default | Argv passed directly to the process. `generic` has no default, so it is required there. | | `args` | `[]` | Appended to `command`. Handy with typed services that supply the command. | | `container` | type fallback, then job container | Aliases resolve through `srtslurm.yaml` like every other container key. | @@ -299,7 +299,7 @@ environment its process needs; the launch path is shared by every kind. Register | `generic` | none (required) | `after_frontend` | `false` | Launches exactly what you wrote. | | `etcd` | `/configs/etcd` from the job container, advertising the node's IP | `infra` | `true` | Implied by the Dynamo frontend. Placement `head`, `infra`, or `dedicated`; supports `external`. Fresh data dir on node-local `/tmp` each job. | | `nats` | `/configs/nats-server -js` from the job container | `infra` | `true` | Implied by the Dynamo frontend. `options.max_payload_mb` writes a server config. Same placements as etcd; supports `external`. | -| `mooncake-master` | `mooncake_master` with the RPC, HTTP metadata, and metrics ports srtctl owns | `before_workers` | `true` | Declared by name; see [Mooncake KV Store](mooncake-kv-store.md). `args` are appended; `options.store_config` is the vLLM connector JSON. Container falls back to the job container. Supports `dedicated` and `external`. | +| `mooncake-master` | `mooncake_master` with the RPC, HTTP metadata, and metrics ports srtctl owns | `before_workers` | `true` | Declared by name; see [Mooncake KV Store](mooncake-kv-store.md). `args` are appended; `options.store_config` is the vLLM or TRT-LLM connector JSON (a TRT-LLM role layers its own keys in `roles..mooncake_store_config`). Container falls back to the job container. Supports `dedicated` and `external`. | | `dcgm-exporter` | `dcgm-exporter --collect-interval= --address :9401` in `nvcr.io/nvidia/k8s/dcgm-exporter` | `after_frontend` | `false` | Implied on worker nodes while tachometer runs. Shell-less (distroless image). `options`: `port`, `collect_interval_ms`. | | `node-exporter` | `/bin/node_exporter` with the cpu, infiniband, and meminfo collectors on 9101 in `quay.io/prometheus/node-exporter` | `after_frontend` | `false` | Implied on worker nodes while tachometer runs. Shell-less. `options`: `port`. | | `process-exporter` | `configs/process-exporter -config.path /process-exporter.yml -web.listen-address=:9256 -threads=true ...` on the bare node | `after_frontend` | `false` | Implied on every allocated node (`placement.node: all`) while tachometer runs. Host-native from the static binary `make setup` installs; skipped with a warning when it is missing. A declared `container` switches to the image's `/bin/process-exporter` with the group file under `/logs`. `options`: `port`, `binary`. | diff --git a/examples/README.md b/examples/README.md index 0766368cf..608b7d990 100644 --- a/examples/README.md +++ b/examples/README.md @@ -8,7 +8,7 @@ Small starting points, one per frontend and topology. The matrix examples serve | --- | --- | --- | --- | | SGLang | `sglang/dynamo-agg.yaml`, `sglang/dynamo-disagg.yaml` | `sglang/sglang-router-agg.yaml`, `sglang/sglang-router-disagg.yaml`, `sglang/smg-disagg.yaml` (SMG) | `sglang/sglang-direct-agg.yaml` | | vLLM | `vllm/dynamo-agg.yaml`, `vllm/dynamo-disagg.yaml` | `vllm/vllm-router-agg.yaml`, `vllm/vllm-router-disagg.yaml`, `vllm/vllm-router-moriio-disagg.yaml` (ROCm, MoRI-IO discovery), `vllm/smg-agg.yaml`, `vllm/smg-dep16.yaml`, `vllm/smg-disagg-grpc.yaml` (SMG; P/D over gRPC with NIXL), `vllm/llm-d-agg.yaml`, `vllm/llm-d-disagg.yaml` (llm-d EPP behind Envoy; P/D through the llm-d sidecar), `vllm/llm-d-dp-ranks.yaml` (llm-d routing to every external-LB DP rank by its KV-cache events) | `vllm/vllm-direct-agg.yaml` | -| TRT-LLM | `trtllm/dynamo-agg.yaml`, `trtllm/dynamo-disagg.yaml`, `trtllm/dynamo-disagg-multinode.yaml` (TP16 decode across two 8-GPU nodes: one srun, one host per rank under `--distribution=arbitrary`) | `trtllm/trtllm-serve-disagg.yaml`, `trtllm/smg-agg.yaml` (SMG) | `trtllm/trtllm-serve-agg.yaml` | +| TRT-LLM | `trtllm/dynamo-agg.yaml`, `trtllm/dynamo-disagg.yaml`, `trtllm/dynamo-disagg-multinode.yaml` (TP16 decode across two 8-GPU nodes: one srun, one host per rank under `--distribution=arbitrary`), `trtllm/dynamo-disagg-mooncake.yaml` (Mooncake Store pool) | `trtllm/trtllm-serve-disagg.yaml`, `trtllm/trtllm-serve-disagg-mooncake.yaml` (Mooncake Store pool), `trtllm/smg-agg.yaml` (SMG) | `trtllm/trtllm-serve-agg.yaml` | | TokenSpeed | `tokenspeed/dynamo-agg.yaml`, `tokenspeed/dynamo-disagg.yaml` | `tokenspeed/smg-agg.yaml`, `tokenspeed/smg-disagg.yaml` (SMG, gRPC engines) | | | Mocker | `mocker/dynamo-agg.yaml` | | | diff --git a/examples/trtllm/dynamo-disagg-mooncake.yaml b/examples/trtllm/dynamo-disagg-mooncake.yaml new file mode 100644 index 000000000..45d4c63d2 --- /dev/null +++ b/examples/trtllm/dynamo-disagg-mooncake.yaml @@ -0,0 +1,97 @@ +# TRT-LLM disaggregated prefill/decode with a shared Mooncake Store pool, behind +# the Dynamo frontend. +# +# Same pool as trtllm-serve-disagg-mooncake.yaml: the `mooncake-master` service runs +# the upstream `mooncake_master` on the infra node, and srtctl gives each role's +# workers its Mooncake client config through MOONCAKE_CONFIG_PATH. `dynamo.trtllm` +# provisions no pool itself, so that variable is how its ranks find the master; no +# wrapper around the worker is needed. Needs a TRT-LLM build with the refined +# Mooncake connector API (NVIDIA/TensorRT-LLM#19706). See docs/mooncake-kv-store.md. +# +# Aliases (`qwen3-0.6b`, `trtllm`) resolve through srtslurm.yaml; see examples/README.md. + +schema: 2 +name: "qwen3-0.6b-trtllm-dynamo-disagg-mooncake" + +slurm: + time_limit: "00:30:00" + +model: + path: "qwen3-0.6b" + container: "trtllm" + precision: "bf16" + +resources: + gpu_type: "h100" + gpus_per_node: 8 +dynamo: + install: false + +frontend: + type: dynamo + enable_multiple_frontends: false + +engine: + type: trtllm + served_model_name: "Qwen/Qwen3-0.6B" +roles: + prefill: + nodes: 1 + workers: 1 + gpus: 1 + env: + PYTHONUNBUFFERED: "1" + TLLM_LOG_LEVEL: "INFO" + args: + backend: "pytorch" + tensor_parallel_size: 1 + max_batch_size: 8 + max_num_tokens: 8192 + max_seq_len: 4096 + trust_remote_code: true + disable_overlap_scheduler: true + kv_cache_config: + free_gpu_memory_fraction: 0.5 + use_kv_cache_manager_v2: true + cache_transceiver_config: # KV cache manager v2 needs the Python NIXL transceiver + backend: "NIXL" + transceiver_runtime: "PYTHON" + kv_connector_config: + connector: mooncake-store + decode: + nodes: colocate + workers: 1 + gpus: 1 + env: + PYTHONUNBUFFERED: "1" + TLLM_LOG_LEVEL: "INFO" + args: + backend: "pytorch" + tensor_parallel_size: 1 + max_batch_size: 64 + max_num_tokens: 4096 + max_seq_len: 4096 + trust_remote_code: true + kv_cache_config: + free_gpu_memory_fraction: 0.5 + use_kv_cache_manager_v2: true + cache_transceiver_config: # KV cache manager v2 needs the Python NIXL transceiver + backend: "NIXL" + transceiver_runtime: "PYTHON" + kv_connector_config: + connector: mooncake-store + +services: + - name: mooncake-master + type: mooncake-master + options: + store_config: # model_key defaults to the served model name + protocol: rdma + global_segment_size: 8GiB + stage_through_host: true + +benchmark: + type: "sa-bench" + isl: 128 + osl: 128 + concurrencies: "4x8" diff --git a/examples/trtllm/trtllm-serve-disagg-mooncake.yaml b/examples/trtllm/trtllm-serve-disagg-mooncake.yaml new file mode 100644 index 000000000..e2f259ab0 --- /dev/null +++ b/examples/trtllm/trtllm-serve-disagg-mooncake.yaml @@ -0,0 +1,100 @@ +# TRT-LLM disaggregated prefill/decode with a shared Mooncake Store pool, fronted +# by `trtllm-serve disaggregated`. +# +# The `mooncake-master` service runs the upstream `mooncake_master` on the infra +# node, the same master SGLang and vLLM recipes use. Both roles select TRT-LLM's +# `mooncake-store` KV connector with no `mooncake_store` block: srtctl renders each +# role's Mooncake client config into /logs/mooncake_store_config_.json and +# sets MOONCAKE_CONFIG_PATH on that role's workers, which TRT-LLM reads instead. +# The context server reads and writes the pool (role: both); the generation server +# only lends its segment (role: capacity), and `roles.decode.mooncake_store_config` +# gives it a larger one. model_key defaults to the served model name. No +# `trtllm-serve mooncake_master` or donor process runs. Needs a TRT-LLM build with the refined Mooncake connector +# API (NVIDIA/TensorRT-LLM#19706). See docs/mooncake-kv-store.md. +# +# Aliases (`qwen3-0.6b`, `trtllm`) resolve through srtslurm.yaml; see examples/README.md. + +schema: 2 +name: "qwen3-0.6b-trtllm-serve-disagg-mooncake" + +slurm: + time_limit: "00:30:00" + +model: + path: "qwen3-0.6b" + container: "trtllm" + precision: "bf16" + +resources: + gpu_type: "h100" + gpus_per_node: 8 +frontend: + type: trtllm_serve + enable_multiple_frontends: false + +engine: + type: trtllm + served_model_name: "Qwen/Qwen3-0.6B" +roles: + prefill: + nodes: 1 + workers: 1 + gpus: 1 + env: + PYTHONUNBUFFERED: "1" + TLLM_LOG_LEVEL: "INFO" + args: + backend: "pytorch" + tensor_parallel_size: 1 + max_batch_size: 8 + max_num_tokens: 8192 + max_seq_len: 4096 + trust_remote_code: true + disable_overlap_scheduler: true + kv_cache_config: + free_gpu_memory_fraction: 0.5 + use_kv_cache_manager_v2: true + cache_transceiver_config: # KV cache manager v2 needs the Python NIXL transceiver + backend: "NIXL" + transceiver_runtime: "PYTHON" + kv_connector_config: + connector: mooncake-store + decode: + nodes: colocate + workers: 1 + gpus: 1 + env: + PYTHONUNBUFFERED: "1" + TLLM_LOG_LEVEL: "INFO" + args: + backend: "pytorch" + tensor_parallel_size: 1 + max_batch_size: 64 + max_num_tokens: 4096 + max_seq_len: 4096 + trust_remote_code: true + kv_cache_config: + free_gpu_memory_fraction: 0.5 + use_kv_cache_manager_v2: true + cache_transceiver_config: # KV cache manager v2 needs the Python NIXL transceiver + backend: "NIXL" + transceiver_runtime: "PYTHON" + kv_connector_config: + connector: mooncake-store + mooncake_store_config: # this role's client keys, over the master's store_config + global_segment_size: 16GiB # decode ranks lend more host memory + +services: + - name: mooncake-master + type: mooncake-master + options: + store_config: # TRT-LLM Mooncake client config keys, shared by every role + protocol: rdma + global_segment_size: 8GiB # per rank; binary units only ("GB" is refused) + stage_through_host: true # no GPUDirect RDMA needed for the KV pools + +benchmark: + type: "sa-bench" + isl: 128 + osl: 128 + concurrencies: "4x8" diff --git a/src/srtctl/backends/AGENTS.md b/src/srtctl/backends/AGENTS.md index 43039352b..9a9a1d5c3 100644 --- a/src/srtctl/backends/AGENTS.md +++ b/src/srtctl/backends/AGENTS.md @@ -40,6 +40,7 @@ Rules for `src/srtctl/backends/`. Every consumer asks a backend through `Backend - The master is the `mooncake-master` service (`services/`); `services/normalize.py` maps a declared entry onto the internal `backend.mooncake_kv_store` field before schema load, and `engine.mooncake_kv_store` sets that field directly. - srtslurm stamps `MOONCAKE_MASTER`, `MOONCAKE_TE_META_DATA_SERVER`, and `MOONCAKE_LOCAL_HOSTNAME` on every worker; `MOONCAKE_LOCAL_HOSTNAME` is the worker's own IP on `runtime.network_interface`. A value in a role's `env` pins the NIC; `MOONCAKE_MASTER` is never set by hand. - vLLM reads its store config from JSON: `store_config` is rendered into the file `MOONCAKE_CONFIG_PATH` names. +- TRT-LLM's `mooncake-store` connector reads JSON too, one file per role: `TRTLLMBackend.build_mooncake_store_config(mode, ...)` renders `store_config` plus `roles..mooncake_store_config` (read through `role_mooncake_store_config`) over a per-role `role` default and a `model_key` of the served model name, `Backend.mooncake_store_configs` lists the files the orchestrator writes, and `get_mooncake_worker_env(..., mode)` sets that role's `MOONCAKE_CONFIG_PATH` on the srun task so every MPI rank inherits it. Roles select the connector with `kv_connector_config: {connector: mooncake-store}` and no `mooncake_store` block; `trtllm-serve mooncake_master` and donors are not used. - SGLang disaggregated recipes must set `disaggregation-transfer-backend: mooncake` in the prefill and decode `args`; the validator rejects a master without it, because workers would silently fall back to the default transport. - Consumers read `backend.mooncake_kv_store` / `backend.get_mooncake_worker_env(...)`; a backend without Mooncake returns `None` / `{}`. diff --git a/src/srtctl/backends/__init__.py b/src/srtctl/backends/__init__.py index a4544cace..4c364cc11 100644 --- a/src/srtctl/backends/__init__.py +++ b/src/srtctl/backends/__init__.py @@ -15,7 +15,7 @@ from .sglang import MooncakeKVStoreConfig, SGLangBackend from .tilert import TileRTBackend from .tokenspeed import TokenSpeedBackend -from .trtllm import TRTLLMBackend +from .trtllm import TRTLLMBackend, TRTLLMMooncakeKVStoreConfig from .vllm import VLLMBackend, VLLMFailoverConfig, VLLMMooncakeKVStoreConfig # Union type for all backend configs @@ -39,6 +39,7 @@ "SrunConfig", # TRTLLM "TRTLLMBackend", + "TRTLLMMooncakeKVStoreConfig", # TileRT "TileRTBackend", # TokenSpeed diff --git a/src/srtctl/backends/base.py b/src/srtctl/backends/base.py index 2d0c83517..8cf8eb1db 100644 --- a/src/srtctl/backends/base.py +++ b/src/srtctl/backends/base.py @@ -19,6 +19,7 @@ from pathlib import Path from srtctl.backends.sglang import MooncakeKVStoreConfig + from srtctl.backends.trtllm import TRTLLMMooncakeKVStoreConfig from srtctl.backends.vllm import VLLMFailoverConfig, VLLMMooncakeKVStoreConfig from srtctl.core.runtime import RuntimeContext from srtctl.core.schema import ProfilingConfig @@ -99,6 +100,9 @@ def extra_args(self) -> Sequence[str]: ... @property def kv_events(self) -> bool | Mapping[str, Any] | None: ... + @property + def mooncake_store_config(self) -> Mapping[str, Any]: ... + class BoundRolesField(fields.Field): """Marshmallow field for an engine's ``roles``: bound by SrtConfig, never read from a recipe, never dumped.""" @@ -127,6 +131,12 @@ def role_env(roles: Mapping[str, RoleSettings], mode: str) -> dict[str, str]: return dict(role.env) if role is not None else {} +def role_mooncake_store_config(roles: Mapping[str, RoleSettings], mode: str) -> dict[str, Any]: + """``roles..mooncake_store_config`` for a worker mode, as a fresh dict; empty when the role is absent.""" + role = role_for_mode(roles, mode) + return dict(role.mooncake_store_config) if role is not None else {} + + def role_kv_events(roles: Mapping[str, RoleSettings], mode: str, defaults: Mapping[str, Any]) -> dict[str, Any] | None: """``roles..kv_events`` for a worker mode over ``defaults``; None when the role publishes none.""" role = role_for_mode(roles, mode) @@ -166,7 +176,9 @@ def type(self) -> str: raise NotImplementedError @property - def mooncake_kv_store(self) -> "MooncakeKVStoreConfig | VLLMMooncakeKVStoreConfig | None": + def mooncake_kv_store( + self, + ) -> "MooncakeKVStoreConfig | VLLMMooncakeKVStoreConfig | TRTLLMMooncakeKVStoreConfig | None": """The recipe's Mooncake KV store block, or None when the engine has none. Set, it implies the mooncake-master service and the MOONCAKE_* worker @@ -277,7 +289,9 @@ def get_process_environment(self, process: "Process") -> dict[str, str]: """ return {} - def get_mooncake_worker_env(self, infra_node_ip: str, local_hostname: str) -> dict[str, str]: + def get_mooncake_worker_env( + self, infra_node_ip: str, local_hostname: str, mode: WorkerMode | None = None + ) -> dict[str, str]: """MOONCAKE_* environment for a worker; empty when mooncake_kv_store is None. - MOONCAKE_LOCAL_HOSTNAME defaults to the worker's resolved IP, but the @@ -291,6 +305,7 @@ def get_mooncake_worker_env(self, infra_node_ip: str, local_hostname: str) -> di infra_node_ip: Resolved IP of the infra node where mooncake_master runs. local_hostname: Resolved IP of the worker's own node, for peer-to-peer transfers. Defaults to the worker's primary network interface IP. + mode: The worker's mode, for engines whose client config differs by role. """ store = self.mooncake_kv_store if store is None: @@ -302,6 +317,17 @@ def get_mooncake_worker_env(self, infra_node_ip: str, local_hostname: str) -> di "MOONCAKE_TE_META_DATA_SERVER": f"http://{infra_node_ip}:{MOONCAKE_HTTP_METADATA_PORT}/metadata", } + def mooncake_store_configs(self, infra_node_ip: str, served_model_name: str) -> dict[str, dict[str, Any]]: + """Mooncake client configs to write into log_dir before workers start, by file name. + + log_dir is mounted at /logs in every worker, where get_mooncake_worker_env + points MOONCAKE_CONFIG_PATH. Empty for an engine that reads its Mooncake + settings from the MOONCAKE_* environment instead (SGLang). + ``served_model_name`` is the name the workers serve, for engines that key the + pool by the model. + """ + return {} + def get_failover_environment(self, process: "Process", job_id: str) -> dict[str, str]: """Shadow engine recovery environment for a worker; empty when failover is None.""" return {} diff --git a/src/srtctl/backends/trtllm.py b/src/srtctl/backends/trtllm.py index e12f0e9b0..8d8a687f6 100644 --- a/src/srtctl/backends/trtllm.py +++ b/src/srtctl/backends/trtllm.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import builtins +import math +import re import uuid from collections.abc import Mapping, Sequence from dataclasses import field, replace @@ -19,9 +21,10 @@ WorkerMode, role_env, role_for_mode, + role_mooncake_store_config, ) from srtctl.backends.sidecar import build_sidecar_launch_command, get_dynamo_sidecar_config, sidecar_grpc_port -from srtctl.ports import DYN_SYSTEM_PORT_BASE, TRTLLM_DIST_INIT_PORTS +from srtctl.ports import DYN_SYSTEM_PORT_BASE, MOONCAKE_MASTER_PORT, TRTLLM_DIST_INIT_PORTS if TYPE_CHECKING: from srtctl.backends.base import SrunConfig @@ -30,6 +33,163 @@ from srtctl.core.topology import Endpoint, NodePortAllocator, Process +# TRT-LLM's Mooncake Store KV connector: the preset name and the module it resolves to. +# TRT-LLM recognizes the connector by the resolved module (``registry.uses_connector``); +# the preset fills ``connector_module`` only when the config leaves it unset. +TRTLLM_MOONCAKE_STORE_CONNECTOR = "mooncake-store" +TRTLLM_MOONCAKE_STORE_MODULE = "tensorrt_llm._torch.pyexecutor.connectors.mooncake_store" + +# A role's ``role`` in its Mooncake client config when the recipe does not set one. +# Context servers read and write the pool; generation servers only lend their +# segment, since prompt KV reaches them over the cache transceiver. +TRTLLM_MOONCAKE_DEFAULT_ROLES: dict[str, str] = {"prefill": "both", "agg": "both", "decode": "capacity"} +TRTLLM_MOONCAKE_STORE_ROLES = frozenset({"both", "producer", "consumer", "capacity"}) + +# Client config keys every server of one pool must agree on, so ``roles..mooncake_store_config`` +# cannot set them: they stay in the ``mooncake-master`` service's ``options.store_config``. +TRTLLM_MOONCAKE_POOL_KEYS: tuple[str, ...] = ( + "model_key", + "namespace", + "metadata_server", + "protocol", + "master_server_address", +) + +# Keys of the ``kv_connector_config.mooncake_store`` block that the client config does +# not read; TRT-LLM would drop them silently. +_TRTLLM_MOONCAKE_BLOCK_ONLY_KEYS: dict[str, str] = { + "segment_size": "the client config calls it global_segment_size", + "pool": "srtctl points every role at the mooncake-master service", + "run_dir": "srtctl hands each role its config through MOONCAKE_CONFIG_PATH", + "master_timeout": "srtctl starts the master before the workers", +} + +# TRT-LLM parses sizes in its Mooncake client config with ``parse_size(..., +# strict_units=True)``: byte counts or binary units only. ``GB``/``MB`` are refused +# because vLLM's reader of the same file scales them by 1024 and TRT-LLM's by 1000. +_TRTLLM_MOONCAKE_SIZE_RE = re.compile(r"^\s*([0-9]+(?:\.[0-9]+)?)\s*([a-zA-Z]*)\s*$") +_TRTLLM_MOONCAKE_SIZE_UNITS: dict[str, int] = { + "": 1, + "b": 1, + "kib": 1 << 10, + "mib": 1 << 20, + "gib": 1 << 30, + "tib": 1 << 40, +} + + +def trtllm_mooncake_store_config_filename(mode: WorkerMode) -> str: + """The client config file srtslurm writes into log_dir (mounted at /logs) for one worker mode.""" + return f"mooncake_store_config_{mode}.json" + + +def _trtllm_mooncake_size(key: str, value: Any) -> int | str: + """``value`` in bytes as TRT-LLM's ``parse_size(strict_units=True)`` reads it, or why it cannot.""" + if isinstance(value, bool) or not isinstance(value, int | float | str): + return f"{key} must be a byte count or a size such as '16GiB', got {value!r}" + if not isinstance(value, str): + if not math.isfinite(value): + return f"{key} must be a finite size, got {value!r}" + return int(value) + match = _TRTLLM_MOONCAKE_SIZE_RE.match(value) + scale = _TRTLLM_MOONCAKE_SIZE_UNITS.get(match.group(2).lower()) if match is not None else None + if match is None or scale is None: + return f"{key} {value!r} must be a byte count or use a binary unit (KiB, MiB, GiB, TiB); 'GB' is refused" + return int(float(match.group(1)) * scale) + + +def trtllm_mooncake_block_only_key_problem(keys: Mapping[str, Any]) -> str | None: + """Why a key of ``kv_connector_config.mooncake_store`` cannot go in a client config, or None.""" + for key, why in _TRTLLM_MOONCAKE_BLOCK_ONLY_KEYS.items(): + if key in keys: + hint = " (set global_segment_size)" if key == "segment_size" else "; drop it" + return f"{key} belongs to the mooncake_store block and the client config ignores it: {why}{hint}" + return None + + +def trtllm_mooncake_role_key_problem(keys: Mapping[str, Any]) -> str | None: + """Why a key cannot go in ``roles..mooncake_store_config``, or None.""" + pool_wide = [key for key in TRTLLM_MOONCAKE_POOL_KEYS if key in keys] + if pool_wide: + return ( + f"{', '.join(pool_wide)} must be the same for every server of the pool; " + "set it in the mooncake-master service's options.store_config" + ) + return None + + +def trtllm_mooncake_store_config_problem(payload: Mapping[str, Any]) -> str | None: + """Why TRT-LLM would refuse a rendered Mooncake client config, or None. + + Mirrors ``MooncakeStoreConnectorConfig.from_file`` and its ``__post_init__``, which + reject these in every rank after the model has loaded, so the recipe fails at load. + """ + model_key = payload.get("model_key") + if not isinstance(model_key, str) or not model_key.strip(): + return ( + f"model_key must be a non-empty string, got {model_key!r}: TRT-LLM keys the pool by it. " + "Leave it unset to use the served model name" + ) + role = payload.get("role") + if str(role).strip().lower() not in TRTLLM_MOONCAKE_STORE_ROLES: + return f"role must be one of {sorted(TRTLLM_MOONCAKE_STORE_ROLES)}, got {role!r}" + if isinstance(payload.get("stage_through_host"), str): + # TRT-LLM takes bool() of it, so the string "false" would mean true. + return f"stage_through_host must be true or false, not the string {payload['stage_through_host']!r}" + for key, minimum in (("global_segment_size", 0), ("local_buffer_size", 1)): + if key in payload: + size = _trtllm_mooncake_size(key, payload[key]) + if isinstance(size, str): + return size + if size < minimum: + return f"{key} must be at least {minimum} bytes, got {payload[key]!r}" + if "transfer_batch_size" in payload: + # TRT-LLM reads it with int(), so "64" and 64.0 are both 64. + batch = payload["transfer_batch_size"] + try: + count = int(batch) + except (TypeError, ValueError, OverflowError): + count = 0 + if isinstance(batch, bool) or count <= 0: + return f"transfer_batch_size must be a positive integer, got {batch!r}" + return None + + +@dataclass(frozen=True) +class TRTLLMMooncakeKVStoreConfig: + """Mooncake Store pool for TRT-LLM's ``mooncake-store`` KV connector. + + Filled from a declared ``mooncake-master`` service (``services/normalize.py``). + srtslurm launches the upstream ``mooncake_master`` on the infra node, the same + master SGLang and vLLM use, and renders one Mooncake client config for every + role whose ``kv_connector_config`` resolves to the Mooncake Store connector into + ``/logs/mooncake_store_config_.json``. That role's workers get it as + ``MOONCAKE_CONFIG_PATH``, which TRT-LLM's connector reads before anything + else, so ``trtllm-serve`` and ``dynamo.trtllm`` (which provisions no pool of its + own) join the pool the same way, without ``trtllm-serve mooncake_master``. + + Each config is ``{"role": , "model_key": , + **store_config, **roles..mooncake_store_config, + "master_server_address": ":8700"}``. The role default is ``both`` + for prefill and agg and ``capacity`` for decode. + + Attributes: + container: Container image for the master; defaults to the job container. + env: Environment injected on every worker, e.g. in-process Mooncake ``MC_*`` knobs. + master_extra_args: Extra arguments appended to the ``mooncake_master`` command. + store_config: Keys of TRT-LLM's Mooncake client config shared by every role + (``protocol``, ``device_name``, ``global_segment_size``, ...); a role + overrides the per-server ones in ``roles..mooncake_store_config``. + """ + + container: str | None = None + env: dict[str, str] = field(default_factory=dict) + master_extra_args: list[str] = field(default_factory=list) + store_config: dict[str, Any] | None = None + + Schema: ClassVar[builtins.type[Schema]] = Schema + + # Log lines that mean the engine behind a TRT-LLM worker step is gone while the # step itself may stay up. ``trtllm-llmapi-launch`` runs the engine as a child of # the rank-0 task and prints ``Rank Task exit code: `` when that child @@ -135,6 +295,10 @@ class TRTLLMBackend(Backend): # engine does not recognise. served_model_name: str | None = None + # Mooncake Store pool for the `mooncake-store` KV connector; a declared + # mooncake-master service maps here. Never written on `engine`. + mooncake_kv_store: TRTLLMMooncakeKVStoreConfig | None = None + # Publish TRT-LLM engine metrics without enabling KV-cache events. # Requires a Dynamo build supporting --publish-metrics; set False to omit # the flag for older builds. Native trtllm-serve and sidecars are unaffected. @@ -286,6 +450,68 @@ def get_process_environment(self, process: "Process") -> dict[str, str]: env.pop(key) return env + def uses_mooncake_store(self, mode: WorkerMode) -> bool: + """Whether this mode's engine config resolves to the Mooncake Store KV connector. + + The rule TRT-LLM applies: compare the resolved module, which the ``mooncake-store`` + preset supplies only when ``connector_module`` is unset. + """ + connector = self.get_config_for_mode(mode).get("kv_connector_config") + if not isinstance(connector, Mapping): + return False + module = connector.get("connector_module") + if module is None and connector.get("connector") == TRTLLM_MOONCAKE_STORE_CONNECTOR: + module = TRTLLM_MOONCAKE_STORE_MODULE + return module == TRTLLM_MOONCAKE_STORE_MODULE + + def mooncake_store_modes(self) -> tuple[WorkerMode, ...]: + """The worker modes whose engine config selects the ``mooncake-store`` KV connector.""" + modes: tuple[WorkerMode, ...] = ("prefill", "decode", "agg") + return tuple(mode for mode in modes if self.uses_mooncake_store(mode)) + + def build_mooncake_store_config( + self, mode: WorkerMode, infra_node_ip: str, served_model_name: str + ) -> dict[str, Any]: + """The Mooncake client config a ``mode`` worker reads from ``MOONCAKE_CONFIG_PATH``. + + Layers, later wins: the role default and ``model_key`` (the served model name; + the pool belongs to this job alone), the service's ``options.store_config``, then + ``roles..mooncake_store_config``. Keys pass through to TRT-LLM's reader, + which supplies its own defaults. ``master_server_address`` is always the master + srtslurm launches; a value set by hand is ignored, as for vLLM. + """ + payload: dict[str, Any] = {"role": TRTLLM_MOONCAKE_DEFAULT_ROLES[mode], "model_key": served_model_name} + store = self.mooncake_kv_store + if store is not None: + payload.update(store.store_config or {}) + payload.update(role_mooncake_store_config(self.roles, mode)) + payload["master_server_address"] = f"{infra_node_ip}:{MOONCAKE_MASTER_PORT}" + return payload + + def mooncake_store_configs(self, infra_node_ip: str, served_model_name: str) -> dict[str, dict[str, Any]]: + """One client config per role that uses the ``mooncake-store`` connector.""" + if self.mooncake_kv_store is None: + return {} + return { + trtllm_mooncake_store_config_filename(mode): self.build_mooncake_store_config( + mode, infra_node_ip, served_model_name + ) + for mode in self.mooncake_store_modes() + } + + def get_mooncake_worker_env( + self, infra_node_ip: str, local_hostname: str, mode: WorkerMode | None = None + ) -> dict[str, str]: + """The shared MOONCAKE_* environment plus this role's ``MOONCAKE_CONFIG_PATH``. + + Set on the srun task, so every rank ``trtllm-llmapi-launch`` starts inherits + it; a path exported from inside the engine process would not reach them. + """ + env = super().get_mooncake_worker_env(infra_node_ip, local_hostname, mode) + if env and mode is not None and self.uses_mooncake_store(mode): + env["MOONCAKE_CONFIG_PATH"] = f"/logs/{trtllm_mooncake_store_config_filename(mode)}" + return env + def allocate_endpoints( self, num_prefill: int, diff --git a/src/srtctl/backends/vllm.py b/src/srtctl/backends/vllm.py index 284e3b388..e3982f5a4 100644 --- a/src/srtctl/backends/vllm.py +++ b/src/srtctl/backends/vllm.py @@ -599,19 +599,27 @@ def get_process_environment(self, process: Process) -> dict[str, str]: env["VLLM_PORT"] = str(process.vllm_scan_port) return env - def get_mooncake_worker_env(self, infra_node_ip: str, local_hostname: str) -> dict[str, str]: + def get_mooncake_worker_env( + self, infra_node_ip: str, local_hostname: str, mode: WorkerMode | None = None + ) -> dict[str, str]: """The shared MOONCAKE_* environment plus ``MOONCAKE_CONFIG_PATH``. ``MOONCAKE_CONFIG_PATH`` points to the JSON file srtslurm writes at job start (mounted into the container at ``/logs``). vLLM's ``MooncakeStoreConnector`` requires this — it does not read the - ``MOONCAKE_*`` env vars directly. + ``MOONCAKE_*`` env vars directly. Every mode reads the same file. """ - env = super().get_mooncake_worker_env(infra_node_ip, local_hostname) + env = super().get_mooncake_worker_env(infra_node_ip, local_hostname, mode) if env: env["MOONCAKE_CONFIG_PATH"] = MOONCAKE_STORE_CONFIG_CONTAINER_PATH return env + def mooncake_store_configs(self, infra_node_ip: str, served_model_name: str) -> dict[str, dict[str, Any]]: + """The shared store config; process-local device configs are written per process.""" + if self.mooncake_kv_store is None: + return {} + return {MOONCAKE_STORE_CONFIG_FILENAME: self.build_mooncake_store_config(infra_node_ip)} + def build_mooncake_store_config(self, infra_node_ip: str) -> dict[str, Any]: """Build the JSON payload for vLLM's ``MooncakeStoreConfig.load_from_env()``. diff --git a/src/srtctl/cli/do_sweep.py b/src/srtctl/cli/do_sweep.py index 73ea2a7da..6810f4265 100644 --- a/src/srtctl/cli/do_sweep.py +++ b/src/srtctl/cli/do_sweep.py @@ -24,7 +24,7 @@ from dataclasses import dataclass from pathlib import Path -from srtctl.backends.vllm import MOONCAKE_STORE_CONFIG_FILENAME, VLLMBackend +from srtctl.backends.vllm import VLLMBackend from srtctl.cli.mixins import ( BenchmarkStageMixin, FrontendStageMixin, @@ -130,18 +130,21 @@ def start_head_infrastructure(self, registry: ProcessRegistry) -> None: self.start_services("infra", registry) def _write_mooncake_store_config(self) -> None: - """vLLM's MooncakeStoreConnector reads its config from a JSON file, not env. + """vLLM's and TRT-LLM's Mooncake connectors read their config from a JSON file, not env. Written into log_dir (mounted at /logs in every worker) before workers start, pointing at the Mooncake master on the infra node. """ backend = self.config.backend + served_model_name = backend.get_served_model_name(self.runtime.model_path.name) + for filename, store_cfg in backend.mooncake_store_configs( + self.runtime.infra_node_ip, served_model_name + ).items(): + store_cfg_path = self.runtime.log_dir / filename + store_cfg_path.write_text(json.dumps(store_cfg, indent=2)) + logger.info("Wrote mooncake_store_config to %s: %s", store_cfg_path, store_cfg) if not isinstance(backend, VLLMBackend) or backend.mooncake_kv_store is None: return - store_cfg = backend.build_mooncake_store_config(self.runtime.infra_node_ip) - store_cfg_path = self.runtime.log_dir / MOONCAKE_STORE_CONFIG_FILENAME - store_cfg_path.write_text(json.dumps(store_cfg, indent=2)) - logger.info("Wrote mooncake_store_config to %s: %s", store_cfg_path, store_cfg) if not backend.mooncake_kv_store.device_names_by_gpu: return # Render only GPU subsets actually launched, rather than all 2**N subsets. diff --git a/src/srtctl/cli/mixins/worker_stage.py b/src/srtctl/cli/mixins/worker_stage.py index d3c13be5b..e4903679d 100644 --- a/src/srtctl/cli/mixins/worker_stage.py +++ b/src/srtctl/cli/mixins/worker_stage.py @@ -419,7 +419,7 @@ def __missing__(self, key: str) -> str: local_hostname = env_to_set.get("MOONCAKE_LOCAL_HOSTNAME") or get_hostname_ip( process.node, self.runtime.network_interface ) - env_to_set.update(backend.get_mooncake_worker_env(self.runtime.infra_node_ip, local_hostname)) + env_to_set.update(backend.get_mooncake_worker_env(self.runtime.infra_node_ip, local_hostname, mode)) self._apply_mooncake_process_config(process, env_to_set) @@ -647,7 +647,7 @@ def start_endpoint_worker(self, endpoint_processes: list["Process"], *, attempt: local_hostname = env_to_set.get("MOONCAKE_LOCAL_HOSTNAME") or get_hostname_ip( leader.node, self.runtime.network_interface ) - env_to_set.update(backend.get_mooncake_worker_env(self.runtime.infra_node_ip, local_hostname)) + env_to_set.update(backend.get_mooncake_worker_env(self.runtime.infra_node_ip, local_hostname, mode)) # Add profiling environment variables after the worker environment. if profiling.enabled and profiling_selects_process: diff --git a/src/srtctl/cli/submit.py b/src/srtctl/cli/submit.py index 1d263c1e8..7a02e6c17 100755 --- a/src/srtctl/cli/submit.py +++ b/src/srtctl/cli/submit.py @@ -790,6 +790,21 @@ def show_config_details(config: SrtConfig) -> None: "MOONCAKE_CONFIG_PATH", "/logs/mooncake_store_config.json (auto)", ) + elif isinstance(backend, TRTLLMBackend): + # One client config per role using the mooncake-store connector; the + # role's workers read it through MOONCAKE_CONFIG_PATH. model_key defaults + # to the served model name, which needs the staged model when not configured. + served_model_name = backend.get_served_model_name("") + for mode in backend.mooncake_store_modes(): + path = backend.get_mooncake_worker_env("", "", mode)["MOONCAKE_CONFIG_PATH"] + details.add_row("mooncake", f"MOONCAKE_CONFIG_PATH ({mode})", f"{path} (auto)") + details.add_row( + "mooncake", + f"store_config ({mode})", + json.dumps( + backend.build_mooncake_store_config(mode, "", served_model_name), indent=2 + ), + ) console.print(Panel(details, border_style="blue")) diff --git a/src/srtctl/core/schema.py b/src/srtctl/core/schema.py index 317bc6786..936f85dbf 100755 --- a/src/srtctl/core/schema.py +++ b/src/srtctl/core/schema.py @@ -50,6 +50,11 @@ VLLMMooncakeKVStoreConfig, ) from srtctl.backends.base import RoleSettings, WorkerMode +from srtctl.backends.trtllm import ( + trtllm_mooncake_block_only_key_problem, + trtllm_mooncake_role_key_problem, + trtllm_mooncake_store_config_problem, +) from srtctl.core.formatting import ( FormattablePath, FormattablePathField, @@ -824,7 +829,7 @@ class RoleConfig(RoleSettings): worker, its environment and engine arguments, and optionally its own engine and image. ``SrtConfig.topology`` derives the per-role counts the launch path reads; ``SrtConfig.backend`` binds the roles onto the engine, which reads ``env`` / ``args`` / - ``extra_args`` / ``kv_events`` from them. + ``extra_args`` / ``kv_events`` / ``mooncake_store_config`` from them. """ # Nodes reserved for this role. `colocate` (decode only) reserves none and packs the @@ -851,6 +856,9 @@ class RoleConfig(RoleSettings): container: str | None = None # `true` for the default ZMQ publisher, or a mapping with `publisher` / `topic`. kv_events: bool | dict[str, Any] | None = None + # TRT-LLM only: Mooncake client config keys for this role's workers (`role`, `global_segment_size`, + # ...), layered over the `mooncake-master` service's `options.store_config`. Pool-wide keys stay there. + mooncake_store_config: dict[str, Any] = field(default_factory=dict) # Run the native engine with a Dynamo sidecar (turns on `dynamo.sidecar`); every role must agree. sidecar: bool | None = field( default=None, @@ -3290,6 +3298,10 @@ def _validate_mooncake_kv_store(self): ``MooncakeConnector``), the master we launch is unused and workers fall back to the default transport — almost never what the user intends. """ + self._validate_role_mooncake_store_config() + if isinstance(self.backend, TRTLLMBackend): + self._validate_trtllm_mooncake_store(self.backend) + return mooncake_cfg = self.backend.mooncake_kv_store if mooncake_cfg is None: return @@ -3352,6 +3364,117 @@ def _vllm_has_mooncake(mode_cfg: dict | None) -> bool: "for you." ) + def _validate_role_mooncake_store_config(self) -> None: + """``roles..mooncake_store_config`` is read by a TRT-LLM role's engine only.""" + for role, spec in self.roles.items(): + if spec.mooncake_store_config and not isinstance(self.role_backends.get(role, self.backend), TRTLLMBackend): + raise ValidationError( + f"roles.{role}.mooncake_store_config is read by TRT-LLM only. For SGLang, set per-role Mooncake " + "settings as MOONCAKE_* variables in roles..env; for vLLM, in the mooncake-master " + "service's options.store_config." + ) + + def _validate_trtllm_mooncake_store(self, backend: TRTLLMBackend) -> None: + """TRT-LLM's Mooncake Store connector against the master srtslurm launches. + + With a ``mooncake-master`` service, srtslurm renders each connector role's + client config and points ``MOONCAKE_CONFIG_PATH`` at it, which TRT-LLM reads + before its ``kv_connector_config.mooncake_store`` block. Everything checked + here would otherwise fail in every rank after the model has loaded. + """ + store = backend.mooncake_kv_store + modes = backend.mooncake_store_modes() + role_keys = { + role: dict(spec.mooncake_store_config) for role, spec in self.roles.items() if spec.mooncake_store_config + } + if store is None: + if role_keys: + raise ValidationError( + f"roles.{min(role_keys)}.mooncake_store_config needs a mooncake-master service: srtslurm writes " + "it into the client config it renders for the master's pool" + ) + self._validate_trtllm_own_mooncake_pool(backend, modes) + return + if any(s.enabled and s.type == "mooncake-master" and s.external for s in self.services): + raise ValidationError( + "an external mooncake-master is not supported for TRT-LLM: srtslurm renders each role's " + "master_server_address from the infra node, where nothing would be listening" + ) + if not modes: + raise ValidationError( + "a mooncake-master service is configured but no roles..args.kv_connector_config resolves to " + "the Mooncake Store connector (connector: mooncake-store, with connector_module unset or " + "tensorrt_llm._torch.pyexecutor.connectors.mooncake_store), so no TRT-LLM worker would use the master." + ) + for mode in modes: + if backend.get_config_for_mode(mode)["kv_connector_config"].get("mooncake_store") is not None: + raise ValidationError( + f"roles.{mode}.args.kv_connector_config.mooncake_store is set, but with a mooncake-master " + "service srtslurm gives this role its Mooncake client config through MOONCAKE_CONFIG_PATH, " + "which TRT-LLM reads instead of that block. Remove the block and put its settings in the " + "service's options.store_config or roles..mooncake_store_config (segment_size is " + "global_segment_size there; pool, run_dir, and master_timeout have no equivalent)." + ) + unused = sorted(set(role_keys) - set(modes)) + if unused: + raise ValidationError( + f"roles.{unused[0]}.mooncake_store_config is set, but only {list(modes)} use the Mooncake Store " + "connector" + ) + sources = {"mooncake-master options.store_config": store.store_config or {}} + sources.update({f"roles.{role}.mooncake_store_config": keys for role, keys in role_keys.items()}) + for where, keys in sources.items(): + problem = trtllm_mooncake_block_only_key_problem(keys) + if problem is None and where.startswith("roles."): + problem = trtllm_mooncake_role_key_problem(keys) + if problem is not None: + raise ValidationError(f"{where}: {problem}") + served_model_name = backend.get_served_model_name("") + for mode in modes: + problem = trtllm_mooncake_store_config_problem( + backend.build_mooncake_store_config(mode, "", served_model_name) + ) + if problem is not None: + raise ValidationError(f"Mooncake client config for roles.{mode}: {problem}") + + def _validate_trtllm_own_mooncake_pool(self, backend: TRTLLMBackend, modes: tuple[WorkerMode, ...]) -> None: + """Connector roles without a mooncake-master service must find a client config themselves. + + Every rank runs under trtllm-llmapi-launch, so a MOONCAKE_CONFIG_PATH that trtllm-serve + exports never reaches it: ranks read the config trtllm-serve rendered back from + ``mooncake_store.run_dir``, which one server claims for itself. ``dynamo.trtllm`` + renders none; a Dynamo sidecar engine is ``tensorrt_llm.commands.serve`` and does. + """ + from srtctl.frontends import get_frontend + + workers = { + "prefill": self.topology.num_prefill, + "decode": self.topology.num_decode, + "agg": self.topology.num_agg, + } + run_dirs: dict[str, WorkerMode] = {} + for mode in modes: + if "MOONCAKE_CONFIG_PATH" in {**self.environment, **backend.get_environment_for_mode(mode)}: + continue + block = backend.get_config_for_mode(mode)["kv_connector_config"].get("mooncake_store") + run_dir = block.get("run_dir") if isinstance(block, Mapping) else None + serve_renders = self.dynamo.sidecar or get_frontend(self.frontend.type).worker_launch == "direct" + if not run_dir or not serve_renders: + raise ValidationError( + f"roles.{mode}.args.kv_connector_config selects the mooncake-store connector, but its workers " + "would not find a Mooncake client config. Declare a mooncake-master service: srtslurm then " + "launches the master and gives this role its config through MOONCAKE_CONFIG_PATH. (Without " + f"one, set MOONCAKE_CONFIG_PATH in roles.{mode}.env, or, when trtllm-serve runs the engine, " + "a mooncake_store block with run_dir.)" + ) + if workers[mode] > 1 or run_dir in run_dirs: + raise ValidationError( + f"roles.{mode}.args.kv_connector_config.mooncake_store.run_dir {run_dir!r} would be shared by " + "several servers, and TRT-LLM lets one live server own a run_dir. Give each role its own " + "run_dir and one worker, or declare a mooncake-master service." + ) + run_dirs[run_dir] = mode + def _profiling_worker_ranks(self, mode: WorkerMode) -> set[int]: """Derive selectable physical ranks from the configured worker layout.""" from srtctl.core.topology import Endpoint diff --git a/src/srtctl/core/schema_docs.py b/src/srtctl/core/schema_docs.py index 1b097851d..45c10ebed 100644 --- a/src/srtctl/core/schema_docs.py +++ b/src/srtctl/core/schema_docs.py @@ -53,6 +53,7 @@ VLLMBackend, ) from srtctl.backends.sglang import MooncakeKVStoreConfig +from srtctl.backends.trtllm import TRTLLMMooncakeKVStoreConfig from srtctl.backends.vllm import VLLMMooncakeKVStoreConfig from srtctl.benchmarks import SHARED_BENCHMARK_FIELDS, get_runner_class, list_benchmarks from srtctl.core.config import LEGACY_TOP_LEVEL_KEYS @@ -386,6 +387,7 @@ def _present(cls: type, keys: frozenset[str]) -> frozenset[str]: INTERNAL_CLASSES: frozenset[type] = frozenset( { MooncakeKVStoreConfig, + TRTLLMMooncakeKVStoreConfig, VLLMMooncakeKVStoreConfig, } ) diff --git a/src/srtctl/services/config.py b/src/srtctl/services/config.py index 19ed89a72..8064d2968 100644 --- a/src/srtctl/services/config.py +++ b/src/srtctl/services/config.py @@ -326,8 +326,8 @@ class ServiceConfig: ``mooncake-master``): use this already-running endpoint and launch nothing; the URL is what the job's processes are pointed at. options: Kind-specific settings (``nats``: ``max_payload_mb``; - ``mooncake-master``: ``store_config`` and ``device_names_by_gpu`` - for vLLM). Unknown keys are rejected by the kind. + ``mooncake-master``: ``store_config`` for vLLM and TRT-LLM, + ``device_names_by_gpu`` for vLLM). Unknown keys are rejected by the kind. metrics: Prometheus endpoints this service serves: one mapping or a list of ``{port, path, nodes, name}`` (``path`` defaults to ``/metrics``, ``nodes`` to ``all``). Tachometer scrapes each on every diff --git a/src/srtctl/services/implicit.py b/src/srtctl/services/implicit.py index 8b504f87f..9e68b4b21 100644 --- a/src/srtctl/services/implicit.py +++ b/src/srtctl/services/implicit.py @@ -17,6 +17,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING +from srtctl.backends.trtllm import TRTLLMMooncakeKVStoreConfig from srtctl.backends.vllm import VLLMBackend, VLLMMooncakeKVStoreConfig from srtctl.ports import ETCD_CLIENT_PORT, NATS_PORT from srtctl.services.config import ServiceConfig, ServicePlacementConfig @@ -145,6 +146,9 @@ def implied_services(config: SrtConfig) -> list[EffectiveService]: options["store_config"] = dict(mooncake_cfg.store_config) if mooncake_cfg.device_names_by_gpu: options["device_names_by_gpu"] = list(mooncake_cfg.device_names_by_gpu) + elif isinstance(mooncake_cfg, TRTLLMMooncakeKVStoreConfig): + if mooncake_cfg.store_config: + options["store_config"] = dict(mooncake_cfg.store_config) implied.append( EffectiveService( ServiceConfig( diff --git a/src/srtctl/services/mooncake_master.py b/src/srtctl/services/mooncake_master.py index 693e806de..56b04171f 100644 --- a/src/srtctl/services/mooncake_master.py +++ b/src/srtctl/services/mooncake_master.py @@ -7,9 +7,12 @@ either way srtctl injects ``MOONCAKE_MASTER``, ``MOONCAKE_TE_META_DATA_SERVER``, and ``MOONCAKE_LOCAL_HOSTNAME`` into every worker (``expand_services`` maps a declared entry onto the internal ``engine.mooncake_kv_store`` field so the -engine-side validation and env injection read one field). Runs on the infra -node, before workers, with the -embedded HTTP metadata server and the metrics endpoint on, all three ports gated. +engine-side validation and env injection read one field). vLLM and TRT-LLM +workers also get ``MOONCAKE_CONFIG_PATH``, a client config srtctl renders from +``options.store_config`` (TRT-LLM: one per role, with ``roles..mooncake_store_config``). +Runs on the infra node, before workers, with the embedded HTTP metadata server +and the metrics endpoint on, all three ports gated. TRT-LLM's +``trtllm-serve mooncake_master`` wraps this same binary, so it is not needed. """ from __future__ import annotations @@ -51,7 +54,8 @@ class MooncakeMasterService(ServiceKind): default_readiness_ports = (MOONCAKE_MASTER_PORT, MOONCAKE_HTTP_METADATA_PORT, MOONCAKE_METRICS_PORT) supports_dedicated = True supports_external = True - option_keys = ("store_config", "device_names_by_gpu") # vLLM: worker-side Mooncake JSON configuration + # Worker-side Mooncake client JSON: store_config (vLLM and TRT-LLM), device_names_by_gpu (vLLM). + option_keys = ("store_config", "device_names_by_gpu") def build_command(self, service: ServiceConfig, ctx: ServiceLaunchContext) -> list[str]: if service.command is not None: diff --git a/tests/snapshots/launch/trtllm__dynamo-disagg-mooncake.txt b/tests/snapshots/launch/trtllm__dynamo-disagg-mooncake.txt new file mode 100644 index 000000000..215ad5a12 --- /dev/null +++ b/tests/snapshots/launch/trtllm__dynamo-disagg-mooncake.txt @@ -0,0 +1,249 @@ +# examples/trtllm/dynamo-disagg-mooncake.yaml +# exit_code: 0 +# srun calls: 13 + +## service_etcd +nodelist: node-01 +output: /logs/service_etcd.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +use_bash_wrapper: True +bash_preamble: | + rm -rf /tmp/etcd && mkdir -p /tmp/etcd +command: | + /configs/etcd --data-dir /tmp/etcd --listen-client-urls http://0.0.0.0:2379 --advertise-client-urls http://127.0.0.1:2379 + +## service_mooncake-master +nodelist: node-01 +output: /logs/service_mooncake-master.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: True +command: | + mooncake_master --port=8700 --enable_http_metadata_server=true --http_metadata_server_port=8701 --eviction_high_watermark_ratio=0.9 --default_kv_lease_ttl=10000 --rpc_thread_num=16 --enable_metric_reporting=true --metrics_port=8702 + +## prefill_0_node-01 +nodes: 1 +ntasks: 1 +nodelist: node-01 +output: /logs/node-01_prefill_w0.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 0 + DYN_LOG: info,dynamo_runtime::pipeline::network::ingress::push_handler=warn + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + DYN_SYSTEM_PORT: 7500 + ETCD_ENDPOINTS: http://127.0.0.1:2379 + HEAD_NODE_IP: 127.0.0.1 + MASTER_ADDR: 127.0.0.1 + MASTER_PORT: 29500 + MOONCAKE_CONFIG_PATH: /logs/mooncake_store_config_prefill.json + MOONCAKE_LOCAL_HOSTNAME: 127.0.0.1 + MOONCAKE_MASTER: 127.0.0.1:8700 + MOONCAKE_TE_META_DATA_SERVER: http://127.0.0.1:8701/metadata + PYTHONUNBUFFERED: 1 + TLLM_LOG_LEVEL: INFO + TRTLLM_EPLB_SHM_NAME: +srun_options: + kill-on-bad-exit: 1 + ntasks-per-node: 1 +mpi: pmix +oversubscribe: True +cpu_bind: verbose,none +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + trtllm-llmapi-launch python3 -m dynamo.trtllm --model-path /model --served-model-name Qwen/Qwen3-0.6B --disaggregation-mode prefill --extra-engine-args /logs/trtllm_config_prefill.yaml --request-plane tcp --publish-metrics + +## decode_0_node-02 +nodes: 1 +ntasks: 1 +nodelist: node-02 +output: /logs/node-02_decode_w0.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 0 + DYN_LOG: info,dynamo_runtime::pipeline::network::ingress::push_handler=warn + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + DYN_SYSTEM_PORT: 7501 + ETCD_ENDPOINTS: http://127.0.0.1:2379 + HEAD_NODE_IP: 127.0.0.1 + MASTER_ADDR: 127.0.0.1 + MASTER_PORT: 29501 + MOONCAKE_CONFIG_PATH: /logs/mooncake_store_config_decode.json + MOONCAKE_LOCAL_HOSTNAME: 127.0.0.1 + MOONCAKE_MASTER: 127.0.0.1:8700 + MOONCAKE_TE_META_DATA_SERVER: http://127.0.0.1:8701/metadata + PYTHONUNBUFFERED: 1 + TLLM_LOG_LEVEL: INFO + TRTLLM_EPLB_SHM_NAME: +srun_options: + kill-on-bad-exit: 1 + ntasks-per-node: 1 +mpi: pmix +oversubscribe: True +cpu_bind: verbose,none +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + trtllm-llmapi-launch python3 -m dynamo.trtllm --model-path /model --served-model-name Qwen/Qwen3-0.6B --disaggregation-mode decode --extra-engine-args /logs/trtllm_config_decode.yaml --request-plane tcp --publish-metrics + +## frontend_0 +nodelist: node-01 +output: /logs/node-01_frontend_0.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + DYN_REQUEST_PLANE: tcp + DYN_SKIP_SGLANG_LOG_FORMATTING: 1 + ETCD_ENDPOINTS: http://127.0.0.1:2379 +mpi: pmix +command: | + python3 -m dynamo.frontend --http-port=8000 + +## service_dcgm-exporter_node-01 +nodelist: node-01 +output: /logs/service_dcgm-exporter_node-01.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-02 +nodelist: node-02 +output: /logs/service_dcgm-exporter_node-02.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-03 +nodelist: node-03 +output: /logs/service_dcgm-exporter_node-03.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_node-exporter_node-01 +nodelist: node-01 +output: /logs/service_node-exporter_node-01.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-02 +nodelist: node-02 +output: /logs/service_node-exporter_node-02.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-03 +nodelist: node-03 +output: /logs/service_node-exporter_node-03.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## tachometer +nodelist: node-01 +output: /logs/tachometer.out +srun_export_env: + POLARS_MAX_THREADS: 4 +use_bash_wrapper: False +command: | + tachometer-scraper --config /logs/tachometer_config.toml --local-dir /logs/tachometer/local --sync-interval 120 + +## +nodelist: node-01 +output: /logs/benchmark.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + SRTCTL_FRONTEND_TYPE: dynamo + SRT_FRONTEND_HOST: 127.0.0.1 + SRT_FRONTEND_PORT: 8000 +command: | + bash /srtctl-benchmarks/sa-bench/bench.sh http://localhost:8000 128 128 4x8 inf /model Qwen/Qwen3-0.6B true 2 1 1 0.8 10 2 '' true random '' false diff --git a/tests/snapshots/launch/trtllm__trtllm-serve-disagg-mooncake.txt b/tests/snapshots/launch/trtllm__trtllm-serve-disagg-mooncake.txt new file mode 100644 index 000000000..3f565080a --- /dev/null +++ b/tests/snapshots/launch/trtllm__trtllm-serve-disagg-mooncake.txt @@ -0,0 +1,220 @@ +# examples/trtllm/trtllm-serve-disagg-mooncake.yaml +# exit_code: 0 +# srun calls: 12 + +## service_mooncake-master +nodelist: node-01 +output: /logs/service_mooncake-master.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: True +command: | + mooncake_master --port=8700 --enable_http_metadata_server=true --http_metadata_server_port=8701 --eviction_high_watermark_ratio=0.9 --default_kv_lease_ttl=10000 --rpc_thread_num=16 --enable_metric_reporting=true --metrics_port=8702 + +## prefill_0_node-01 +nodes: 1 +ntasks: 1 +nodelist: node-01 +output: /logs/node-01_prefill_w0.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 0 + HEAD_NODE_IP: 127.0.0.1 + MASTER_ADDR: 127.0.0.1 + MASTER_PORT: 29500 + MOONCAKE_CONFIG_PATH: /logs/mooncake_store_config_prefill.json + MOONCAKE_LOCAL_HOSTNAME: 127.0.0.1 + MOONCAKE_MASTER: 127.0.0.1:8700 + MOONCAKE_TE_META_DATA_SERVER: http://127.0.0.1:8701/metadata + PYTHONUNBUFFERED: 1 + TLLM_LOG_LEVEL: INFO + TRTLLM_EPLB_SHM_NAME: +srun_options: + kill-on-bad-exit: 1 + ntasks-per-node: 1 +mpi: pmix +oversubscribe: True +cpu_bind: verbose,none +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + trtllm-llmapi-launch trtllm-serve /model --host 0.0.0.0 --port 6100 --tensor_parallel_size 1 --config /logs/trtllm_config_prefill.yaml --served_model_name Qwen/Qwen3-0.6B + +## decode_0_node-02 +nodes: 1 +ntasks: 1 +nodelist: node-02 +output: /logs/node-02_decode_w0.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + CUDA_VISIBLE_DEVICES: 0 + HEAD_NODE_IP: 127.0.0.1 + MASTER_ADDR: 127.0.0.1 + MASTER_PORT: 29501 + MOONCAKE_CONFIG_PATH: /logs/mooncake_store_config_decode.json + MOONCAKE_LOCAL_HOSTNAME: 127.0.0.1 + MOONCAKE_MASTER: 127.0.0.1:8700 + MOONCAKE_TE_META_DATA_SERVER: http://127.0.0.1:8701/metadata + PYTHONUNBUFFERED: 1 + TLLM_LOG_LEVEL: INFO + TRTLLM_EPLB_SHM_NAME: +srun_options: + kill-on-bad-exit: 1 + ntasks-per-node: 1 +mpi: pmix +oversubscribe: True +cpu_bind: verbose,none +bash_preamble: | + ( python3 <(cat <<'__FINGERPRINT_EOF__' + + __FINGERPRINT_EOF__ + ) || true ) +command: | + trtllm-llmapi-launch trtllm-serve /model --host 0.0.0.0 --port 6100 --tensor_parallel_size 1 --config /logs/trtllm_config_decode.yaml --served_model_name Qwen/Qwen3-0.6B + +## trtllm_serve_orchestrator +nodelist: node-01 +output: /logs/node-01_frontend_0.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +mpi: pmix +command: | + trtllm-serve disaggregated --config /logs/ser.yaml + +## service_dcgm-exporter_node-01 +nodelist: node-01 +output: /logs/service_dcgm-exporter_node-01.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-02 +nodelist: node-02 +output: /logs/service_dcgm-exporter_node-02.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_dcgm-exporter_node-03 +nodelist: node-03 +output: /logs/service_dcgm-exporter_node-03.out +container_image: nvcr.io#nvidia/k8s/dcgm-exporter:3.3.9-3.6.1-ubuntu22.04 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + dcgm-exporter --collect-interval=1000 --address :9401 + +## service_node-exporter_node-01 +nodelist: node-01 +output: /logs/service_node-exporter_node-01.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-02 +nodelist: node-02 +output: /logs/service_node-exporter_node-02.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## service_node-exporter_node-03 +nodelist: node-03 +output: /logs/service_node-exporter_node-03.out +container_image: quay.io#prometheus/node-exporter:v1.8.2 +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +srun_export_env: + ETCD_ENDPOINTS: http://127.0.0.1:2379 +use_bash_wrapper: False +command: | + /bin/node_exporter --web.listen-address=:9101 --collector.disable-defaults --collector.cpu --collector.infiniband --collector.meminfo --collector.processes --collector.stat --collector.vmstat --collector.pressure --collector.meminfo_numa '--collector.vmstat.fields=^(oom_kill|pgpg|pswp|pgsteal|pg.*fault).*' + +## tachometer +nodelist: node-01 +output: /logs/tachometer.out +srun_export_env: + POLARS_MAX_THREADS: 4 +use_bash_wrapper: False +command: | + tachometer-scraper --config /logs/tachometer_config.toml --local-dir /logs/tachometer/local --sync-interval 120 + +## +nodelist: node-01 +output: /logs/benchmark.out +container_image: trtllm +container_mounts: + /src/srtctl/benchmarks/scripts: /srtctl-benchmarks + /src/srtctl/runtime_scripts: /srtctl-runtime + : /model + /logs: /logs +env_to_set: + SRTCTL_FRONTEND_TYPE: trtllm_serve + SRT_FRONTEND_HOST: 127.0.0.1 + SRT_FRONTEND_PORT: 8000 +command: | + bash /srtctl-benchmarks/sa-bench/bench.sh http://localhost:8000 128 128 4x8 inf /model Qwen/Qwen3-0.6B true 2 1 1 0.8 10 2 '' true random '' false diff --git a/tests/test_dry_run.py b/tests/test_dry_run.py index b7f7df978..9f29656b3 100644 --- a/tests/test_dry_run.py +++ b/tests/test_dry_run.py @@ -762,6 +762,36 @@ def test_vllm_mooncake_store_config_in_dry_run(self, capsys): assert "P2PHANDSHAKE" in output assert "100GB" in output + def test_trtllm_mooncake_role_configs_in_dry_run(self, capsys): + """Each TRT-LLM role using the mooncake-store connector shows its client config and path.""" + connector = {"kv_connector_config": {"connector": "mooncake-store"}} + config = _make_config( + { + "engine": "trtllm", + "roles": { + "prefill": {"args": connector}, + "decode": {"args": connector, "mooncake_store_config": {"global_segment_size": "16GiB"}}, + }, + "services": [ + { + "name": "mooncake-master", + "type": "mooncake-master", + "options": {"store_config": {"global_segment_size": "8GiB"}}, + } + ], + } + ) + show_config_details(config) + output = capsys.readouterr().out + assert "MOONCAKE_CONFIG_PATH" in output + assert "store_config (prefill)" in output + assert "store_config (decode)" in output + assert '"capacity"' in output + assert "16GiB" in output + # model_key defaults to the served model name, unknown before the model is staged + # (the table wraps the placeholder, so match its start). + assert '"model_key": " dict[str, Any]: + """The example recipe for ``frontend``, with the model and container resolvable without srtslurm.yaml.""" + name = "dynamo-disagg-mooncake.yaml" if frontend == "dynamo" else "trtllm-serve-disagg-mooncake.yaml" + recipe = yaml.safe_load((EXAMPLES / name).read_text()) + recipe["model"].update(path="hf:Qwen/Qwen3-0.6B", container="trtllm.sqsh") + return recipe + + +def _master(recipe: dict[str, Any]) -> dict[str, Any]: + return next(service for service in recipe["services"] if service["type"] == "mooncake-master") + + +def _load(recipe: dict[str, Any], tmp_path: Path) -> SrtConfig: + path = tmp_path / "recipe.yaml" + path.write_text(yaml.safe_dump(recipe)) + return SrtConfig.from_yaml(path) + + +def _sweep(recipe: dict[str, Any], tmp_path: Path) -> tuple[list[dict[str, Any]], Path]: + path = tmp_path / "recipe.yaml" + path.write_text(yaml.safe_dump(recipe)) + launches: list[dict[str, Any]] = [] + exit_code = run_mock_sweep( + config_path=path, + output_dir=tmp_path / "outputs" / "4301", + job_id="4301", + options=MockOptions(child_duration_s=0.05, phase_pause_s=0.01, nodelist=("node-01",), on_srun=launches.append), + ) + assert exit_code == 0 + return launches, tmp_path / "outputs" + + +def _written(outputs: Path, filename: str) -> dict[str, Any]: + (path,) = outputs.rglob(filename) + return json.loads(path.read_text()) + + +@pytest.mark.parametrize("frontend", ["trtllm_serve", "dynamo"]) +def test_workers_join_the_native_master_through_their_role_config(tmp_path: Path, frontend: str) -> None: + """The native master runs before workers; each role reads its own rendered client config.""" + launches, outputs = _sweep(_recipe(frontend), tmp_path) + + steps = [launch.get("step_name", "") for launch in launches] + master = steps.index("service_mooncake-master") + workers = { + mode: next(launch for launch in launches if launch.get("step_name", "").startswith(f"{mode}_")) + for mode in ("prefill", "decode") + } + assert master < min(steps.index(worker["step_name"]) for worker in workers.values()) + assert launches[master]["command"][:2] == ["mooncake_master", "--port=8700"] + # Nothing goes through TRT-LLM's own pool provisioning or a donor. + assert not any( + "mooncake_donor" in launch["command"] or "mooncake_master" in launch["command"][1:] for launch in launches + ) + + for mode, role in (("prefill", "both"), ("decode", "capacity")): + worker = workers[mode] + assert worker["env_to_set"]["MOONCAKE_CONFIG_PATH"] == f"/logs/mooncake_store_config_{mode}.json" + launcher = worker["command"].index("trtllm-llmapi-launch") + entry = worker["command"][launcher + 1 : launcher + 4] + assert entry[0] == "trtllm-serve" if frontend == "trtllm_serve" else entry == ["python3", "-m", "dynamo.trtllm"] + written = _written(outputs, f"mooncake_store_config_{mode}.json") + assert written["role"] == role + assert written["model_key"] == "Qwen/Qwen3-0.6B" + assert written["stage_through_host"] is True + assert written["master_server_address"].endswith(":8700") + assert written["master_server_address"] == worker["env_to_set"]["MOONCAKE_MASTER"] + # The engine config selects the connector and carries no mooncake_store block. + engine_yaml = yaml.safe_load(next(outputs.rglob(f"trtllm_config_{mode}.yaml")).read_text()) + assert engine_yaml["kv_connector_config"] == CONNECTOR + + +def test_role_mooncake_store_config_layers_over_the_shared_keys(tmp_path: Path) -> None: + _, outputs = _sweep(_recipe(), tmp_path) + assert _written(outputs, "mooncake_store_config_prefill.json")["global_segment_size"] == "8GiB" + assert _written(outputs, "mooncake_store_config_decode.json")["global_segment_size"] == "16GiB" + + +def test_model_key_defaults_to_the_served_model_name(tmp_path: Path) -> None: + """The pool belongs to this job alone, so the name the workers serve is a safe key; a recipe value wins.""" + recipe = _recipe() + del recipe["engine"]["served_model_name"] + launches, outputs = _sweep(recipe, tmp_path) + prefill = next(launch for launch in launches if launch.get("step_name", "").startswith("prefill_")) + model_arg = prefill["command"][prefill["command"].index("trtllm-serve") + 1] + for mode in ("prefill", "decode"): + assert _written(outputs, f"mooncake_store_config_{mode}.json")["model_key"] == Path(model_arg).name + + _master(recipe)["options"]["store_config"]["model_key"] = "qwen3-pool" + backend = _load(recipe, tmp_path).backend + assert isinstance(backend, TRTLLMBackend) + assert backend.build_mooncake_store_config("decode", "10.1.1.1", "served")["model_key"] == "qwen3-pool" + + +@pytest.mark.parametrize("engine", ["sglang", "vllm"]) +def test_role_mooncake_store_config_is_read_by_trtllm_only(tmp_path: Path, engine: str) -> None: + recipe = yaml.safe_load((EXAMPLES.parent / engine / "dynamo-disagg.yaml").read_text()) + recipe["model"].update(path="hf:Qwen/Qwen3-0.6B", container=f"{engine}.sqsh") + _load(recipe, tmp_path) + recipe["roles"]["decode"]["mooncake_store_config"] = {"global_segment_size": "16GiB"} + with pytest.raises(ValidationError, match="roles.decode.mooncake_store_config is read by TRT-LLM only"): + _load(recipe, tmp_path) + + +def test_role_mooncake_store_config_needs_a_master(tmp_path: Path) -> None: + """Without a mooncake-master service srtslurm renders no client config, so the per-role keys would be dropped.""" + recipe = _recipe() + recipe["services"] = [] + for mode in ("prefill", "decode"): + recipe["roles"][mode]["env"]["MOONCAKE_CONFIG_PATH"] = "/data/mooncake.json" + with pytest.raises(ValidationError, match="roles.decode.mooncake_store_config needs a mooncake-master service"): + _load(recipe, tmp_path) + + +def test_role_without_the_connector_gets_no_client_config(tmp_path: Path) -> None: + recipe = _recipe() + del recipe["roles"]["decode"]["args"]["kv_connector_config"] + del recipe["roles"]["decode"]["mooncake_store_config"] + launches, outputs = _sweep(recipe, tmp_path) + decode = next(launch for launch in launches if launch.get("step_name", "").startswith("decode_")) + assert "MOONCAKE_CONFIG_PATH" not in decode["env_to_set"] + assert not list(outputs.rglob("mooncake_store_config_decode.json")) + assert _written(outputs, "mooncake_store_config_prefill.json")["role"] == "both" + + +def test_render_defaults_role_by_mode_and_owns_the_master_address(tmp_path: Path) -> None: + recipe = _recipe() + _master(recipe)["options"]["store_config"]["master_server_address"] = "10.0.0.9:1234" + recipe["roles"]["decode"]["mooncake_store_config"]["role"] = "both" + backend = _load(recipe, tmp_path).backend + assert isinstance(backend, TRTLLMBackend) + assert backend.mooncake_store_modes() == ("prefill", "decode") + prefill = backend.build_mooncake_store_config("prefill", "10.1.1.1", "served") + assert prefill["role"] == "both" + assert prefill["master_server_address"] == "10.1.1.1:8700" + assert backend.build_mooncake_store_config("decode", "10.1.1.1", "served")["role"] == "both" + assert set(backend.mooncake_store_configs("10.1.1.1", "served")) == { + "mooncake_store_config_prefill.json", + "mooncake_store_config_decode.json", + } + + +def test_engines_without_a_client_config_file_write_none() -> None: + assert SGLangBackend().mooncake_store_configs("10.1.1.1", "served") == {} + assert TRTLLMBackend().mooncake_store_configs("10.1.1.1", "served") == {} + + +@pytest.mark.parametrize( + ("mutate", "message"), + [ + (lambda r: [r["roles"][m]["args"].pop("kv_connector_config") for m in ("prefill", "decode")], "no roles"), + ( + lambda r: r["roles"]["prefill"]["args"]["kv_connector_config"].update( + mooncake_store={"pool": "file:///logs/pool.json", "model_key": "m"} + ), + "mooncake_store is set", + ), + (lambda r: _master(r)["options"]["store_config"].update(model_key=""), "model_key must be a non-empty"), + (lambda r: _master(r)["options"]["store_config"].update(global_segment_size="16GB"), "'GB' is refused"), + (lambda r: _master(r)["options"]["store_config"].update(local_buffer_size=0), "at least 1 bytes"), + (lambda r: _master(r)["options"]["store_config"].update(global_segment_size=-1), "at least 0 bytes"), + (lambda r: _master(r)["options"]["store_config"].update(transfer_batch_size=0), "positive integer"), + (lambda r: _master(r)["options"]["store_config"].update(stage_through_host="false"), "not the string"), + (lambda r: _master(r)["options"]["store_config"].update(segment_size="160GiB"), "set global_segment_size"), + (lambda r: _master(r)["options"]["store_config"].update(pool="file:///logs/pool.json"), "drop it"), + (lambda r: _master(r).update(external="10.0.0.5:8700"), "external mooncake-master"), + (lambda r: r["roles"]["decode"]["mooncake_store_config"].update(role="donor"), "role must be"), + ( + lambda r: r["roles"]["decode"]["args"].pop("kv_connector_config"), + r"roles.decode.mooncake_store_config is set, but only \['prefill'\]", + ), + ( + lambda r: r["roles"]["decode"]["mooncake_store_config"].update(protocol="tcp"), + "roles.decode.mooncake_store_config: protocol must be the same for every server of the pool", + ), + ( + lambda r: r["roles"]["decode"]["mooncake_store_config"].update(model_key="other"), + "model_key must be the same for every server of the pool", + ), + ], +) +def test_invalid_recipes_fail_at_load(tmp_path: Path, mutate, message: str) -> None: + recipe = _recipe() + mutate(recipe) + with pytest.raises(ValidationError, match=message): + _load(recipe, tmp_path) + + +def _own_pool(frontend: str, pool: dict[str, Any]) -> dict[str, Any]: + """The example without the mooncake-master service, each role naming its pool through ``pool``.""" + recipe = _recipe(frontend) + recipe["services"] = [] + for mode in ("prefill", "decode"): + recipe["roles"][mode].pop("mooncake_store_config", None) + for key, value in pool.items(): + recipe["roles"][mode][key] = {**recipe["roles"][mode][key], **copy.deepcopy(value)} + return recipe + + +BLOCK = {"pool": "10.0.0.5:50051", "model_key": "m"} + + +@pytest.mark.parametrize( + ("frontend", "pool"), + [ + ("trtllm_serve", {"env": {"MOONCAKE_CONFIG_PATH": "/data/mooncake.json"}}), + ("dynamo", {"env": {"MOONCAKE_CONFIG_PATH": "/data/mooncake.json"}}), + ( + "trtllm_serve", + {"args": {"kv_connector_config": {**CONNECTOR, "mooncake_store": {**BLOCK, "run_dir": "/d"}}}}, + ), + ], +) +def test_connector_with_a_pool_of_its_own_needs_no_master(tmp_path: Path, frontend: str, pool: dict[str, Any]) -> None: + """A role whose ranks can find a client config of their own is left alone.""" + recipe = _own_pool(frontend, pool) + if "args" in pool: # one server owns a run_dir, so each role needs its own + recipe["roles"]["decode"]["args"]["kv_connector_config"]["mooncake_store"]["run_dir"] = "/d-decode" + assert _load(recipe, tmp_path).backend.mooncake_kv_store is None + + +@pytest.mark.parametrize( + ("frontend", "pool"), + [ + ("trtllm_serve", {}), + # The ranks run under trtllm-llmapi-launch and never see the path trtllm-serve exports. + ("trtllm_serve", {"args": {"kv_connector_config": {**CONNECTOR, "mooncake_store": BLOCK}}}), + # dynamo.trtllm renders no client config from the block at all. + ("dynamo", {"args": {"kv_connector_config": {**CONNECTOR, "mooncake_store": {**BLOCK, "run_dir": "/d"}}}}), + ], +) +def test_connector_without_a_reachable_pool_fails_at_load(tmp_path: Path, frontend: str, pool: dict[str, Any]) -> None: + with pytest.raises(ValidationError, match="would not find a Mooncake client config"): + _load(_own_pool(frontend, pool), tmp_path) + + +def test_connector_named_by_module_is_recognized(tmp_path: Path) -> None: + """TRT-LLM resolves the connector by module, so the explicit module form gets a client config too.""" + recipe = _recipe() + recipe["roles"]["decode"]["args"]["kv_connector_config"] = { + "connector_module": "tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", + "connector_scheduler_class": "MooncakeStoreConnectorScheduler", + "connector_worker_class": "MooncakeStoreConnectorWorker", + } + backend = _load(recipe, tmp_path).backend + assert isinstance(backend, TRTLLMBackend) + assert backend.mooncake_store_modes() == ("prefill", "decode") + # A preset name with another module resolves to that module, as in TRT-LLM. + recipe["roles"]["decode"]["args"]["kv_connector_config"] = {**CONNECTOR, "connector_module": "my.connector"} + del recipe["roles"]["decode"]["mooncake_store_config"] + assert _load(recipe, tmp_path).backend.mooncake_store_modes() == ("prefill",) + + +def test_role_and_flags_are_read_as_trtllm_reads_them(tmp_path: Path) -> None: + recipe = _recipe() + _master(recipe)["options"]["store_config"]["stage_through_host"] = 1 + recipe["roles"]["decode"]["mooncake_store_config"]["role"] = " Capacity " + assert _load(recipe, tmp_path).backend.mooncake_kv_store is not None + + +def test_block_only_keys_are_reported_where_they_were_written(tmp_path: Path) -> None: + for where, mutate in ( + ("options.store_config:", lambda r: _master(r)["options"]["store_config"].update(master_timeout=900)), + ( + "roles.decode.mooncake_store_config:", + lambda r: r["roles"]["decode"]["mooncake_store_config"].update(run_dir="/logs/mc"), + ), + ): + recipe = _recipe() + mutate(recipe) + with pytest.raises(ValidationError, match=where): + _load(recipe, tmp_path) + + +@pytest.mark.parametrize(("value", "ok"), [("64", True), (64.0, True), ("lots", False), (True, False)]) +def test_transfer_batch_size_is_read_with_int(tmp_path: Path, value: Any, ok: bool) -> None: + recipe = _recipe() + _master(recipe)["options"]["store_config"]["transfer_batch_size"] = value + if ok: + _load(recipe, tmp_path) + else: + with pytest.raises(ValidationError, match="positive integer"): + _load(recipe, tmp_path) + + +def test_non_finite_size_names_the_key(tmp_path: Path) -> None: + recipe = _recipe() + _master(recipe)["options"]["store_config"]["global_segment_size"] = float("inf") + with pytest.raises(ValidationError, match="global_segment_size must be a finite size"): + _load(recipe, tmp_path) + + +def test_disabled_external_master_entry_is_ignored(tmp_path: Path) -> None: + recipe = _recipe() + recipe["services"].append( + {"name": "mooncake-master-ext", "type": "mooncake-master", "enabled": False, "external": "10.0.0.5:8700"} + ) + assert _load(recipe, tmp_path).backend.mooncake_kv_store is not None + + +def test_run_dir_shared_by_several_servers_fails_at_load(tmp_path: Path) -> None: + block = {**CONNECTOR, "mooncake_store": {**BLOCK, "run_dir": "/logs/mc"}} + recipe = _own_pool("trtllm_serve", {"args": {"kv_connector_config": block}}) + with pytest.raises(ValidationError, match="would be shared by several servers"): + _load(recipe, tmp_path) + recipe = _own_pool("trtllm_serve", {"args": {"kv_connector_config": block}}) + del recipe["roles"]["decode"]["args"]["kv_connector_config"] + recipe["roles"]["prefill"]["workers"] = 2 + with pytest.raises(ValidationError, match="would be shared by several servers"): + _load(recipe, tmp_path) + + +def test_services_only_trtllm_recipe_still_loads(tmp_path: Path) -> None: + """No connector role, no master, no frontend: nothing to check, and frontend.type none must not be looked up.""" + recipe = yaml.safe_load((EXAMPLES.parent / "features" / "torchrun-pool.yaml").read_text()) + recipe["engine"] = {"type": "trtllm"} + _load(recipe, tmp_path) + + +def test_dynamo_sidecar_engine_provisions_its_own_pool(tmp_path: Path) -> None: + """A Dynamo sidecar runs the engine as tensorrt_llm.commands.serve, which renders the block's config.""" + recipe = _recipe("dynamo") + recipe["services"] = [] + agg = recipe["roles"].pop("prefill") + del recipe["roles"]["decode"] + agg["args"]["kv_connector_config"] = {**CONNECTOR, "mooncake_store": {**BLOCK, "run_dir": "/logs/mc"}} + recipe["roles"] = {"agg": agg} + with pytest.raises(ValidationError, match="would not find a Mooncake client config"): + _load(recipe, tmp_path) + recipe["dynamo"]["sidecar"] = True + assert _load(recipe, tmp_path).backend.mooncake_store_modes() == ("agg",)