diff --git a/docs/mooncake-kv-store.md b/docs/mooncake-kv-store.md index 8f5bbf29e..a6ead07dc 100644 --- a/docs/mooncake-kv-store.md +++ b/docs/mooncake-kv-store.md @@ -5,31 +5,34 @@ For nodes with a known physical-GPU-to-HCA mapping, opt in through the v2 `mooncake-master` service options: ```yaml +roles: + decode: + mooncake_store_config: + metadata_server: P2PHANDSHAKE + global_segment_size: 150GB + local_buffer_size: 4GB + protocol: rdma services: - name: mooncake-master type: mooncake-master options: device_names_by_gpu: [mlx5_0, mlx5_1, mlx5_2, mlx5_3] - store_config: - metadata_server: P2PHANDSHAKE - global_segment_size: 150GB - local_buffer_size: 4GB - protocol: rdma ``` The same fields are supported under `engine.mooncake_kv_store` (or legacy `backend.mooncake_kv_store`), but do not combine that form with a declared `mooncake-master` service. `srtctl migrate` moves the legacy mapping into the service's `options.device_names_by_gpu` without changing its meaning. -This mapping is indexed by **physical GPU ID on each node**, not by the CUDA ordinal inside a container. A process on GPU 2 receives `device_name: mlx5_2`; two TP2 processes on GPUs `[0,1]` and `[2,3]` receive disjoint device subsets. Only the configs needed by the allocated processes are generated in `/logs`, and each worker receives its own `MOONCAKE_CONFIG_PATH`. Other store settings, including segment capacity and the managed master address, are preserved. When the map is nonempty, it overrides `store_config.device_name` in each process-local JSON; the original shared configuration is not mutated. The map must cover every physical GPU on the node; repeated HCA names are allowed for machines where multiple GPUs share an HCA. Without the option, or with an empty list, the existing shared JSON behavior is unchanged. Config loading (including dry-run) rejects wrong-length maps and entries that are empty, contain whitespace, or contain a comma, before allocating nodes. +This mapping is indexed by **physical GPU ID on each node**, not by the CUDA ordinal inside a container. A process on GPU 2 receives `device_name: mlx5_2`; two TP2 processes on GPUs `[0,1]` and `[2,3]` receive disjoint device subsets. Only the configs needed by the allocated processes are generated in `/logs`, and each worker receives its own `MOONCAKE_CONFIG_PATH`. Other store settings, including segment capacity and the managed master address, are preserved. When the map is nonempty, it overrides the role's `device_name` in each process-local JSON (`/logs/mooncake_store_config__gpu.json`); the role's own configuration is not mutated. The map must cover every physical GPU on the node; repeated HCA names are allowed for machines where multiple GPUs share an HCA. Without the option, or with an empty list, each role's JSON is used unchanged. Config loading (including dry-run) rejects wrong-length maps and entries that are empty, contain whitespace, or contain a comma, before allocating nodes. 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 +85,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: ... # deprecated: use roles..mooncake_store_config # placement.node: dedicated # share the reserved infra node with etcd/nats # external: 10.0.0.5:8700 # a master that already runs; nothing launches ``` @@ -133,7 +136,7 @@ services: ## Quick Start (vLLM) -vLLM's `MooncakeStoreConnector` reads its configuration from a JSON file pointed to by `MOONCAKE_CONFIG_PATH` rather than directly from env vars, so the vLLM master entry takes an extra `options.store_config:` section that srtslurm renders into that JSON at job start: +vLLM's `MooncakeStoreConnector` reads its configuration from a JSON file pointed to by `MOONCAKE_CONFIG_PATH` rather than directly from env vars, so each role that uses it takes a `mooncake_store_config:` section that srtslurm renders into that role's JSON at job start: ```yaml engine: vllm @@ -144,30 +147,92 @@ roles: MC_STORE_CLIENT_METRIC: "1" args: kv-transfer-config: '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_connector_extra_config":{"load_async":true}}]}}' + mooncake_store_config: # rendered into this role's MOONCAKE_CONFIG_PATH JSON + metadata_server: "P2PHANDSHAKE" + global_segment_size: "100GB" + local_buffer_size: "4GB" + protocol: "rdma" + device_name: "mlx5_0,mlx5_1" decode: env: MC_ENABLE_DEST_DEVICE_AFFINITY: "1" MC_STORE_CLIENT_METRIC: "1" args: kv-transfer-config: '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false}},{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_connector_extra_config":{"load_async":true}}]}}' + mooncake_store_config: + metadata_server: "P2PHANDSHAKE" + global_segment_size: "100GB" + local_buffer_size: "4GB" + protocol: "rdma" + device_name: "mlx5_0,mlx5_1" services: - name: mooncake-master type: mooncake-master - options: - store_config: # rendered into the MOONCAKE_CONFIG_PATH JSON - metadata_server: "P2PHANDSHAKE" - global_segment_size: "100GB" - local_buffer_size: "4GB" - protocol: "rdma" - device_name: "mlx5_0,mlx5_1" ``` Real-world production form: `MultiConnector` wraps `NixlConnector` (P2P transfer between prefill and decode) **and** `MooncakeStoreConnector` (shared store for cross-instance reuse), both with `kv_role: "kv_both"` so prefill and decode workers run identical connector stacks. srtslurm's validator accepts `MooncakeStoreConnector` standalone or wrapped in `MultiConnector`. -srtslurm stamps `MOONCAKE_MASTER`, `MOONCAKE_TE_META_DATA_SERVER`, `MOONCAKE_LOCAL_HOSTNAME`, and `MOONCAKE_CONFIG_PATH` on every worker; you supply the rest. `master_server_address` in `store_config` is also auto-filled from the infra node IP and ignored if set by hand. +srtslurm stamps `MOONCAKE_MASTER`, `MOONCAKE_TE_META_DATA_SERVER`, and `MOONCAKE_LOCAL_HOSTNAME` on every worker, and `MOONCAKE_CONFIG_PATH=/logs/mooncake_store_config_.json` on the workers of each role whose `kv-transfer-config` uses Mooncake; you supply the rest. `master_server_address` is also auto-filled from the infra node IP and ignored if set by hand. `metadata_server` and `protocol` must have the same value in every role. + +Older recipes set these keys once, in the `mooncake-master` service's `options.store_config`. That still works: the keys apply under every role's own `mooncake_store_config`, and srtslurm logs a deprecation warning. 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 + mooncake_store_config: # this role's Mooncake client config + 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 + 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: + protocol: rdma # keys the whole pool shares match across roles + device_name: "mlx5_0,mlx5_1" + global_segment_size: 32GiB # decode ranks lend more + stage_through_host: true +services: + - name: mooncake-master + type: mooncake-master +``` + +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), then `roles..mooncake_store_config` (over the deprecated service `options.store_config`, if any); `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 `model_key` set in every role overrides it. Keys every server of the pool must agree on (`model_key`, `namespace`, `metadata_server`, `protocol`) must have the same value in every 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 `mooncake_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 +241,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 | Points to the JSON file srtslurm renders from the role's `mooncake_store_config`, one per role: `/logs/mooncake_store_config_.json`. Mounted under `/logs` in every worker. | +| `master_server_address` in the client 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 every role's `mooncake_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 @@ -228,23 +296,28 @@ roles: MC_TE_METRIC: "0" # default 0 (disabled) args: kv-transfer-config: '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}' + mooncake_store_config: # rendered into this role's MOONCAKE_CONFIG_PATH JSON + metadata_server: "P2PHANDSHAKE" + global_segment_size: "100GB" + local_buffer_size: "4GB" + protocol: "rdma" + device_name: "mlx5_0,mlx5_1" decode: env: MC_ENABLE_DEST_DEVICE_AFFINITY: "1" args: kv-transfer-config: '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}' + mooncake_store_config: + metadata_server: "P2PHANDSHAKE" + global_segment_size: "200GB" # decode ranks lend more + local_buffer_size: "4GB" + protocol: "rdma" + device_name: "mlx5_0,mlx5_1" services: - name: mooncake-master type: mooncake-master container: ... # optional, default: job container args: [] # optional, appended to mooncake_master - options: - store_config: # required for vLLM; rendered into MOONCAKE_CONFIG_PATH JSON - metadata_server: "P2PHANDSHAKE" - global_segment_size: "100GB" - local_buffer_size: "4GB" - protocol: "rdma" - device_name: "mlx5_0,mlx5_1" ``` ### Fields @@ -253,13 +326,17 @@ 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]`, deprecated): The client config keys every role used to share. They still apply under each role's `roles..mooncake_store_config`, and srtslurm logs a deprecation warning; move them to the roles. On the roles: +- **`roles..mooncake_store_config`** (vLLM and TensorRT-LLM, `dict[str, Any]`, optional): The Mooncake client config for that role's workers, rendered as JSON into the file `MOONCAKE_CONFIG_PATH` names. The role must use Mooncake (vLLM: its `kv-transfer-config`; TensorRT-LLM: `kv_connector_config.connector: mooncake-store`), and a `mooncake-master` service must exist. Keys every client of the pool must agree on have to match across roles (vLLM: `metadata_server`, `protocol`; TensorRT-LLM: also `model_key` and `namespace`). SGLang takes per-role Mooncake settings in `roles..env` instead. + - For **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. + - 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. + - `master_server_address` is auto-filled and any user value is ignored. - **`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. + - 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 the role's `mooncake_store_config`), not from these env vars. - Setting `MOONCAKE_MASTER`, `MOONCAKE_TE_META_DATA_SERVER`, or `MOONCAKE_CONFIG_PATH` here is a no-op (srtslurm always wins). Mooncake `v0.3.11+` adds NoF SSD-tier eviction flags. Enable them only in recipes using a compatible Mooncake image: @@ -323,6 +400,16 @@ 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 every engine, `roles..mooncake_store_config` needs a role whose workers read a Mooncake client config file (refused on SGLang roles and on vLLM or TensorRT-LLM roles without a Mooncake connector) and a `mooncake-master` service, and the keys the pool shares must match across roles. + +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). +- `roles..mooncake_store_config` and the deprecated `options.store_config` must not carry `mooncake_store`-block keys (`segment_size`, `pool`, `run_dir`, `master_timeout`), which the client config would ignore. +- 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..7587078fd 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] | `{}` | Mooncake client config for this role's workers (vLLM, TRT-LLM): the keys srtslurm writes into the JSON file MOONCAKE_CONFIG_PATH names (`protocol`, `global_segment_size`, ...). Keys the whole pool shares must match across roles. The mooncake-master service's deprecated `options.store_config` still applies underneath. | | `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..9b3b8d036 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": "Mooncake client config for this role's workers (vLLM, TRT-LLM): the keys srtslurm writes into the JSON file MOONCAKE_CONFIG_PATH names (`protocol`, `global_segment_size`, ...). Keys the whole pool shares must match across roles. The mooncake-master service's deprecated `options.store_config` still applies underneath.", + "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..0b208e536 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`: `device_names_by_gpu` (vLLM), and the deprecated `store_config` (each role's `mooncake_store_config` replaces it). | | `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; Each vLLM or TRT-LLM role's connector JSON is its `roles..mooncake_store_config`; the deprecated `options.store_config` still applies underneath. 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..5d308f3d1 --- /dev/null +++ b/examples/trtllm/dynamo-disagg-mooncake.yaml @@ -0,0 +1,100 @@ +# 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 + mooncake_store_config: + protocol: rdma + global_segment_size: 8GiB + stage_through_host: true + 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: + protocol: rdma + global_segment_size: 8GiB + stage_through_host: true + +services: + - name: mooncake-master # model_key defaults to the served model name + type: mooncake-master + +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..4d01c09c9 --- /dev/null +++ b/examples/trtllm/trtllm-serve-disagg-mooncake.yaml @@ -0,0 +1,101 @@ +# 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. +# Each role's `mooncake_store_config` is that role's Mooncake client config. The +# context server reads and writes the pool (role: both); the generation server only +# lends its segment (role: capacity) and lends more. 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 + mooncake_store_config: # this role's Mooncake client config + protocol: rdma + global_segment_size: 8GiB # host memory each prefill rank lends; binary units only + stage_through_host: true # no GPUDirect RDMA needed for the KV pools + 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: + protocol: rdma # keys the whole pool shares must match across roles + global_segment_size: 16GiB # decode ranks lend more host memory + stage_through_host: true + +services: + - name: mooncake-master + type: mooncake-master + +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..61fb4dee9 100644 --- a/src/srtctl/backends/AGENTS.md +++ b/src/srtctl/backends/AGENTS.md @@ -39,9 +39,10 @@ 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. +- vLLM reads its store config from JSON, one file per role whose `kv-transfer-config` uses Mooncake: `VLLMBackend.build_mooncake_store_config(mode, ...)` renders `mooncake_store_config_for_mode(mode)`, and `get_mooncake_worker_env(..., mode)` points that role's `MOONCAKE_CONFIG_PATH` at it. +- TRT-LLM's `mooncake-store` connector reads JSON too, one file per role: `TRTLLMBackend.build_mooncake_store_config(mode, ...)` renders `mooncake_store_config_for_mode(mode)` (`roles..mooncake_store_config` over the deprecated service `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` / `{}`. +- Consumers read `backend.mooncake_kv_store` / `backend.get_mooncake_worker_env(...)`; a backend without Mooncake returns `None` / `{}`. Per-role client config goes through `backend.mooncake_store_modes()` and `backend.mooncake_store_config_for_mode(mode)` (empty for SGLang); `Backend.mooncake_pool_keys` lists the keys the validator checks match across roles. ## Shadow engine recovery (vLLM, `engine.failover`) 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 8548d6b73..ded6b3a09 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 @@ -88,6 +89,7 @@ class RoleSettings(ABC): args: Mapping[str, Any] extra_args: Sequence[str] kv_events: bool | Mapping[str, Any] | None + mooncake_store_config: Mapping[str, Any] class BoundRolesField(fields.Field): @@ -117,6 +119,17 @@ 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 mooncake_store_config_filename(mode: WorkerMode) -> str: + """The Mooncake client config file srtslurm writes into log_dir (mounted at /logs) for one worker mode.""" + return f"mooncake_store_config_{mode}.json" + + 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) @@ -148,6 +161,9 @@ class Backend(ABC): #: Environment that turns on the engine's own NVTX ranges; the automatic nsys #: preset (``observability.nsys``) sets it on every worker it wraps. nvtx_environment: ClassVar[Mapping[str, str]] = {} + #: Mooncake client config keys every server of one pool must agree on. Roles set + #: their client config separately, so the validator checks these match across roles. + mooncake_pool_keys: ClassVar[tuple[str, ...]] = () @property @abstractmethod @@ -156,7 +172,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 @@ -264,7 +282,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 @@ -278,6 +298,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: @@ -289,6 +310,25 @@ 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 mooncake_store_modes(self) -> tuple[WorkerMode, ...]: + """Worker modes whose workers read a Mooncake client config file; empty for an engine that reads none.""" + return () + + def mooncake_store_config_for_mode(self, mode: WorkerMode) -> dict[str, Any]: + """The Mooncake client config keys the recipe sets for ``mode``'s workers; empty for an engine that reads none.""" + 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..184153ffb 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 @@ -17,11 +19,13 @@ ProcessGroup, RoleSettings, WorkerMode, + mooncake_store_config_filename, 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 +34,135 @@ 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"}) + +# 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_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_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": , + **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: Deprecated: the ``mooncake-master`` service's ``options.store_config``. + Its keys still apply under every role's ``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 @@ -96,6 +229,8 @@ class TRTLLMBackend(Backend): # return_perf_metrics is true); its /metrics route is JSON iteration stats. prometheus_metrics_path: ClassVar[str] = "/prometheus/metrics" nvtx_environment: ClassVar[Mapping[str, str]] = {"TLLM_LLMAPI_ENABLE_NVTX": "1", "TLLM_PROFILE_LOG_RANKS": "all"} + # TRT-LLM keys its pool by model_key and namespace; protocol and metadata server are the transport. + mooncake_pool_keys: ClassVar[tuple[str, ...]] = ("model_key", "namespace", "metadata_server", "protocol") # The roles this engine runs (`roles.` of the recipe), bound by SrtConfig and # never written on `engine:`. Per-role env and args (the engine YAML) are read from @@ -135,6 +270,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 +425,71 @@ 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 mooncake_store_config_for_mode(self, mode: WorkerMode) -> dict[str, Any]: + """``roles..mooncake_store_config`` over the deprecated service ``options.store_config``.""" + store = self.mooncake_kv_store + base = dict(store.store_config or {}) if store is not None else {} + return {**base, **role_mooncake_store_config(self.roles, 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), then ``mooncake_store_config_for_mode``. 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} + payload.update(self.mooncake_store_config_for_mode(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 { + 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/{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..c4394d7e1 100644 --- a/src/srtctl/backends/vllm.py +++ b/src/srtctl/backends/vllm.py @@ -33,8 +33,10 @@ ProcessGroup, RoleSettings, WorkerMode, + mooncake_store_config_filename, role_args, role_kv_events, + role_mooncake_store_config, ) from srtctl.backends.sidecar import build_sidecar_launch_command, get_dynamo_sidecar_config, sidecar_grpc_port from srtctl.ports import ( @@ -153,13 +155,6 @@ def find_vllm_orchestration_recipe_flags(backend: VLLMBackend) -> list[tuple[str return findings -# Filename for the mooncake-store JSON config srtslurm writes to log_dir at job -# start. log_dir is mounted into every worker at /logs, so workers read the JSON -# from MOONCAKE_STORE_CONFIG_CONTAINER_PATH. -MOONCAKE_STORE_CONFIG_FILENAME = "mooncake_store_config.json" -MOONCAKE_STORE_CONFIG_CONTAINER_PATH = f"/logs/{MOONCAKE_STORE_CONFIG_FILENAME}" - - @dataclass(frozen=True) class VLLMMooncakeKVStoreConfig: """Mooncake KV store config for the vLLM backend. @@ -171,12 +166,14 @@ class VLLMMooncakeKVStoreConfig: MOONCAKE_MASTER = :8700 MOONCAKE_TE_META_DATA_SERVER = http://:8701/metadata MOONCAKE_LOCAL_HOSTNAME = - MOONCAKE_CONFIG_PATH = /logs/mooncake_store_config.json + MOONCAKE_CONFIG_PATH = /logs/mooncake_store_config_.json The JSON file referenced by ``MOONCAKE_CONFIG_PATH`` is generated by - srtslurm at job start (see ``do_sweep._write_mooncake_store_config``) from - ``store_config`` below. vLLM's ``MooncakeStoreConnector`` reads this - file via ``MooncakeStoreConfig.load_from_env()``. + srtslurm at job start (see ``do_sweep._write_mooncake_store_config``), one + per role whose ``kv-transfer-config`` uses Mooncake, from that role's + ``roles..mooncake_store_config`` over the deprecated ``store_config`` + below. vLLM's ``MooncakeStoreConnector`` reads this file via + ``MooncakeStoreConfig.load_from_env()``. ``env:`` is injected on every vLLM worker (alongside the auto-stamped ``MOONCAKE_*`` vars above), not on the standalone ``mooncake_master`` @@ -208,7 +205,8 @@ class VLLMMooncakeKVStoreConfig: container: str | None = None env: dict[str, str] = field(default_factory=dict) master_extra_args: list[str] = field(default_factory=list) - # ``store_config`` values are JSON-serialized into MOONCAKE_CONFIG_PATH and + # Deprecated: set these keys in roles..mooncake_store_config, where they layer over + # this. Values are JSON-serialized into MOONCAKE_CONFIG_PATH and # parsed by vLLM's ``MooncakeStoreConfig`` dataclass — fields are a mix of # str (e.g. ``protocol``), int (e.g. ``port``), and human-readable sizes # (e.g. ``"4GB"``). Type as ``dict[str, Any]`` to avoid forcing users to @@ -395,6 +393,8 @@ class VLLMBackend(Backend): "VLLM_NVTX_SCOPES_FOR_PROFILING": "1", "VLLM_WORKER_MULTIPROC_METHOD": "spawn", } + # The transport every client of one pool must share. + mooncake_pool_keys: ClassVar[tuple[str, ...]] = ("metadata_server", "protocol") def find_dp_modes(self) -> list[tuple[str, dict[str, Any]]]: """Return the roles whose configured data-parallel size is greater than one, with their args.""" @@ -599,23 +599,55 @@ 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]: - """The shared MOONCAKE_* environment plus ``MOONCAKE_CONFIG_PATH``. + def uses_mooncake_store(self, mode: WorkerMode) -> bool: + """Whether this mode's ``kv-transfer-config`` names a Mooncake connector, at the top or in a MultiConnector.""" + for key, value in self.get_config_for_mode(mode).items(): + if normalize_vllm_config_key(key) != "kv-transfer-config": + continue + text = value if isinstance(value, str) else json.dumps(value) if isinstance(value, Mapping) else "" + if "mooncake" in text.lower(): + return True + return False + + def mooncake_store_modes(self) -> tuple[WorkerMode, ...]: + """The worker modes whose ``kv-transfer-config`` uses Mooncake.""" + modes: tuple[WorkerMode, ...] = ("prefill", "decode", "agg") + return tuple(mode for mode in modes if self.uses_mooncake_store(mode)) + + def mooncake_store_config_for_mode(self, mode: WorkerMode) -> dict[str, Any]: + """``roles..mooncake_store_config`` over the deprecated service ``options.store_config``.""" + store = self.mooncake_kv_store + base = dict(store.store_config or {}) if store is not None else {} + return {**base, **role_mooncake_store_config(self.roles, mode)} + + 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``. ``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. Each role reads its own file. """ - env = super().get_mooncake_worker_env(infra_node_ip, local_hostname) - if env: - env["MOONCAKE_CONFIG_PATH"] = MOONCAKE_STORE_CONFIG_CONTAINER_PATH + 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/{mooncake_store_config_filename(mode)}" return env - def build_mooncake_store_config(self, infra_node_ip: str) -> dict[str, Any]: + def mooncake_store_configs(self, infra_node_ip: str, served_model_name: str) -> dict[str, dict[str, Any]]: + """One store config per role that uses Mooncake; process-local device configs are written per process.""" + if self.mooncake_kv_store is None: + return {} + return { + mooncake_store_config_filename(mode): self.build_mooncake_store_config(mode, infra_node_ip) + for mode in self.mooncake_store_modes() + } + + def build_mooncake_store_config(self, mode: WorkerMode, infra_node_ip: str) -> dict[str, Any]: """Build the JSON payload for vLLM's ``MooncakeStoreConfig.load_from_env()``. - Pass-through of ``mooncake_kv_store.store_config`` with + Pass-through of ``mooncake_store_config_for_mode(mode)`` with ``master_server_address`` force-overridden to the infra node IP (the user can't know that at config time). srtslurm intentionally does not provide defaults for the other fields — values like @@ -624,9 +656,7 @@ def build_mooncake_store_config(self, infra_node_ip: str) -> dict[str, Any]: failing loudly. Set them explicitly in YAML; see ``docs/mooncake-kv-store.md``. """ - result: dict[str, Any] = {} - if self.mooncake_kv_store is not None and self.mooncake_kv_store.store_config: - result.update(self.mooncake_kv_store.store_config) + result = self.mooncake_store_config_for_mode(mode) result["master_server_address"] = f"{infra_node_ip}:{MOONCAKE_MASTER_PORT}" return result @@ -649,9 +679,10 @@ def build_mooncake_process_config( # Sharing an HCA between physical GPUs is legal. Avoid duplicating its # name when a process spans multiple GPUs mapped to the same HCA. process_devices = list(dict.fromkeys(devices[gpu] for gpu in gpu_ids)) - payload = self.build_mooncake_store_config(infra_node_ip) + mode = process.endpoint_mode + payload = self.build_mooncake_store_config(mode, infra_node_ip) payload["device_name"] = ",".join(process_devices) - filename = "mooncake_store_config_gpu" + "-".join(map(str, gpu_ids)) + ".json" + filename = f"mooncake_store_config_{mode}_gpu" + "-".join(map(str, gpu_ids)) + ".json" return filename, payload def is_grpc_mode(self, mode: WorkerMode) -> bool: 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..d61bdc711 100755 --- a/src/srtctl/cli/submit.py +++ b/src/srtctl/cli/submit.py @@ -38,6 +38,7 @@ from rich.text import Text from srtctl.backends import VLLMBackend, VLLMMooncakeKVStoreConfig +from srtctl.backends.base import mooncake_store_config_filename from srtctl.core.config import ( expand_engine_config_defaults, generate_override_configs, @@ -769,26 +770,25 @@ def show_config_details(config: SrtConfig) -> None: details.add_row("mooncake", "container", mooncake_cfg.container or "") if isinstance(mooncake_cfg, VLLMMooncakeKVStoreConfig) and mooncake_cfg.device_names_by_gpu: details.add_row("mooncake", "device_names_by_gpu", str(mooncake_cfg.device_names_by_gpu)) - details.add_row("mooncake", "process config", "/logs/mooncake_store_config_gpu.json") + details.add_row( + "mooncake", "process config", "/logs/mooncake_store_config__gpu.json" + ) details.add_row("mooncake", "master_port", f"{MOONCAKE_MASTER_PORT} (auto)") if mooncake_cfg.master_extra_args: details.add_row("mooncake", "master_extra_args", shlex.join(mooncake_cfg.master_extra_args)) - if isinstance(backend, VLLMBackend): - # vLLM workers need MOONCAKE_CONFIG_PATH pointing at a JSON file - # — srtslurm writes this at job start. Show the resolved JSON - # so operators can sanity-check protocol/device_name/sizes - # before submitting. infra IP is unknown until allocation, so - # use a placeholder for master_server_address. - store_cfg = backend.build_mooncake_store_config("") - details.add_row( - "mooncake", - "store_config", - json.dumps(store_cfg, indent=2), - ) + # One client config per role whose workers read one (vLLM, TRT-LLM); the role's + # workers get it through MOONCAKE_CONFIG_PATH. The infra IP is unknown until + # allocation, and TRT-LLM's model_key default needs the staged model, so both + # show as placeholders. + served_model_name = backend.get_served_model_name("") + store_configs = backend.mooncake_store_configs("", 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", - "MOONCAKE_CONFIG_PATH", - "/logs/mooncake_store_config.json (auto)", + f"store_config ({mode})", + json.dumps(store_configs[mooncake_store_config_filename(mode)], 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..4198d5d4b 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 ( + TRTLLMMooncakeKVStoreConfig, + trtllm_mooncake_block_only_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,11 @@ 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 + # Mooncake client config for this role's workers (vLLM, TRT-LLM): the keys srtslurm writes into the + # JSON file MOONCAKE_CONFIG_PATH names (`protocol`, `global_segment_size`, ...). Keys the whole pool + # shares must match across roles. The mooncake-master service's deprecated `options.store_config` still + # applies underneath. + 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, @@ -2792,6 +2802,7 @@ def __post_init__(self): self._validate_services_only() self._validate_services() self._warn_dp_launch_mode() + self._warn_mooncake_master_store_config() @cached_property def topology(self) -> Topology: @@ -3290,6 +3301,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 +3367,140 @@ def _vllm_has_mooncake(mode_cfg: dict | None) -> bool: "for you." ) + def _validate_role_mooncake_store_config(self) -> None: + """``roles..mooncake_store_config`` belongs to a role whose workers read a Mooncake client config file. + + Each role carries its own client config, so the keys every client of one pool must + share (``Backend.mooncake_pool_keys``) are checked to match across the roles. + """ + for role, spec in self.roles.items(): + if not spec.mooncake_store_config: + continue + backend = self.role_backends.get(role, self.backend) + if role not in backend.mooncake_store_modes(): + raise ValidationError( + f"roles.{role}.mooncake_store_config is set, but roles.{role} reads no Mooncake client config " + "file: its args select no Mooncake connector, or its engine takes Mooncake settings another way " + "(SGLang: MOONCAKE_* variables in roles..env)" + ) + if backend.mooncake_kv_store is None: + raise ValidationError( + f"roles.{role}.mooncake_store_config needs a mooncake-master service: srtslurm writes it into " + "the client config it renders for the master's pool" + ) + backend = self.backend + modes = backend.mooncake_store_modes() + if backend.mooncake_kv_store is None or len(modes) < 2: + return + for key in backend.mooncake_pool_keys: + values = {mode: backend.mooncake_store_config_for_mode(mode).get(key) for mode in modes} + if len({repr(value) for value in values.values()}) > 1: + raise ValidationError( + f"{key} must be the same for every role in the Mooncake pool, got {values}: set it to one " + "value in each roles..mooncake_store_config" + ) + + def _warn_mooncake_master_store_config(self) -> None: + """``options.store_config`` on the mooncake-master service still applies, but roles carry the client config now.""" + mooncake_cfg = self.backend.mooncake_kv_store + if ( + isinstance(mooncake_cfg, VLLMMooncakeKVStoreConfig | TRTLLMMooncakeKVStoreConfig) + and mooncake_cfg.store_config + ): + logger.warning( + "services[mooncake-master].options.store_config is deprecated: set its keys in " + "roles..mooncake_store_config of each role that uses Mooncake. Until then they apply " + "under every role's own keys" + ) + + 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: + 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 " + "roles..mooncake_store_config (segment_size is global_segment_size there; pool, " + "run_dir, and master_timeout have no equivalent)." + ) + 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 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..8dac51e3c 100644 --- a/tests/test_dry_run.py +++ b/tests/test_dry_run.py @@ -728,7 +728,7 @@ def test_vllm_mooncake_kv_store_dry_run(self, capsys): assert "nof_eviction" in output def test_vllm_mooncake_store_config_in_dry_run(self, capsys): - """vLLM store_config + MOONCAKE_CONFIG_PATH appear in the dry-run extensions panel.""" + """The deprecated service store_config still reaches each role's config and MOONCAKE_CONFIG_PATH in the dry run.""" kv_cfg = '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}' mooncake_env = {"MOONCAKE_PROTOCOL": "rdma"} config = _make_config( @@ -758,10 +758,41 @@ def test_vllm_mooncake_store_config_in_dry_run(self, capsys): show_config_details(config) output = capsys.readouterr().out assert "MOONCAKE_CONFIG_PATH" in output - assert "/logs/mooncake_store_config.json" in output + assert "store_config (prefill)" in output + assert "store_config (decode)" in output 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": ".mooncake_store_config: each role's Mooncake client config, for vLLM, without Slurm or GPUs.""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Any + +import pytest +import yaml +from marshmallow import ValidationError + +from srtctl.backends import VLLMBackend +from srtctl.core.schema import SrtConfig +from srtctl.mock import MockOptions, run_mock_sweep + +EXAMPLES = Path(__file__).resolve().parents[1] / "examples" +KV_TRANSFER_CONFIG = json.dumps( + { + "kv_connector": "MultiConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "connectors": [ + {"kv_connector": "NixlConnector", "kv_role": "kv_both"}, + {"kv_connector": "MooncakeStoreConnector", "kv_role": "kv_both"}, + ] + }, + } +) + + +def _recipe() -> dict[str, Any]: + """examples/vllm/dynamo-disagg.yaml on a Mooncake pool, each role carrying its own client config.""" + recipe = yaml.safe_load((EXAMPLES / "vllm" / "dynamo-disagg.yaml").read_text()) + recipe["model"].update(path="hf:Qwen/Qwen3-0.6B", container="vllm.sqsh") + recipe["engine"] = {"type": "vllm"} + for mode, size in (("prefill", "8GB"), ("decode", "16GB")): + role = recipe["roles"][mode] + role["args"]["kv-transfer-config"] = KV_TRANSFER_CONFIG + role["mooncake_store_config"] = {"protocol": "rdma", "metadata_server": "P2PHANDSHAKE", "global_segment_size": size} + recipe["services"] = [{"name": "mooncake-master", "type": "mooncake-master"}] + return recipe + + +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 test_vllm_roles_carry_their_own_client_config(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + """Each role's workers read the file rendered from that role's keys; nothing is deprecated.""" + path = tmp_path / "recipe.yaml" + path.write_text(yaml.safe_dump(_recipe())) + launches: list[dict[str, Any]] = [] + with caplog.at_level(logging.WARNING): + exit_code = run_mock_sweep( + config_path=path, + output_dir=tmp_path / "outputs" / "4302", + job_id="4302", + options=MockOptions(child_duration_s=0.05, phase_pause_s=0.01, nodelist=("node-01",), on_srun=launches.append), + ) + assert exit_code == 0 + assert "options.store_config is deprecated" not in caplog.text + for mode, size in (("prefill", "8GB"), ("decode", "16GB")): + (written,) = (tmp_path / "outputs").rglob(f"mooncake_store_config_{mode}.json") + config = json.loads(written.read_text()) + assert config["global_segment_size"] == size + assert config["protocol"] == "rdma" + assert config["master_server_address"].endswith(":8700") + worker = next(launch for launch in launches if launch.get("step_name", "").startswith(f"{mode}_")) + assert worker["env_to_set"]["MOONCAKE_CONFIG_PATH"] == f"/logs/mooncake_store_config_{mode}.json" + assert not list((tmp_path / "outputs").rglob("mooncake_store_config.json")) + + +def test_master_store_config_still_applies_under_each_role_and_warns( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """The service's options.store_config is deprecated but keeps working: a role's own keys win over it.""" + recipe = _recipe() + recipe["services"][0]["options"] = {"store_config": {"local_buffer_size": "4GB", "global_segment_size": "2GB"}} + del recipe["roles"]["prefill"]["mooncake_store_config"]["global_segment_size"] + with caplog.at_level(logging.WARNING): + backend = _load(recipe, tmp_path).backend + assert "options.store_config is deprecated" in caplog.text + assert isinstance(backend, VLLMBackend) + prefill = backend.build_mooncake_store_config("prefill", "10.1.1.1") + decode = backend.build_mooncake_store_config("decode", "10.1.1.1") + assert (prefill["global_segment_size"], decode["global_segment_size"]) == ("2GB", "16GB") + assert prefill["local_buffer_size"] == decode["local_buffer_size"] == "4GB" + + +@pytest.mark.parametrize("key", ["protocol", "metadata_server"]) +def test_keys_the_pool_shares_must_match_across_roles(tmp_path: Path, key: str) -> None: + recipe = _recipe() + recipe["roles"]["decode"]["mooncake_store_config"][key] = "tcp" + with pytest.raises(ValidationError, match=f"{key} must be the same for every role in the Mooncake pool"): + _load(recipe, tmp_path) + + +def test_a_role_without_a_mooncake_connector_takes_no_client_config(tmp_path: Path) -> None: + recipe = _recipe() + recipe["roles"]["decode"]["args"]["kv-transfer-config"] = json.dumps({"kv_connector": "NixlConnector"}) + with pytest.raises(ValidationError, match="roles.decode reads no Mooncake client config file"): + _load(recipe, tmp_path) + del recipe["roles"]["decode"]["mooncake_store_config"] + backend = _load(recipe, tmp_path).backend + assert backend.mooncake_store_modes() == ("prefill",) diff --git a/tests/test_trtllm_mooncake.py b/tests/test_trtllm_mooncake.py new file mode 100644 index 000000000..a62762d82 --- /dev/null +++ b/tests/test_trtllm_mooncake.py @@ -0,0 +1,377 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TRT-LLM's mooncake-store KV connector on the native mooncake-master service, without Slurm or GPUs.""" + +from __future__ import annotations + +import copy +import json +from pathlib import Path +from typing import Any + +import pytest +import yaml +from marshmallow import ValidationError + +from srtctl.backends import SGLangBackend, TRTLLMBackend +from srtctl.core.schema import SrtConfig +from srtctl.mock import MockOptions, run_mock_sweep + +EXAMPLES = Path(__file__).resolve().parents[1] / "examples" / "trtllm" +CONNECTOR = {"connector": "mooncake-store"} + + +def _recipe(frontend: str = "trtllm_serve") -> 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 _role(recipe: dict[str, Any], mode: str = "prefill") -> dict[str, Any]: + """A role's Mooncake client config in the recipe.""" + return recipe["roles"][mode]["mooncake_store_config"] + + +def _both(recipe: dict[str, Any], **keys: Any) -> None: + """Set ``keys`` in every connector role's client config (pool-wide keys must match).""" + for mode in ("prefill", "decode"): + _role(recipe, mode).update(keys) + + +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_each_role_writes_its_own_client_config(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 + + _both(recipe, 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_needs_a_role_that_reads_one(tmp_path: Path, engine: str) -> None: + """SGLang reads no client config file; a vLLM role without a Mooncake connector does not either.""" + 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 reads no Mooncake client config file"): + _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="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() + _role(recipe)["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"), r["roles"][m].pop("mooncake_store_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: _both(r, model_key=""), "model_key must be a non-empty"), + (lambda r: _role(r).update(global_segment_size="16GB"), "'GB' is refused"), + (lambda r: _role(r).update(local_buffer_size=0), "at least 1 bytes"), + (lambda r: _role(r).update(global_segment_size=-1), "at least 0 bytes"), + (lambda r: _role(r).update(transfer_batch_size=0), "positive integer"), + (lambda r: _role(r).update(stage_through_host="false"), "not the string"), + (lambda r: _role(r).update(segment_size="160GiB"), "set global_segment_size"), + (lambda r: _role(r).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"), + "roles.decode.mooncake_store_config is set, but roles.decode reads no Mooncake client config file", + ), + ( + lambda r: r["roles"]["decode"]["mooncake_store_config"].update(protocol="tcp"), + "protocol must be the same for every role in the Mooncake pool", + ), + ( + lambda r: r["roles"]["decode"]["mooncake_store_config"].update(model_key="other"), + "model_key must be the same for every role in the Mooncake 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() + _role(recipe)["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).setdefault("options", {}).setdefault("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() + _role(recipe)["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() + _role(recipe)["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") + agg.pop("mooncake_store_config") # the pool is the engine's own, not the master's + 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",)