diff --git a/requirements.txt b/requirements.txt index d3945937aef6..308cfe9b386b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -116,3 +116,4 @@ cache-dit>=1.3.5 librosa msgpack uvloop>=0.19.0 +mooncake-transfer-engine-cuda13==0.3.13 diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py new file mode 100644 index 000000000000..1bf4602e2088 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""A Mooncake distributed store to back a KV cache connector. + +The store is a shared CPU memory pool addressed by content, so a prefix computed +by one engine can be replayed by another, which regular block reuse cannot do +because it never leaves the instance that computed it. + +This is a different component from the Mooncake transfer engine that the C++ +cache transceiver uses for disaggregated prefill/decode handoff: that moves KV +point to point between two known peers, while this one publishes pages into a +pool addressed by content. The two compose, so a context server can write pages +here and still hand off over NIXL. + +`master.py` brings a pool up: it resolves or launches the `mooncake_master` +and renders the client config the workers read. Capacity comes only from +processes that open a store handle, which in a disaggregated deployment is the +context servers alone, so `donor.py` lends a node's memory to the pool without +giving it a connector. `trtllm-serve mooncake_master` and `mooncake_donor` +expose both. Both need the Mooncake Python bindings, which `tensorrt-llm` pulls +in as `mooncake-transfer-engine-cuda13`. + +`keys.py` and `staging.py` hold what the store side shares with the connector +that moves pages in and out of the pool: how a block of tokens becomes a store +key, and how pages reach the fabric on hosts without GPUDirect RDMA. The +connector itself, and the `LlmArgs` surface that selects it, land with the KV +cache manager V2 support it depends on. +""" + +from .config import MooncakeStoreConnectorConfig, StoreRole, parse_size +from .donor import DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment +from .master import ( + PoolSpec, + local_address, + master_timeout, + provision_pool, + resolve_master_address, + running_master, + wait_for_master, + write_client_config, +) + +__all__ = [ + "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", + "MooncakeStoreConnectorConfig", + "PoolSpec", + "StoreRole", + "donate_segment", + "local_address", + "master_timeout", + "parse_size", + "provision_pool", + "resolve_master_address", + "running_master", + "wait_for_master", + "write_client_config", +] diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py new file mode 100644 index 000000000000..b73caf63182e --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Configuration for the Mooncake store KV cache connector. + +Topology settings are read from the JSON file named by `MOONCAKE_CONFIG_PATH`, +the same file and environment variable the vLLM Mooncake store connector uses, +so one deployment can point both engines at the same pool. + +`KvCacheConnectorConfig` carries no free-form dictionary, so the two settings +that are TensorRT-LLM's rather than Mooncake's, the read/write role and the key +prefix, are also taken from the environment. +""" + +import json +import os +import re +from dataclasses import dataclass +from enum import Enum +from typing import Any, Optional + +__all__ = [ + "CLIENT_CONFIG_NAME", + "CONFIG_PATH_ENV", + "MooncakeStoreConnectorConfig", + "ROLE_ENV", + "RUN_DIR_ENV", + "STAGE_THROUGH_HOST_ENV", + "StoreRole", + "provisioned_config_path", +] + +CONFIG_PATH_ENV = "MOONCAKE_CONFIG_PATH" +#: Where a server keeps the client config it renders and the master's log. Set +#: it to keep them after shutdown; otherwise they live in a temporary directory. +RUN_DIR_ENV = "TRTLLM_MOONCAKE_RUN_DIR" +#: Name the rendered client config takes in the run directory. +CLIENT_CONFIG_NAME = "mooncake.json" +ROLE_ENV = "TRTLLM_MOONCAKE_STORE_ROLE" +CACHE_PREFIX_ENV = "TRTLLM_MOONCAKE_STORE_PREFIX" +MODEL_KEY_ENV = "TRTLLM_MOONCAKE_STORE_MODEL_KEY" +STAGE_THROUGH_HOST_ENV = "TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST" + +DEFAULT_GLOBAL_SEGMENT_SIZE = 3355443200 +DEFAULT_LOCAL_BUFFER_SIZE = 1073741824 +DEFAULT_CACHE_PREFIX = "trtllm" +DEFAULT_STAGING_BUFFER_SIZE = 536870912 +#: Mooncake's own peer-to-peer handshake, which keeps a separate metadata +#: process out of the deployment. Nothing else is a sensible fallback: an empty +#: connstring is not one of the forms `store.setup` accepts, so a config that +#: leaves the field out means this rather than meaning no metadata service. +DEFAULT_METADATA_SERVER = "P2PHANDSHAKE" + +_TRUE = {"1", "true", "yes", "on"} +_FALSE = {"0", "false", "no", "off"} + +_SIZE_UNITS = { + "": 1, + "b": 1, + "k": 1000, + "kb": 1000, + "m": 1000**2, + "mb": 1000**2, + "g": 1000**3, + "gb": 1000**3, + "t": 1000**4, + "tb": 1000**4, + "kib": 1024, + "mib": 1024**2, + "gib": 1024**3, + "tib": 1024**4, +} +_SIZE_RE = re.compile(r"^\s*([0-9]+(?:\.[0-9]+)?)\s*([a-zA-Z]*)\s*$") + + +class StoreRole(Enum): + """Which directions of traffic this engine is allowed to drive. + + A disaggregated deployment typically runs context servers as `both` and + leaves generation servers unconfigured: generated tokens are rarely a reused + prefix, so writing them costs bandwidth for no hit rate. + """ + + PRODUCER = "producer" + CONSUMER = "consumer" + BOTH = "both" + + @property + def loads(self) -> bool: + """Whether this role reads previously stored KV back onto the GPU.""" + return self is not StoreRole.PRODUCER + + @property + def saves(self) -> bool: + """Whether this role writes newly computed KV into the store.""" + return self is not StoreRole.CONSUMER + + +def parse_size(value: Any) -> int: + """Accept either a byte count or a suffixed string such as `"4GiB"`.""" + if isinstance(value, bool): + raise ValueError(f"expected a size, got {value!r}") + if isinstance(value, int): + return value + if isinstance(value, float): + return int(value) + match = _SIZE_RE.match(str(value)) + if match is None: + raise ValueError(f"cannot parse size {value!r}") + magnitude, unit = match.groups() + scale = _SIZE_UNITS.get(unit.lower()) + if scale is None: + raise ValueError(f"unknown size unit {unit!r} in {value!r}") + return int(float(magnitude) * scale) + + +def provisioned_config_path() -> Optional[str]: + """The client config a server on this node rendered, if there is one. + + `provision_pool` writes one and exports `MOONCAKE_CONFIG_PATH`, which the + ranks the LLM constructor spawns inherit. Ranks an external launcher + started, one task per rank, were already running by then and never see it, + so they read the config back from the run directory instead. + + Only possible when the deployment named that directory, since it otherwise + defaults to a per-process temporary one that no other rank could read. + """ + run_dir = os.getenv(RUN_DIR_ENV) + if not run_dir: + return None + path = os.path.join(run_dir, CLIENT_CONFIG_NAME) + return path if os.path.exists(path) else None + + +@dataclass(frozen=True) +class MooncakeStoreConnectorConfig: + """Everything needed to open a store handle and name keys in it.""" + + master_server_address: str + metadata_server: str = DEFAULT_METADATA_SERVER + protocol: str = "rdma" + device_name: str = "" + global_segment_size: int = DEFAULT_GLOBAL_SEGMENT_SIZE + local_buffer_size: int = DEFAULT_LOCAL_BUFFER_SIZE + local_hostname: Optional[str] = None + tenant_id: Optional[str] = None + role: StoreRole = StoreRole.BOTH + cache_prefix: str = DEFAULT_CACHE_PREFIX + #: Identity the keys are namespaced by. Two engines share cache only when + #: they agree on this, and two that disagree about what it names read each + #: other's pages, so it is required rather than defaulted. See + #: :meth:`resolve_model_key`. + model_key: Optional[str] = None + #: How many page keys go into one store call. Bounds the size of a single + #: RPC without bounding how much a request may transfer. + transfer_batch_size: int = 64 + #: Pass pages through a pinned host buffer instead of registering the KV + #: pools with Mooncake. Costs a copy each way, but works without GPUDirect + #: RDMA, which registering device memory requires. + stage_through_host: bool = False + #: Ceiling on the pinned allocation per direction when staging. Slots are + #: sized from the layout's largest page, so this caps how many pages may be + #: in flight rather than how large one may be. + staging_buffer_bytes: int = DEFAULT_STAGING_BUFFER_SIZE + + def __post_init__(self) -> None: + """Reject settings that would fail later, inside a transfer.""" + if not self.master_server_address: + raise ValueError("master_server_address is required") + if self.local_buffer_size <= 0: + raise ValueError("local_buffer_size must be > 0") + if self.global_segment_size < 0: + raise ValueError("global_segment_size must be >= 0") + if self.transfer_batch_size <= 0: + raise ValueError("transfer_batch_size must be > 0") + if self.stage_through_host and self.staging_buffer_bytes <= 0: + raise ValueError("staging_buffer_bytes must be > 0 when staging is on") + + @staticmethod + def from_file(path: str) -> "MooncakeStoreConnectorConfig": + """Read the topology from a vLLM-compatible Mooncake JSON config.""" + with open(path) as handle: + raw = json.load(handle) + return MooncakeStoreConnectorConfig( + master_server_address=raw.get("master_server_address", ""), + metadata_server=raw.get("metadata_server") or DEFAULT_METADATA_SERVER, + protocol=raw.get("protocol", "rdma"), + device_name=raw.get("device_name", ""), + global_segment_size=parse_size( + raw.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE) + ), + local_buffer_size=parse_size(raw.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)), + local_hostname=raw.get("local_hostname") or None, + tenant_id=raw.get("tenant_id") or None, + role=StoreRole(str(raw.get("role", StoreRole.BOTH.value)).strip().lower()), + cache_prefix=str(raw.get("cache_prefix", DEFAULT_CACHE_PREFIX)), + model_key=raw.get("model_key") or None, + transfer_batch_size=int(raw.get("transfer_batch_size", 64)), + stage_through_host=bool(raw.get("stage_through_host", False)), + staging_buffer_bytes=parse_size( + raw.get("staging_buffer_bytes", DEFAULT_STAGING_BUFFER_SIZE) + ), + ) + + @staticmethod + def from_env() -> "MooncakeStoreConnectorConfig": + """Load the JSON config, then apply the TensorRT-LLM env overrides.""" + path = os.getenv(CONFIG_PATH_ENV) or provisioned_config_path() + if not path: + raise ValueError( + f"The mooncake-store connector needs {CONFIG_PATH_ENV} set to a " + "Mooncake JSON config (metadata_server, master_server_address, " + "protocol, device_name, global_segment_size, local_buffer_size), " + "or kv_connector_config.mooncake_store set so the server renders " + f"one, into ${RUN_DIR_ENV} if this rank was started by the " + "launcher rather than spawned by the server." + ) + config = MooncakeStoreConnectorConfig.from_file(path) + return config.with_env_overrides() + + def with_env_overrides(self) -> "MooncakeStoreConnectorConfig": + """Apply `TRTLLM_MOONCAKE_STORE_*` on top of the file's settings.""" + import dataclasses + + updates: dict[str, Any] = {} + role = os.getenv(ROLE_ENV) + if role: + try: + updates["role"] = StoreRole(role.strip().lower()) + except ValueError as exc: + known = ", ".join(member.value for member in StoreRole) + raise ValueError(f"{ROLE_ENV}={role!r} is not one of: {known}") from exc + prefix = os.getenv(CACHE_PREFIX_ENV) + if prefix: + updates["cache_prefix"] = prefix + model_key = os.getenv(MODEL_KEY_ENV) + if model_key: + updates["model_key"] = model_key + staging = os.getenv(STAGE_THROUGH_HOST_ENV) + if staging: + normalized = staging.strip().lower() + if normalized in _TRUE: + updates["stage_through_host"] = True + elif normalized in _FALSE: + updates["stage_through_host"] = False + else: + known = ", ".join(sorted(_TRUE | _FALSE)) + raise ValueError( + f"{STAGE_THROUGH_HOST_ENV}={staging!r} is not a boolean; use one of: {known}" + ) + return dataclasses.replace(self, **updates) if updates else self + + def resolve_model_key(self, model: Any) -> str: + """The model identity to namespace keys by. + + Deliberately has no default. Deriving one from the model path would + make two checkpoints that happen to share a directory name, such as + `org-a/model` and `org-b/model` or two revisions mounted alike, agree + on a namespace while disagreeing on what the pages mean, and each would + read the other's KV as its own. + """ + if self.model_key: + return self.model_key + raise ValueError( + f"The mooncake-store connector needs a model key to namespace its " + f"pool keys by, and there is no safe default: set " + f"kv_connector_config.mooncake_store.model_key, the model_key field " + f"of the Mooncake JSON config, or ${MODEL_KEY_ENV}. Give it a value " + f"that separates this checkpoint from any other an engine sharing " + f"the pool might load, rather than one derived from {model!r}." + ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py new file mode 100644 index 000000000000..071390ca286d --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Put a node's host memory into a Mooncake pool without reading or writing it. + +Pool capacity comes only from processes that open a store handle: `setup` +registers `global_segment_size` bytes of the caller's host memory and the master +then places blocks in it. In a disaggregated deployment only the context servers +configure the connector, so the pool is entirely prefill-node memory, which +overlaps what TensorRT-LLM's own host offload already does. + +Running `trtllm-serve mooncake_donor` on a generation node puts that node's +memory into the same pool, so prefill writes blocks that land on decode-side +DRAM. The generation engine stays free of any connector and keeps its single +cache transceiver for the prefill-to-decode handoff. + +Donation is not a `StoreRole`. The roles describe an engine's traffic and none +of them means "contribute memory only", so capacity and traffic stay separate +concerns and a donor holds a store handle of its own. + +The memory is charged to the donating process, so size it together with that +node's `kv_cache_config.host_cache_size`. +""" + +import contextlib +import time +from typing import Iterator, Optional + +from tensorrt_llm.logger import logger + +from .config import DEFAULT_METADATA_SERVER +from .master import local_address + +__all__ = [ + "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", + "donate_segment", +] + +#: A donor never transfers, but `setup` rejects a zero-sized transfer buffer. +DEFAULT_DONOR_LOCAL_BUFFER_SIZE = 64 * 1024**2 + + +@contextlib.contextmanager +def donate_segment( + master_server_address: str, + segment_size: int, + protocol: str = "rdma", + device_name: str = "", + metadata_server: str = DEFAULT_METADATA_SERVER, + local_buffer_size: int = DEFAULT_DONOR_LOCAL_BUFFER_SIZE, + hostname: Optional[str] = None, +) -> Iterator[str]: + """Hold `segment_size` bytes of this node's memory in the pool. + + Yields the host the segment is registered under, which is how the master + and the engines reading from it identify the capacity. + + Dropping the store handle unmounts the segment and the master starts + reporting the blocks that lived in it as lost, so the caller must stay + inside this context for as long as the capacity is meant to exist. + """ + try: + from mooncake.store import MooncakeDistributedStore + except ImportError as exc: + raise ImportError( + "Donating memory needs the Mooncake Python bindings, which " + "`tensorrt-llm` pulls in as `mooncake-transfer-engine-cuda13`, so " + "reinstall that package if this environment dropped it. The C++ " + f"transfer engine is a separate component: {exc}" + ) from exc + + host = hostname or local_address() + donated = f"{segment_size / 1024**3:.1f}GiB" + # Byte counts are spelled out next to the human-readable form. A misparsed + # size string otherwise surfaces only as a pool that evicts far too eagerly. + logger.info( + f"mooncake-store: lending memory to the pool at {master_server_address} " + f"as capacity only, no reads or writes: host={host} " + f"segment_size={donated} ({segment_size} bytes) " + f"protocol={protocol} device={device_name or '(none)'} " + f"metadata_server={metadata_server} " + f"local_buffer_size={local_buffer_size} bytes" + ) + + store = MooncakeDistributedStore() + started = time.monotonic() + status = store.setup( + host, + metadata_server, + segment_size, + local_buffer_size, + protocol, + device_name, + master_server_address, + ) + elapsed = time.monotonic() - started + if status != 0: + raise RuntimeError( + f"Mooncake store.setup failed with status {status} after " + f"{elapsed:.1f}s, so no memory was lent to the pool. The master at " + f"{master_server_address} must already be accepting connections; " + f"protocol={protocol!r} with device={device_name or '(none)'!r} " + f"must be usable from {host}; and this node must have " + f"{donated} of memory to spare, which it does not if its own " + "kv_cache_config.host_cache_size has already claimed it." + ) + + logger.info( + f"mooncake-store: {donated} of {host} is now part of the pool, " + f"registered in {elapsed:.1f}s; the master at {master_server_address} " + "can place blocks here from now on" + ) + try: + yield host + finally: + # The segment stays mounted while anything references the handle. + del store + logger.info( + f"mooncake-store: withdrew the {donated} lent from {host}; the " + "master will report blocks that lived there as lost" + ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py new file mode 100644 index 000000000000..56609877843d --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Block identity and store key naming for the Mooncake store connector. + +`KVCacheManagerV2` exposes no block hashes to a connector, since `RequestData` +reports them empty, so content identity is derived here instead. The chain is +the standard one: a block's hash covers its own tokens *and* every token before +it, so a key can only be reused by a request whose prefix is byte-identical. + +A key is `/`. The namespace pins down everything that +would make the stored bytes mean something different: the model, the shard that +produced them, the layer group inside that shard, the tokens each page holds and +how many bytes a page is. Anything that changes those reads as a cache miss +rather than as garbage. +""" + +import hashlib +from dataclasses import dataclass +from typing import List, Optional, Sequence + +__all__ = [ + "BlockHashChain", + "KeyNamespace", + "HASH_DIGEST_BYTES", +] + +#: 128 bits. Collisions decide whether one request reads another's KV, so the +#: digest is sized to make that negligible over any realistic cache lifetime, +#: while staying half the width of a full blake2b digest in every key. +HASH_DIGEST_BYTES = 16 + + +def _digest(*parts: bytes) -> bytes: + hasher = hashlib.blake2b(digest_size=HASH_DIGEST_BYTES) + for part in parts: + hasher.update(part) + return hasher.digest() + + +class BlockHashChain: + """Rolling hashes of a request's full blocks, one entry per block ordinal. + + Extended in place as a request's token list grows, so generation steps cost + one digest per newly completed block rather than a rehash of the prompt. + """ + + def __init__(self, tokens_per_block: int, cache_salt: Optional[str] = None): + if tokens_per_block <= 0: + raise ValueError(f"tokens_per_block must be > 0, got {tokens_per_block}") + self._tokens_per_block = int(tokens_per_block) + # The salt seeds the chain rather than being mixed into every block, so + # a request carrying a different salt diverges from the first block on. + salt_bytes = b"" if cache_salt is None else str(cache_salt).encode() + self._seed = _digest(b"salt", salt_bytes) + self._hashes: List[bytes] = [] + + @property + def tokens_per_block(self) -> int: + """Tokens covered by each entry in the chain.""" + return self._tokens_per_block + + @property + def hashes(self) -> Sequence[bytes]: + """Hashes computed so far, indexed by block ordinal.""" + return self._hashes + + def extend(self, tokens: Sequence[int]) -> Sequence[bytes]: + """Grow the chain to cover every full block of `tokens`. + + Args: + tokens: The request's complete token list, prompt first. Must be an + extension of what was passed previously; a request's tokens only + ever grow, so a shorter list means the caller mixed up requests. + + Returns: + The full chain, indexed by block ordinal. + """ + num_full_blocks = len(tokens) // self._tokens_per_block + if num_full_blocks < len(self._hashes): + raise ValueError( + f"token list shrank from {len(self._hashes)} to {num_full_blocks} " + "full blocks; a hash chain belongs to exactly one request" + ) + for ordinal in range(len(self._hashes), num_full_blocks): + start = ordinal * self._tokens_per_block + block = tokens[start : start + self._tokens_per_block] + parent = self._hashes[-1] if self._hashes else self._seed + # Fixed-width little-endian token ids: a delimiter-free encoding + # would let two different token sequences serialize identically. + payload = b"".join(int(token).to_bytes(8, "little", signed=True) for token in block) + self._hashes.append(_digest(parent, payload)) + return self._hashes + + +@dataclass(frozen=True) +class KeyNamespace: + """The part of a store key that is fixed for one shard and layer group.""" + + cache_prefix: str + model_key: str + #: Global rank of the shard whose KV these bytes are, and the world size it + #: was produced under. Both are needed: rank 3 of 8 holds different heads + #: than rank 3 of 4. + rank: int + world_size: int + layer_group_id: int + tokens_per_block: int + bytes_per_page: int + + @property + def prefix(self) -> str: + """The literal string every key in this namespace starts with.""" + return ( + f"{self.cache_prefix}/{self.model_key}" + f"/w{self.world_size}r{self.rank}" + f"/lg{self.layer_group_id}" + f"/t{self.tokens_per_block}b{self.bytes_per_page}" + ) + + def key(self, block_hash: bytes) -> str: + """The store key holding one page of this namespace.""" + return f"{self.prefix}/{block_hash.hex()}" diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py new file mode 100644 index 000000000000..e94b0bc37350 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -0,0 +1,609 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Bring a Mooncake store's pool up and tell its workers where to find it. + +The connector needs two things that are not the engine's to produce: a reachable +`mooncake_master`, and a JSON client config named by `MOONCAKE_CONFIG_PATH` that +points every worker at it. + +`provision_pool` produces both. It resolves the master, either launching one or +checking that the configured one answers, renders the client config, and exports +`MOONCAKE_CONFIG_PATH`. Everything it started is torn down when the context +exits. + +A master launched by the process that provisions lives and dies with it, so it +suits one engine talking to its own pool. Several engines sharing a pool, or a +pool meant to survive a restart, want `trtllm-serve mooncake_master`, which +holds one for as long as that command runs and publishes its address for the +rest of the deployment to read back. +""" + +import contextlib +import json +import os +import shutil +import socket +import subprocess # nosec B404 +import tempfile +import time +from dataclasses import dataclass, fields +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union + +from tensorrt_llm.logger import logger + +from .config import CLIENT_CONFIG_NAME, CONFIG_PATH_ENV, DEFAULT_METADATA_SERVER, RUN_DIR_ENV + +__all__ = [ + "PoolSpec", + "local_address", + "master_timeout", + "provision_pool", + "resolve_master_address", + "running_master", + "wait_for_master", + "write_client_config", +] + +#: Override the binary that `launch_master` runs. +MASTER_BINARY_ENV = "TRTLLM_MOONCAKE_MASTER_BINARY" +#: How long to wait for a master to accept connections, in seconds. +MASTER_TIMEOUT_ENV = "TRTLLM_MOONCAKE_MASTER_TIMEOUT" +DEFAULT_MASTER_BINARY = "mooncake_master" +DEFAULT_MASTER_TIMEOUT = 60.0 +MASTER_LOG_NAME = "mooncake_master.log" +#: Name a launched master's address is always published under in the run +#: directory, so even a run that named no address file records its pool. +MASTER_ADDRESS_NAME = "master.addr" +#: Prefix that makes `master_server_address` name a file holding the address +#: rather than the address itself. +ADDRESS_FILE_SCHEME = "file://" +#: Lines of the master's log to quote when startup fails, since its last words +#: (a port in use, a bad flag) are usually the whole diagnosis. +LOG_TAIL_LINES = 20 + + +@dataclass(frozen=True) +class PoolSpec: + """A pool to provision: whose master owns it, and what its workers are told. + + Deliberately internal. The settings below end up in the Mooncake client + config, whose schema is vLLM's and is the contract the connector reads; + this is only the shape `provision_pool` takes them in. Keeping it out of + `LlmArgs` leaves the user-facing spelling to be settled alongside the + connector that gives these fields meaning. + + Exactly one of `master_server_address` and `launch_master` says where the + master is, which `provision_pool` enforces rather than this class, since a + caller may fill the two in from different places. + """ + + #: `host:port`, or `file://` naming a file that holds one. The file + #: is how to reach a master whose host a scheduler chose. + master_server_address: Optional[str] = None + #: Start a `mooncake_master` and use it, instead of joining one. + launch_master: bool = False + #: Extra path to publish a launched master's address to, beyond the copy + #: always written to the run directory. + master_address_file: Optional[str] = None + master_port: int = 50051 + master_metrics_port: int = 9004 + master_eviction_ratio: float = 0.05 + metadata_server: str = DEFAULT_METADATA_SERVER + #: `rdma` or `tcp`. TCP is for bring-up only. + protocol: str = "rdma" + #: RDMA device to transfer over, from `ibv_devinfo`. One config serves a + #: whole pool, so naming a device here names it for every node; left empty + #: each client discovers its own, which is what differing HCAs need. + device_name: str = "" + #: Host memory each worker process contributes. Pool capacity is this times + #: the number of processes that open a store handle. + global_segment_size: Union[int, str] = "16GiB" + #: Per-process Mooncake transfer buffer, not pool capacity. + local_buffer_size: Union[int, str] = "1GiB" + transfer_batch_size: int = 64 + #: Left unset so the rendered config omits the key and the connector's own + #: default applies, rather than carrying a second copy of it here. + cache_prefix: Optional[str] = None + #: Identity the keys are namespaced by. Engines sharing a pool read each + #: other's pages exactly when they agree on this, so there is no safe + #: default; see `MooncakeStoreConnectorConfig.resolve_model_key`. + model_key: Optional[str] = None + #: Copy pages through a pinned host buffer instead of registering the KV + #: pools, for hosts without GPUDirect RDMA. + stage_through_host: bool = False + staging_buffer_bytes: Optional[Union[int, str]] = None + + @staticmethod + def from_json(raw: Dict[str, Any], **overrides: Any) -> "PoolSpec": + """Build a spec from a Mooncake client config, plus what a caller adds. + + Lets a deployment describe its pool in the same JSON the connector + reads, rather than in a second format, and lets the CLI layer on the + parts that are the command's rather than the pool's. + """ + known = {f.name for f in fields(PoolSpec)} + settings = {key: value for key, value in raw.items() if key in known} + settings.update({key: value for key, value in overrides.items() if value is not None}) + return PoolSpec(**settings) + + +def _log_tail(path: str, lines: int = LOG_TAIL_LINES) -> str: + """The end of the master's log, ready to append to a failure message.""" + try: + with open(path, errors="replace") as handle: + tail = handle.read().splitlines()[-lines:] + except OSError as exc: + return f" Its log at {path} could not be read: {exc}." + if not tail: + return ( + f" Its log at {path} is empty, which usually means it failed " + "before glog opened; check that the binary runs at all." + ) + quoted = "\n ".join(tail) + return f" The last {len(tail)} lines of {path}:\n {quoted}" + + +def local_address() -> str: + """The address this host is known by inside the pool. + + Uses the same derivation as the connector worker's own hostname, so the + master and the segments registering with it agree on which host they are on. + """ + try: + return socket.gethostbyname(socket.gethostname()) + except OSError: + return "127.0.0.1" + + +def master_timeout() -> float: + raw = os.getenv(MASTER_TIMEOUT_ENV) + if not raw: + return DEFAULT_MASTER_TIMEOUT + try: + timeout = float(raw) + except ValueError as exc: + raise ValueError(f"{MASTER_TIMEOUT_ENV}={raw!r} is not a number") from exc + if timeout <= 0: + raise ValueError(f"{MASTER_TIMEOUT_ENV}={raw!r} must be > 0") + return timeout + + +def _split_address(address: str) -> Optional[Tuple[str, int]]: + """Split `host:port`, or return `None` if it is not in that form.""" + host, separator, port = address.rpartition(":") + if not separator or not port.isdigit(): + return None + return host.strip("[]"), int(port) + + +def resolve_master_address(address: str, timeout: float) -> str: + """Read a `file://` address through, and pass anything else along. + + A master with its own lifetime runs on whichever host its scheduler gave + it, which is not known when the worker configs are written. Naming the file + it publishes to keeps the address out of both the config and the launch + script: `trtllm-serve mooncake_master --address-file` writes it, every + worker's `master_server_address` names the same path, and the wait here + doubles as the wait for the master to exist at all. + """ + if not address.startswith(ADDRESS_FILE_SCHEME): + return address + + path = address[len(ADDRESS_FILE_SCHEME) :] + started = time.monotonic() + deadline = started + timeout + announced = started + logger.info(f"mooncake-store: reading the master's address from {path}") + while True: + try: + published = open(path).read().strip() + except FileNotFoundError: + published = "" + if published: + logger.info(f"mooncake-store: {path} names the master at {published}") + return published + now = time.monotonic() + if now - announced >= 5.0: + announced = now + # Waiting on a master in another job is normal here, so say so + # rather than letting the wait look like a hang. + logger.info( + f"mooncake-store: no master address in {path} yet " + f"({now - started:.0f}s of {timeout:g}s); waiting for the " + "master to start and publish it" + ) + if now >= deadline: + raise TimeoutError( + f"No Mooncake master address appeared in {path} within " + f"{timeout:g}s. Start one with 'trtllm-serve mooncake_master " + f"--address-file {path}', or name a reachable host:port in " + f"master_server_address. Raise {MASTER_TIMEOUT_ENV} if the " + "master is only slow to start." + ) + time.sleep(0.5) + + +def _wait_until_accepting( + host: str, + port: int, + timeout: float, + process: Optional[subprocess.Popen] = None, + log_path: Optional[str] = None, +) -> float: + """Block until the master accepts connections, and say how long it took. + + A worker that opens its store handle before the master is listening fails + outright, so the ordering has to wait on the port rather than on the + presence of a process. When the master is ours, its exit is checked first + each pass, so a master that died is reported as such rather than as a + timeout. + + The wait is narrated as it happens, since silence here is + indistinguishable from a hang elsewhere in bringup. + """ + started = time.monotonic() + deadline = started + timeout + announced = started + while True: + if process is not None and (code := process.poll()) is not None: + raise RuntimeError( + f"mooncake_master exited with code {code} after " + f"{time.monotonic() - started:.1f}s, before it accepted " + f"connections on {host}:{port}." + f"{_log_tail(log_path) if log_path else ''}" + ) + try: + with socket.create_connection((host, port), timeout=1.0): + return time.monotonic() - started + except OSError as exc: + last_error = exc + now = time.monotonic() + if now >= deadline: + raise TimeoutError( + f"The Mooncake master at {host}:{port} did not accept " + f"connections within {timeout:g}s ({last_error}). Raise " + f"{MASTER_TIMEOUT_ENV} if it is only slow to start." + f"{_log_tail(log_path) if log_path else ''}" + ) + if now - announced >= 5.0: + announced = now + logger.info( + f"mooncake-store: still waiting for the master at {host}:{port}" + f" ({now - started:.0f}s of {timeout:g}s, {last_error})" + ) + time.sleep(0.5) + + +def wait_for_master(master_address: str, timeout: Optional[float] = None) -> Optional[float]: + """Block until the master at `master_address` accepts connections. + + Reaching a master that is not there otherwise fails deep inside + `store.setup`, in every rank, after the model has loaded, as a bare status + code. One socket beforehand turns that into a line naming the address. + + Returns how long it took, or `None` if the address was not in `host:port` + form and could not be checked. + """ + timeout = master_timeout() if timeout is None else timeout + endpoint = _split_address(master_address) + if endpoint is None: + logger.warning( + f"mooncake-store: cannot parse master_server_address=" + f"{master_address!r} as host:port, so its reachability is left " + "for the workers to discover." + ) + return None + elapsed = _wait_until_accepting(*endpoint, timeout) + logger.info(f"mooncake-store: the master at {master_address} answered in {elapsed:.1f}s") + return elapsed + + +def _client_config(pool: PoolSpec, master_address: str) -> Dict[str, Any]: + """Render the Mooncake client config for a pool. + + The schema is vLLM's, so one pool can serve both engines. `role` is written + as `both` because the file describes the pool; the directions of traffic a + given process drives come from its own `TRTLLM_MOONCAKE_STORE_ROLE`. + """ + config: Dict[str, Any] = { + "metadata_server": pool.metadata_server, + "master_server_address": master_address, + "protocol": pool.protocol, + "device_name": pool.device_name, + "global_segment_size": pool.global_segment_size, + "local_buffer_size": pool.local_buffer_size, + "role": "both", + "transfer_batch_size": pool.transfer_batch_size, + "stage_through_host": pool.stage_through_host, + } + if pool.cache_prefix is not None: + config["cache_prefix"] = pool.cache_prefix + if pool.model_key is not None: + config["model_key"] = pool.model_key + # Left out when unset so the connector's own default applies instead of a + # second copy of it here. + if pool.staging_buffer_bytes is not None: + config["staging_buffer_bytes"] = pool.staging_buffer_bytes + return config + + +def write_client_config(pool: PoolSpec, master_address: str, run_dir: str) -> str: + """Render `pool`'s client config into `run_dir` and return its path. + + Workers reach it either through `MOONCAKE_CONFIG_PATH`, which + `provision_pool` exports, or by looking in `$TRTLLM_MOONCAKE_RUN_DIR` when + an external launcher started them too early to inherit it; see + `provisioned_config_path`. Writing it next to the master's address file is + what lets `trtllm-serve mooncake_master --run_dir` describe a whole pool + without any engine having to. + """ + os.makedirs(run_dir, exist_ok=True) + config_path = os.path.join(run_dir, CLIENT_CONFIG_NAME) + config = _client_config(pool, master_address) + with open(config_path, "w") as handle: + json.dump(config, handle, indent=2) + logger.info(f"mooncake-store: wrote {config_path} ({json.dumps(config, sort_keys=True)})") + return config_path + + +@dataclass +class LaunchedMaster: + """A `mooncake_master` owned by this process.""" + + process: subprocess.Popen + address: str + log_path: str + + def stop(self, timeout: float = 10.0) -> None: + if self.process.poll() is not None: + return + self.process.terminate() + try: + self.process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + self.process.kill() + self.process.wait() + + +def _launch_master(pool: PoolSpec, run_dir: str) -> LaunchedMaster: + """Start a master on this host and wait for it to answer.""" + binary = os.getenv(MASTER_BINARY_ENV) or DEFAULT_MASTER_BINARY + resolved = shutil.which(binary) + if resolved is None: + raise FileNotFoundError( + f"{binary!r} is not on PATH, so launch_master cannot start a " + "Mooncake master. It comes from the Mooncake Python client, which " + "`tensorrt-llm` pulls in as `mooncake-transfer-engine-cuda13`, so " + "reinstall that package if this environment dropped it. Otherwise " + f"point {MASTER_BINARY_ENV} at the binary, or drop launch_master " + "and set master_server_address to a master you run yourself." + ) + + host = local_address() + log_path = os.path.join(run_dir, MASTER_LOG_NAME) + + # glog writes to files under /tmp unless redirected, so without + # GLOG_logtostderr the log opened below stays empty. GLOG_v=1 adds the + # per-RPC lines showing segments registering and keys moving, which is the + # only view of the pool's own side of the conversation short of scraping + # the metrics port. + env = dict(os.environ, GLOG_logtostderr="1") + env.setdefault("GLOG_v", "1") + command = [ + resolved, + f"--rpc_port={pool.master_port}", + f"--metrics_port={pool.master_metrics_port}", + f"--eviction_ratio={pool.master_eviction_ratio}", + ] + + logger.info(f"mooncake-store: starting {' '.join(command)} on {host}") + with open(log_path, "wb") as log_file: + process = subprocess.Popen( # nosec B603 + command, env=env, stdout=log_file, stderr=subprocess.STDOUT + ) + master = LaunchedMaster( + process=process, address=f"{host}:{pool.master_port}", log_path=log_path + ) + logger.info( + f"mooncake-store: master pid={process.pid} logging to {log_path} " + f"(GLOG_v={env['GLOG_v']}); waiting for it to accept connections" + ) + try: + elapsed = _wait_until_accepting( + host, pool.master_port, master_timeout(), process=process, log_path=log_path + ) + except BaseException: + master.stop() + raise + + logger.info( + f"mooncake-store: master ready at {master.address} after {elapsed:.1f}s " + f"(metrics http://{host}:{pool.master_metrics_port}, log {log_path})" + ) + return master + + +@contextlib.contextmanager +def _published_address(address: str, paths: Sequence[str]) -> Iterator[None]: + """Write `address` to every path for the life of the context. + + Publishing is how anything else finds this master: a donor or a second + server names the path as `file://` in `master_server_address`. + Retracting on the way out matters as much as writing, since an address that + outlives its master sends the next run's workers to a dead port. + """ + for path in paths: + directory = os.path.dirname(path) + if directory: + os.makedirs(directory, exist_ok=True) + # Renamed into place so a reader never sees a partial address. + staging = f"{path}.partial" + with open(staging, "w") as handle: + handle.write(f"{address}\n") + os.replace(staging, path) + logger.info(f"mooncake-store: published master {address} to {path}") + try: + yield + finally: + for path in paths: + with contextlib.suppress(OSError): + os.remove(path) + logger.info(f"mooncake-store: withdrew the master address at {path}") + + +def _address_files(run_dir: str, extra: Optional[str] = None) -> List[str]: + """Where a master this process starts should publish its address. + + Always the run directory, plus wherever the deployment asked for. + """ + paths = [os.path.join(run_dir, MASTER_ADDRESS_NAME)] + if extra and os.path.abspath(extra) not in {os.path.abspath(p) for p in paths}: + paths.append(extra) + return paths + + +@contextlib.contextmanager +def running_master( + pool: PoolSpec, run_dir: str, address_file: Optional[str] = None +) -> Iterator[LaunchedMaster]: + """Run a master whose lifetime is this process's rather than an engine's. + + `provision_pool` covers the server that owns its pool. Several engines on + one pool, or a pool that has to survive a restart, need the master + somewhere that is not any of them. + + `address_file` receives `host:port` once the master answers, so workers can + name the file instead of an address nobody knows until the scheduler has + placed this process. One is written to `run_dir` either way. + """ + os.makedirs(run_dir, exist_ok=True) + master = _launch_master(pool, run_dir) + try: + with _published_address(master.address, _address_files(run_dir, address_file)): + yield master + finally: + master.stop() + logger.info(f"mooncake-store: master at {master.address} stopped") + + +@contextlib.contextmanager +def provision_pool(pool: PoolSpec, run_dir: Optional[str] = None) -> Iterator[Optional[str]]: + """Make `pool` reachable and name it in this process's environment. + + Yields the path of the client config written, or `None` when an inherited + `MOONCAKE_CONFIG_PATH` was left in charge. + + Args: + pool: The pool to bring up. + run_dir: Where to write the client config and the master's log. + Defaults to `TRTLLM_MOONCAKE_RUN_DIR`, else a temporary directory + that is removed on exit. + """ + _check_exactly_one_master(pool) + + inherited = os.getenv(CONFIG_PATH_ENV) + if inherited: + logger.info( + f"mooncake-store: {CONFIG_PATH_ENV}={inherited} is already set, so " + "the pool it names is used as is and nothing is provisioned here." + ) + yield None + return + + keep_run_dir = bool(run_dir or os.getenv(RUN_DIR_ENV)) + run_dir = run_dir or os.getenv(RUN_DIR_ENV) or tempfile.mkdtemp(prefix="trtllm-mooncake-") + os.makedirs(run_dir, exist_ok=True) + if keep_run_dir: + logger.info(f"mooncake-store: provisioning the pool, run directory {run_dir}") + else: + logger.info( + f"mooncake-store: provisioning the pool in {run_dir}, which is " + f"removed at shutdown along with the master's log; set " + f"{RUN_DIR_ENV} to keep them" + ) + + master: Optional[LaunchedMaster] = None + exported = False + with contextlib.ExitStack() as stack: + try: + if pool.launch_master: + master = _launch_master(pool, run_dir) + master_address = master.address + # Published even when only this server uses it, since that is + # how its donors reach it. + stack.enter_context( + _published_address( + master_address, _address_files(run_dir, pool.master_address_file) + ) + ) + else: + master_address = resolve_master_address( + pool.master_server_address, master_timeout() + ) + wait_for_master(master_address) + logger.info(f"mooncake-store: using the master at {master_address}") + + config_path = write_client_config(pool, master_address, run_dir) + # Inherited by the ranks the LLM constructor spawns. Ranks an + # external launcher started were already running, so they read the + # config out of the run directory instead; see + # provisioned_config_path. + os.environ[CONFIG_PATH_ENV] = config_path + exported = True + logger.info(f"mooncake-store: {CONFIG_PATH_ENV}={config_path}") + # Capacity is what explains a low hit rate, so state the + # arithmetic instead of leaving it to be derived later. + logger.info( + "mooncake-store: each rank that opens a handle will contribute " + f"global_segment_size={pool.global_segment_size} to the pool; " + "total capacity is that times the number of such ranks, plus " + "whatever any donor lends" + ) + yield config_path + finally: + if exported: + os.environ.pop(CONFIG_PATH_ENV, None) + if master is not None: + master.stop() + logger.info(f"mooncake-store: master at {master.address} stopped") + if not keep_run_dir: + shutil.rmtree(run_dir, ignore_errors=True) + + +def _check_exactly_one_master(pool: PoolSpec) -> None: + """Reject a pool that names no master, or two. + + Checked here rather than on `PoolSpec` because a caller may fill the two + fields in from different places, such as a command line and a config file, + and is entitled to an incomplete spec until it has read both. + """ + if pool.launch_master and pool.master_server_address: + raise ValueError( + "mooncake-store: set either launch_master or master_server_address, " + "not both. launch_master starts a master here; " + "master_server_address joins an existing pool." + ) + if not pool.launch_master and not pool.master_server_address: + raise ValueError( + "mooncake-store: needs a master. Set master_server_address to join " + "an existing pool, or launch_master to start one." + ) + if pool.master_address_file and not pool.launch_master: + raise ValueError( + "mooncake-store: master_address_file publishes the address of a " + "master started here, so it needs launch_master. To read an " + "address a master elsewhere published, set master_server_address " + "to file://." + ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py new file mode 100644 index 000000000000..f1eb2db7ee14 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The per-iteration work list the scheduler hands the workers. + +Instances are broadcast from rank 0 to every worker, so these carry only plain +data: a page's identity (its block hash) and where that page currently lives on +this rank (a layer group and a page slot index). Deliberately no store keys: +each worker prefixes its own rank namespace, so one broadcast serves all shards. +""" + +from dataclasses import dataclass, field +from typing import List + +__all__ = ["MooncakeStoreMetadata", "PageTransfer", "RequestTransfers"] + + +@dataclass +class PageTransfer: + """One page of one layer group, to move in either direction.""" + + #: Content identity from `BlockHashChain`; names the key, not the location. + block_hash: bytes + layer_group_id: int + #: Page slot index within `layer_group_id`, as reported by + #: `RequestData.new_block_ids_by_layer_group`. + page_index: int + + +@dataclass +class RequestTransfers: + """Pages belonging to one request, kept together for save bookkeeping. + + The worker owes `get_finished` an answer per request, so a save's owner has + to survive the trip from scheduler to worker. + """ + + request_id: int + pages: List[PageTransfer] = field(default_factory=list) + + +@dataclass +class MooncakeStoreMetadata: + """Loads to perform before the next forward pass, saves to start after it.""" + + loads: List[RequestTransfers] = field(default_factory=list) + saves: List[RequestTransfers] = field(default_factory=list) + + def __bool__(self) -> bool: + """Whether there is any work at all this iteration.""" + return bool(self.loads or self.saves) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py new file mode 100644 index 000000000000..98e3b61a0a62 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Pinned host slots that stand in for GPU pages when the pool cannot reach them. + +The connector's default path registers the KV pools themselves with Mooncake, so +the store reads and writes device memory directly. That needs the HCA to be able +to pin GPU pages, which means GPUDirect RDMA via `nvidia_peermem` or dma-buf. +Where that is unavailable, `ibv_reg_mr` fails on every pool range and the +connector cannot start. + +Staging trades a copy for that dependency. Mooncake is given a pinned host +buffer instead of the pools, and each page passes through a slot in it: gathered +from its device regions before a write, scattered back to them after a read. The +store then only ever registers host memory. + +A slot holds the page's regions concatenated in region order, which is exactly +the payload the zero-copy path produces from the same regions. The stored bytes +are therefore identical either way, so a pool written by one path is readable by +the other, including by another engine sharing the pool. + +Copies go through `cudaMemcpyAsync` rather than the batched Triton kernel in +`disaggregation/native/bounce/gather_scatter.py`. That kernel is the better tool +for device-to-device gather, but here one side is host memory, which the copy +engines move over the host link by DMA. +""" + +from typing import List, Optional, Sequence, Tuple + +import torch + +try: + from cuda.bindings import runtime as cudart +except ImportError: + from cuda import cudart + +from tensorrt_llm._utils import CUASSERT +from tensorrt_llm.logger import logger + +__all__ = ["HostStagingPool", "plan_slot_geometry", "sync_stream"] + +#: Stated explicitly rather than inferred from the pointers, which would be +#: wrong for a host pointer outside the unified address space. +_DEVICE_TO_HOST = cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost +_HOST_TO_DEVICE = cudart.cudaMemcpyKind.cudaMemcpyHostToDevice + + +def _memcpy_async(dst: int, src: int, size: int, kind, stream: int) -> None: + """One asynchronous copy between a device page and a host slot.""" + status = cudart.cudaMemcpyAsync(int(dst), int(src), int(size), kind, stream)[0] + if status == cudart.cudaError_t.cudaSuccess: + return + # Raised here rather than through CUASSERT so the operands appear in the + # message; a bare cudaErrorInvalidValue names no cause. + device = torch.cuda.current_device() if torch.cuda.is_available() else None + raise RuntimeError( + f"cudaMemcpyAsync failed with {status} staging a KV page: " + f"dst={int(dst):#x} src={int(src):#x} size={size} " + f"stream={int(stream):#x} current_device={device}. An invalid value here " + "is usually a stream created on a different device than the pages, which " + "happens when a thread issues the copy without inheriting the rank's " + "device, since torch's current device is thread-local." + ) + + +def sync_stream(stream: int) -> None: + """Wait for a stream's copies to finish, given its raw handle.""" + CUASSERT(cudart.cudaStreamSynchronize(stream)) + + +def plan_slot_geometry( + max_bytes_per_page: int, + transfer_batch_size: int, + budget_bytes: int, +) -> Tuple[int, int]: + """Choose how many pages may be staged at once, and how wide a slot is. + + A slot has to hold the largest page any layer group produces, so the page size + is a floor on the allocation: a budget below one page is raised to one rather + than refused, since the alternative is not starting. + + Args: + max_bytes_per_page: Largest page payload across layer groups. + transfer_batch_size: Pages the connector puts in one store call. There is + no point staging more than that. + budget_bytes: Ceiling on this pool's pinned allocation. + + Returns: + Slot width in bytes, and the number of slots. + """ + if max_bytes_per_page <= 0: + raise ValueError(f"max_bytes_per_page must be > 0, got {max_bytes_per_page}") + if transfer_batch_size <= 0: + raise ValueError(f"transfer_batch_size must be > 0, got {transfer_batch_size}") + + affordable = budget_bytes // max_bytes_per_page + num_slots = max(1, min(transfer_batch_size, affordable)) + return max_bytes_per_page, num_slots + + +class HostStagingPool: + """A registered pinned buffer, sliced into per-page slots. + + One pool serves one direction. Loads run on the executor thread and saves + on the connector's background thread, so sharing slots between them would + need a lock on the transfer path for no benefit. + """ + + def __init__( + self, + *, + slot_bytes: int, + num_slots: int, + store, + label: str, + ): + self._slot_bytes = int(slot_bytes) + self._num_slots = int(num_slots) + self._store = store + self._label = label + + # Page-locking is a correctness requirement here rather than a + # copy-speed preference: this memory is handed to the store to register. + pin = torch.cuda.is_available() + self._buffer = torch.empty( + self._slot_bytes * self._num_slots, dtype=torch.uint8, pin_memory=pin + ) + self._base = int(self._buffer.data_ptr()) + + status = store.register_buffer(self._base, self._buffer.numel()) + if status != 0: + raise RuntimeError( + f"MooncakeDistributedStore.register_buffer failed with status " + f"{status} for the {label} host staging buffer at " + f"[{self._base:#x}, {self._base + self._buffer.numel():#x}). Host " + f"memory registration failing points at the pool or the fabric " + f"rather than at GPUDirect, which is what staging avoids." + ) + logger.info( + f"mooncake-store {label} staging: {self._num_slots} slots x " + f"{self._slot_bytes} B = {self._buffer.numel() / 1024**2:.1f} MiB pinned " + f"(pinned={pin})" + ) + + def close(self) -> None: + """Hand the registration back before the buffer is freed. + + The store registers an address range rather than the tensor, so a buffer + freed while it still holds one leaves the fabric able to reach memory the + allocator has since handed out again. Called once the connector's pending + transfers have drained. Idempotent, and a failed unregistration keeps the + buffer rather than freeing memory the store may still reach. + """ + if self._buffer is None: + return + status = self._store.unregister_buffer(self._base) + if status != 0: + raise RuntimeError( + f"MooncakeDistributedStore.unregister_buffer failed with status " + f"{status} for the {self._label} host staging buffer at " + f"{self._base:#x}. The buffer is kept alive rather than freed." + ) + self._buffer = None + + @property + def num_slots(self) -> int: + """Pages this pool can hold at once.""" + return self._num_slots + + @property + def slot_bytes(self) -> int: + """Capacity of one slot.""" + return self._slot_bytes + + def slot_address(self, index: int) -> int: + """Address of slot `index`.""" + if not 0 <= index < self._num_slots: + raise IndexError(f"slot {index} out of range [0, {self._num_slots})") + return self._base + index * self._slot_bytes + + def _check_fits(self, total: int) -> None: + if total > self._slot_bytes: + raise ValueError( + f"a {total} B page does not fit the {self._slot_bytes} B " + f"{self._label} staging slot; the pool was sized from the layout's " + "largest page, so this means the layout changed after registration" + ) + + def gather( + self, + index: int, + addresses: Sequence[int], + sizes: Sequence[int], + stream: int, + ) -> Tuple[int, int]: + """Copy one page's device regions into slot `index`, concatenated. + + Args: + index: Slot to fill. + addresses: Device addresses of the page's regions, in region order. + sizes: Byte counts matching `addresses`. + stream: CUDA stream handle the copies are issued on. + + Returns: + The slot's address and the total bytes written, ready to hand to the + store as a single buffer. + """ + total = sum(sizes) + self._check_fits(total) + destination = self.slot_address(index) + offset = 0 + for address, size in zip(addresses, sizes, strict=True): + _memcpy_async(destination + offset, address, size, _DEVICE_TO_HOST, stream) + offset += size + return destination, total + + def scatter( + self, + index: int, + addresses: Sequence[int], + sizes: Sequence[int], + stream: int, + ) -> None: + """Copy slot `index` back out to one page's device regions. + + The inverse of :meth:`gather`, walking the regions in the same order so + the split matches the concatenation the slot holds. + """ + self._check_fits(sum(sizes)) + source = self.slot_address(index) + offset = 0 + for address, size in zip(addresses, sizes, strict=True): + _memcpy_async(address, source + offset, size, _HOST_TO_DEVICE, stream) + offset += size + + def reserve(self, total: int) -> None: + """Assert a page of `total` bytes is stageable, without copying.""" + self._check_fits(total) + + +def stage_batch_for_put( + pool: HostStagingPool, + addresses: Sequence[Sequence[int]], + sizes: Sequence[Sequence[int]], + stream: int, +) -> Tuple[List[List[int]], List[List[int]]]: + """Gather a batch of device pages into slots and describe them for the store. + + Args: + pool: Slots to stage through. The batch must not exceed its slot count. + addresses: Per-page device region addresses. + sizes: Per-page device region sizes. + stream: Stream the copies are issued on. The caller must synchronize it + before the store reads the slots. + + Returns: + Per-page address and size lists, each a single staged buffer. + """ + if len(addresses) > pool.num_slots: + raise ValueError(f"batch of {len(addresses)} pages exceeds {pool.num_slots} staging slots") + staged_addresses: List[List[int]] = [] + staged_sizes: List[List[int]] = [] + for index, (page_addresses, page_sizes) in enumerate(zip(addresses, sizes, strict=True)): + slot, total = pool.gather(index, page_addresses, page_sizes, stream) + staged_addresses.append([slot]) + staged_sizes.append([total]) + return staged_addresses, staged_sizes + + +def describe_batch_for_get( + pool: HostStagingPool, + sizes: Sequence[Sequence[int]], +) -> Tuple[List[List[int]], List[List[int]]]: + """Describe slots for the store to read a batch into, before scattering. + + Unlike the put direction there is nothing to copy first: the slots are the + destination, and :func:`unstage_batch_after_get` moves the bytes on once the + store has filled them. + """ + if len(sizes) > pool.num_slots: + raise ValueError(f"batch of {len(sizes)} pages exceeds {pool.num_slots} staging slots") + staged_addresses: List[List[int]] = [] + staged_sizes: List[List[int]] = [] + for index, page_sizes in enumerate(sizes): + total = sum(page_sizes) + pool.reserve(total) + staged_addresses.append([pool.slot_address(index)]) + staged_sizes.append([total]) + return staged_addresses, staged_sizes + + +def unstage_batch_after_get( + pool: HostStagingPool, + addresses: Sequence[Sequence[int]], + sizes: Sequence[Sequence[int]], + stream: int, + only: Optional[Sequence[int]] = None, +) -> None: + """Scatter filled slots back to their device pages. + + Args: + pool: Slots the store just wrote into. + addresses: Per-page device region addresses. + sizes: Per-page device region sizes. + stream: Stream the copies are issued on. The caller must synchronize it + before the pages are read. + only: Slot indices to scatter. Defaults to all of them; a caller that + knows some reads failed passes the rest so a failed page is not + written over its device slot with whatever the slot held. + """ + indices = range(len(addresses)) if only is None else only + for index in indices: + pool.scatter(index, addresses[index], sizes[index], stream) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 371d7ed85327..e4e619c367af 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3775,6 +3775,45 @@ def resume_request(self, req: LlmRequest) -> bool: return False return self._resume_and_restore(req.py_request_id, kv_cache) + # ---- preemption ---- + # + # Suspension only unpins pages; the eviction controller then migrates them + # one cache level down. With GPU as the last level a suspended page stays + # `HELD`, which `CacheLevelManager.is_evictable` refuses to evict, so + # suspension frees nothing and the scheduler has no way out of a full pool. + # + # Preemption is the fallback for that case. It gives the pages up instead + # of parking them, which costs a re-prefill but always works. + + @property + def has_cache_tier_below_gpu(self) -> bool: + """True when a suspended page has somewhere to be evicted to.""" + return len(self.impl.cache_tier_list) > 1 + + def preempt_request(self, req: LlmRequest) -> bool: + """Give up *req*'s KV cache so its pages can be reclaimed. + + Unlike :meth:`suspend_request` this does not keep the pages. Closing + the request's `_KVCache` returns its committed blocks to the radix tree + as reusable prefix and leaves their pages `DROPPABLE`, which is + evictable at every level, unlike `HELD`. The data is not thrown away: + it stays resident and locally matchable until something else needs the + space. + + The request is reset to context state by the caller and re-prefills + whatever it can no longer match. + + Returns whether the pages were released. Callers must check: pages + still being read out of, such as by a connector save in flight, are + not released until that completes. + """ + self._release_preempted(req) + return True + + def _release_preempted(self, req: LlmRequest) -> None: + self.free_resources(req) + req.py_num_connector_matched_tokens = 0 + # ---- prepare_resources ---- def _life_cycle_by_layer_group(self) -> List[AttnLifeCycle]: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 3fc930fa8a48..01eaf0d10e1f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -745,6 +745,11 @@ def __init__( self.resource_manager, should_store_blocks=self.enable_disagg_partial_reuse_store) + # Wire the transfer manager into the V2 scheduler's deadlock detector. + if hasattr(self.scheduler, "set_async_transfer_manager"): + self.scheduler.set_async_transfer_manager( + self.async_transfer_manager) + # Router is built after async_transfer_manager so KVCacheAwareADPRouter # can receive the transfer-manager reference at construction time. self.adp_router: ADPRouter = ADPRouter.create( diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 2929fe386200..c11210f6b5ed 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -15,7 +15,8 @@ import enum import os -from typing import Optional +from collections import Counter +from typing import Callable, Optional from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, ContextChunkingPolicy from tensorrt_llm.logger import logger @@ -170,6 +171,7 @@ def __init__( enable_recompute_pause: bool = True, ) -> None: self.max_num_tokens = max_num_tokens + self._stalled_schedules = 0 self.max_num_requests = ( scheduler_capacity if scheduler_capacity is not None else max_batch_size ) @@ -255,6 +257,20 @@ def __init__( os.environ.get("TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN", "0") == "1" ) + # Registered by PyExecutor; see set_async_transfer_manager. + self._async_transfer_manager = None + + def set_async_transfer_manager(self, mgr) -> None: + """Register the AsyncTransferManager the deadlock detector consults. + + A finished disaggregated context sender leaves active_requests once its + response is emitted, but the transfer manager still owns it and its + pinned KV pages until the send lands. Without this reference the + detector cannot see those pages and may report a false deadlock on a + context server whose pool is full of in-flight sends. + """ + self._async_transfer_manager = mgr + @property def scheduling_state_range( self, @@ -443,6 +459,27 @@ def _schedule_loop(self, active_requests, inflight_request_ids): req_it += 1 + # Requests whose pages were given up during this pass. A victim that is + # itself a started context request still sits in pending_ctx, so + # re-admitting it would spend the pages its own preemption released. + preempted_ids: set[int] = set() + + def preempt_for_pages(req: LlmRequest) -> bool: + """Free pages for `req` by giving up one started request. + + A success ends the phase 2 loop, reserving the pages for the + request that paid a re-prefill for them. Letting a later context + request take them instead would leave `req` to preempt again on + the next pass, repeating without ever admitting it. The cost is + one iteration of admission. + """ + protected = {r.py_request_id for r in scheduled_gen} + protected.update(r.py_request_id for r in scheduled_ctx) + protected.add(req.py_request_id) + return self._try_preempt_for_pages( + requests_list, protected, inflight_request_ids, recompute_paused, preempted_ids + ) + # --- Phase 2: schedule deferred context / encoder requests --- # Generation PEFT pages are now fully committed in the budget. # @@ -456,6 +493,11 @@ def _schedule_loop(self, active_requests, inflight_request_ids): contributed_blocks = self._collect_contributed_blocks( requests_list, pending_ctx, inflight_request_ids ) + # A deferral behind an in-flight contributor leaves nothing on any of + # the scheduled lists, so the deadlock detector below would read the + # iteration as a stall even though the contributor is running. Deferring + # is itself the progress in that case. + deferred_behind_contributor = False for req in pending_ctx: if budget.requests_full: @@ -468,6 +510,8 @@ def _schedule_loop(self, active_requests, inflight_request_ids): and not self._has_context_chunk_budget(budget) ): continue + if req.py_request_id in preempted_ids: + continue # Probe context requests before peft_pages_needed and before # _try_schedule_context so that a deferral costs nothing: KV pages # are allocated inline, so a skip decided after prepare_context @@ -488,6 +532,7 @@ def _schedule_loop(self, active_requests, inflight_request_ids): f"Deferring context request {req.py_request_id}: its first new " "block is already contributed by a request that runs this iteration" ) + deferred_behind_contributor = True continue peft_pages = budget.peft_pages_needed(req) if peft_pages is None: @@ -499,7 +544,9 @@ def _schedule_loop(self, active_requests, inflight_request_ids): scheduled_encoder.append(req) budget.commit(req, tokens, peft_pages) else: - action, tokens, chunking_flag = self._try_schedule_context(req, budget) + action, tokens, chunking_flag = self._try_schedule_context( + req, budget, preempt_for_pages + ) if action is ScheduleAction.STOP: break if action is ScheduleAction.SKIP: @@ -521,50 +568,21 @@ def _schedule_loop(self, active_requests, inflight_request_ids): if first_new_block is not None: contributed_blocks.add(first_new_block) - # Deadlock detection: if generation requests exist but none were - # scheduled and none were evicted, no forward pass will run and no - # KV cache pages will ever be freed — the scheduler will spin - # forever. This typically happens when the KV cache pool is exhausted - # and no secondary cache tier is available for suspend/resume. - if not scheduled_gen and not scheduled_ctx: - num_gen_candidates = sum( - 1 - for r in active_requests - if r.is_generation_in_progress_state - and not r.is_generation_to_complete_state - and r.request_id not in inflight_request_ids - and not self._has_pending_connector_load(r) - ) - if ( - num_gen_candidates > 0 - and not evicted - and not recompute_paused - and not inflight_request_ids - and not any(self._has_pending_connector_load(r) for r in active_requests) - ): - # A connector rejects every tier below GPU at bring-up - # (`PyExecutor._reject_non_gpu_cache_tiers`), so offering those - # two settings there is advice the user cannot act on. - if getattr(self.kv_cache_manager, "kv_connector_manager", None) is not None: - remedy = ( - "A KV connector is attached, which requires a GPU-only cache, " - "so no secondary tier can be configured. Increase " - "kv_cache_config.max_tokens or kv_cache_config." - "free_gpu_memory_fraction, or lower max_num_tokens to hand " - "memory back to the KV pool." - ) - else: - remedy = ( - "Configure kv_cache_config.host_cache_size or " - "kv_cache_config.disk_cache_size, or increase " - "kv_cache_config.max_tokens." - ) - raise RuntimeError( - f"V2 scheduler deadlock: {num_gen_candidates} generation " - f"request(s) active but none could be scheduled or " - f"evicted or recompute-paused. KV cache pool is likely exhausted with no " - f"secondary cache tier for suspend/resume offload. {remedy}" - ) + self._detect_deadlock( + active_requests, + inflight_request_ids, + pending_ctx, + preempted_ids, + made_progress=bool( + scheduled_gen + or scheduled_ctx + or scheduled_encoder + or disagg_candidates + or evicted + or recompute_paused + or deferred_behind_contributor + ), + ) return ( scheduled_encoder, @@ -734,7 +752,10 @@ def _try_schedule_disagg_gen_init( return ScheduleAction.SCHEDULED, 0 def _try_schedule_context( - self, req: LlmRequest, budget: BudgetTracker + self, + req: LlmRequest, + budget: BudgetTracker, + preempt_for_pages: Callable[[LlmRequest], bool], ) -> tuple[ScheduleAction, int, bool]: """Try to schedule a context request (chunked or non-chunked). @@ -747,9 +768,9 @@ def _try_schedule_context( """ first_chunk = req.is_first_context_chunk if self.chunking_enabled: - result = self._try_schedule_context_chunked(req, budget) + result = self._try_schedule_context_chunked(req, budget, preempt_for_pages) else: - result = self._try_schedule_context_full(req, budget) + result = self._try_schedule_context_full(req, budget, preempt_for_pages) if first_chunk and result[0] is not ScheduleAction.SCHEDULED: # Failed admission must not retain prefix-reuse holds. Suspension @@ -766,7 +787,10 @@ def _try_schedule_context( return result def _try_schedule_context_full( - self, req: LlmRequest, budget: BudgetTracker + self, + req: LlmRequest, + budget: BudgetTracker, + preempt_for_pages: Callable[[LlmRequest], bool], ) -> tuple[ScheduleAction, int, bool]: """Try to schedule a non-chunked context request. @@ -805,6 +829,12 @@ def _try_schedule_context_full( # V2 resizes KV cache directly in the scheduler (no separate # prepareResources for main cache), so include draft tokens. if not self._try_allocate_context(req, context_tokens + draft_len): + # Out of pages. Give up one started request so this one can + # proceed, and retry next iteration: a failed resize leaves a + # first chunk suspended, so the retry has to go back through + # prepare_context to resume it. + if preempt_for_pages(req): + return ScheduleAction.STOP, 0, False return ScheduleAction.SKIP, 0, False cross_action = self._try_schedule_cross_context(req) @@ -825,7 +855,10 @@ def _has_context_chunk_budget(self, budget: BudgetTracker) -> bool: ) def _try_schedule_context_chunked( - self, req: LlmRequest, budget: BudgetTracker + self, + req: LlmRequest, + budget: BudgetTracker, + preempt_for_pages: Callable[[LlmRequest], bool], ) -> tuple[ScheduleAction, int, bool]: """FCFS interleaved chunking for a single context request. @@ -883,6 +916,9 @@ def _try_schedule_context_chunked( chunk_size = (chunk_size // self.chunk_unit_size) * self.chunk_unit_size if chunk_size <= 0: + # Out of token budget rather than out of pages, so releasing pages + # would not help; the next iteration gets a fresh budget. Not + # suspended either, to avoid pathological suspend/resume cycles. return ScheduleAction.SKIP, 0, False chunk_size = self._align_chunk_to_mm_block( @@ -910,6 +946,9 @@ def _try_schedule_context_chunked( # V2 resizes KV cache directly in the scheduler, so include # draft tokens for last chunk. if not self._try_allocate_context(req, resize_tokens): + # Out of pages, as in _try_schedule_context_full. + if preempt_for_pages(req): + return ScheduleAction.STOP, 0, False return ScheduleAction.SKIP, 0, False cross_action = self._try_schedule_cross_context(req) @@ -1375,6 +1414,246 @@ def _suspend_request(self, req: LlmRequest) -> None: def _clear_request_runtime_state(self, req: LlmRequest) -> None: req.py_batch_idx = None + def _try_preempt_for_pages( + self, + requests_list: RequestList, + protected_ids: set[int], + inflight_request_ids: set[int], + recompute_paused: RequestList, + preempted_ids: set[int], + ) -> bool: + """Release one started request's KV cache so another can allocate. + + The fallback for a pool that suspension cannot drain; see + `KVCacheManagerV2.preempt_request`. With a cache tier below GPU, + suspension is cheaper and keeps the pages, so that path is left alone. + + The victim leaves on `recompute_paused`, the same channel the generation + side uses, because a re-prefill needs more teardown than the KV cache: + the executor frees the request's remaining resources, its sequence slot + included, and `reset_for_recompute` rewrites the prompt and resyncs the + Python-side mirrors of it. Pausing the request here instead would leave + the slot owned by SeqSlotManager while `py_seq_slot` is None, which + asserts on the next schedule. + + Returns True when pages became available in this iteration. + """ + if self.kv_cache_manager.has_cache_tier_below_gpu: + return False + # A disaggregated generation worker received its context KV rather + # than computing it, so it cannot replay a prefill at all. + if not self.enable_recompute_pause: + return False + + # Newest first, so the requests closest to completing keep their + # pages and the pool drains instead of thrashing. + for i in range(len(requests_list) - 1, -1, -1): + victim = requests_list[i] + if victim.py_request_id in protected_ids: + continue + if not self._is_recompute_pause_candidate(victim, inflight_request_ids): + continue + if not self.kv_cache_manager.is_request_active(victim.py_request_id): + continue + + self.kv_cache_manager.preempt_request(victim) + logger.debug( + f"[V2Scheduler] Preempting request {victim.py_request_id} " + f"(state={victim.state.name})" + ) + self._clear_request_runtime_state(victim) + if self.draft_kv_cache_manager is not None: + self.draft_kv_cache_manager.free_resources(victim) + recompute_paused.append(victim) + preempted_ids.add(victim.py_request_id) + return True + + return False + + # Consecutive scheduling passes that reclaimed nothing before this counts + # as a deadlock. A stalled pass costs ~2ms, so it trips within seconds, + # while transient one-iteration deferrals (multimodal chunk alignment, + # PEFT budget, IndexMapper slots) clear long before. + _DEADLOCK_STALL_ITERS = 1000 + + # States in which an in-flight KV transfer still owns pages it is about to + # release: a context-only request sending its cache after prefill, one + # whose send landed and which the executor has yet to reap, and a + # generation request receiving a cache. + _TRANSFER_HOLDING_STATE_VALUES = frozenset( + { + LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS.value, + LlmRequestState.DISAGG_CONTEXT_COMPLETE.value, + LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS.value, + } + ) + + def _detect_deadlock( + self, + active_requests: RequestList, + inflight_request_ids: set[int], + pending_ctx: RequestList, + preempted_ids: set[int], + made_progress: bool, + ) -> None: + """Fail loudly when no request can be scheduled or reclaimed. + + Without this the executor spins at full speed while scheduling + nothing, which looks healthy to the hang detector and to `/health` + while the job burns its wall clock. Context candidates count alongside + generation ones because a disaggregated prefill server has no + generation requests at all. + """ + if made_progress: + if self._stalled_schedules: + logger.debug( + f"[V2Scheduler] Stall cleared after {self._stalled_schedules} pass(es); " + "scheduling resumed." + ) + self._stalled_schedules = 0 + return + + num_gen_candidates = sum( + 1 + for r in active_requests + if r.is_generation_in_progress_state + and not r.is_generation_to_complete_state + and r.request_id not in inflight_request_ids + and not self._has_pending_connector_load(r) + ) + num_ctx_candidates = sum( + 1 + for r in pending_ctx + if r.py_request_id not in preempted_ids and r.request_id not in inflight_request_ids + ) + if num_gen_candidates == 0 and num_ctx_candidates == 0: + # Legitimately idle: nothing to schedule. + if self._stalled_schedules: + logger.debug( + f"[V2Scheduler] Stall cleared after {self._stalled_schedules} pass(es); " + "no candidates remain." + ) + self._stalled_schedules = 0 + return + + # A connector load in flight releases its pages when it lands, so a + # pass that reclaims nothing while one is outstanding is not a stall. + if any(self._has_pending_connector_load(r) for r in active_requests): + self._stalled_schedules = 0 + return + + # Waiting on a transfer is not a deadlock: the pages come back when it + # lands. None of those states are schedulable, so a context server + # whose pool is full of pending sends shows no progress here at all. A + # send that never lands is the transfer layer's timeout to report. + # + # The active_requests scan covers senders still on the list and the + # generation-side receiver. The transfer manager covers finished + # senders that have already left active_requests (see + # set_async_transfer_manager). + transfer_holding = any( + req.state_value in self._TRANSFER_HOLDING_STATE_VALUES for req in active_requests + ) or ( + self._async_transfer_manager is not None + and self._async_transfer_manager.has_any_inflight_requests() + ) + if transfer_holding: + if self._stalled_schedules: + logger.debug( + f"[V2Scheduler] {num_gen_candidates} generation and " + f"{num_ctx_candidates} context request(s) cannot allocate, but KV " + "transfers in flight will release pages; stall cleared." + ) + self._stalled_schedules = 0 + return + + self._stalled_schedules += 1 + if self._stalled_schedules < self._DEADLOCK_STALL_ITERS: + # Warn as the stall builds so the raise is not a surprise. A quarter + # of the threshold keeps this to a handful of lines beforehand. + if self._stalled_schedules == 1: + logger.debug( + "[V2Scheduler] Scheduling stall started: " + f"{self._summarize_stall(active_requests, pending_ctx, inflight_request_ids, preempted_ids)}" + ) + elif self._stalled_schedules % max(1, self._DEADLOCK_STALL_ITERS // 4) == 0: + logger.warning( + f"[V2Scheduler] Stalled {self._stalled_schedules}/" + f"{self._DEADLOCK_STALL_ITERS} consecutive passes: " + f"{self._summarize_stall(active_requests, pending_ctx, inflight_request_ids, preempted_ids)}" + ) + return + + # A connector rejects every tier below GPU at bring-up + # (`PyExecutor._reject_non_gpu_cache_tiers`), so offering host_cache_size + # there is advice the user cannot act on. + if getattr(self.kv_cache_manager, "kv_connector_manager", None) is not None: + remedy = ( + "A KV connector is attached, which requires a GPU-only cache, " + "so no secondary tier can be configured. Increase " + "kv_cache_config.max_tokens or " + "kv_cache_config.free_gpu_memory_fraction, or lower " + "max_num_tokens to hand memory back to the KV pool." + ) + else: + remedy = ( + "Configure kv_cache_config.host_cache_size, increase " + "kv_cache_config.max_tokens, or lower max_batch_size." + ) + # Logged before the raise so the diagnostic survives even if the + # exception is caught or truncated upstream. + logger.warning( + f"[V2Scheduler] Declaring deadlock after {self._stalled_schedules} passes: " + f"{self._summarize_stall(active_requests, pending_ctx, inflight_request_ids, preempted_ids)} " + f"{remedy}" + ) + raise RuntimeError( + f"V2 scheduler deadlock: {num_gen_candidates} generation and " + f"{num_ctx_candidates} context request(s) active but none could " + f"be scheduled, suspended or preempted in " + f"{self._stalled_schedules} consecutive attempts. The KV cache " + f"pool is likely exhausted. {remedy}" + ) + + def _summarize_stall( + self, + active_requests: RequestList, + pending_ctx: RequestList, + inflight_request_ids: set[int], + preempted_ids: set[int], + ) -> str: + """One-line breakdown of a stalled pass for the deadlock logs. + + Re-derives the blocked candidates so the per-pass detector stays cheap: + this runs only on the rare logging branches. + """ + gen = [ + r + for r in active_requests + if r.is_generation_in_progress_state + and not r.is_generation_to_complete_state + and r.request_id not in inflight_request_ids + ] + ctx = [ + r + for r in pending_ctx + if r.py_request_id not in preempted_ids and r.request_id not in inflight_request_ids + ] + states = Counter(LlmRequestState(r.state_value).name for r in gen + ctx) + sample_ids = [r.py_request_id for r in (gen + ctx)[:5]] + has_connector = getattr(self.kv_cache_manager, "kv_connector_manager", None) is not None + transfers_in_flight = ( + self._async_transfer_manager is not None + and self._async_transfer_manager.has_any_inflight_requests() + ) + return ( + f"{len(gen)} generation and {len(ctx)} context request(s) blocked; " + f"states={dict(states)}; sample_ids={sample_ids}; " + f"kv_connector={has_connector}; " + f"cache_tier_below_gpu={self.kv_cache_manager.has_cache_tier_below_gpu}; " + f"transfers_in_flight={transfers_in_flight}" + ) + def _free_kv_caches(self, req: LlmRequest) -> None: self.kv_cache_manager.free_resources(req) if self.draft_kv_cache_manager is not None: diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py new file mode 100644 index 000000000000..47b3b8e32fcd --- /dev/null +++ b/tensorrt_llm/commands/mooncake.py @@ -0,0 +1,415 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The two pieces of a Mooncake pool that are not any one engine's to run. + +`mooncake_master` owns a pool for as long as the command runs, and with +`--config` also writes the client config its workers read into `--run_dir`. +Pointing every server at that directory with `$TRTLLM_MOONCAKE_RUN_DIR`, or at +the file with `$MOONCAKE_CONFIG_PATH`, is how a pool is described today. + +`mooncake_donor` lends a node's host memory to a pool it does not otherwise +use, which is how capacity comes from nodes whose engines have no connector. +""" + +import contextlib +import json +import os +import signal +import tempfile +import time +from typing import Optional + +import click + +import tensorrt_llm.usage as usage +from tensorrt_llm.commands import _telemetry as _command_telemetry +from tensorrt_llm.logger import logger +from tensorrt_llm.usage.config import UsageContext + +#: How long a donor with heartbeats turned off sleeps between wakeups. Only +#: the signal that ends the command interrupts it, so the value is arbitrary. +_IDLE_POLL_SECONDS = 60.0 + +#: Both commands sit in the telemetry-aware `trtllm-serve` group, so they have +#: to offer the same opt-out the group documents. +_telemetry_option = click.option( + "--telemetry/--no-telemetry", + default=True, + help="Enable or disable anonymous usage telemetry collection.", +) + + +def _apply_cli_telemetry(telemetry: bool) -> None: + """Honor --no-telemetry for a command that reads no config of its own. + + The group already applies the flag it finds in argv, so this matters when + the command is reached through its callback rather than through the CLI. + """ + if telemetry: + return + usage.apply_usage_session_config( + {"disabled": True}, + default_usage_context=UsageContext.CLI_SERVE.value, + component="server", + lifecycle_phase="config_validation", + ) + + +@contextlib.contextmanager +def _signal_handoff(): + """Turn SIGINT and SIGTERM into the exit the telemetry boundary expects. + + Both commands hold a resource, a child process or a mounted segment, whose + release is in a `finally`. Default SIGTERM handling would skip it, leaving + the master unreaped or the pool advertising memory that has gone. + + `raise_signal_exit` unwinds those context managers and carries the signal + number out to `trtllm-serve`, which is what reports the exit as a signal + rather than as a clean one. + + Wrap this around the resource so the log below follows the release. + """ + for received in (signal.SIGINT, signal.SIGTERM): + signal.signal(received, _command_telemetry.raise_signal_exit) + try: + yield + except _command_telemetry.SignalExit as stopping: + # Logged here rather than in the handler, which must not take the + # logging lock. + logger.info(f"mooncake-store: signal {stopping.signal_number} received, shut down") + raise + + +@contextlib.contextmanager +def _announce_ready(ready_file: Optional[str], host: str, segment_size: int): + """Hold a readiness file for exactly as long as the segment is mounted. + + Launchers gate prefill on this file existing, so it has to go when the + capacity does. Left behind, it would clear that gate for a pool the donor + no longer contributes to, and prefill would start writing into a pool + short of the memory the launcher waited for. + """ + if ready_file is None: + yield + return + + with open(ready_file, "w") as handle: + handle.write(f"{host} {segment_size}\n") + logger.info( + f"mooncake-store: announced this segment in {ready_file}, so a " + "launcher waiting on the pool's capacity can proceed" + ) + try: + yield + finally: + # A launcher that consumes the file by removing it has already had the + # signal, so its absence here is a success and not worth failing the + # shutdown over. + with contextlib.suppress(OSError): + os.unlink(ready_file) + logger.info( + f"mooncake-store: withdrew the announcement in {ready_file}; a " + "launcher waiting on this pool's capacity will wait again" + ) + + +@click.command("mooncake_master") +@click.option( + "--rpc_port", + type=int, + default=50051, + show_default=True, + help="Port the store clients reach the master on.", +) +@click.option( + "--metrics_port", + type=int, + default=9004, + show_default=True, + help="Prometheus port. Pool occupancy and eviction are read " + "from here or from the master's log.", +) +@click.option( + "--eviction_ratio", + type=float, + default=0.05, + show_default=True, + help="Fraction of the pool freed per eviction pass.", +) +@click.option( + "--address_file", + type=str, + default=None, + help="File to publish 'host:port' to once the master answers. " + "Workers name it as master_server_address: file://, which " + "is how they reach a master whose host the scheduler chose. " + "Removed on exit so a stale address is never dialed.", +) +@click.option( + "--run_dir", + type=str, + default=None, + help="Where to keep the master's log. Defaults to " + "$TRTLLM_MOONCAKE_RUN_DIR, else a temporary directory.", +) +@click.option( + "--config", + type=str, + default=None, + help="Mooncake JSON config describing the pool. When given, a copy " + "naming this master is written to --run_dir as the client config " + "every worker reads, so no external script has to render one. " + "Workers find it by setting $TRTLLM_MOONCAKE_RUN_DIR to that " + "directory, or $MOONCAKE_CONFIG_PATH to the file.", +) +@click.option( + "--heartbeat_seconds", + type=int, + default=300, + show_default=True, + help="Interval between liveness lines. 0 disables them.", +) +@_telemetry_option +def mooncake_master( + rpc_port: int, + metrics_port: int, + eviction_ratio: float, + address_file: Optional[str], + run_dir: Optional[str], + config: Optional[str], + heartbeat_seconds: int, + telemetry: bool, +): + """Run a mooncake_master for as long as this command runs.""" + _apply_cli_telemetry(telemetry) + + # Imported lazily so other subcommands and --help do not pay for the + # connector package. + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( + PoolSpec, + running_master, + write_client_config, + ) + + raw = {} + if config: + with open(config) as handle: + raw = json.load(handle) + # This command is the master, so whichever one the config names is not + # the one being described here. + raw.pop("master_server_address", None) + + pool = PoolSpec.from_json( + raw, + launch_master=True, + master_port=rpc_port, + master_metrics_port=metrics_port, + master_eviction_ratio=eviction_ratio, + ) + run_dir = ( + run_dir + or os.getenv("TRTLLM_MOONCAKE_RUN_DIR") + or tempfile.mkdtemp(prefix="trtllm-mooncake-master-") + ) + + with _signal_handoff(), running_master(pool, run_dir, address_file=address_file) as master: + if config: + write_client_config(pool, master.address, run_dir) + logger.info( + f"mooncake-store: this master owns the pool until this command " + f"stops; address {master.address}, log {master.log_path}, metrics " + f"http://{master.address.rsplit(':', 1)[0]}:{metrics_port}/metrics" + ) + started = time.monotonic() + announced = started + while True: + if (code := master.process.poll()) is not None: + # The pool is gone once the master dies, and every client is + # about to start failing. + raise click.ClickException( + f"mooncake_master exited with code {code}. See {master.log_path}" + ) + time.sleep(1.0) + now = time.monotonic() + # Distinguishes a dead master from a dead fabric. + if heartbeat_seconds > 0 and now - announced >= heartbeat_seconds: + announced = now + logger.info( + f"mooncake-store: master at {master.address} alive after " + f"{(now - started) / 60:.0f}m" + ) + + +@click.command("mooncake_donor") +@click.option( + "--master_server_address", + type=str, + default=None, + help="Master to join, as host:port or file:// naming a " + "file that holds one. Defaults to the master_server_address in " + "--config.", +) +@click.option( + "--segment_size", + type=str, + default="32GiB", + show_default=True, + help="Host memory to contribute from this node. Deliberately " + "separate from a config's global_segment_size, which is sized " + "for an engine worker rather than a node lending what it can " + "spare.", +) +@click.option( + "--config", + type=str, + default=None, + help="Mooncake JSON config describing the pool, for the " + "settings not given here. Defaults to $MOONCAKE_CONFIG_PATH.", +) +@click.option( + "--protocol", + type=str, + default=None, + help="Transport, 'rdma' or 'tcp'. Defaults to --config's, else rdma.", +) +@click.option( + "--device_name", + type=str, + default=None, + help="RDMA device, from ibv_devinfo. Defaults to --config's.", +) +@click.option( + "--metadata_server", + type=str, + default=None, + help="Mooncake metadata service. Defaults to --config's, else P2PHANDSHAKE.", +) +@click.option( + "--local_buffer_size", + type=str, + default=None, + help="Mooncake transfer buffer for this process. Deliberately " + "separate from a config's local_buffer_size, which is sized for " + "an engine worker: a donor never transfers, and only needs one " + "because setup rejects a zero-sized buffer. Defaults to 64MiB.", +) +@click.option( + "--ready_file", + type=str, + default=None, + help="File to create once the segment is mounted, for launchers " + "that must not let prefill start writing before the pool has " + "this capacity.", +) +@click.option( + "--heartbeat_seconds", + type=int, + default=300, + show_default=True, + help="Interval between liveness lines. 0 disables them.", +) +@_telemetry_option +def mooncake_donor( + master_server_address: Optional[str], + segment_size: str, + config: Optional[str], + protocol: Optional[str], + device_name: Optional[str], + metadata_server: Optional[str], + local_buffer_size: Optional[str], + ready_file: Optional[str], + heartbeat_seconds: int, + telemetry: bool, +): + """Lend this node's host memory to a Mooncake pool, for as long as it runs. + + Running this on the generation nodes puts their memory into the pool while + leaving those engines connector-free. + """ + _apply_cli_telemetry(telemetry) + + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( + DEFAULT_DONOR_LOCAL_BUFFER_SIZE, + donate_segment, + master_timeout, + parse_size, + resolve_master_address, + wait_for_master, + ) + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + CONFIG_PATH_ENV, + DEFAULT_METADATA_SERVER, + ) + + raw = {} + config = config or os.getenv(CONFIG_PATH_ENV) + if config: + with open(config) as handle: + raw = json.load(handle) + + master = master_server_address or raw.get("master_server_address", "") + if not master: + raise click.UsageError( + "No master to join. Pass --master_server_address, or a --config " + f"naming one (or set {CONFIG_PATH_ENV})." + ) + + def size_option(name: str, value: str) -> int: + """Parse a size option, reporting a bad one as a usage error.""" + try: + return parse_size(value) + except ValueError as exc: + raise click.UsageError(f"{name}: {exc}") from exc + + donating = size_option("--segment_size", segment_size) + # None means the option was left off. An empty string was passed, so it goes + # to parse_size and is rejected rather than silently taking the default. + buffer_size = ( + DEFAULT_DONOR_LOCAL_BUFFER_SIZE + if local_buffer_size is None + else size_option("--local_buffer_size", local_buffer_size) + ) + resolved = resolve_master_address(master, master_timeout()) + wait_for_master(resolved) + + with ( + _signal_handoff(), + donate_segment( + resolved, + donating, + protocol=protocol or raw.get("protocol") or "rdma", + device_name=device_name or raw.get("device_name", "") or "", + metadata_server=( + metadata_server or raw.get("metadata_server") or DEFAULT_METADATA_SERVER + ), + local_buffer_size=buffer_size, + ) as host, + # After `as host`, so the announcement names the segment that mounted + # and is withdrawn before the segment is. + _announce_ready(ready_file, host, donating), + ): + # Idle by design: a put or get here would make this node a traffic + # client, which is what donation exists to avoid. + started = time.monotonic() + while True: + if heartbeat_seconds <= 0: + time.sleep(_IDLE_POLL_SECONDS) + continue + time.sleep(heartbeat_seconds) + logger.info( + f"mooncake-store: {host} still lending " + f"{donating / 1024**3:.1f}GiB to the pool at {master} " + f"after {(time.monotonic() - started) / 60:.0f}m" + ) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 2bf9ae8f9a5a..f94e7ffc5688 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -47,6 +47,7 @@ from tensorrt_llm._utils import mpi_rank, set_prometheus_multiproc_dir from tensorrt_llm.commands import _telemetry as _command_telemetry from tensorrt_llm.commands._serve_stability import stability_option +from tensorrt_llm.commands.mooncake import mooncake_donor, mooncake_master from tensorrt_llm.commands.utils import (collect_explicit_cli_keys, get_is_diffusion_only_model) from tensorrt_llm.executor.utils import MAX_NUM_FRONTENDS, LlmLauncherEnvs @@ -2745,7 +2746,11 @@ def resolve_command(self, ctx, args): "disaggregated": disaggregated, "disaggregated_mpi_worker": disaggregated_mpi_worker, "mm_embedding_serve": serve_encoder, - "embeddings": serve_embedding + "embeddings": serve_embedding, + # The parts of a Mooncake pool that cannot belong to a server, for + # deployments where a pool outlives or spans them. + "mooncake_master": mooncake_master, + "mooncake_donor": mooncake_donor, }) if __name__ == "__main__": diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index b35205380b5f..d41067842c82 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -41,6 +41,10 @@ l0_a10: - unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_kv_cache_layout.py - unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py + - unittest/_torch/executor/test_mooncake_store_cli.py + - unittest/_torch/executor/test_mooncake_store_common.py + - unittest/_torch/executor/test_mooncake_store_donor.py + - unittest/_torch/executor/test_mooncake_store_master.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/executor/test_resource_manager.py - unittest/_torch/executor/test_seq_slot_sizing.py diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 7a941a960796..cc663b218a4d 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -38,6 +38,8 @@ CONTEXT_INIT = LlmRequestState.CONTEXT_INIT.value # 10 GEN_IN_PROGRESS = LlmRequestState.GENERATION_IN_PROGRESS.value # 13 GEN_TO_COMPLETE = LlmRequestState.GENERATION_TO_COMPLETE.value # 14 +DISAGG_CTX_TRANS_IN_PROGRESS = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS.value # 21 +DISAGG_CTX_COMPLETE = LlmRequestState.DISAGG_CONTEXT_COMPLETE.value # 22 # --------------------------------------------------------------------------- @@ -57,6 +59,7 @@ def make_gen_request( req.lora_task_id = lora_task_id req.is_context_init_state = False req.is_generation_in_progress_state = True + req.is_generation_to_complete_state = False req.is_first_context_chunk = is_first_context_chunk req.py_encoder_output_ready_event = None req.py_multimodal_data = None @@ -172,6 +175,7 @@ def make_kv_cache_manager( block_reuse_policy=BlockReusePolicy.PER_REQUEST, first_new_block_fn=None, is_vswa=False, + has_cache_tier_below_gpu=True, ): mgr = Mock() mgr.kv_connector_manager = None @@ -216,6 +220,9 @@ def suspend_request(req): mgr.suspend_request.side_effect = suspend_request mgr.is_request_active.side_effect = lambda req_id: mgr.kv_cache_map[req_id].is_active + # The default here has a cache tier below GPU, which leaves preemption off. + mgr.has_cache_tier_below_gpu = has_cache_tier_below_gpu + mgr.preempt_request.side_effect = lambda req: True return mgr @@ -1285,6 +1292,473 @@ def test_self_eviction_no_started_in_range(self): assert len(out.context_requests) == 0 +# =========================================================================== +# Preemption (context side, no cache tier below GPU) +# =========================================================================== + + +def _out_of_pages_for(request_id): + """resize_context that only fails for *request_id*.""" + return lambda req, n: req.py_request_id != request_id + + +#: What the executor passes to reset_for_recompute. Its choice, not the +#: scheduler's, so the exact value does not matter here. +UNBOUNDED_MAX_INPUT_LEN = 0x7FFFFFFF + + +class TestContextPreemption: + """Releasing a started request's pages when suspension cannot help. + + With GPU as the last cache level a suspended page stays HELD and + unevictable, so suspension frees nothing. These tests cover the fallback + that gives the pages up instead. + """ + + def test_out_of_pages_preempts_started_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + reqs = [make_ctx_request(0, 100), victim] + + out = sched.schedule_request(reqs, set()) + + mgr.preempt_request.assert_called_once_with(victim) + assert ids(out.recompute_paused_requests) == [99] + # Deferred to the next iteration: a failed resize leaves the first + # chunk suspended, so the retry has to go back through + # prepare_context. + assert len(out.context_requests) == 0 + + def test_released_victim_is_left_for_the_executor_to_reset(self): + """The rest of the teardown belongs to the recompute-pause path.""" + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + victim.py_batch_idx = 7 + + out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + victim.pause.assert_not_called() + victim.reset_for_recompute.assert_not_called() + assert ids(out.recompute_paused_requests) == [99] + assert victim.py_batch_idx is None + + def test_preemption_releases_the_draft_pool_too(self): + """The draft pool mirrors the target, so a half-released victim leaks.""" + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + draft_mgr = make_kv_cache_manager() + sched = make_scheduler(mgr, max_num_tokens=1000, draft_kv_cache_manager=draft_mgr) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + mgr.preempt_request.assert_called_once_with(victim) + draft_mgr.free_resources.assert_called_once_with(victim) + assert ids(out.recompute_paused_requests) == [99] + + def test_blocked_request_is_admitted_after_the_executor_recomputes(self): + """The whole handoff, not just the scheduler's half of it. + + A preemption is only worth anything if the blocked request gets in on a + later pass, and that depends on the executor freeing the victim and + resetting it for recompute in between. + """ + out_of_pages = {0} + + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: req.py_request_id not in out_of_pages, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + blocked = make_ctx_request(0, 100) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + first = sched.schedule_request([blocked, victim], set()) + assert ids(first.context_requests) == [] + assert ids(first.recompute_paused_requests) == [99] + + # Stand in for the executor: free the victim's resources, reset it for + # recompute, and let the pages it gave up satisfy the blocked request. + for req in first.recompute_paused_requests: + mgr.free_resources(req) + req.reset_for_recompute(UNBOUNDED_MAX_INPUT_LEN) + req.is_first_context_chunk = True + out_of_pages.clear() + + second = sched.schedule_request([blocked, victim], set()) + + assert 0 in ids(second.context_requests) + victim.reset_for_recompute.assert_called_once() + + def test_disagg_generation_worker_never_preempts(self): + """It received its context KV, so it cannot replay a prefill.""" + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000, enable_recompute_pause=False) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + mgr.preempt_request.assert_not_called() + + def test_preempted_victim_not_scheduled_in_the_same_pass(self): + """Re-admitting the victim would spend the pages it just released.""" + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + assert ids(out.context_requests) == [] + + def test_freed_pages_are_reserved_for_the_request_that_preempted(self): + """No later context request may spend them first.""" + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + blocked = make_ctx_request(0, 100) + # Fits in what the preemption releases, and is only behind `blocked` + # in arrival order. + later = make_ctx_request(1, 10) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([blocked, later, victim], set()) + + assert ids(out.recompute_paused_requests) == [99] + assert ids(out.context_requests) == [] + + def test_a_failed_preemption_still_lets_the_pass_continue(self): + """With nothing to give up there is no capacity to reserve.""" + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + # Only first chunks, so there is no preemption victim among them. + reqs = [make_ctx_request(0, 100), make_ctx_request(1, 10)] + + out = sched.schedule_request(reqs, set()) + + mgr.preempt_request.assert_not_called() + assert ids(out.context_requests) == [1] + + def test_generation_scheduled_before_a_preemption_is_unaffected(self): + """Phase 1 has already committed, so ending phase 2 costs it nothing.""" + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(1), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + blocked = make_ctx_request(1, 100) + later = make_ctx_request(2, 10) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_gen_request(0), blocked, later, victim], set()) + + assert ids(out.generation_requests) == [0] + assert ids(out.recompute_paused_requests) == [99] + assert ids(out.context_requests) == [] + + def test_skipped_when_a_cache_tier_exists_below_gpu(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=True, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 100), victim], set()) + + mgr.preempt_request.assert_not_called() + # Suspension is cheaper and keeps the pages, so the request is + # simply skipped. + assert ids(out.context_requests) == [99] + + def test_never_preempts_a_scheduled_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(1), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + # gen0 is scheduled in phase 1; ctx1 then runs out of pages and must + # not take the pages out from under it. + reqs = [make_gen_request(0), make_ctx_request(1, 100)] + + out = sched.schedule_request(reqs, set()) + + assert ids(out.generation_requests) == [0] + mgr.preempt_request.assert_not_called() + + def test_never_preempts_itself(self): + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + req = make_ctx_request(0, 100, is_first_context_chunk=False) + + sched.schedule_request([req], set()) + + mgr.preempt_request.assert_not_called() + + def test_never_preempts_an_inflight_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + sched.schedule_request([make_ctx_request(0, 100), victim], {99}) + + mgr.preempt_request.assert_not_called() + + def test_never_preempts_a_first_chunk_or_suspended_request(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + # First chunk: holds no pages worth taking. + first_chunk = make_ctx_request(98, 100, is_first_context_chunk=True) + # Already suspended: preempting it frees nothing extra. + suspended = make_ctx_request(99, 100, is_first_context_chunk=False) + mgr.kv_cache_map[suspended.py_request_id].is_active = False + + sched.schedule_request([make_ctx_request(0, 100), first_chunk, suspended], set()) + + mgr.preempt_request.assert_not_called() + + def test_chunked_context_out_of_pages_preempts(self): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler( + mgr, + max_num_tokens=1000, + ctx_chunk_config=(ContextChunkingPolicy.FIRST_COME_FIRST_SERVED, 64), + ) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([make_ctx_request(0, 500), victim], set()) + + mgr.preempt_request.assert_called_once_with(victim) + assert ids(out.recompute_paused_requests) == [99] + + +# =========================================================================== +# Deadlock detection +# =========================================================================== + + +class TestDeadlockDetection: + """The scheduler must fail loudly rather than spin scheduling nothing. + + A stalled pass costs a couple of milliseconds, so an undetected stall + burns a job's whole wall clock. + """ + + def test_raises_after_repeated_stalls_with_context_candidates(self): + """A prefill-only worker has no generation requests to count.""" + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + reqs = [make_ctx_request(0, 100, is_first_context_chunk=False)] + + for _ in range(2): + sched.schedule_request(reqs, set()) + with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): + sched.schedule_request(reqs, set()) + + def test_raises_after_repeated_stalls_with_generation_candidates(self): + mgr = make_kv_cache_manager(try_allocate_generation_fn=lambda req: False, can_evict=True) + sched = make_scheduler(mgr, max_num_tokens=100) + sched._DEADLOCK_STALL_ITERS = 3 + # Self-eviction suspends it on the first pass, which counts as + # progress. Afterwards it is inactive and nothing can be reclaimed. + reqs = [make_gen_request(0)] + + for _ in range(3): + sched.schedule_request(reqs, set()) + with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): + sched.schedule_request(reqs, set()) + + def test_transient_stall_does_not_raise(self): + """One bad iteration is normal; the counter has to reset.""" + fail = [True] + + def resize_fn(req, n): + return not fail[0] + + mgr = make_kv_cache_manager(resize_context_fn=resize_fn, has_cache_tier_below_gpu=False) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + reqs = [make_ctx_request(0, 100, is_first_context_chunk=False)] + + for _ in range(10): + sched.schedule_request(reqs, set()) + fail[0] = not fail[0] + + def test_idle_scheduler_never_raises(self): + mgr = make_kv_cache_manager() + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 2 + + for _ in range(5): + out = sched.schedule_request([], set()) + assert len(out.context_requests) == 0 + + def test_all_candidates_inflight_never_raises(self): + """Requests in the PP pipeline are progressing, just not here.""" + mgr = make_kv_cache_manager(resize_context_fn=lambda req, n: False) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 2 + reqs = [make_ctx_request(0, 100)] + + for _ in range(5): + sched.schedule_request(reqs, {0}) + + @pytest.mark.parametrize( + "holder_state", + [DISAGG_CTX_TRANS_IN_PROGRESS, DISAGG_CTX_COMPLETE, DISAGG_GEN_TRANS_IN_PROGRESS], + ) + def test_a_pending_transfer_holding_pages_is_not_a_deadlock(self, holder_state): + """A context server's pool can be full of sends that have not landed. + + Those requests are past the schedulable states, so they never reach + the scheduled lists, and the context requests they block cannot + allocate. The transfer's own timeout covers a send that never lands. + """ + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + reqs = [ + make_ctx_request(0, 100, is_first_context_chunk=False), + make_filtered_request(1, state_value=holder_state), + ] + + for _ in range(sched._DEADLOCK_STALL_ITERS + 2): + sched.schedule_request(reqs, set()) + + assert sched._stalled_schedules == 0 + + def test_a_deadlock_behind_a_finished_transfer_is_still_reported(self): + """The reprieve lasts only as long as the transfer does.""" + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + blocked = make_ctx_request(0, 100, is_first_context_chunk=False) + sending = make_filtered_request(1, state_value=DISAGG_CTX_TRANS_IN_PROGRESS) + + for _ in range(5): + sched.schedule_request([blocked, sending], set()) + + # The send landed and the executor reaped it, but the pool is still + # full, so the stall is now the scheduler's to report. + for _ in range(sched._DEADLOCK_STALL_ITERS - 1): + sched.schedule_request([blocked], set()) + with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): + sched.schedule_request([blocked], set()) + + def test_inflight_send_that_left_active_list_is_not_a_deadlock(self): + """A sender off the active list but still in flight is not a deadlock. + + The sender never appears in active_requests, so only the transfer + manager can tell the detector its pinned pages are coming back. + """ + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + transfer_mgr = Mock() + transfer_mgr.has_any_inflight_requests.return_value = True + sched.set_async_transfer_manager(transfer_mgr) + # Only the blocked request is on the active list; the sender has left. + blocked = make_ctx_request(0, 100, is_first_context_chunk=False) + + for _ in range(sched._DEADLOCK_STALL_ITERS + 2): + sched.schedule_request([blocked], set()) + + assert sched._stalled_schedules == 0 + + def test_blocked_request_schedules_after_send_completes(self): + """Once the off-list send lands, its pages free the blocked request.""" + pages_held = [True] + + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: not pages_held[0], + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + transfer_mgr = Mock() + transfer_mgr.has_any_inflight_requests.return_value = True + sched.set_async_transfer_manager(transfer_mgr) + blocked = make_ctx_request(0, 100, is_first_context_chunk=False) + + for _ in range(sched._DEADLOCK_STALL_ITERS + 2): + out = sched.schedule_request([blocked], set()) + assert len(out.context_requests) == 0 + assert sched._stalled_schedules == 0 + + # The send lands: the transfer manager drops it and its pages return. + transfer_mgr.has_any_inflight_requests.return_value = False + pages_held[0] = False + out = sched.schedule_request([blocked], set()) + assert ids(out.context_requests) == [0] + + def test_finished_send_still_deadlocks_with_transfer_manager_wired(self): + """The transfer-manager reprieve ends the moment no send is in flight. + + With the manager wired but reporting nothing in flight, a still-full + pool must deadlock exactly as it does without a manager attached. + """ + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, + has_cache_tier_below_gpu=False, + ) + sched = make_scheduler(mgr, max_num_tokens=1000) + sched._DEADLOCK_STALL_ITERS = 3 + transfer_mgr = Mock() + transfer_mgr.has_any_inflight_requests.return_value = False + sched.set_async_transfer_manager(transfer_mgr) + blocked = make_ctx_request(0, 100, is_first_context_chunk=False) + + for _ in range(sched._DEADLOCK_STALL_ITERS - 1): + sched.schedule_request([blocked], set()) + with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): + sched.schedule_request([blocked], set()) + + # =========================================================================== # PEFT / LoRA # =========================================================================== @@ -3399,6 +3873,24 @@ def test_deferral_never_empties_the_batch_alone(self): assert ids(out.context_requests) == [] assert mgr.resize_context.call_count == 2 + def test_deferral_behind_an_inflight_contributor_is_not_a_stall(self): + """Deferring is progress when the contributor is the one in flight. + + Nothing reaches a scheduled list on such a pass, but the duplicate + stays in pending_ctx and still counts as a candidate, so without an + explicit signal the detector reads a working engine as hung. + """ + mgr = self._keyed_manager({5: b"blockA", 1: b"blockA"}) + sched = make_scheduler(mgr, max_num_tokens=10000) + sched._DEADLOCK_STALL_ITERS = 3 + contributor = make_ctx_request(5, 100, is_first_context_chunk=False) + duplicate = make_ctx_request(1, 500) + reqs = [duplicate, contributor] + + for _ in range(sched._DEADLOCK_STALL_ITERS + 1): + out = sched.schedule_request(reqs, {5}) + assert ids(out.context_requests) == [] + def test_registered_contributor_cannot_be_evicted(self): """Eviction and deferral cannot collide. diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index fbb26fedafd0..1f3b77e7b170 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -1397,6 +1397,7 @@ class _ContextRequest: return_perf_metrics: bool = False context_current_position: int = 0 py_connector_served_position: int = 0 + py_num_connector_matched_tokens: int = 0 prepopulated_prompt: tuple[int, int] | None = None multimodal_hashes: None = None multimodal_positions: None = None @@ -1595,6 +1596,73 @@ def _run_context( _update_context_resources(manager, batch) +def _try_run_context(manager: KVCacheManagerV2, request: _ContextRequest) -> bool: + """Run a context request all the way through, reporting whether it fit.""" + batch = _prepare_context_resources(manager, request) + if not manager.prepare_context(request): + return False + request.context_remaining_length = request.prompt_len - request.context_current_position + if not manager.resize_context(request, num_tokens=request.context_remaining_length): + return False + request.context_current_position = request.prompt_len + request.context_remaining_length = 0 + _update_context_resources(manager, batch) + return True + + +def test_preempt_request_gives_a_full_pool_back_its_pages( + manager: KVCacheManagerV2, +) -> None: + """Preemption is the only way out of a full GPU-only pool. + + With GPU as the last cache level a suspended page stays HELD, which the + eviction controller refuses to move, so suspension frees nothing. This + covers the release itself: the pages returning, and the connector's + matched-token count dropping. + """ + victim = _ContextRequest(1, list(range(MAX_SEQ_LEN)), MAX_SEQ_LEN, "conv-1") + victim.py_num_connector_matched_tokens = TOKENS_PER_BLOCK + started: list[_ContextRequest] = [victim] + + try: + assert _try_run_context(manager, victim) + assert manager.is_request_active(victim.py_request_id) + + # The pool's capacity in requests follows from the fixture's layout and + # windows, so find it rather than hard-code it. Distinct tokens per + # request keep block reuse from hiding the pressure. + blocked = None + for request_id in range(2, 12): + candidate = _ContextRequest( + request_id, + [request_id * 1000 + i for i in range(MAX_SEQ_LEN)], + MAX_SEQ_LEN, + f"conv-{request_id}", + ) + started.append(candidate) + if not _try_run_context(manager, candidate): + blocked = candidate + break + assert blocked is not None, "the pool never filled, so there is nothing to preempt for" + + # Its own partial allocation goes back first, as the scheduler's retry + # does. What is missing is the victim's pages. + manager.free_resources(blocked) + + assert manager.preempt_request(victim) is True + assert not manager.is_request_active(victim.py_request_id) + assert victim.py_request_id not in manager.kv_cache_map + # A replay recomputes from its own reuse match, so a stale count would + # skip a prefix the pool no longer holds for it. + assert victim.py_num_connector_matched_tokens == 0 + + blocked.context_current_position = 0 + assert _try_run_context(manager, blocked) + finally: + for request in started: + _free_if_active(manager, request) + + def test_per_conversation_policy_delays_commit_until_last_context_chunk( manager: KVCacheManagerV2, ) -> None: diff --git a/tests/unittest/_torch/executor/test_mooncake_store_cli.py b/tests/unittest/_torch/executor/test_mooncake_store_cli.py new file mode 100644 index 000000000000..1f9e6a0f5c59 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_cli.py @@ -0,0 +1,502 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the two Mooncake pool commands as `trtllm-serve` runs them. + +Runs without a Mooncake installation and without a GPU. The resource each +command holds is a context manager that records its own release, and the wait +the command would otherwise spend idling is where the signal under test is +delivered to this process. +""" + +import contextlib +import json +import os +import signal +import time +from types import SimpleNamespace + +import click +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors import mooncake_store +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + CONFIG_PATH_ENV, + DEFAULT_METADATA_SERVER, +) +from tensorrt_llm.commands import _telemetry +from tensorrt_llm.commands import mooncake as mooncake_commands +from tensorrt_llm.usage.config import UsageContext + +HANDLED_SIGNALS = [signal.SIGINT, signal.SIGTERM] +COMMANDS = ["mooncake_master", "mooncake_donor"] + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + """No ambient pool config: these tests drive the commands from their flags.""" + monkeypatch.delenv(CONFIG_PATH_ENV, raising=False) + + +@pytest.fixture(autouse=True) +def restore_signal_handlers(): + """Both commands install handlers process-wide, so put pytest's back.""" + installed = {number: signal.getsignal(number) for number in HANDLED_SIGNALS} + yield + for number, handler in installed.items(): + signal.signal(number, handler) + + +@pytest.fixture(autouse=True) +def terminal_outcomes(monkeypatch): + """Collect what the shared boundary reports, without a usage session. + + Autouse because a command reaching the real session would apply this + process's opt-out for the rest of the run. + """ + outcomes = [] + monkeypatch.setattr( + _telemetry.usage, "report_exit", lambda outcome, **_kwargs: outcomes.append(outcome) + ) + monkeypatch.setattr(_telemetry.usage, "apply_usage_session_config", lambda *a, **k: True) + monkeypatch.setattr(_telemetry.usage, "set_lifecycle_phase", lambda *a, **k: None) + monkeypatch.setattr(_telemetry.usage, "get_observed_signal", lambda: 0) + monkeypatch.setattr(_telemetry.usage, "record_observed_signal", lambda *a, **k: None) + return outcomes + + +class Resource: + """Stands in for the master process or the mounted segment. + + Records its release, so a test can tell a signal that unwound the command + from one that ended it with the `finally` never reached, and what it was + asked for, so a test can check what the command resolved before asking. + """ + + def __init__(self, held): + self.held = held + self.released = False + #: `None` until the resource is taken, which distinguishes a command + #: that gave up beforehand from one that took it. + self.taken_with = None + + @contextlib.contextmanager + def holding(self, *args, **_kwargs): + self.taken_with = args + try: + yield self.held + finally: + self.released = True + + +class CommandUnderTest: + """One command, its stubbed resource, and the arguments that reach it.""" + + def __init__(self, resource: Resource, args: list[str]): + self.resource = resource + self.args = args + + def run(self, *extra_args: str) -> None: + group = _telemetry.TelemetryGroup( + name="trtllm-serve", + telemetry_usage_context=UsageContext.CLI_SERVE, + telemetry_component="server", + commands={ + "mooncake_master": mooncake_commands.mooncake_master, + "mooncake_donor": mooncake_commands.mooncake_donor, + }, + ) + group.main( + args=[*self.args, *extra_args], + prog_name="trtllm-serve", + standalone_mode=False, + ) + + +@pytest.fixture +def master(monkeypatch, tmp_path) -> CommandUnderTest: + resource = Resource( + SimpleNamespace( + address="127.0.0.1:50051", + log_path=str(tmp_path / "master.log"), + # Alive, so the command idles rather than reporting a dead master. + process=SimpleNamespace(poll=lambda: None), + ) + ) + monkeypatch.setattr(mooncake_store, "running_master", resource.holding) + return CommandUnderTest(resource, ["mooncake_master", "--run_dir", str(tmp_path)]) + + +@pytest.fixture +def donor(monkeypatch) -> CommandUnderTest: + resource = Resource("10.0.0.1:12345") + monkeypatch.setattr(mooncake_store, "donate_segment", resource.holding) + monkeypatch.setattr(mooncake_store, "resolve_master_address", lambda address, _timeout: address) + monkeypatch.setattr(mooncake_store, "wait_for_master", lambda _address: None) + return CommandUnderTest( + resource, + ["mooncake_donor", "--master_server_address", "127.0.0.1:50051", "--segment_size", "1GiB"], + ) + + +@pytest.fixture +def command(request, master, donor) -> CommandUnderTest: + return {"mooncake_master": master, "mooncake_donor": donor}[request.param] + + +def signal_on_idle(monkeypatch, number: int, on_idle=None) -> None: + """Deliver `number` to this process the first time a command idles. + + The handler runs at the next bytecode boundary in the main thread, so the + loop here only has to give the interpreter one. It bounds the wait rather + than blocking, so a handler that never fires fails the test rather than + hanging it. + + `on_idle` runs before the signal, which is the only point at which a + command's resources are all held. Anything a command releases on its way + out has to be observed from there rather than after `run` returns. + """ + + def sleep(_seconds): + if on_idle is not None: + on_idle() + os.kill(os.getpid(), number) + for _ in range(100): + time.sleep(0.01) + + monkeypatch.setattr( + mooncake_commands, "time", SimpleNamespace(sleep=sleep, monotonic=time.monotonic) + ) + + +@pytest.mark.parametrize("number", HANDLED_SIGNALS) +@pytest.mark.parametrize("command", COMMANDS, indirect=True) +def test_a_signal_is_reported_once_and_still_releases_what_was_held( + monkeypatch, terminal_outcomes, command, number +): + """The signal reaches the boundary, and the resource is still given up. + + A command that returned normally instead would be reported as a clean + exit before any model was loaded. + """ + signal_on_idle(monkeypatch, number) + + with pytest.raises(_telemetry.SignalExit) as stopped: + command.run() + + assert stopped.value.signal_number == number + assert command.resource.released + + assert len(terminal_outcomes) == 1 + outcome = terminal_outcomes[0] + assert outcome.termination_kind == "signal" + assert outcome.signal_number == number + assert outcome.exit_code == 128 + number + + +@pytest.mark.parametrize("flag", ["--telemetry", "--no-telemetry"]) +@pytest.mark.parametrize("command", COMMANDS, indirect=True) +def test_the_opt_out_the_group_documents_is_accepted(monkeypatch, command, flag): + """Without the option Click rejects the flag as a usage error instead.""" + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + command.run(flag) + + assert command.resource.released + + +# ---- what the donor resolves before it offers anything ---- + + +@pytest.fixture +def segment(monkeypatch) -> Resource: + """Stub only the segment, so the command's address handling still runs.""" + resource = Resource("10.0.0.1:12345") + monkeypatch.setattr(mooncake_store, "donate_segment", resource.holding) + return resource + + +def donor_running(resource: Resource, *args: str) -> CommandUnderTest: + return CommandUnderTest(resource, ["mooncake_donor", *args]) + + +def test_a_donor_resolves_a_published_master_address_before_joining(monkeypatch, segment, tmp_path): + """What lets a generation node name a path instead of a scheduler's choice.""" + address_file = tmp_path / "master.addr" + address_file.write_text("10.0.0.9:50051\n") + probed = [] + monkeypatch.setattr(mooncake_store, "wait_for_master", probed.append) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running( + segment, + "--master_server_address", + f"file://{address_file}", + "--segment_size", + "2GiB", + ).run() + + # The resolved address is both what was probed and what Mooncake is given: + # Mooncake cannot dial a file:// URL. + assert probed == ["10.0.0.9:50051"] + assert segment.taken_with[0] == "10.0.0.9:50051" + # A size string reaching Mooncake unparsed would be a segment of nothing. + assert segment.taken_with[1] == 2 * 1024**3 + + +def test_a_donor_reports_an_unreachable_master_before_offering_the_segment(monkeypatch, segment): + """Otherwise this is a status code from setup, with no address in it.""" + + def refuse(address): + raise TimeoutError(f"The Mooncake master at {address} did not accept connections") + + monkeypatch.setattr(mooncake_store, "resolve_master_address", lambda address, _t: address) + monkeypatch.setattr(mooncake_store, "wait_for_master", refuse) + + with pytest.raises(TimeoutError, match="10.0.0.1:50051"): + donor_running(segment, "--master_server_address", "10.0.0.1:50051").run() + + assert segment.taken_with is None + + +@pytest.mark.parametrize("size", ["", "16 GB!"]) +def test_a_donor_rejects_a_size_it_cannot_parse(monkeypatch, segment, size): + """Caught as a usage error rather than as a segment of some other size.""" + monkeypatch.setattr(mooncake_store, "resolve_master_address", lambda address, _t: address) + monkeypatch.setattr(mooncake_store, "wait_for_master", lambda _address: None) + + with pytest.raises(click.UsageError, match="--segment_size"): + donor_running( + segment, "--master_server_address", "10.0.0.1:50051", "--segment_size", size + ).run() + + assert segment.taken_with is None + + +# ---- how the donor resolves its flags against a config ---- + + +class RecordingDonation: + """Records the arguments the donor resolves before donate_segment is called. + + The command idles inside the donation context, so a test breaks the loop + with a signal once the resolved arguments and the ready file are in hand. + """ + + def __init__(self, host: str = "10.0.0.1:12345"): + self.host = host + self.resolved: dict = {} + + @contextlib.contextmanager + def holding(self, master, segment, **kwargs): + self.resolved = dict(master=master, segment=segment, **kwargs) + yield self.host + + +@pytest.fixture +def donation(monkeypatch) -> RecordingDonation: + """Drive the donor to the point of donation, capturing what it resolved.""" + recorder = RecordingDonation() + monkeypatch.setattr(mooncake_store, "donate_segment", recorder.holding) + monkeypatch.setattr(mooncake_store, "resolve_master_address", lambda address, _t: address) + monkeypatch.setattr(mooncake_store, "wait_for_master", lambda _address: None) + return recorder + + +def _config_file(tmp_path, **entries) -> str: + path = tmp_path / "pool.json" + path.write_text(json.dumps(entries)) + return str(path) + + +def test_the_donor_uses_its_documented_defaults_without_a_config(monkeypatch, donation): + """With only the required master, transport and buffer defaults apply.""" + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running(donation, "--master_server_address", "10.0.0.1:50051").run() + + assert donation.resolved["protocol"] == "rdma" + assert donation.resolved["device_name"] == "" + assert donation.resolved["metadata_server"] == DEFAULT_METADATA_SERVER + assert donation.resolved["local_buffer_size"] == mooncake_store.DEFAULT_DONOR_LOCAL_BUFFER_SIZE + + +def test_the_donor_resolves_every_option_from_the_cli(monkeypatch, donation): + """Flags alone describe the donation, with no config on hand.""" + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running( + donation, + "--master_server_address", + "10.0.0.1:50051", + "--segment_size", + "2GiB", + "--protocol", + "tcp", + "--device_name", + "mlx5_0", + "--metadata_server", + "http://cli:8080/metadata", + ).run() + + assert donation.resolved["master"] == "10.0.0.1:50051" + assert donation.resolved["segment"] == 2 * 1024**3 + assert donation.resolved["protocol"] == "tcp" + assert donation.resolved["device_name"] == "mlx5_0" + assert donation.resolved["metadata_server"] == "http://cli:8080/metadata" + + +def test_the_donor_falls_back_to_the_config_for_unset_options(monkeypatch, donation, tmp_path): + """A config names the master and the transport when no flag overrides it.""" + config = _config_file( + tmp_path, + master_server_address="10.0.0.9:50051", + protocol="tcp", + device_name="mlx5_1", + metadata_server="http://config:8080/metadata", + ) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running(donation, "--config", config, "--segment_size", "1GiB").run() + + assert donation.resolved["master"] == "10.0.0.9:50051" + assert donation.resolved["protocol"] == "tcp" + assert donation.resolved["device_name"] == "mlx5_1" + assert donation.resolved["metadata_server"] == "http://config:8080/metadata" + + +def test_the_donor_prefers_the_cli_over_the_config(monkeypatch, donation, tmp_path): + """A flag wins over the same setting in the config.""" + config = _config_file( + tmp_path, + master_server_address="10.0.0.9:50051", + protocol="tcp", + device_name="mlx5_1", + metadata_server="http://config:8080/metadata", + ) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running( + donation, + "--config", + config, + "--master_server_address", + "10.0.0.1:50051", + "--segment_size", + "1GiB", + "--protocol", + "rdma", + "--device_name", + "mlx5_0", + "--metadata_server", + "http://cli:8080/metadata", + ).run() + + assert donation.resolved["master"] == "10.0.0.1:50051" + assert donation.resolved["protocol"] == "rdma" + assert donation.resolved["device_name"] == "mlx5_0" + assert donation.resolved["metadata_server"] == "http://cli:8080/metadata" + + +def test_the_donor_defaults_the_protocol_to_rdma_when_the_config_value_is_null( + monkeypatch, donation, tmp_path +): + """A null protocol in the config must not reach setup unnormalized.""" + config = _config_file(tmp_path, master_server_address="10.0.0.9:50051", protocol=None) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + donor_running(donation, "--config", config, "--segment_size", "1GiB").run() + + assert donation.resolved["protocol"] == "rdma" + + +def test_the_donors_ready_file_lasts_exactly_as_long_as_the_segment( + monkeypatch, donation, tmp_path +): + """A launcher waits on this file before letting prefill fill the pool. + + So the file has to appear once the segment is mounted and go when it is + withdrawn. Outliving the segment would clear that gate for capacity this + donor no longer contributes, and prefill would start writing into a pool + short of the memory the launcher waited for. + """ + ready_file = tmp_path / "ready" + announced = [] + signal_on_idle( + monkeypatch, + signal.SIGTERM, + on_idle=lambda: announced.append(ready_file.read_text()), + ) + + with pytest.raises(_telemetry.SignalExit): + donor_running( + donation, + "--master_server_address", + "10.0.0.1:50051", + "--segment_size", + "2GiB", + "--ready_file", + str(ready_file), + ).run() + + assert announced == [f"{donation.host} {2 * 1024**3}\n"] + assert not ready_file.exists() + + +# ---- describing the pool from the master's own run directory ---- + + +def test_the_master_writes_the_client_config_it_was_given(monkeypatch, master, tmp_path): + """Lets workers be pointed at the run directory instead of a hand-written file.""" + config = tmp_path / "pool.json" + config.write_text( + json.dumps({"protocol": "tcp", "global_segment_size": "8GiB", "model_key": "m"}) + ) + written = {} + monkeypatch.setattr( + mooncake_store, + "write_client_config", + lambda pool, address, run_dir: written.update(pool=pool, address=address), + ) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + master.run("--config", str(config)) + + # Written only once the master answers, and naming that master rather than + # whichever one the config was copied from. + assert written["address"] == "127.0.0.1:50051" + assert written["pool"].protocol == "tcp" + assert written["pool"].global_segment_size == "8GiB" + assert written["pool"].launch_master is True + assert written["pool"].master_server_address is None + + +def test_the_master_writes_no_client_config_unless_asked(monkeypatch, master): + """A deployment that renders its own must not have it overwritten.""" + calls = [] + monkeypatch.setattr(mooncake_store, "write_client_config", lambda *args: calls.append(args)) + signal_on_idle(monkeypatch, signal.SIGTERM) + + with pytest.raises(_telemetry.SignalExit): + master.run() + + assert calls == [] diff --git a/tests/unittest/_torch/executor/test_mooncake_store_common.py b/tests/unittest/_torch/executor/test_mooncake_store_common.py new file mode 100644 index 000000000000..0347ca026eb7 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_common.py @@ -0,0 +1,525 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the Mooncake store pieces every deployment shares. + +Covers how a block of tokens becomes a store key, how the JSON config is read, +and the pinned host slots pages pass through where GPUDirect RDMA is +unavailable. + +Runs without a Mooncake installation: the store handle is replaced by an +in-process fake that records what it was handed. +""" + +import json + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import staging as staging_module +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + MooncakeStoreConnectorConfig, + StoreRole, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.keys import ( + BlockHashChain, + KeyNamespace, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.metadata import ( + MooncakeStoreMetadata, + PageTransfer, + RequestTransfers, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.staging import ( + HostStagingPool, + describe_batch_for_get, + plan_slot_geometry, + stage_batch_for_put, + unstage_batch_after_get, +) + +TOKENS_PER_BLOCK = 4 + +#: Device addresses the staging tests gather from and scatter back to. Never +#: dereferenced: the copies are recorded rather than issued. +PAGE_ADDRESSES = [0xA000, 0xB000] +PAGE_SIZES = [64, 128] +PAGE_BYTES = sum(PAGE_SIZES) + + +# ---- fixtures and fakes ---- + + +class FakeStore: + """Records buffer registrations, nothing more. + + Staging only ever asks the store to register its host buffer; the transfer + calls belong to the connector. + """ + + def __init__(self, register_status=0, unregister_status=0): + self.registered = [] + self.unregistered = [] + self._register_status = register_status + self._unregister_status = unregister_status + + def register_buffer(self, address, size): + self.registered.append((address, size)) + return self._register_status + + def unregister_buffer(self, address): + self.unregistered.append(address) + return self._unregister_status + + +@pytest.fixture +def store_config(tmp_path, monkeypatch): + path = tmp_path / "mooncake.json" + path.write_text( + json.dumps( + { + "metadata_server": "http://127.0.0.1:8080/metadata", + "master_server_address": "127.0.0.1:50051", + "protocol": "tcp", + "device_name": "", + "global_segment_size": "1GiB", + "local_buffer_size": "256MiB", + "model_key": "test-model", + } + ) + ) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(path)) + monkeypatch.delenv("TRTLLM_MOONCAKE_STORE_ROLE", raising=False) + monkeypatch.delenv("TRTLLM_MOONCAKE_STORE_PREFIX", raising=False) + monkeypatch.delenv("TRTLLM_MOONCAKE_STORE_MODEL_KEY", raising=False) + monkeypatch.delenv("TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST", raising=False) + return path + + +@pytest.fixture +def staged_copies(monkeypatch): + """Record the copies staging would issue, instead of running them.""" + copies = [] + monkeypatch.setattr( + staging_module, + "_memcpy_async", + lambda dst, src, size, kind, stream: copies.append((int(dst), int(src), int(size))), + ) + return copies + + +def make_pool(*, slot_bytes=PAGE_BYTES, num_slots=4, store=None): + return HostStagingPool( + slot_bytes=slot_bytes, + num_slots=num_slots, + store=store if store is not None else FakeStore(), + label="save", + ) + + +# ---- keys ---- + + +def test_hash_chain_is_deterministic_and_prefix_sensitive(): + tokens = list(range(3 * TOKENS_PER_BLOCK)) + first = list(BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)) + second = list(BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)) + assert first == second + + # Changing a token in block 0 must change every hash after it, which is what + # makes a key safe to share: a hit implies the whole prefix matched. + altered = list(tokens) + altered[0] += 1 + changed = list(BlockHashChain(TOKENS_PER_BLOCK).extend(altered)) + assert all(a != b for a, b in zip(first, changed)) + + +def test_hash_chain_ignores_partial_trailing_block(): + full = list(range(2 * TOKENS_PER_BLOCK)) + chain = BlockHashChain(TOKENS_PER_BLOCK) + assert len(chain.extend(full)) == 2 + assert len(chain.extend(full + [99])) == 2 + + +def test_hash_chain_extends_incrementally(): + tokens = list(range(4 * TOKENS_PER_BLOCK)) + incremental = BlockHashChain(TOKENS_PER_BLOCK) + for end in range(0, len(tokens) + 1, TOKENS_PER_BLOCK): + incremental.extend(tokens[:end]) + assert list(incremental.hashes) == list(BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)) + + +def test_hash_chain_separates_cache_salts(): + tokens = list(range(TOKENS_PER_BLOCK)) + unsalted = BlockHashChain(TOKENS_PER_BLOCK).extend(tokens) + salted = BlockHashChain(TOKENS_PER_BLOCK, cache_salt="tenant-a").extend(tokens) + other = BlockHashChain(TOKENS_PER_BLOCK, cache_salt="tenant-b").extend(tokens) + assert unsalted[0] != salted[0] != other[0] + assert salted[0] != other[0] + + +def test_hash_chain_rejects_shrinking_token_list(): + chain = BlockHashChain(TOKENS_PER_BLOCK) + chain.extend(list(range(2 * TOKENS_PER_BLOCK))) + with pytest.raises(ValueError, match="shrank"): + chain.extend(list(range(TOKENS_PER_BLOCK))) + + +def test_key_namespace_separates_every_dimension(): + base = dict( + cache_prefix="trtllm", + model_key="m", + rank=0, + world_size=2, + layer_group_id=0, + tokens_per_block=32, + bytes_per_page=1024, + ) + block_hash = b"\x01" * 16 + reference = KeyNamespace(**base).key(block_hash) + for field, value in [ + ("cache_prefix", "other"), + ("model_key", "n"), + ("rank", 1), + ("world_size", 4), + ("layer_group_id", 1), + ("tokens_per_block", 64), + ("bytes_per_page", 2048), + ]: + assert KeyNamespace(**{**base, field: value}).key(block_hash) != reference + + +# ---- config ---- + + +def test_config_reads_sizes_and_staging_from_the_json(store_config): + """Sizes arrive as unit strings, and staging is off until the JSON asks.""" + config = MooncakeStoreConnectorConfig.from_env() + assert config.global_segment_size == 1024**3 + assert config.local_buffer_size == 256 * 1024**2 + assert config.role is StoreRole.BOTH + assert config.resolve_model_key("/models/ignored") == "test-model" + assert config.stage_through_host is False + + raw = json.loads(store_config.read_text()) + raw["stage_through_host"] = True + raw["staging_buffer_bytes"] = "256MiB" + store_config.write_text(json.dumps(raw)) + + config = MooncakeStoreConnectorConfig.from_env() + assert config.stage_through_host is True + assert config.staging_buffer_bytes == 256 * 1024**2 + + +def test_config_role_comes_from_environment(store_config, monkeypatch): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "producer") + config = MooncakeStoreConnectorConfig.from_env() + assert config.role is StoreRole.PRODUCER + assert config.role.saves and not config.role.loads + + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "consumer") + config = MooncakeStoreConnectorConfig.from_env() + assert config.role.loads and not config.role.saves + + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_ROLE", "nonsense") + with pytest.raises(ValueError, match="TRTLLM_MOONCAKE_STORE_ROLE"): + MooncakeStoreConnectorConfig.from_env() + + +def test_config_requires_the_env_var(monkeypatch): + monkeypatch.delenv("MOONCAKE_CONFIG_PATH", raising=False) + monkeypatch.delenv("TRTLLM_MOONCAKE_RUN_DIR", raising=False) + with pytest.raises(ValueError, match="MOONCAKE_CONFIG_PATH"): + MooncakeStoreConnectorConfig.from_env() + + +def test_config_falls_back_to_the_run_directory(tmp_path, monkeypatch): + # A rank an external launcher started was already running when its leader + # provisioned the pool, so it never inherited the exported path and reads + # the rendered config out of the shared run directory instead. + monkeypatch.delenv("MOONCAKE_CONFIG_PATH", raising=False) + monkeypatch.setenv("TRTLLM_MOONCAKE_RUN_DIR", str(tmp_path)) + (tmp_path / "mooncake.json").write_text( + json.dumps({"master_server_address": "10.0.0.1:50051", "global_segment_size": "8GiB"}) + ) + + config = MooncakeStoreConnectorConfig.from_env() + + assert config.master_server_address == "10.0.0.1:50051" + assert config.global_segment_size == 8 * 1024**3 + + +def test_config_run_directory_without_a_rendered_config_still_asks(tmp_path, monkeypatch): + # An empty run directory means no leader provisioned anything, which is a + # missing pool rather than a default one. + monkeypatch.delenv("MOONCAKE_CONFIG_PATH", raising=False) + monkeypatch.setenv("TRTLLM_MOONCAKE_RUN_DIR", str(tmp_path)) + with pytest.raises(ValueError, match="MOONCAKE_CONFIG_PATH"): + MooncakeStoreConnectorConfig.from_env() + + +def test_config_env_var_wins_over_the_run_directory(tmp_path, monkeypatch): + # An externally managed pool stays reachable, since the run directory is + # only consulted when nothing was passed in. + named = tmp_path / "external.json" + named.write_text(json.dumps({"master_server_address": "external:50051"})) + (tmp_path / "mooncake.json").write_text( + json.dumps({"master_server_address": "provisioned:50051"}) + ) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(named)) + monkeypatch.setenv("TRTLLM_MOONCAKE_RUN_DIR", str(tmp_path)) + + assert MooncakeStoreConnectorConfig.from_env().master_server_address == "external:50051" + + +@pytest.mark.parametrize("named", [{}, {"metadata_server": ""}], ids=["omitted", "empty"]) +def test_config_metadata_server_falls_back_to_the_handshake(tmp_path, monkeypatch, named): + # No metadata service means Mooncake's peer-to-peer handshake. An empty + # connstring is not one of the forms setup accepts, so leaving the field + # out of a hand-written config must not reach it. + path = tmp_path / "metadata.json" + path.write_text(json.dumps({"master_server_address": "127.0.0.1:50051", **named})) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(path)) + assert MooncakeStoreConnectorConfig.from_env().metadata_server == "P2PHANDSHAKE" + + +def test_config_requires_an_explicit_model_key(store_config, tmp_path, monkeypatch): + # A default derived from the path would let org-a/model and org-b/model + # agree on a namespace while disagreeing on what the pages mean. + path = tmp_path / "no_model_key.json" + path.write_text(json.dumps({"master_server_address": "127.0.0.1:50051"})) + monkeypatch.setenv("MOONCAKE_CONFIG_PATH", str(path)) + config = MooncakeStoreConnectorConfig.from_env() + with pytest.raises(ValueError, match="needs a model key"): + config.resolve_model_key("/models/MiniMax-M3/") + + +@pytest.mark.parametrize( + "value,expected", [("1", True), ("true", True), ("on", True), ("0", False), ("off", False)] +) +def test_config_staging_env_override(store_config, monkeypatch, value, expected): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST", value) + assert MooncakeStoreConnectorConfig.from_env().stage_through_host is expected + + +def test_config_namespace_env_overrides(store_config, monkeypatch): + # Both feed KeyNamespace, which decides whether two engines share cache. + # Getting either wrong silently loses every hit, or lets one engine read + # another's pages. + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_PREFIX", "tenant-a") + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_MODEL_KEY", "llama-3.1-8b@rev7") + + config = MooncakeStoreConnectorConfig.from_env() + + assert config.cache_prefix == "tenant-a" + # Overrides the JSON's model_key rather than only supplying a missing one. + assert config.resolve_model_key("/models/test-model") == "llama-3.1-8b@rev7" + + +def test_config_rejects_a_non_boolean_staging_env(store_config, monkeypatch): + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_STAGE_THROUGH_HOST", "sometimes") + with pytest.raises(ValueError, match="not a boolean"): + MooncakeStoreConnectorConfig.from_env() + + +# The connector itself, the registry preset that selects it, and the startup +# gates on parallelism and index-V caching land with the KV cache manager V2 +# support they depend on. + + +# ---- host staging ---- + + +@pytest.mark.parametrize( + "page_bytes,batch,budget,expected_slots", + [ + (1024, 64, 1 << 20, 64), # budget is ample: the full batch stages + (1024, 64, 8 * 1024, 8), # budget binds before the batch does + (1024, 8, 1 << 20, 8), # batch binds before the budget does + (1024, 64, 1024, 1), # exactly one page fits + (1024, 64, 1, 1), # below one page, raised to one rather than refused + ], +) +def test_plan_slot_geometry(page_bytes, batch, budget, expected_slots): + slot_bytes, num_slots = plan_slot_geometry(page_bytes, batch, budget) + # A slot always holds a whole page: the budget bounds the count, not the width. + assert slot_bytes == page_bytes + assert num_slots == expected_slots + + +@pytest.mark.parametrize("bad", [(0, 8, 1024), (-1, 8, 1024), (1024, 0, 1024)]) +def test_plan_slot_geometry_rejects_degenerate_inputs(bad): + with pytest.raises(ValueError): + plan_slot_geometry(*bad) + + +def test_staging_pool_registers_its_whole_buffer_once(): + """One registration covering every slot, since the store addresses bytes.""" + store = FakeStore() + pool = make_pool(slot_bytes=256, num_slots=4, store=store) + + assert pool.slot_bytes == 256 + assert pool.num_slots == 4 + assert store.registered == [(pool.slot_address(0), 256 * 4)] + + +def test_staging_pool_refuses_to_start_when_registration_fails(): + # Host memory that cannot be registered is the one failure staging exists to + # rule out, so it has to be loud rather than fall back to the pools. + with pytest.raises(RuntimeError, match="register_buffer failed"): + make_pool(store=FakeStore(register_status=-1)) + + +def test_staging_pool_close_unregisters_before_the_buffer_goes(): + store = FakeStore() + pool = make_pool(slot_bytes=256, num_slots=4, store=store) + base = pool.slot_address(0) + + pool.close() + pool.close() + + assert store.unregistered == [base] + + +def test_staging_pool_close_keeps_the_buffer_when_unregistration_fails(): + store = FakeStore(unregister_status=-1) + pool = make_pool(slot_bytes=256, num_slots=4, store=store) + + with pytest.raises(RuntimeError, match="unregister_buffer failed"): + pool.close() + + assert pool._buffer is not None + + +def test_staging_pool_slots_are_contiguous_and_bounds_checked(): + pool = make_pool(slot_bytes=256, num_slots=4) + base = pool.slot_address(0) + + assert [pool.slot_address(i) for i in range(4)] == [base + 256 * i for i in range(4)] + with pytest.raises(IndexError): + pool.slot_address(4) + with pytest.raises(IndexError): + pool.slot_address(-1) + + +def test_staging_pool_gather_concatenates_regions_in_region_order(staged_copies): + """The slot layout is what the zero-copy path would have written.""" + pool = make_pool(num_slots=4) + slot = pool.slot_address(1) + + address, total = pool.gather(1, PAGE_ADDRESSES, PAGE_SIZES, stream=0) + + assert (address, total) == (slot, PAGE_BYTES) + assert staged_copies == [ + (slot, PAGE_ADDRESSES[0], PAGE_SIZES[0]), + (slot + PAGE_SIZES[0], PAGE_ADDRESSES[1], PAGE_SIZES[1]), + ] + + +def test_staging_pool_scatter_is_the_inverse_of_gather(staged_copies): + pool = make_pool(num_slots=4) + slot = pool.slot_address(1) + + pool.scatter(1, PAGE_ADDRESSES, PAGE_SIZES, stream=0) + + # Same split in the same order, with source and destination swapped. + assert staged_copies == [ + (PAGE_ADDRESSES[0], slot, PAGE_SIZES[0]), + (PAGE_ADDRESSES[1], slot + PAGE_SIZES[0], PAGE_SIZES[1]), + ] + + +@pytest.mark.parametrize("operation", ["gather", "scatter", "reserve"]) +def test_staging_pool_rejects_a_page_wider_than_a_slot(staged_copies, operation): + # The pool is sized from the layout's largest page, so an overflow means the + # layout changed after registration rather than a bad argument. + pool = make_pool(slot_bytes=PAGE_BYTES - 1) + + with pytest.raises(ValueError, match="does not fit"): + if operation == "reserve": + pool.reserve(PAGE_BYTES) + else: + getattr(pool, operation)(0, PAGE_ADDRESSES, PAGE_SIZES, stream=0) + + +def test_stage_batch_for_put_hands_the_store_one_buffer_per_page(staged_copies): + pool = make_pool(num_slots=4) + addresses, sizes = stage_batch_for_put( + pool, [PAGE_ADDRESSES, PAGE_ADDRESSES], [PAGE_SIZES, PAGE_SIZES], stream=0 + ) + + # Two regions per page collapse into the one slot the store reads. + assert addresses == [[pool.slot_address(0)], [pool.slot_address(1)]] + assert sizes == [[PAGE_BYTES], [PAGE_BYTES]] + assert len(staged_copies) == 4 + + +def test_describe_batch_for_get_points_at_slots_without_copying(staged_copies): + """Nothing to gather in this direction: the slots are the destination.""" + pool = make_pool(num_slots=4) + + addresses, sizes = describe_batch_for_get(pool, [PAGE_SIZES, PAGE_SIZES]) + + assert addresses == [[pool.slot_address(0)], [pool.slot_address(1)]] + assert sizes == [[PAGE_BYTES], [PAGE_BYTES]] + assert staged_copies == [] + + +@pytest.mark.parametrize("helper", ["put", "get"]) +def test_batch_helpers_reject_a_batch_wider_than_the_pool(staged_copies, helper): + pool = make_pool(num_slots=1) + batch_sizes = [PAGE_SIZES, PAGE_SIZES] + + with pytest.raises(ValueError, match="exceeds 1 staging slots"): + if helper == "put": + stage_batch_for_put(pool, [PAGE_ADDRESSES, PAGE_ADDRESSES], batch_sizes, stream=0) + else: + describe_batch_for_get(pool, batch_sizes) + + +def test_unstage_batch_after_get_scatters_every_page_by_default(staged_copies): + pool = make_pool(num_slots=4) + pages = [[0xA000, 0xB000], [0xC000, 0xD000]] + + unstage_batch_after_get(pool, pages, [PAGE_SIZES, PAGE_SIZES], stream=0) + + assert [copy[0] for copy in staged_copies] == [0xA000, 0xB000, 0xC000, 0xD000] + + +def test_unstage_batch_after_get_leaves_the_pages_not_asked_for_alone(staged_copies): + """A page whose read failed keeps its device bytes rather than slot garbage.""" + pool = make_pool(num_slots=4) + pages = [[0xA000, 0xB000], [0xC000, 0xD000]] + + unstage_batch_after_get(pool, pages, [PAGE_SIZES, PAGE_SIZES], stream=0, only=[1]) + + assert [copy[0] for copy in staged_copies] == [0xC000, 0xD000] + + +# ---- metadata ---- + + +def test_metadata_is_falsy_until_there_is_work(): + """The worker skips the iteration entirely on an empty work list.""" + assert not MooncakeStoreMetadata() + assert MooncakeStoreMetadata(loads=[RequestTransfers(request_id=1)]) + assert MooncakeStoreMetadata(saves=[RequestTransfers(request_id=1)]) + + +def test_request_transfers_do_not_share_a_page_list(): + first = RequestTransfers(request_id=1) + second = RequestTransfers(request_id=2) + + first.pages.append(PageTransfer(block_hash=b"\x01" * 16, layer_group_id=0, page_index=3)) + + assert second.pages == [] diff --git a/tests/unittest/_torch/executor/test_mooncake_store_donor.py b/tests/unittest/_torch/executor/test_mooncake_store_donor.py new file mode 100644 index 000000000000..36cbc3aded57 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for lending a node's host memory to a Mooncake pool. + +Runs without a Mooncake installation and without a GPU. The store is a fake +recording what `setup` was called with, since the contract under test is what +the donor asks Mooncake for and how long it holds it. +""" + +import sys +from types import ModuleType + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import donor as donor_module +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.donor import ( + DEFAULT_DONOR_LOCAL_BUFFER_SIZE, + donate_segment, +) + +GIB = 1024**3 + + +class FakeStore: + """The slice of `MooncakeDistributedStore` a donor drives.""" + + instances = [] + + def __init__(self): + self.setup_args = None + self.status = 0 + FakeStore.instances.append(self) + + def setup(self, *args): + self.setup_args = args + return self.status + + +@pytest.fixture +def fake_bindings(monkeypatch): + """Stand in for `mooncake.store`, which is not installed here.""" + FakeStore.instances = [] + package = ModuleType("mooncake") + store = ModuleType("mooncake.store") + store.MooncakeDistributedStore = FakeStore + package.store = store + monkeypatch.setitem(sys.modules, "mooncake", package) + monkeypatch.setitem(sys.modules, "mooncake.store", store) + return FakeStore + + +@pytest.fixture +def failing_bindings(fake_bindings): + """Bindings whose `setup` refuses, as an unreachable master would.""" + + class Refusing(fake_bindings): + def setup(self, *args): + super().setup(*args) + return 7 + + sys.modules["mooncake.store"].MooncakeDistributedStore = Refusing + return Refusing + + +def test_a_donor_registers_the_segment_it_was_asked_for(fake_bindings): + """Mooncake's `setup` is positional, so all seven arguments are pinned.""" + donation = donate_segment( + "10.0.0.1:50051", + 32 * GIB, + protocol="rdma", + device_name="mlx5_0", + metadata_server="P2PHANDSHAKE", + hostname="10.0.0.5", + ) + + with donation as host: + assert host == "10.0.0.5" + ( + registered_host, + metadata_server, + segment_size, + local_buffer_size, + protocol, + device_name, + master, + ) = fake_bindings.instances[0].setup_args + + assert registered_host == "10.0.0.5" + assert metadata_server == "P2PHANDSHAKE" + assert segment_size == 32 * GIB + assert protocol == "rdma" + assert device_name == "mlx5_0" + assert master == "10.0.0.1:50051" + assert local_buffer_size == DEFAULT_DONOR_LOCAL_BUFFER_SIZE + + +def test_a_donor_that_cannot_join_says_which_master_it_could_not_reach(failing_bindings): + with pytest.raises(RuntimeError, match="status 7"): + with donate_segment("10.0.0.1:50051", GIB, hostname="10.0.0.5"): + pytest.fail("donation should not have yielded") + + +def test_a_donor_given_no_host_registers_under_the_pool_s_view_of_this_node( + fake_bindings, monkeypatch +): + """The master and the segments registering with it must agree on the host.""" + monkeypatch.setattr(donor_module, "local_address", lambda: "10.1.2.3") + + with donate_segment("10.0.0.1:50051", GIB) as host: + assert host == "10.1.2.3" + assert fake_bindings.instances[0].setup_args[0] == "10.1.2.3" + + +def test_missing_bindings_are_reported_as_the_separate_component_they_are(monkeypatch): + """The container's C++ transfer engine is not these Python bindings.""" + monkeypatch.setitem(sys.modules, "mooncake", None) + monkeypatch.setitem(sys.modules, "mooncake.store", None) + + with pytest.raises(ImportError, match=r"mooncake-transfer-engine-cuda13"): + with donate_segment("10.0.0.1:50051", GIB): + pytest.fail("donation should not have yielded") + + +# The address a donor joins is resolved and probed by the `mooncake_donor` +# command before it gets here; see test_mooncake_store_cli.py. diff --git a/tests/unittest/_torch/executor/test_mooncake_store_master.py b/tests/unittest/_torch/executor/test_mooncake_store_master.py new file mode 100644 index 000000000000..9eb27056a037 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -0,0 +1,630 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for provisioning a Mooncake store pool during server bringup. + +Runs without a Mooncake installation and without a GPU. A master this process +launches is a fake standing in for `Popen` that opens the RPC port, which is +all the readiness handshake ever observes. A master someone else runs is a +plain socket. +""" + +import json +import os +import shutil +import socket +import subprocess +import threading +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import master as master_module +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + CONFIG_PATH_ENV, + MooncakeStoreConnectorConfig, + provisioned_config_path, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.master import ( + PoolSpec, + provision_pool, + resolve_master_address, +) + + +def free_port() -> int: + with socket.socket() as probe: + probe.bind(("", 0)) + return probe.getsockname()[1] + + +class FakeMasterProcess: + """The slice of `Popen` that launching a master actually drives. + + `listen_on` makes it answer on that port, which is what a real master does + last and what the readiness wait keys off. `exit_code` makes it a master + that failed to start. + """ + + def __init__(self, command, env, listen_on=None, exit_code=None): + self.command = command + self.env = env + self.pid = 4242 + self.terminated = False + self.killed = False + self._exit_code = exit_code + self._listener = None + if listen_on is not None: + self._listener = socket.socket() + self._listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self._listener.bind(("", listen_on)) + self._listener.listen(8) + + def poll(self): + return self._exit_code + + def terminate(self): + self.terminated = True + self._exit_code = -15 + if self._listener is not None: + self._listener.close() + self._listener = None + + def wait(self, timeout=None): + return self._exit_code + + def kill(self): + self.killed = True + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + """No ambient pool: these tests are about what provisioning does itself.""" + for name in ( + CONFIG_PATH_ENV, + master_module.MASTER_BINARY_ENV, + master_module.MASTER_TIMEOUT_ENV, + master_module.RUN_DIR_ENV, + ): + monkeypatch.delenv(name, raising=False) + # Nothing here is slow to start, so a wait that runs long is a failure + # rather than something that needs more time. + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "10") + + +@pytest.fixture +def fake_master(monkeypatch): + """Replace the master binary and its process with in-process fakes. + + Returns a callable that arms the fake. Once provisioning has run, the + launched instance is available as `.process` for inspection. + """ + + class Launcher: + def __init__(self): + self.process = None + + def arm(self, listen_on=None, exit_code=None, log_text=None): + def popen(command, env=None, stdout=None, **_kwargs): + # A real master writes its log through this handle. + if log_text is not None and stdout is not None: + stdout.write(log_text.encode()) + stdout.flush() + self.process = FakeMasterProcess( + command, env, listen_on=listen_on, exit_code=exit_code + ) + return self.process + + # Swap the modules as this module sees them rather than patching + # attributes on the shared stdlib ones. + monkeypatch.setattr( + master_module, + "shutil", + SimpleNamespace(which=lambda name: f"/opt/bin/{name}", rmtree=shutil.rmtree), + ) + monkeypatch.setattr( + master_module, + "subprocess", + SimpleNamespace( + Popen=popen, + STDOUT=subprocess.STDOUT, + TimeoutExpired=subprocess.TimeoutExpired, + ), + ) + + return Launcher() + + +@pytest.fixture +def running_master(): + """A socket standing in for a master someone else is running.""" + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen(8) + try: + yield f"127.0.0.1:{listener.getsockname()[1]}" + finally: + listener.close() + + +# ---- configuration ---- + + +@pytest.mark.parametrize( + "kwargs, message", + [ + (dict(launch_master=True, master_server_address="host:50051"), "not both"), + (dict(), "needs a master"), + # master_address_file only writes an address; reading one is + # master_server_address, so publishing without launching is incoherent. + ( + dict(master_server_address="host:50051", master_address_file="/shared/master.addr"), + "needs launch_master", + ), + ], + ids=["two_masters", "no_master", "publishing_without_launching"], +) +def test_provisioning_needs_exactly_one_master(kwargs, message): + """Rejected before anything is launched, dialed, or written.""" + with pytest.raises(ValueError, match=message): + with provision_pool(PoolSpec(**kwargs)): + pytest.fail("provisioning should not have yielded") + assert CONFIG_PATH_ENV not in os.environ + + +# ---- the rendered client config ---- + + +def test_client_config_is_what_the_connector_reads_back(tmp_path): + """The generated JSON has to survive the connector's own parser.""" + pool = PoolSpec( + master_server_address="10.0.0.1:50051", + protocol="rdma", + device_name="mlx5_0", + global_segment_size="64GiB", + local_buffer_size="4GiB", + cache_prefix="trtllm-m3", + model_key="minimax-m3@rev7", + stage_through_host=True, + transfer_batch_size=32, + ) + path = tmp_path / "mooncake.json" + path.write_text(json.dumps(master_module._client_config(pool, "10.0.0.1:50051"))) + + parsed = MooncakeStoreConnectorConfig.from_file(str(path)) + assert parsed.master_server_address == "10.0.0.1:50051" + assert parsed.metadata_server == "P2PHANDSHAKE" + assert parsed.protocol == "rdma" + assert parsed.device_name == "mlx5_0" + assert parsed.global_segment_size == 64 * 1024**3 + assert parsed.local_buffer_size == 4 * 1024**3 + assert parsed.cache_prefix == "trtllm-m3" + # Reaches the ranks the server spawns, which otherwise have no model key + # and would refuse to name a page. + assert parsed.resolve_model_key("/models/ignored") == "minimax-m3@rev7" + assert parsed.stage_through_host is True + assert parsed.transfer_batch_size == 32 + + +def test_client_config_omits_the_fields_the_pool_left_unset(): + """An absent key leaves the connector its own default; a null would not.""" + pool = PoolSpec(master_server_address="host:50051") + written = master_module._client_config(pool, "host:50051") + assert "cache_prefix" not in written + assert "model_key" not in written + assert "staging_buffer_bytes" not in written + + +def test_a_written_client_config_is_where_the_ranks_look_for_it(tmp_path, monkeypatch): + """A rank an external launcher started never inherits MOONCAKE_CONFIG_PATH. + + Writing the config into the run directory is what lets + 'trtllm-serve mooncake_master --run_dir' describe a pool to servers it + did not spawn. + """ + pool = PoolSpec(master_server_address="10.0.0.1:50051", model_key="m") + + written = master_module.write_client_config(pool, "10.0.0.1:50051", str(tmp_path)) + + monkeypatch.setenv(master_module.RUN_DIR_ENV, str(tmp_path)) + assert provisioned_config_path() == written + assert MooncakeStoreConnectorConfig.from_file(written).model_key == "m" + + +def test_a_spec_takes_the_pool_settings_from_the_json_the_connector_reads(): + """One schema for the pool, rather than a second spelling of it.""" + pool = PoolSpec.from_json( + { + "protocol": "tcp", + "global_segment_size": "8GiB", + "model_key": "m", + # Not a PoolSpec field; the client config carries settings that + # only the connector reads, and they must not break the parse. + "role": "both", + }, + launch_master=True, + ) + + assert pool.protocol == "tcp" + assert pool.global_segment_size == "8GiB" + assert pool.model_key == "m" + assert pool.launch_master is True + + +@pytest.mark.parametrize( + "address, expected", + [ + ("host:50051", ("host", 50051)), + ("[::1]:50051", ("::1", 50051)), + ("unix:///var/run/mooncake", None), + ("host", None), + ], +) +def test_master_addresses_are_split_or_declined(address, expected): + assert master_module._split_address(address) == expected + + +# ---- provisioning against a master someone else runs ---- + + +def test_provisioning_points_the_workers_at_a_running_master(running_master): + pool = PoolSpec(master_server_address=running_master) + + with provision_pool(pool) as config_path: + # The workers are spawned inside this window and are told about the + # pool through the environment, so both have to hold while it is open. + assert os.environ[CONFIG_PATH_ENV] == config_path + written = json.loads(open(config_path).read()) + assert written["master_server_address"] == running_master + + assert CONFIG_PATH_ENV not in os.environ + assert not os.path.exists(config_path) + + +def test_a_staging_buffer_can_be_sized_where_staging_is_turned_on(running_master): + """Undersizing it silently shrinks the transfer batch, so it must be settable.""" + pool = PoolSpec( + master_server_address=running_master, + stage_through_host=True, + staging_buffer_bytes="4GiB", + ) + + with provision_pool(pool) as config_path: + written = json.loads(open(config_path).read()) + assert written["stage_through_host"] is True + assert written["staging_buffer_bytes"] == "4GiB" + + +def test_provisioning_fails_before_the_model_loads_if_the_master_is_absent(monkeypatch): + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") + pool = PoolSpec(master_server_address=f"127.0.0.1:{free_port()}") + + with pytest.raises(TimeoutError, match="did not accept connections"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + assert CONFIG_PATH_ENV not in os.environ + + +def test_an_unparseable_master_address_is_left_to_the_workers(): + """Not every address is host:port, so an unprobeable one passes through.""" + pool = PoolSpec(master_server_address="unix:///var/run/mooncake") + + with provision_pool(pool) as config_path: + written = json.loads(open(config_path).read()) + assert written["master_server_address"] == "unix:///var/run/mooncake" + + +def test_an_inherited_config_path_wins(monkeypatch, tmp_path): + """An externally managed pool names itself this way, so provisioning defers.""" + harness_config = tmp_path / "harness.json" + harness_config.write_text("{}") + monkeypatch.setenv(CONFIG_PATH_ENV, str(harness_config)) + pool = PoolSpec(launch_master=True) + + with provision_pool(pool) as config_path: + assert config_path is None + assert os.environ[CONFIG_PATH_ENV] == str(harness_config) + + assert os.environ[CONFIG_PATH_ENV] == str(harness_config) + + +# ---- provisioning with a master of our own ---- + + +def test_a_launched_master_is_named_in_the_config_and_stopped_on_exit(fake_master): + port = free_port() + fake_master.arm(listen_on=port) + pool = PoolSpec(launch_master=True, master_port=port) + + with provision_pool(pool) as config_path: + written = json.loads(open(config_path).read()) + host, _, named_port = written["master_server_address"].rpartition(":") + assert int(named_port) == port + # The address in the config is all a worker on another host gets, so + # it has to be dialable. + with socket.create_connection((host, port), timeout=5): + pass + + assert fake_master.process.terminated + assert not fake_master.process.killed + + +def test_a_launched_master_gets_the_flags_and_logging_it_needs(fake_master): + port = free_port() + fake_master.arm(listen_on=port) + pool = PoolSpec( + launch_master=True, + master_port=port, + master_metrics_port=free_port(), + master_eviction_ratio=0.1, + ) + + with provision_pool(pool): + command = fake_master.process.command + assert command[0].endswith("mooncake_master") + assert f"--rpc_port={port}" in command + assert f"--metrics_port={pool.master_metrics_port}" in command + assert "--eviction_ratio=0.1" in command + # Without these the master logs to a file under /tmp and the log the + # run directory holds stays empty. + assert fake_master.process.env["GLOG_logtostderr"] == "1" + assert fake_master.process.env["GLOG_v"] == "1" + + +def test_a_master_that_dies_during_startup_says_so(fake_master): + fake_master.arm(exit_code=3) + pool = PoolSpec(launch_master=True, master_port=free_port()) + + with pytest.raises(RuntimeError, match="exited with code 3"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + assert CONFIG_PATH_ENV not in os.environ + + +def test_a_master_that_never_listens_times_out(monkeypatch, fake_master): + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") + fake_master.arm() + pool = PoolSpec(launch_master=True, master_port=free_port()) + + with pytest.raises(TimeoutError, match="did not accept connections"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + assert fake_master.process.terminated + + +def test_a_missing_master_binary_names_the_alternatives(monkeypatch): + monkeypatch.setattr( + master_module, + "shutil", + SimpleNamespace(which=lambda _name: None, rmtree=shutil.rmtree), + ) + pool = PoolSpec(launch_master=True) + + with pytest.raises(FileNotFoundError, match="master_server_address"): + with provision_pool(pool): + pytest.fail("provisioning should not have yielded") + + +def test_a_run_dir_keeps_the_master_log_and_the_config(fake_master, tmp_path): + port = free_port() + fake_master.arm(listen_on=port) + run_dir = tmp_path / "pool" + pool = PoolSpec(launch_master=True, master_port=port) + + with provision_pool(pool, run_dir=str(run_dir)) as config_path: + assert config_path == str(run_dir / master_module.CLIENT_CONFIG_NAME) + + # An explicit run directory outlives the run that filled it, since the + # master's log is where pool occupancy and eviction are read from. + assert (run_dir / master_module.MASTER_LOG_NAME).exists() + assert (run_dir / master_module.CLIENT_CONFIG_NAME).exists() + + +# ---- reaching a master whose host nobody knew in advance ---- + + +@pytest.mark.parametrize("address", ["10.0.0.1:50051", "unix:///var/run/mooncake"]) +def test_an_address_that_is_not_a_file_passes_through(address): + assert resolve_master_address(address, timeout=1.0) == address + + +def test_a_published_address_is_read_from_the_file_that_names_it(tmp_path): + published = tmp_path / "master.addr" + published.write_text("10.0.0.7:50051\n") + + assert resolve_master_address(f"file://{published}", timeout=1.0) == "10.0.0.7:50051" + + +def test_an_address_not_published_yet_is_waited_for(tmp_path): + """Master and workers are started together; neither one orders the other.""" + published = tmp_path / "master.addr" + threading.Timer(0.5, published.write_text, ["10.0.0.9:50051\n"]).start() + + assert resolve_master_address(f"file://{published}", timeout=10.0) == "10.0.0.9:50051" + + +def test_an_empty_address_file_is_not_taken_for_an_address(tmp_path): + """An existing file is not the same as a published address.""" + published = tmp_path / "master.addr" + published.write_text("") + + with pytest.raises(TimeoutError, match="No Mooncake master address"): + resolve_master_address(f"file://{published}", timeout=1.0) + + +def test_an_unpublished_address_names_the_command_that_publishes_it(tmp_path): + with pytest.raises(TimeoutError, match="--address-file"): + resolve_master_address(f"file://{tmp_path / 'absent'}", timeout=1.0) + + +# ---- a master with a lifetime of its own ---- + + +def test_a_standalone_master_publishes_an_address_that_can_be_dialed(fake_master, tmp_path): + port = free_port() + fake_master.arm(listen_on=port) + address_file = tmp_path / "master.addr" + pool = PoolSpec(launch_master=True, master_port=port) + + with master_module.running_master( + pool, str(tmp_path / "run"), address_file=str(address_file) + ) as master: + assert resolve_master_address(f"file://{address_file}", timeout=5.0) == master.address + host, _, named_port = master.address.rpartition(":") + assert int(named_port) == port + with socket.create_connection((host, port), timeout=5): + pass + + +def test_a_stopped_master_leaves_no_address_behind(fake_master, tmp_path): + """A stale address would send the next run's workers to a dead port.""" + port = free_port() + fake_master.arm(listen_on=port) + address_file = tmp_path / "master.addr" + pool = PoolSpec(launch_master=True, master_port=port) + + with master_module.running_master(pool, str(tmp_path / "run"), address_file=str(address_file)): + assert address_file.exists() + + assert not address_file.exists() + assert fake_master.process.terminated + + +def test_a_standalone_master_keeps_its_log(fake_master, tmp_path): + """A standalone master outlives the servers that used it, so its log is kept.""" + port = free_port() + fake_master.arm(listen_on=port) + run_dir = tmp_path / "run" + pool = PoolSpec(launch_master=True, master_port=port) + + with master_module.running_master(pool, str(run_dir)): + pass + + assert (run_dir / master_module.MASTER_LOG_NAME).exists() + + +def test_provisioning_joins_a_master_it_was_never_given_the_address_of(fake_master, tmp_path): + """The address file is how workers reach a master no config names a host for.""" + port = free_port() + fake_master.arm(listen_on=port) + address_file = tmp_path / "master.addr" + standalone = PoolSpec(launch_master=True, master_port=port) + worker = PoolSpec(master_server_address=f"file://{address_file}") + + with master_module.running_master( + standalone, str(tmp_path / "run"), address_file=str(address_file) + ) as master: + with provision_pool(worker) as config_path: + # Mooncake cannot dial a file:// URL, so what reaches the workers + # has to be the address it resolved to. + written = json.loads(open(config_path).read()) + assert written["master_server_address"] == master.address + + +# ---- a master a server launched, made findable ---- + + +def test_a_launched_master_publishes_where_its_run_left_its_logs(fake_master, tmp_path): + """A finished run's logs still say which pool it used.""" + port = free_port() + fake_master.arm(listen_on=port) + run_dir = tmp_path / "run" + pool = PoolSpec(launch_master=True, master_port=port) + + with provision_pool(pool, run_dir=str(run_dir)): + address = (run_dir / master_module.MASTER_ADDRESS_NAME).read_text().strip() + assert address.endswith(f":{port}") + + assert not (run_dir / master_module.MASTER_ADDRESS_NAME).exists() + + +def test_a_launched_master_can_be_published_where_the_donors_look(fake_master, tmp_path): + """This is what lets a server that launched its own master have donors.""" + port = free_port() + fake_master.arm(listen_on=port) + shared = tmp_path / "shared" / "master.addr" + pool = PoolSpec(launch_master=True, master_port=port, master_address_file=str(shared)) + + with provision_pool(pool, run_dir=str(tmp_path / "run")): + assert resolve_master_address(f"file://{shared}", timeout=5.0).endswith(f":{port}") + + # Retracted, so the next run's donors wait for a live master rather than + # joining a pool that no longer exists. + assert not shared.exists() + + +def test_a_half_written_address_is_never_read(tmp_path): + """A reader sees the whole address or nothing, never a prefix of one.""" + target = tmp_path / "master.addr" + + with master_module._published_address("10.0.0.7:50051", [str(target)]): + assert not (tmp_path / "master.addr.partial").exists() + assert target.read_text().strip() == "10.0.0.7:50051" + + +# ---- saying why bringup is stuck ---- + + +def test_an_absent_master_is_named_rather_than_left_to_store_setup(monkeypatch): + """Otherwise the failure is a bare status code in every rank, after loading.""" + monkeypatch.setenv(master_module.MASTER_TIMEOUT_ENV, "1") + address = f"127.0.0.1:{free_port()}" + + with pytest.raises(TimeoutError, match=address): + master_module.wait_for_master(address) + + +def test_an_address_of_a_shape_we_cannot_probe_is_not_fatal(): + """Mooncake may accept addresses this cannot dial; leave them to it.""" + assert master_module.wait_for_master("unix:///var/run/mooncake") is None + + +def test_a_master_that_died_starting_is_reported_with_its_last_words(fake_master, tmp_path): + """The reason is in the master's log, which is only read if the error quotes it.""" + run_dir = tmp_path / "run" + fake_master.arm(exit_code=1, log_text="E0903 bind(50051) failed: Address already in use\n") + pool = PoolSpec(launch_master=True, master_port=free_port()) + + with pytest.raises(RuntimeError, match="Address already in use"): + with provision_pool(pool, run_dir=str(run_dir)): + pytest.fail("provisioning should not have yielded") + + +# ---- which fabric a client transfers over ---- + + +def test_the_pool_config_leaves_the_fabric_to_each_client(fake_master, tmp_path): + """One config describes a whole pool, so it must not name one node's HCAs. + + Mooncake discovers devices locally on each client when `device_name` is + empty. Resolving a device while writing the config would instead freeze + whatever the provisioning host happens to have into a file every node + reads, which is wrong as soon as their HCAs differ. + """ + port = free_port() + fake_master.arm(listen_on=port) + pool = PoolSpec(launch_master=True, master_port=port, protocol="rdma") + + with provision_pool(pool, run_dir=str(tmp_path / "run")) as config_path: + assert json.loads(open(config_path).read())["device_name"] == "" + + +def test_a_named_device_is_what_the_workers_are_told(fake_master, tmp_path): + """Naming one is still how a deployment overrides that discovery.""" + port = free_port() + fake_master.arm(listen_on=port) + pool = PoolSpec(launch_master=True, master_port=port, protocol="rdma", device_name="mlx5_7") + + with provision_pool(pool, run_dir=str(tmp_path / "run")) as config_path: + assert json.loads(open(config_path).read())["device_name"] == "mlx5_7"