diff --git a/docker/common/install_mooncake.sh b/docker/common/install_mooncake.sh index 0935c7f31a2f..a145b3756f62 100644 --- a/docker/common/install_mooncake.sh +++ b/docker/common/install_mooncake.sh @@ -53,3 +53,60 @@ cd ../.. rm -rf Mooncake echo "export LD_LIBRARY_PATH=${MOONCAKE_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> "${ENV}" + +# The source build above provides only the C++ transfer engine, which is what +# the cache transceiver links against. MooncakeDistributedStore, the shared CPU +# pool behind the mooncake-store KV cache connector, comes from the Python +# wheel instead, for two reasons. +# +# First, `make install` emits a `mooncake` Python package that omits +# libmooncake_store.so, so importing mooncake.store raises ImportError. It has +# to be removed wherever it landed, and where that is depends on the +# environment: mooncake-integration/CMakeLists.txt picks its install directory +# as the first sys.path entry whose name merely contains "packages". +# +# - With nvidia-cutlass-dsl installed, that is +# nvidia_cutlass_dsl/dsl_packages, which nvidia_cutlass_dsl_packages.pth +# puts at sys.path[0], so it shadows anything pip installs. CUTLASS DSL +# does not reference `mooncake`, so removing the package is safe. +# - Without it, the package lands in dist-packages and collides with the +# wheel: CMake writes store.cpython-312-x86_64-linux-gnu.so, the wheel +# writes store.so, and importlib prefers the interpreter-tagged suffix, so +# the broken extension wins even after pip reports success. +# +# Remove the directory outright rather than trying to identify leftovers, since +# pip overwrites __init__.py in the collision case and leaves no marker to key +# on. +python3 - <<'PY' +import os +import shutil +import sys +import sysconfig + +paths = sysconfig.get_paths() +for entry in list(sys.path) + [paths["purelib"], paths["platlib"]]: + if not entry: + continue + package = os.path.join(entry, "mooncake") + if os.path.isdir(package): + print(f"removing CMake-generated mooncake package: {package}") + shutil.rmtree(package, ignore_errors=True) +PY + +# Second, the `mooncake-transfer-engine` wheel is built against CUDA 12 while +# these images ship CUDA 13 only, so its extensions cannot resolve +# libcudart.so.12. `mooncake-transfer-engine-cuda13` is the same project built +# for CUDA 13. It is versioned independently, with releases starting at 0.3.9, +# so it cannot track MOONCAKE_VERSION above. The store client only has to agree +# with the mooncake_master it connects to, and this wheel supplies both. +MOONCAKE_WHEEL_VERSION="0.3.13" +pip3 install --no-cache-dir "mooncake-transfer-engine-cuda13==${MOONCAKE_WHEEL_VERSION}" + +# Fail the build rather than ship an image whose import is broken. +python3 - <<'PY' +from mooncake.store import MooncakeDistributedStore +import mooncake.store + +MooncakeDistributedStore() +print(f"mooncake.store OK: {mooncake.store.__file__}") +PY diff --git a/scripts/attribution/scan/metadata/mooncake.yml b/scripts/attribution/scan/metadata/mooncake.yml index c8d51e0f81b8..b1e6dab2c141 100644 --- a/scripts/attribution/scan/metadata/mooncake.yml +++ b/scripts/attribution/scan/metadata/mooncake.yml @@ -1,5 +1,8 @@ name: mooncake -description: Mooncake transfer engine for distributed KV cache +description: Mooncake transfer engine and distributed store for distributed KV cache source: container directory_matches: - /usr/local/Mooncake +- mooncake +basename_matches: +- mooncake_transfer_engine 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..be28d0590594 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/__init__.py @@ -0,0 +1,71 @@ +# 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. + +The pool is described in `KvCacheConnectorConfig.mooncake_store`, which lets +`trtllm-serve` provision it during bringup so no external script has to; see +`master.py`. 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. Both need the +Mooncake Python bindings (`pip install mooncake-transfer-engine`). + +`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. +`connector.py` holds the connector classes, which are placeholders. +""" + +from .config import MooncakeStoreConnectorConfig, StoreRole, parse_size +from .connector import MooncakeStoreConnectorScheduler, MooncakeStoreConnectorWorker +from .donor import DEFAULT_DONOR_LOCAL_BUFFER_SIZE, donate_segment, maybe_donate_segment +from .master import ( + local_address, + master_timeout, + maybe_provision_pool, + provision_pool, + resolve_device_name, + resolve_master_address, + running_master, + wait_for_master, +) + +__all__ = [ + "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", + "MooncakeStoreConnectorConfig", + "MooncakeStoreConnectorScheduler", + "MooncakeStoreConnectorWorker", + "StoreRole", + "donate_segment", + "local_address", + "master_timeout", + "maybe_donate_segment", + "maybe_provision_pool", + "parse_size", + "provision_pool", + "resolve_device_name", + "resolve_master_address", + "running_master", + "wait_for_master", +] 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/connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/connector.py new file mode 100644 index 000000000000..2e5dc2e66cc0 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/connector.py @@ -0,0 +1,54 @@ +# 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. +"""Placeholders for the connector that moves KV pages in and out of the pool. + +The store side stands on its own: a pool can be provisioned, lent memory, and +shared with an engine elsewhere without a connector here. Moving pages needs the +KV cache layout description, so neither class is implemented here. + +`CONNECTOR_REGISTRY` names both and `py_executor_creator` resolves them by name, +so they exist to turn `connector: mooncake-store` into a clear refusal instead +of an `AttributeError` raised once the pool is already provisioned. + +Neither subclasses `KvCacheConnectorWorker` or `KvCacheConnectorScheduler`, +whose `__init_subclass__` checks a method set there is nothing here to satisfy. +""" + +__all__ = ["MooncakeStoreConnectorScheduler", "MooncakeStoreConnectorWorker"] + +_UNAVAILABLE = ( + "The mooncake-store KV cache connector is not available in this build, so an " + "engine here cannot read or write the pool. The pool itself is: " + "'trtllm-serve mooncake_master' owns one, 'trtllm-serve mooncake_donor' and " + "the mooncake_donation setting lend it host memory, and an engine that does " + "have the connector can share it. Drop kv_connector_config.connector to " + "serve without one." +) + + +class _Unavailable: + """Reports what this build can and cannot do with a Mooncake pool.""" + + def __init__(self, llm_args): + del llm_args + raise NotImplementedError(_UNAVAILABLE) + + +class MooncakeStoreConnectorWorker(_Unavailable): + """Moves this rank's KV pages between its cache and the pool.""" + + +class MooncakeStoreConnectorScheduler(_Unavailable): + """Decides which pages a request loads from the pool and saves to it.""" 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..4fe167a5bdaf --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/donor.py @@ -0,0 +1,171 @@ +# 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. + +Donating alongside a generation server 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 Any, Iterator, Optional + +from tensorrt_llm.logger import logger + +from .config import DEFAULT_METADATA_SERVER, parse_size +from .master import ( + local_address, + master_timeout, + resolve_device_name, + resolve_master_address, + wait_for_master, +) + +__all__ = [ + "DEFAULT_DONOR_LOCAL_BUFFER_SIZE", + "donate_segment", + "maybe_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 " + "(`pip install mooncake-transfer-engine`). The C++ transfer engine " + f"in the container is a different 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" + ) + + +@contextlib.contextmanager +def maybe_donate_segment(donation: Any) -> Iterator[Optional[str]]: + """Lend memory for this process's lifetime if the config asked to. + + Args: + donation: A `MooncakeDonationConfig`, or `None` to do nothing, so + callers need no condition of their own. + + Yields the host the segment is registered under, or `None`. + """ + if donation is None: + yield None + return + + # Bringup blocks here on a master that may belong to a different job, so + # name the address before waiting on it. + logger.info( + "mooncake-store: mooncake_donation is set, so this server lends host " + f"memory to the pool at {donation.master_server_address} without using " + "it; resolving the master now" + ) + master_address = resolve_master_address(donation.master_server_address, master_timeout()) + wait_for_master(master_address) + with donate_segment( + master_server_address=master_address, + segment_size=parse_size(donation.segment_size), + protocol=donation.protocol, + device_name=resolve_device_name(donation.protocol, donation.device_name), + metadata_server=donation.metadata_server, + ) as host: + yield host 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..93f1f4843182 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -0,0 +1,598 @@ +# 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 the Mooncake store's pool up as part of a server's own startup. + +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` does that work inside the serving process. It resolves the +master, either launching one here or checking that the configured one answers, +renders the client config, and exports `MOONCAKE_CONFIG_PATH`, which reaches the +ranks because the LLM constructor spawns them from this process. Everything it +started is torn down when the context exits. + +A master launched here lives and dies with the server, so it suits one engine +talking to its own pool. Several engines sharing a pool, or a pool meant to +survive a restart, need a master with its own lifetime named by +`master_server_address`. +""" + +import contextlib +import json +import os +import shutil +import socket +import subprocess # nosec B404 +import tempfile +import time +from dataclasses import dataclass +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple + +from tensorrt_llm.logger import logger + +from ..registry import uses_connector +from .config import CLIENT_CONFIG_NAME, CONFIG_PATH_ENV, RUN_DIR_ENV + +__all__ = [ + "local_address", + "maybe_provision_pool", + "master_timeout", + "provision_pool", + "resolve_device_name", + "resolve_master_address", + "running_master", + "wait_for_master", +] + +#: 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 + + +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) + + +#: Where the InfiniBand devices of a host are described. +IB_SYSFS_ROOT = "/sys/class/infiniband" + + +def _highest_rate_ib_devices(sysfs_root: Optional[str] = None) -> List[str]: + """The active InfiniBand devices on the compute fabric, fastest first. + + A node's HCAs are not interchangeable. On GB300 six are exposed, of which + four run at 800Gb/s (two per NUMA node, one per GPU) while the rest share a + PCI device with an Ethernet port and serve storage or management. Taking + every device at the highest rate picks the compute fabric on any node type, + where a hardcoded name would be wrong on the next one. + """ + sysfs_root = sysfs_root or IB_SYSFS_ROOT + rated: Dict[str, int] = {} + try: + devices = sorted(os.listdir(sysfs_root)) + except OSError: + return [] + for device in devices: + port = os.path.join(sysfs_root, device, "ports", "1") + + def attribute(name: str) -> str: + try: + with open(os.path.join(port, name)) as handle: + return handle.read().strip() + except OSError: + return "" + + if attribute("link_layer") != "InfiniBand": + continue + if "ACTIVE" not in attribute("state"): + continue + # "800 Gb/sec (4X XDR)" + rate = attribute("rate").split() + if not rate or not rate[0].isdigit(): + continue + rated[device] = int(rate[0]) + + if not rated: + return [] + fastest = max(rated.values()) + return [device for device, rate in sorted(rated.items()) if rate == fastest] + + +def resolve_device_name(protocol: str, configured: str, sysfs_root: Optional[str] = None) -> str: + """The RDMA devices to transfer over, detected if the config left it open. + + Which HCAs a node has is a property of the node, not of the deployment, so + requiring it in a config would tie that config to one machine type. + Detecting it keeps `protocol: rdma` portable; setting `device_name` + overrides the detection. + """ + if configured or protocol != "rdma": + return configured + detected = _highest_rate_ib_devices(sysfs_root) + if not detected: + logger.warning( + "mooncake-store: protocol is rdma but no active InfiniBand device " + f"was found under {sysfs_root or IB_SYSFS_ROOT}, so device_name is " + "left empty for Mooncake's own discovery. Set device_name to " + "choose explicitly." + ) + return "" + joined = ",".join(detected) + logger.info( + f"mooncake-store: transferring over the fastest active InfiniBand " + f"devices on this host: {joined}" + ) + return joined + + +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: Any, master_address: str, device_name: Optional[str] = None +) -> 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 if device_name is None else 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 + + +@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: Any, 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 ships with the Mooncake runtime, which " + "docker/common/install_mooncake.sh installs. Point " + f"{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: Any, 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: Any, 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: A `MooncakeStoreConfig`. + 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. + """ + inherited = os.getenv(CONFIG_PATH_ENV) + if inherited: + logger.info( + f"mooncake-store: {CONFIG_PATH_ENV}={inherited} is already set, so " + "kv_connector_config.mooncake_store is ignored and the pool it " + "names is used as is." + ) + 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 = os.path.join(run_dir, CLIENT_CONFIG_NAME) + config = _client_config( + pool, master_address, resolve_device_name(pool.protocol, pool.device_name) + ) + with open(config_path, "w") as handle: + json.dump(config, handle, indent=2) + # 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} " + f"({json.dumps(config, sort_keys=True)})" + ) + # Capacity is what explains a low hit rate, so state the + # arithmetic instead of leaving it to be derived later. + logger.info( + "mooncake-store: this server's ranks will each contribute " + f"global_segment_size={pool.global_segment_size} to the pool; " + "total capacity is that times the number of ranks that open a " + "handle, plus whatever any mooncake_donation adds" + ) + 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) + + +@contextlib.contextmanager +def maybe_provision_pool(kv_connector_config: Any) -> Iterator[None]: + """Provision the pool if this deployment asked the server to. + + A no-op for every other connector, and for a `mooncake-store` config that + left `mooncake_store` unset, since such a deployment is told about its pool + through `MOONCAKE_CONFIG_PATH` instead. + """ + if not uses_connector(kv_connector_config, "mooncake-store"): + yield + return + pool = kv_connector_config.mooncake_store + if pool is None: + yield + return + with provision_pool(pool): + yield 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/connectors/mooncake_store/validation.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py new file mode 100644 index 000000000000..4df451e42b5a --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py @@ -0,0 +1,60 @@ +# 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. +"""Startup gates for the Mooncake store connector. + +Every rejection here is a configuration whose failure mode is a wrong answer +rather than a slow one: KV that gets replayed without all of the state it was +computed with. Beam search, attention data parallelism, host and disk cache +tiers, and Mamba caches are rejected for all connectors in `py_executor`, so +they are not repeated. + +Checks run at construction, before any request is admitted, so a bad deployment +fails at startup instead of after the first cache hit. +""" + +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + +__all__ = ["validate_llm_args"] + + +def validate_llm_args(llm_args: TorchLlmArgs) -> None: + """Reject parallel and model configurations this connector cannot serve.""" + if getattr(llm_args, "context_parallel_size", 1) > 1: + raise NotImplementedError( + "The mooncake-store connector does not support context parallelism. " + "A stored page is keyed by the tokens it holds, but under context " + "parallelism a rank holds a slice of the sequence rather than whole " + "blocks of it, so the same key would name different bytes on " + "different ranks." + ) + + if getattr(llm_args, "pipeline_parallel_size", 1) > 1: + raise NotImplementedError( + "The mooncake-store connector does not support pipeline parallelism. " + "Keys are namespaced per rank, so each stage would store only its own " + "layers and a prefix hit would require every stage to agree; that path " + "is untested. Run with tensor parallelism only." + ) + + sparse_config = getattr(llm_args, "sparse_attention_config", None) + if sparse_config is not None and not getattr(sparse_config, "sparse_disable_index_value", True): + raise NotImplementedError( + "The mooncake-store connector requires " + "sparse_attention_config.sparse_disable_index_value=True. The index-V " + "cache is a plain tensor outside the KV cache manager's paged pools, " + "so it is neither described to the connector nor transferred; a " + "replayed prefix would carry index-K from the store alongside stale " + "index-V. This is the same restriction disaggregated serving applies." + ) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py index 9a00cdadd7fd..c0a40bf480a6 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/registry.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/registry.py @@ -19,6 +19,11 @@ it is resolved at runtime via importlib in py_executor_creator.py. """ +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig + CONNECTOR_REGISTRY: dict[str, dict[str, str]] = { "lmcache": { "connector_module": "lmcache.integration.tensorrt_llm.tensorrt_adapter", @@ -35,4 +40,27 @@ "connector_scheduler_class": "DynamoKVBMConnectorLeader", "connector_worker_class": "DynamoKVBMConnectorWorker", }, + # Both classes currently refuse construction; see the module's docstring. + "mooncake-store": { + "connector_module": "tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", + "connector_scheduler_class": "MooncakeStoreConnectorScheduler", + "connector_worker_class": "MooncakeStoreConnectorWorker", + }, } + + +def uses_connector(kv_connector_config: Optional["KvCacheConnectorConfig"], name: str) -> bool: + """Report whether a connector config resolves to the named preset. + + Compares the resolved module rather than the `connector` field, so a config + that names the module explicitly instead of using the preset is still + recognized. Accepts `None` to save every caller a null check. + """ + if kv_connector_config is None: + return False + preset = CONNECTOR_REGISTRY.get(name) + if preset is None: + raise ValueError( + f"Unknown connector preset: {name!r}. Known presets: {list(CONNECTOR_REGISTRY)}" + ) + return kv_connector_config.connector_module == preset["connector_module"] 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 2b47f3f5982d..193a5809fe68 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 @@ -3728,6 +3728,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_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 6ece1cc72c3c..6d2ac676f422 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -46,6 +46,7 @@ resolve_cache_transceiver_config, uses_vswa_kv_cache_layout) from .connectors.kv_cache_connector import KvCacheConnectorManager +from .connectors.registry import uses_connector from .dwdp import DwdpManager, get_global_dwdp_manager from .guided_decoder import CapturableGuidedDecoder, GuidedDecoder from .hang_diagnostics import monitor_executor_initialization @@ -394,6 +395,19 @@ def _create_py_executor_impl( kv_cache_config.enable_block_reuse = False kv_cache_config.enable_partial_reuse = False + # Must happen before the KV cache manager is built, since the manager reads + # enable_partial_reuse to construct its block pools. + if (kv_cache_config.enable_partial_reuse + and uses_connector(kv_connector_config, "mooncake-store")): + logger.warning( + "Disabling partial reuse: it is not usable with the mooncake-store " + "connector. The store is addressed by whole blocks, so a partial " + "device match leaves the matched length off a block boundary and " + "the connector declines the lookup rather than resume a block from " + "the middle. Partial reuse therefore trades part of one block for " + "every stored block of the remaining prefix.") + kv_cache_config.enable_partial_reuse = False + # The tokenizer is stripped from MPI kwargs in proxy.py to avoid pickle # failures with trust_remote_code models. Reload it from the checkpoint # when guided decoding needs it. diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index d43186fc72b0..e313605df255 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -15,7 +15,7 @@ import enum import os -from typing import Optional +from typing import Callable, Optional from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, ContextChunkingPolicy from tensorrt_llm.logger import logger @@ -170,6 +170,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 ) @@ -443,6 +444,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 +478,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 +495,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 +517,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 +529,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,48 +553,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 - ) - if ( - num_gen_candidates > 0 - and not evicted - and not recompute_paused - and not inflight_request_ids - ): - # 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, @@ -733,7 +738,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). @@ -754,9 +762,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 @@ -773,7 +781,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. @@ -812,6 +823,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) @@ -832,7 +849,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. @@ -890,6 +910,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( @@ -917,6 +940,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) @@ -1380,6 +1406,153 @@ 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: + 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 + ) + 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. + 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. + if any(req.state_value in self._TRANSFER_HOLDING_STATE_VALUES for req in active_requests): + self._stalled_schedules = 0 + return + + self._stalled_schedules += 1 + if self._stalled_schedules < self._DEADLOCK_STALL_ITERS: + 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." + ) + 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 _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..5aa6c7d05bf6 --- /dev/null +++ b/tensorrt_llm/commands/mooncake.py @@ -0,0 +1,365 @@ +# 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 outlive any one engine. + +A server that owns its pool needs neither: it describes the pool in +`kv_connector_config.mooncake_store` and `trtllm-serve` provisions it during +bringup. These commands exist for the pools it cannot own, such as one shared by +several engines, one that has to survive a restart, or one whose capacity comes +from nodes that run 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 + + +@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( + "--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], + heartbeat_seconds: int, + telemetry: bool, +): + """Run a mooncake_master for as long as this command runs. + + A single server with a pool of its own should set + `mooncake_store.launch_master` instead. + """ + _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 running_master + from tensorrt_llm.llmapi.llm_args import MooncakeStoreConfig + + pool = MooncakeStoreConfig( + 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: + 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", "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, + ): + if ready_file: + with open(ready_file, "w") as handle: + handle.write(f"{host} {donating}\n") + logger.info( + f"mooncake-store: announced this segment in " + f"{ready_file}, so a launcher waiting on the pool's " + "capacity can proceed" + ) + + # 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 b4007ceea6c1..5bd0f74e0142 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -18,8 +18,8 @@ from importlib.util import find_spec from pathlib import Path from types import FrameType -from typing import (TYPE_CHECKING, Any, Dict, NamedTuple, NoReturn, Optional, - Sequence, Set) +from typing import (TYPE_CHECKING, Any, Dict, Iterator, NamedTuple, NoReturn, + Optional, Sequence, Set) import click import torch @@ -33,6 +33,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 @@ -44,7 +45,9 @@ parse_disagg_config_file, parse_metadata_server_config_file, validate_config_bool) -from tensorrt_llm.llmapi.llm_args import MultimodalConfig, TorchLlmArgs +from tensorrt_llm.llmapi.llm_args import (KvCacheConnectorConfig, + MooncakeDonationConfig, + MultimodalConfig, TorchLlmArgs) from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr, split_mpi_env from tensorrt_llm.llmapi.reasoning_parser import (ReasoningParserFactory, @@ -611,6 +614,49 @@ def _terminate_attached_frontends(children: list) -> None: child.kill() +@contextlib.contextmanager +def _provision_kv_cache_pool(llm_args: dict, + owns_engine: bool = True) -> Iterator[None]: + """Bring up the shared cache this server needs or feeds, for its lifetime. + + A connector backed by a cluster-wide pool needs that pool reachable before + any rank opens a handle, and the ranks are spawned by the LLM constructor, + so this wraps the construction. A deployment that provisions the pool + externally is detected and left alone. + + A server may also lend the pool host memory without using it, which is how + a pool spans nodes whose engines have no connector. That segment has to be + mounted before traffic arrives and held for as long as the pages placed in + it are expected to be there. + + Only the process that owns the engine does either. An attached frontend + re-execs this command line but shares the launcher's executor, so it would + otherwise stand up a second pool and lend a second segment. + """ + if not owns_engine: + yield + return + + from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import ( + maybe_donate_segment, maybe_provision_pool) + + connector_config = llm_args.get("kv_connector_config") + if isinstance(connector_config, dict): + # A YAML config section arrives unvalidated, and the pool has to be + # described before the LLM constructor would coerce it. Hand the + # validated model on so it is not parsed twice. + connector_config = KvCacheConnectorConfig(**connector_config) + llm_args["kv_connector_config"] = connector_config + + donation = llm_args.get("mooncake_donation") + if isinstance(donation, dict): + donation = MooncakeDonationConfig(**donation) + llm_args["mooncake_donation"] = donation + + with maybe_provision_pool(connector_config), maybe_donate_segment(donation): + yield + + def launch_server( host: str, port: int, @@ -692,54 +738,56 @@ def launch_server( # until uvicorn takes it over, so no one can steal the port in between. _publish_bound_address(report_addr, host, port) - if backend == 'pytorch': - llm_args.pop("build_config", None) - llm = PyTorchLLM(**llm_args) - elif backend == '_autodeploy': - from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM - - # AutoDeploy does not support build_config - llm_args.pop("build_config", None) - llm = AutoDeployLLM(**llm_args) - else: - raise click.BadParameter( - f"{backend} is not a known backend, check help for available options.", - param_hint="backend") - - # The finally below is the cleanup boundary for the attached - # frontends: it must cover everything from their spawn through - # server construction, middleware registration, and runtime, or a - # failure in between leaks the child processes. - frontend_children = [] - try: - if multi_frontend.is_launcher: - frontend_children = _spawn_attached_frontends( - llm, multi_frontend.num_frontends) - - server = OpenAIServer( - generator=llm, - model=model, - tool_parser=tool_parser, - server_role=server_role, - metadata_server_cfg=metadata_server_cfg, - disagg_cluster_config=disagg_cluster_config, - multimodal_server_config=multimodal_server_config, - chat_template=chat_template, - allow_request_chat_template=allow_request_chat_template, - input_processor_workers=num_input_processor_workers, - media_load_workers=num_media_load_workers, - internal_disagg_auth_key=internal_disagg_auth_key) - _apply_fastapi_middlewares(server.app, middleware) - - # Optionally disable GC (default: not disabled) - if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": - gc.disable() - - _signal_frontend_ready(multi_frontend) - uvloop.run(server(host, port, sockets=[s])) - finally: - if frontend_children: - _terminate_attached_frontends(frontend_children) + with _provision_kv_cache_pool( + llm_args, owns_engine=not multi_frontend.is_attached_frontend): + if backend == 'pytorch': + llm_args.pop("build_config", None) + llm = PyTorchLLM(**llm_args) + elif backend == '_autodeploy': + from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM + + # AutoDeploy does not support build_config + llm_args.pop("build_config", None) + llm = AutoDeployLLM(**llm_args) + else: + raise click.BadParameter( + f"{backend} is not a known backend, check help for available options.", + param_hint="backend") + + # The finally below is the cleanup boundary for the attached + # frontends: it must cover everything from their spawn through + # server construction, middleware registration, and runtime, or a + # failure in between leaks the child processes. + frontend_children = [] + try: + if multi_frontend.is_launcher: + frontend_children = _spawn_attached_frontends( + llm, multi_frontend.num_frontends) + + server = OpenAIServer( + generator=llm, + model=model, + tool_parser=tool_parser, + server_role=server_role, + metadata_server_cfg=metadata_server_cfg, + disagg_cluster_config=disagg_cluster_config, + multimodal_server_config=multimodal_server_config, + chat_template=chat_template, + allow_request_chat_template=allow_request_chat_template, + input_processor_workers=num_input_processor_workers, + media_load_workers=num_media_load_workers, + internal_disagg_auth_key=internal_disagg_auth_key) + _apply_fastapi_middlewares(server.app, middleware) + + # Optionally disable GC (default: not disabled) + if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": + gc.disable() + + _signal_frontend_ready(multi_frontend) + uvloop.run(server(host, port, sockets=[s])) + finally: + if frontend_children: + _terminate_attached_frontends(frontend_children) def launch_mm_encoder_server( @@ -1560,10 +1608,16 @@ def _serve_llm(): "https://buf.build/gen/python " "\"tensorrt_llm[openengine]\"`.") from error - launch_grpc_server(host, - port, - llm_args, - served_model_name=served_model_name) + # launch_smg_server provisions from inside itself; OpenEngine's + # server is an optional package this repo does not own, so the + # pool is brought up around it here instead. Either way the + # context has to outlive engine construction, which happens + # inside the launch call. + with _provision_kv_cache_pool(llm_args): + launch_grpc_server(host, + port, + llm_args, + served_model_name=served_model_name) else: # Default: launch OpenAI HTTP server launch_server( @@ -2728,7 +2782,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/tensorrt_llm/grpc/smg/server.py b/tensorrt_llm/grpc/smg/server.py index 608baa47428f..5351e7dc9e3d 100644 --- a/tensorrt_llm/grpc/smg/server.py +++ b/tensorrt_llm/grpc/smg/server.py @@ -131,7 +131,12 @@ def signal_handler() -> None: logger.info("LLM engine stopped") logger.info("Shutdown complete") - uvloop.run(serve_grpc_async()) + # Imported here rather than at module scope: tensorrt_llm.commands.serve + # reaches into this module to launch the server. + from tensorrt_llm.commands.serve import _provision_kv_cache_pool + + with _provision_kv_cache_pool(llm_args): + uvloop.run(serve_grpc_async()) __all__ = ["launch_smg_server"] diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 040dab09a0bb..bf6a0951e88b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2302,6 +2302,181 @@ def num_capture_layers(self) -> int: return 0 +class MooncakeStoreConfig(StrictBaseModel): + """The Mooncake store pool the `mooncake-store` connector should join. + + Describes the pool: which master owns it, how workers reach it, and how + much memory each contributes. A worker's own relationship to the pool, such + as its read/write role and key namespace, stays in the + `TRTLLM_MOONCAKE_STORE_*` environment variables, since that is per process + while this is per deployment. + + Setting this makes `trtllm-serve` render the Mooncake client config and + export `MOONCAKE_CONFIG_PATH` itself. An inherited `MOONCAKE_CONFIG_PATH` + still wins, so an externally managed pool stays reachable. + + Fields opt out of telemetry because they size and address one site's pool + rather than saying which features are in use; `kv_connector_config.connector` + already records that the store is on. + """ + master_server_address: Optional[str] = Field( + None, + description="Address of an already-running mooncake_master, as " + "host:port or file:// naming a file that holds one. The file is " + "how to reach a master whose host a scheduler chose, since " + "'trtllm-serve mooncake_master --address_file' publishes it there " + "once it answers. Mutually exclusive with launch_master.") + launch_master: bool = Field( + False, + telemetry=False, + description="Start a mooncake_master in this server's process group " + "and use it. The pool then dies with the server, so this is only " + "correct for a single engine: several engines sharing a pool, or a " + "pool that must outlive a restart, need master_server_address.") + master_address_file: Optional[str] = Field( + None, + telemetry=False, + description="Where a master started by launch_master should publish " + "its host:port, for donors and other servers to read back as " + "file://. One is always written to the run directory; set this " + "to put a second copy somewhere the rest of the deployment already " + "names, such as a shared filesystem. Removed when the master stops.") + master_port: int = Field( + 50051, + telemetry=False, + description="RPC port for a master started by launch_master.") + master_metrics_port: int = Field( + 9004, + telemetry=False, + description="Prometheus port for a master started by launch_master.") + master_eviction_ratio: float = Field( + 0.05, + telemetry=False, + description="Fraction of the pool a master started by launch_master " + "frees per eviction pass.") + metadata_server: str = Field( + "P2PHANDSHAKE", + description="Mooncake metadata service. P2PHANDSHAKE keeps a separate " + "metadata process out of the deployment.") + protocol: str = Field( + "rdma", + description="Transport for page traffic: 'rdma' or 'tcp'. " + "TCP is for bring-up only; it invalidates performance conclusions.") + device_name: str = Field( + "", + description="RDMA device to transfer over, from ibv_devinfo. Empty " + "with protocol 'tcp'.") + global_segment_size: Union[int, str] = Field( + "16GiB", + telemetry=False, + description="Host memory each worker process contributes to the pool. " + "Pool capacity is this times the number of processes that open a " + "store handle, so a prefill-only connector gives a prefill-only pool.") + local_buffer_size: Union[int, str] = Field( + "1GiB", + telemetry=False, + description="Per-process Mooncake transfer buffer, not pool capacity.") + transfer_batch_size: int = Field(64, + telemetry=False, + description="Page keys per store call.") + cache_prefix: Optional[str] = Field( + None, + description="Key namespace for the pool. Bump it after any change to " + "page layout or contents. Defaults to 'trtllm'.") + model_key: Optional[str] = Field( + None, + telemetry=False, + description="Identity the keys are namespaced by. Required whenever a " + "pool is described here, because engines sharing a pool read each " + "other's pages exactly when they agree on this. Give checkpoints that " + "differ in weights, revision, quantization, or anything else that " + "changes the KV distinct values. TRTLLM_MOONCAKE_STORE_MODEL_KEY " + "overrides it for a single process.") + stage_through_host: bool = Field( + False, + telemetry=False, + description="Copy pages through a pinned host buffer instead of " + "registering the KV pools with Mooncake. Needed where the HCA cannot " + "pin GPU pages (no GPUDirect RDMA); costs a copy each way.") + staging_buffer_bytes: Optional[Union[int, str]] = Field( + None, + telemetry=False, + description="Size of the buffer stage_through_host copies through, " + "per process. Pages move transfer_batch_size at a time, and a buffer " + "that cannot hold that many reduces the batch instead of failing, so " + "undersizing it costs throughput quietly. Defaults to the connector's " + "own 512MiB.") + + @model_validator(mode="after") + def _require_exactly_one_master(self) -> "MooncakeStoreConfig": + if self.launch_master and self.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 self.launch_master and not self.master_server_address: + raise ValueError( + "mooncake_store: needs a master. Set master_server_address to " + "join an existing pool, or launch_master: true to start one " + "for this server alone.") + if self.master_address_file and not self.launch_master: + raise ValueError( + "mooncake_store: master_address_file publishes the address of " + "a master this server starts, so it needs launch_master: " + "true. To read an address a master elsewhere published, set " + "master_server_address: file://.") + return self + + +class MooncakeDonationConfig(StrictBaseModel): + """Host memory this server lends to a Mooncake pool it does not use. + + Pool capacity comes only from processes that open a store handle, which in + a disaggregated deployment is the context servers alone. Setting this on + the generation servers puts their memory into the same pool, so prefill + writes blocks that land on decode-side DRAM while the generation engine + stays free of any connector and keeps its cache transceiver for the + handoff. + + Capacity is kept separate from `kv_connector_config` because attaching a + connector would also start this server reading and writing the store. + + The memory is charged to this process, so size it together with + `kv_cache_config.host_cache_size`. + + Fields opt out of telemetry because they describe one site's pool rather + than which features are in use. + """ + master_server_address: str = Field( + ..., + telemetry=False, + description="Master of the pool to lend memory to, as host:port or " + "file:// naming a file that holds one. The file is what a " + "context server's launch_master publishes, so the generation servers " + "can name a path instead of a host chosen by a scheduler, and they " + "wait for the master rather than having to start after it.") + segment_size: Union[int, str] = Field( + "32GiB", + telemetry=False, + description="Host memory this server contributes. Charged once per " + "server process, not per rank, so a node running several servers " + "contributes this much for each of them.") + protocol: str = Field( + "rdma", + telemetry=False, + description="Transport the pool's traffic reaches this memory over: " + "'rdma' or 'tcp'. Must match the pool's.") + device_name: str = Field( + "", + telemetry=False, + description="RDMA device to serve the segment over, from ibv_devinfo. " + "Empty with protocol 'tcp'.") + metadata_server: str = Field( + "P2PHANDSHAKE", + telemetry=False, + description="Mooncake metadata service. Must match the pool's.") + + class KvCacheConnectorConfig(StrictBaseModel): """Configuration for the KV Cache Connector. @@ -2316,7 +2491,8 @@ class KvCacheConnectorConfig(StrictBaseModel): description="Named connector preset (e.g. 'lmcache'). " "When set, connector_module/scheduler_class/worker_class are " "auto-populated from the preset registry.", - telemetry=TelemetryField.categorical('lmcache', 'lmcache-mp', 'kvbm')) + telemetry=TelemetryField.categorical('lmcache', 'lmcache-mp', 'kvbm', + 'mooncake-store')) connector_module: Optional[str] = Field( None, description= @@ -2331,11 +2507,16 @@ class KvCacheConnectorConfig(StrictBaseModel): description="URL for an external connector server " "(e.g. 'tcp://localhost:5555'). Connectors that run in " "multi-process mode use this to reach the cache server.") + mooncake_store: Optional[MooncakeStoreConfig] = Field( + None, + description="Pool topology for the 'mooncake-store' connector. When " + "set, trtllm-serve provisions the pool during bringup instead of " + "requiring MOONCAKE_CONFIG_PATH from an external script.") @model_validator(mode="after") def _resolve_preset(self) -> "KvCacheConnectorConfig": - from tensorrt_llm._torch.pyexecutor.connectors.registry import \ - CONNECTOR_REGISTRY + from tensorrt_llm._torch.pyexecutor.connectors.registry import ( + CONNECTOR_REGISTRY, uses_connector) if self.connector is not None: preset = CONNECTOR_REGISTRY.get(self.connector) if preset is None: @@ -2353,6 +2534,22 @@ def _resolve_preset(self) -> "KvCacheConnectorConfig": raise ValueError("connector_scheduler_class is required") if self.connector_worker_class is None: raise ValueError("connector_worker_class is required") + if self.mooncake_store is not None and not uses_connector( + self, "mooncake-store"): + raise ValueError( + "mooncake_store describes a Mooncake pool, but this config " + f"resolves to connector_module={self.connector_module!r}. " + "Set connector: mooncake-store, or drop mooncake_store.") + if (self.mooncake_store is not None + and self.mooncake_store.model_key is None): + raise ValueError( + "mooncake_store.model_key is required. Engines sharing a pool " + "read each other's pages exactly when they agree on this, so " + "there is no safe default: deriving one from the model path " + "would give org-a/model and org-b/model, or two revisions " + "mounted under one directory name, the same namespace for " + "different weights. Set it to a value that separates this " + "checkpoint from any other an engine on this pool may load.") return self @@ -6025,6 +6222,17 @@ def is_partial_model_loading(self) -> bool: status="prototype", ) + mooncake_donation: Optional[MooncakeDonationConfig] = Field( + default=None, + description="Host memory to lend to a Mooncake pool this server does " + "not otherwise use. Separate from kv_connector_config because it adds " + "capacity without attaching a connector, which is what lets a " + "generation server hold pages for a pool only prefill reads and " + "writes. Honored by trtllm-serve, which holds the segment for the " + "server's lifetime.", + status="prototype", + ) + mm_encoder_only: bool = Field( default=False, description= diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 257d1ba70c11..07e67f75ce52 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -781,7 +781,8 @@ "allowed_values": [ "lmcache", "lmcache-mp", - "kvbm" + "kvbm", + "mooncake-store" ], "capture_policy": "allowlist|none", "kind": "categorical", diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 8415c70c1dab..fa2745cfe4a5 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_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py index 55c917fd6f9e..d7c298fa2dc3 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py @@ -1395,6 +1395,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 @@ -1593,6 +1594,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/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 4124d7b657c7..fc7b3e4e7ec1 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.tokens_per_block = tokens_per_block @@ -215,6 +219,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 @@ -1284,6 +1291,402 @@ 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()) + + # =========================================================================== # PEFT / LoRA # =========================================================================== @@ -3398,6 +3801,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/test_mooncake_store_cli.py b/tests/unittest/_torch/executor/test_mooncake_store_cli.py new file mode 100644 index 000000000000..3f486fb1375f --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_cli.py @@ -0,0 +1,201 @@ +# 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 os +import signal +import time +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors import mooncake_store +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import CONFIG_PATH_ENV +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. + """ + + def __init__(self, held): + self.held = held + self.released = False + + @contextlib.contextmanager + def holding(self, *_args, **_kwargs): + 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) -> 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. + """ + + def sleep(_seconds): + 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 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..8bcb00d51fb9 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_common.py @@ -0,0 +1,607 @@ +# 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, +which deployments are refused at startup, 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 importlib +import json +from types import SimpleNamespace + +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, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.validation import validate_llm_args +from tensorrt_llm._torch.pyexecutor.connectors.registry import uses_connector +from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig + +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 + + +def make_llm_args(): + return SimpleNamespace( + model="/models/test-model", + kv_cache_config=SimpleNamespace(tokens_per_block=TOKENS_PER_BLOCK), + tensor_parallel_size=1, + pipeline_parallel_size=1, + context_parallel_size=1, + sparse_attention_config=None, + ) + + +@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() + + +# ---- validation ---- + + +@pytest.mark.parametrize( + "field,value,match", + [ + ("context_parallel_size", 2, "context parallelism"), + ("pipeline_parallel_size", 2, "pipeline parallelism"), + ], +) +def test_validate_llm_args_rejects_unsupported_parallelism(field, value, match): + args = make_llm_args() + setattr(args, field, value) + with pytest.raises(NotImplementedError, match=match): + validate_llm_args(args) + + +def test_validate_llm_args_rejects_m3_index_value_cache(): + args = make_llm_args() + args.sparse_attention_config = SimpleNamespace(sparse_disable_index_value=False) + with pytest.raises(NotImplementedError, match="sparse_disable_index_value"): + validate_llm_args(args) + + args.sparse_attention_config = SimpleNamespace(sparse_disable_index_value=True) + validate_llm_args(args) + + +# ---- connector identification ---- +# +# py_executor_creator turns partial reuse off for this connector and finds it +# through uses_connector. Failing to recognize the config would silently cost +# the reuse the store exists to provide. + + +@pytest.mark.parametrize( + "config, expected", + [ + (KvCacheConnectorConfig(connector="mooncake-store"), True), + ( + KvCacheConnectorConfig( + connector_module="tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", + connector_scheduler_class="MooncakeStoreConnectorScheduler", + connector_worker_class="MooncakeStoreConnectorWorker", + ), + True, + ), + (KvCacheConnectorConfig(connector="kvbm"), False), + (None, False), + ], + ids=["preset", "hand_written_module", "another_connector", "no_connector"], +) +def test_uses_connector_recognizes_the_connector_however_it_is_spelled(config, expected): + assert uses_connector(config, "mooncake-store") is expected + + +def test_uses_connector_rejects_an_unknown_preset(): + config = KvCacheConnectorConfig(connector="mooncake-store") + with pytest.raises(ValueError, match="Unknown connector preset"): + uses_connector(config, "mooncake-stroe") + + +def test_the_registered_preset_resolves_to_refusing_classes(): + """py_executor_creator resolves both by name from the package.""" + config = KvCacheConnectorConfig(connector="mooncake-store") + module = importlib.import_module(config.connector_module) + + for class_name in (config.connector_scheduler_class, config.connector_worker_class): + with pytest.raises(NotImplementedError, match="not available in this build"): + getattr(module, class_name)(llm_args=None) + + +# ---- 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..cda334bc385a --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_donor.py @@ -0,0 +1,201 @@ +# 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, + maybe_donate_segment, +) +from tensorrt_llm.llmapi.llm_args import MooncakeDonationConfig + +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 + + +@pytest.mark.parametrize("entry", ["direct", "config"]) +def test_a_donor_registers_the_segment_it_was_asked_for( + entry, fake_bindings, reachable_master, monkeypatch +): + """Both entry paths must reach Mooncake with the same seven setup arguments. + + The config-driven path is what makes a generation server a donor, and it + derives the hostname and parses the size string on the way: a size string + reaching Mooncake unparsed would be a segment of nothing. + """ + if entry == "direct": + expected_host, expected_size, expected_device = "10.0.0.5", 32 * GIB, "mlx5_0" + donation = donate_segment( + "10.0.0.1:50051", + 32 * GIB, + protocol="rdma", + device_name="mlx5_0", + metadata_server="P2PHANDSHAKE", + hostname="10.0.0.5", + ) + else: + monkeypatch.setattr(donor_module, "local_address", lambda: "10.1.2.3") + expected_host, expected_size, expected_device = "10.1.2.3", 320 * GIB, "mlx5_1" + donation = maybe_donate_segment( + MooncakeDonationConfig( + master_server_address="10.0.0.1:50051", + segment_size="320GiB", + protocol="rdma", + device_name="mlx5_1", + ) + ) + + with donation as host: + assert host == expected_host + ( + registered_host, + metadata_server, + segment_size, + local_buffer_size, + protocol, + device_name, + master, + ) = fake_bindings.instances[0].setup_args + + assert registered_host == expected_host + assert metadata_server == "P2PHANDSHAKE" + assert segment_size == expected_size + assert protocol == "rdma" + assert device_name == expected_device + 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="mooncake-transfer-engine"): + with donate_segment("10.0.0.1:50051", GIB): + pytest.fail("donation should not have yielded") + + +@pytest.fixture +def reachable_master(monkeypatch): + """Skip the socket probe: these tests are about what donation asks for.""" + monkeypatch.setattr(donor_module, "wait_for_master", lambda address: 0.0) + + +def test_a_server_that_was_not_asked_lends_nothing(fake_bindings): + """Every deployment that does not lend memory takes this path.""" + with maybe_donate_segment(None) as host: + assert host is None + assert fake_bindings.instances == [] + + +def test_a_published_master_address_is_read_before_joining( + fake_bindings, reachable_master, tmp_path +): + """What lets a generation server 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") + donation = MooncakeDonationConfig( + master_server_address=f"file://{address_file}", + segment_size=GIB, + ) + + with maybe_donate_segment(donation): + assert fake_bindings.instances[0].setup_args[6] == "10.0.0.9:50051" + + +def test_an_unreachable_master_is_reported_before_the_segment_is_offered( + fake_bindings, monkeypatch +): + """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(donor_module, "wait_for_master", refuse) + donation = MooncakeDonationConfig(master_server_address="10.0.0.1:50051") + + with pytest.raises(TimeoutError, match="10.0.0.1:50051"): + with maybe_donate_segment(donation): + pytest.fail("donation should not have yielded") + assert fake_bindings.instances == [] 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..7b96ad99f9f8 --- /dev/null +++ b/tests/unittest/_torch/executor/test_mooncake_store_master.py @@ -0,0 +1,670 @@ +# 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, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.master import ( + maybe_provision_pool, + provision_pool, + resolve_master_address, +) +from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig, MooncakeStoreConfig + + +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_pool_needs_exactly_one_master(kwargs, message): + with pytest.raises(ValueError, match=message): + MooncakeStoreConfig(**kwargs) + + +def test_pool_is_rejected_unless_the_connector_is_mooncake_store(): + """The validator keys off the connector, however that was spelled.""" + with pytest.raises(ValueError, match="mooncake_store describes a Mooncake pool"): + KvCacheConnectorConfig( + connector="lmcache", + mooncake_store=MooncakeStoreConfig(launch_master=True, model_key="m"), + ) + # Naming the module rather than the preset selects the same connector. + KvCacheConnectorConfig( + connector_module="tensorrt_llm._torch.pyexecutor.connectors.mooncake_store", + connector_scheduler_class="MooncakeStoreConnectorScheduler", + connector_worker_class="MooncakeStoreConnectorWorker", + mooncake_store=MooncakeStoreConfig(launch_master=True, model_key="m"), + ) + + +def test_a_described_pool_needs_a_model_key(): + """Two checkpoints that agree on the namespace read each other's pages.""" + with pytest.raises(ValueError, match="mooncake_store.model_key is required"): + KvCacheConnectorConfig( + connector="mooncake-store", + mooncake_store=MooncakeStoreConfig(launch_master=True), + ) + + +# ---- 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 = MooncakeStoreConfig( + 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 = MooncakeStoreConfig(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 + + +@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 = MooncakeStoreConfig(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 = MooncakeStoreConfig( + 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 = MooncakeStoreConfig(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 = MooncakeStoreConfig(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 = MooncakeStoreConfig(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 = MooncakeStoreConfig(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 = MooncakeStoreConfig( + 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 = MooncakeStoreConfig(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 = MooncakeStoreConfig(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 = MooncakeStoreConfig(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 = MooncakeStoreConfig(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() + + +# ---- the entry point servers call ---- + + +@pytest.mark.parametrize( + "config", + [ + KvCacheConnectorConfig(connector="lmcache"), + None, + KvCacheConnectorConfig(connector="mooncake-store"), + ], + ids=["another_connector", "no_connector", "pool_left_undescribed"], +) +def test_provisioning_is_a_no_op_unless_a_pool_is_described(config): + """Without `mooncake_store`, MOONCAKE_CONFIG_PATH is still the only input.""" + with maybe_provision_pool(config): + assert CONFIG_PATH_ENV not in os.environ + + +def test_a_described_pool_is_provisioned(running_master): + config = KvCacheConnectorConfig( + connector="mooncake-store", + mooncake_store=MooncakeStoreConfig( + master_server_address=running_master, model_key="test-model" + ), + ) + with maybe_provision_pool(config): + written = json.loads(open(os.environ[CONFIG_PATH_ENV]).read()) + assert written["master_server_address"] == running_master + assert CONFIG_PATH_ENV not in os.environ + + +# ---- 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 = MooncakeStoreConfig(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 = MooncakeStoreConfig(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 = MooncakeStoreConfig(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 = MooncakeStoreConfig(launch_master=True, master_port=port) + worker = MooncakeStoreConfig(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 = MooncakeStoreConfig(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 = MooncakeStoreConfig( + 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 = MooncakeStoreConfig(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") + + +# ---- choosing the fabric without naming it in a config ---- + + +def fake_hca(root, device, link_layer="InfiniBand", state="4: ACTIVE", rate="800 Gb/sec"): + port = root / device / "ports" / "1" + port.mkdir(parents=True) + (port / "link_layer").write_text(f"{link_layer}\n") + (port / "state").write_text(f"{state}\n") + (port / "rate").write_text(f"{rate}\n") + + +def test_the_compute_fabric_is_picked_over_the_management_adapter(tmp_path): + """A node's HCAs are not interchangeable: only some are the fast fabric.""" + fake_hca(tmp_path, "mlx5_0") + fake_hca(tmp_path, "mlx5_1") + fake_hca(tmp_path, "mlx5_2", rate="400 Gb/sec") + fake_hca(tmp_path, "mlx5_3", state="1: DOWN") + fake_hca(tmp_path, "mlx5_4", link_layer="Ethernet") + + assert ( + master_module.resolve_device_name("rdma", "", sysfs_root=str(tmp_path)) == "mlx5_0,mlx5_1" + ) + + +def test_a_named_device_is_not_second_guessed(tmp_path): + fake_hca(tmp_path, "mlx5_0") + assert master_module.resolve_device_name("rdma", "mlx5_7", sysfs_root=str(tmp_path)) == "mlx5_7" + + +def test_tcp_needs_no_device_and_looks_for_none(tmp_path): + assert master_module.resolve_device_name("tcp", "", sysfs_root=str(tmp_path)) == "" + + +def test_a_node_without_infiniband_is_left_to_mooncake_s_own_discovery(tmp_path): + """Falling back beats failing, since Mooncake may still find a usable device.""" + assert master_module.resolve_device_name("rdma", "", sysfs_root=str(tmp_path / "absent")) == "" + + +def test_the_detected_device_is_what_the_workers_are_told(fake_master, tmp_path, monkeypatch): + sysfs = tmp_path / "sysfs" + fake_hca(sysfs, "mlx5_0") + monkeypatch.setattr(master_module, "IB_SYSFS_ROOT", str(sysfs)) + port = free_port() + fake_master.arm(listen_on=port) + pool = MooncakeStoreConfig(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"] == "mlx5_0" diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 8728888a2868..9ecbbc61e23b 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -291,6 +291,10 @@ methods: annotation: Optional[tensorrt_llm.llmapi.llm_args.KvCacheConnectorConfig] default: null status: prototype + mooncake_donation: + annotation: Optional[tensorrt_llm.llmapi.llm_args.MooncakeDonationConfig] + default: null + status: prototype enable_lm_head_tp_in_adp: annotation: bool default: False