diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1876bedff89e..a7061daca62d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -103,7 +103,6 @@ common-files: &common_files | scripts/test_to_stage_mapping.py | setup.py | tensorrt_llm/__init__.py | - tensorrt_llm/_ray_utils.py | tensorrt_llm/_torch/__init__.py | tensorrt_llm/_torch/attention_backend/__init__.py | tensorrt_llm/_torch/attention_backend/flashinfer.py | @@ -375,8 +374,10 @@ common-files: &common_files | tensorrt_llm/executor/ipc.py | tensorrt_llm/executor/postproc_worker.py | tensorrt_llm/executor/proxy.py | - tensorrt_llm/executor/ray_executor.py | - tensorrt_llm/executor/ray_gpu_worker.py | + tensorrt_llm/executor/ray/executor.py | + tensorrt_llm/executor/ray/gpu_worker.py | + tensorrt_llm/executor/ray/stub.py | + tensorrt_llm/executor/ray/utils.py | tensorrt_llm/executor/request.py | tensorrt_llm/executor/result.py | tensorrt_llm/executor/rpc/__init__.py | @@ -426,7 +427,6 @@ common-files: &common_files | tensorrt_llm/quantization/utils/__init__.py | tensorrt_llm/quantization/utils/fp4_utils.py | tensorrt_llm/quantization/utils/fp8_utils.py | - tensorrt_llm/ray_stub.py | tensorrt_llm/runtime/__init__.py | tensorrt_llm/runtime/memory_pools/__init__.py | tensorrt_llm/scaffolding/__init__.py | @@ -877,7 +877,6 @@ legacy-files: &legacy_files | scripts/test_to_stage_mapping.py | setup.py | tensorrt_llm/__init__.py | - tensorrt_llm/_ray_utils.py | tensorrt_llm/_torch/__init__.py | tensorrt_llm/_torch/attention_backend/__init__.py | tensorrt_llm/_torch/attention_backend/flashinfer.py | @@ -1149,8 +1148,10 @@ legacy-files: &legacy_files | tensorrt_llm/executor/ipc.py | tensorrt_llm/executor/postproc_worker.py | tensorrt_llm/executor/proxy.py | - tensorrt_llm/executor/ray_executor.py | - tensorrt_llm/executor/ray_gpu_worker.py | + tensorrt_llm/executor/ray/executor.py | + tensorrt_llm/executor/ray/gpu_worker.py | + tensorrt_llm/executor/ray/stub.py | + tensorrt_llm/executor/ray/utils.py | tensorrt_llm/executor/request.py | tensorrt_llm/executor/result.py | tensorrt_llm/executor/rpc/__init__.py | @@ -1200,7 +1201,6 @@ legacy-files: &legacy_files | tensorrt_llm/quantization/utils/__init__.py | tensorrt_llm/quantization/utils/fp4_utils.py | tensorrt_llm/quantization/utils/fp8_utils.py | - tensorrt_llm/ray_stub.py | tensorrt_llm/runtime/__init__.py | tensorrt_llm/runtime/memory_pools/__init__.py | tensorrt_llm/scaffolding/__init__.py | diff --git a/docs/source/features/ray-orchestrator.md b/docs/source/features/ray-orchestrator.md index 4984c180b70f..e949febdd3d1 100644 --- a/docs/source/features/ray-orchestrator.md +++ b/docs/source/features/ray-orchestrator.md @@ -37,6 +37,6 @@ Currently available: - Integration with RLHF frameworks, such as [Verl](https://github.com/volcengine/verl) and [NVIDIA NeMo-RL](https://github.com/NVIDIA-NeMo/RL). ## Architecture -This feature introduces new classes such as [RayExecutor](/tensorrt_llm/executor/ray_executor.py) and [RayGPUWorker](/tensorrt_llm/executor/ray_gpu_worker.py) for Ray actor lifecycle management and distributed inference. In Ray mode, collective ops run on [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html) without MPI. We welcome contributions to improve and extend this support. +This feature introduces new classes such as [RayExecutor](/tensorrt_llm/executor/ray/executor.py) and [RayGPUWorker](/tensorrt_llm/executor/ray/gpu_worker.py) for Ray actor lifecycle management and distributed inference. In Ray mode, collective ops run on [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html) without MPI. We welcome contributions to improve and extend this support. ![Ray orchestrator architecture](/docs/source/media/ray_orchestrator_architecture.jpg) diff --git a/examples/ray_orchestrator/README.md b/examples/ray_orchestrator/README.md index f8ab0c402968..107ceba42e7f 100644 --- a/examples/ray_orchestrator/README.md +++ b/examples/ray_orchestrator/README.md @@ -40,7 +40,7 @@ This example is the same as in `/examples/llm-api`, with the only change being ` - Integration with RLHF frameworks, such as [Verl](https://github.com/volcengine/verl) and [NVIDIA Nemo-RL](https://github.com/NVIDIA-NeMo/RL). ## Architecture -This feature introduces new classes such as [RayExecutor](/tensorrt_llm/executor/ray_executor.py) and [RayGPUWorker](/tensorrt_llm/executor/ray_gpu_worker.py) for Ray actor lifecycle management and distributed inference. In Ray mode, collective ops run on [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html) without MPI. We welcome contributions to improve and extend this support. +This feature introduces new classes such as [RayExecutor](/tensorrt_llm/executor/ray/executor.py) and [RayGPUWorker](/tensorrt_llm/executor/ray/gpu_worker.py) for Ray actor lifecycle management and distributed inference. In Ray mode, collective ops run on [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html) without MPI. We welcome contributions to improve and extend this support. ![Ray orchestrator architecture](/docs/source/media/ray_orchestrator_architecture.jpg) diff --git a/legacy-files.txt b/legacy-files.txt index 7007c5129043..b01ca7fbe20c 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -95,7 +95,6 @@ scripts/rename_docker_images.py scripts/test_to_stage_mapping.py setup.py tensorrt_llm/__init__.py -tensorrt_llm/_ray_utils.py tensorrt_llm/_torch/__init__.py tensorrt_llm/_torch/attention_backend/__init__.py tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -367,8 +366,10 @@ tensorrt_llm/executor/executor.py tensorrt_llm/executor/ipc.py tensorrt_llm/executor/postproc_worker.py tensorrt_llm/executor/proxy.py -tensorrt_llm/executor/ray_executor.py -tensorrt_llm/executor/ray_gpu_worker.py +tensorrt_llm/executor/ray/executor.py +tensorrt_llm/executor/ray/gpu_worker.py +tensorrt_llm/executor/ray/stub.py +tensorrt_llm/executor/ray/utils.py tensorrt_llm/executor/request.py tensorrt_llm/executor/result.py tensorrt_llm/executor/rpc/__init__.py @@ -418,7 +419,6 @@ tensorrt_llm/quantization/mode.py tensorrt_llm/quantization/utils/__init__.py tensorrt_llm/quantization/utils/fp4_utils.py tensorrt_llm/quantization/utils/fp8_utils.py -tensorrt_llm/ray_stub.py tensorrt_llm/runtime/__init__.py tensorrt_llm/runtime/memory_pools/__init__.py tensorrt_llm/scaffolding/__init__.py diff --git a/pyproject.toml b/pyproject.toml index df0493ad447a..3a6750deadcb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -152,7 +152,6 @@ exclude = [ "scripts/test_to_stage_mapping.py", "setup.py", "tensorrt_llm/__init__.py", - "tensorrt_llm/_ray_utils.py", "tensorrt_llm/_torch/__init__.py", "tensorrt_llm/_torch/attention_backend/__init__.py", "tensorrt_llm/_torch/attention_backend/flashinfer.py", @@ -424,8 +423,10 @@ exclude = [ "tensorrt_llm/executor/ipc.py", "tensorrt_llm/executor/postproc_worker.py", "tensorrt_llm/executor/proxy.py", - "tensorrt_llm/executor/ray_executor.py", - "tensorrt_llm/executor/ray_gpu_worker.py", + "tensorrt_llm/executor/ray/executor.py", + "tensorrt_llm/executor/ray/gpu_worker.py", + "tensorrt_llm/executor/ray/stub.py", + "tensorrt_llm/executor/ray/utils.py", "tensorrt_llm/executor/request.py", "tensorrt_llm/executor/result.py", "tensorrt_llm/executor/rpc/__init__.py", @@ -475,7 +476,6 @@ exclude = [ "tensorrt_llm/quantization/utils/__init__.py", "tensorrt_llm/quantization/utils/fp4_utils.py", "tensorrt_llm/quantization/utils/fp8_utils.py", - "tensorrt_llm/ray_stub.py", "tensorrt_llm/runtime/__init__.py", "tensorrt_llm/runtime/memory_pools/__init__.py", "tensorrt_llm/scaffolding/__init__.py", diff --git a/ruff-legacy-baseline.json b/ruff-legacy-baseline.json index eae44689bafc..c707e213216b 100644 --- a/ruff-legacy-baseline.json +++ b/ruff-legacy-baseline.json @@ -343,7 +343,7 @@ "D212": 1, "E722": 1 }, - "tensorrt_llm/executor/ray_executor.py": { + "tensorrt_llm/executor/ray/executor.py": { "D205": 1, "D212": 2, "E712": 1 diff --git a/ruff-legacy.toml b/ruff-legacy.toml index b1adcc5e002a..ff58435cf25a 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -112,7 +112,6 @@ include = [ "scripts/test_to_stage_mapping.py", "setup.py", "tensorrt_llm/__init__.py", - "tensorrt_llm/_ray_utils.py", "tensorrt_llm/_torch/__init__.py", "tensorrt_llm/_torch/attention_backend/__init__.py", "tensorrt_llm/_torch/attention_backend/flashinfer.py", @@ -384,8 +383,10 @@ include = [ "tensorrt_llm/executor/ipc.py", "tensorrt_llm/executor/postproc_worker.py", "tensorrt_llm/executor/proxy.py", - "tensorrt_llm/executor/ray_executor.py", - "tensorrt_llm/executor/ray_gpu_worker.py", + "tensorrt_llm/executor/ray/executor.py", + "tensorrt_llm/executor/ray/gpu_worker.py", + "tensorrt_llm/executor/ray/stub.py", + "tensorrt_llm/executor/ray/utils.py", "tensorrt_llm/executor/request.py", "tensorrt_llm/executor/result.py", "tensorrt_llm/executor/rpc/__init__.py", @@ -435,7 +436,6 @@ include = [ "tensorrt_llm/quantization/utils/__init__.py", "tensorrt_llm/quantization/utils/fp4_utils.py", "tensorrt_llm/quantization/utils/fp8_utils.py", - "tensorrt_llm/ray_stub.py", "tensorrt_llm/runtime/__init__.py", "tensorrt_llm/runtime/memory_pools/__init__.py", "tensorrt_llm/scaffolding/__init__.py", diff --git a/tensorrt_llm/_ray_utils.py b/tensorrt_llm/_ray_utils.py index 489dad9cc8f2..433489bfdcf7 100644 --- a/tensorrt_llm/_ray_utils.py +++ b/tensorrt_llm/_ray_utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -12,47 +12,29 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import functools -from contextlib import contextmanager -from typing import Callable, Optional +"""Compatibility shim for ``tensorrt_llm._ray_utils``. -try: - import ray -except ImportError: - import tensorrt_llm.ray_stub as ray +Will be removed once all usages are migrated to +``tensorrt_llm.executor.ray.utils``. +DO NOT ADD ANYTHING TO THIS FILE. +""" -@contextmanager -def unwrap_ray_errors(): - try: - yield - except ray.exceptions.RayTaskError as e: - raise e.as_instanceof_cause() from e +import warnings +from tensorrt_llm.executor.ray.utils import ( # noqa: F401 + control_action_decorator, + unwrap_ray_errors, +) -def control_action_decorator(func: Optional[Callable] = None, - *, - drain: bool = True) -> Callable: - """Wrap a method in the ``control_action`` context manager. +warnings.warn( + "tensorrt_llm._ray_utils has moved to tensorrt_llm.executor.ray.utils " + "and will be removed in a future release.", + FutureWarning, + stacklevel=2, +) - Supports both bare and parameterized forms:: - - @control_action_decorator # drain=True (default) - def shutdown(self): ... - - @control_action_decorator(drain=False) # non-draining variant - def update_weights_via_ipc_zmq(self): ... - """ - - def decorator(f: Callable) -> Callable: - - @functools.wraps(f) - def wrapper(self, *args, **kwargs): - with self.engine.control_action(drain=drain): - return f(self, *args, **kwargs) - - return wrapper - - if func is None: - return decorator - return decorator(func) +__all__ = [ + "control_action_decorator", + "unwrap_ray_errors", +] diff --git a/tensorrt_llm/_torch/distributed/communicator.py b/tensorrt_llm/_torch/distributed/communicator.py index 2e0fc1e0e4c3..1d2ed9e474a5 100644 --- a/tensorrt_llm/_torch/distributed/communicator.py +++ b/tensorrt_llm/_torch/distributed/communicator.py @@ -30,7 +30,7 @@ try: import ray except ModuleNotFoundError: - from tensorrt_llm import ray_stub as ray + from tensorrt_llm.executor.ray import stub as ray class ReduceOp(IntEnum): diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index 98afe8daa44f..3dcfe0df0701 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -480,7 +480,7 @@ def _create_ray_executor( tp_size: int, ): logger.warning(f"Orchestrator is creating Ray executor") - from .ray_executor import RayExecutor + from .ray.executor import RayExecutor return RayExecutor(worker_kwargs, model_world_size=model_world_size, diff --git a/tensorrt_llm/executor/ray/__init__.py b/tensorrt_llm/executor/ray/__init__.py new file mode 100644 index 000000000000..1190c80e50ae --- /dev/null +++ b/tensorrt_llm/executor/ray/__init__.py @@ -0,0 +1,22 @@ +# 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. +"""Ray executor integration. + +This package deliberately re-exports nothing. ``executor`` and ``gpu_worker`` +import ``ray`` at module scope, while ``stub`` is the stand-in used when Ray is +*not* installed -- so a re-export here would make importing the stub require the +very package the stub exists to replace, and it would fail only in environments +without Ray. Import the submodules directly. +""" diff --git a/tensorrt_llm/executor/ray/executor.py b/tensorrt_llm/executor/ray/executor.py new file mode 100644 index 000000000000..6ff1abf666d8 --- /dev/null +++ b/tensorrt_llm/executor/ray/executor.py @@ -0,0 +1,564 @@ +# 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. +import asyncio +import os +import time +from typing import Any, Dict, List, Optional, Tuple + +try: + import ray +except ModuleNotFoundError as e: + e.msg = """Cannot import Ray. Please install 'ray' package to use ray orchestrator""" + raise + +from ray.util.placement_group import (get_current_placement_group, + placement_group) + +try: + # Ray >= 2.55.0 + from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +except ImportError: + # Older Ray re-exported it from ray.util.placement_group + from ray.util.placement_group import PlacementGroupSchedulingStrategy + +from tensorrt_llm._utils import nvtx_range_debug +from tensorrt_llm.executor.ray.utils import unwrap_ray_errors +from tensorrt_llm.logger import logger + +from ...llmapi.utils import logger_debug +from ..executor import GenerationExecutor +from ..postproc_worker import PostprocWorkerConfig +from ..request import GenerationRequest +from ..result import GenerationResult +from ..rpc_proxy_mixin import RpcExecutorMixin +from ..utils import has_event_loop +from .gpu_worker import RayGPUWorker, RayWorkerWrapper + +__all__ = [ + "RayExecutor", +] + + +class RayExecutor(RpcExecutorMixin, GenerationExecutor): + + def __init__(self, + worker_kwargs: Dict, + model_world_size: int, + postproc_worker_config: PostprocWorkerConfig, + is_llm_executor: bool, + tp_size=1): + os.environ['RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES'] = '1' + os.environ["RAY_DEDUP_LOGS"] = "0" # for debug + + super().__init__(model_world_size, postproc_worker_config, + is_llm_executor) + + self.has_start_local_cluser = False + runtime_env = { + "env_vars": { + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1" + } + } + + ray_init_args = { + "include_dashboard": False, + "namespace": "trtllm", + "ignore_reinit_error": True, + "runtime_env": runtime_env + } + + try: + if os.environ.get("TLLM_RAY_FORCE_LOCAL_CLUSTER", "0") != "1": + try: + ray.init(address="auto", **ray_init_args) + logger.info(f"Attached to an existing Ray cluster.") + except ConnectionError: + logger.info(f"Ray cluster not found, starting a new one.") + + if not ray.is_initialized(): + ray.init(**ray_init_args) + self.has_start_local_cluser = True + else: + ray.init(address="local", **ray_init_args) + self.has_start_local_cluser = True + + self.world_size = model_world_size + self.tp_size = tp_size + self.master_address = ray.util.get_node_ip_address() + + self.worker_kwargs = dict( + **worker_kwargs, + postproc_worker_config=postproc_worker_config, + is_llm_executor=is_llm_executor) + + self.init_rpc_executor() + # Inject the generated HMAC key into worker_kwargs for workers + self.worker_kwargs['hmac_key'] = self.hmac_key + self.worker_kwargs['rpc_addr'] = self.rpc_addr + + placement_config = getattr(self.worker_kwargs['llm_args'], + 'ray_placement_config', None) + defer_workers_init = placement_config.defer_workers_init if placement_config else False + + if defer_workers_init: + self.workers = [ + ] # Placeholder, will be initialized in setup_async + self._mainloop_started = False # DO NOT start mainloop until after setup_engine_remote_async is called + else: + if not has_event_loop(): + self.init_workers_sync() + self.setup_engine_remote() + self.setup_mainloop(tasks=[self._fetch_responses_loop_async], + thread_name="ray_executor_main_loop") + + except Exception as e: + self.shutdown() + logger.error(f"Failed to initialize RayExecutor: {e}") + raise e + + def create_workers(self, worker_cls, worker_kwargs): + llm_args = worker_kwargs.get("llm_args") + placement_config = getattr(llm_args, 'ray_placement_config', + None) if llm_args else None + ray_worker_nsight_options = getattr( + llm_args, 'ray_worker_nsight_options', None) if llm_args else None + + # When set to be a fraction, it allows Ray to schedule + # multiple actors on a single GPU for colocate use cases. + num_gpus = float(os.getenv("TRTLLM_RAY_PER_WORKER_GPUS", "1.0")) + if placement_config and placement_config.per_worker_gpu_share is not None: + num_gpus = placement_config.per_worker_gpu_share + + logger.debug(f"{num_gpus=} for each worker.") + + runtime_env = ray.runtime_env.RuntimeEnv() + # Exclude node-local env vars. e.g., The raylet that spawns each worker sets + # RAY_RAYLET_PID to its own PID at exec time. + _NODE_LOCAL_VARS = { + "RAY_RAYLET_PID", + "RAY_NODE_IP_ADDRESS", + } + + runtime_env["env_vars"] = { + k: v + for k, v in os.environ.items() if k not in _NODE_LOCAL_VARS + } + runtime_env["env_vars"].update({ + "TLLM_DISABLE_MPI": "1", + "MASTER_ADDR": self.master_address, # head-IP for NCCL/Gloo + }) + if ray_worker_nsight_options: + runtime_env["nsight"] = ray_worker_nsight_options + + placement_groups, self.bundle_indices = self._get_placement_group( + tp_size=self.tp_size, worker_kwargs=worker_kwargs) + + if isinstance(placement_groups, list): + self.placement_group = None + else: + self.placement_group = placement_groups + + self.workers = [] + for rank in range(self.world_size): + pg = placement_groups[rank] if isinstance( + placement_groups, list) else placement_groups + worker = RayWorkerWrapper.options( + num_gpus=num_gpus, + runtime_env=runtime_env, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=self.bundle_indices[rank], + )).remote(worker_cls, worker_kwargs, self.world_size, rank) + self.workers.append(worker) + + def init_workers_sync(self): + self.create_workers(RayGPUWorker, self.worker_kwargs) + try: + ray.get(self._get_worker_ready_futures()) + except ray.exceptions.ActorDiedError as e: + raise RuntimeError("RayGPUWorker died during initialization") from e + port = self.call_all_ray_workers("setup_tcp_store", + leader_only=True, + async_call=False)[0] + self.call_all_ray_workers("setup_distributed_env_and_worker", + leader_only=False, + async_call=False, + port=port) + + async def init_workers_async(self): + self.create_workers(RayGPUWorker, self.worker_kwargs) + try: + await asyncio.gather(*self._get_worker_ready_futures()) + except ray.exceptions.ActorDiedError as e: + raise RuntimeError("RayGPUWorker died during initialization") from e + port = (await asyncio.gather(*self.call_all_ray_workers( + "setup_tcp_store", leader_only=True, async_call=True)))[0] + await asyncio.gather( + *self.call_all_ray_workers("setup_distributed_env_and_worker", + leader_only=False, + async_call=True, + port=port)) + + @unwrap_ray_errors() + def call_all_ray_workers(self, func: str, leader_only: bool, + async_call: bool, *args, **kwargs): + workers = (self.workers[0], ) if leader_only else self.workers + if async_call: + return [ + getattr(worker, func).remote(*args, **kwargs) + for worker in workers + ] + else: + return ray.get([ + getattr(worker, func).remote(*args, **kwargs) + for worker in workers + ]) + + @unwrap_ray_errors() + def collective_rpc( + self, + method: str, + args: tuple = (), + kwargs: Optional[dict] = None, + non_block: bool = False, + unique_reply_rank: Optional[int] = None, + target_ranks: int | list[int] | None = None) -> list[Any]: + if target_ranks is None: + target_ranks = unique_reply_rank + workers = (self.workers if target_ranks is None else + [self.workers[rank] for rank in target_ranks] if isinstance( + target_ranks, list) else [self.workers[target_ranks]]) + kwargs = kwargs or {} + + refs = [] + for w in workers: + try: + refs.append(getattr(w, method).remote(*args, **kwargs)) + except AttributeError: + # Here worker is the RayWorkerWrapper. + # For extended worker methods, we need to use call_worker_method since + # Ray actor doesn't work with __getattr__ delegation. + refs.append(w.call_worker_method.remote(method, *args, + **kwargs)) + return refs if non_block else ray.get(refs) + + @unwrap_ray_errors() + async def collective_rpc_async( + self, + method: str, + args: tuple = (), + kwargs: Optional[dict] = None, + unique_reply_rank: Optional[int] = None, + target_ranks: int | list[int] | None = None) -> list[Any]: + refs = self.collective_rpc(method, + args, + kwargs, + non_block=True, + unique_reply_rank=unique_reply_rank, + target_ranks=target_ranks) + return await asyncio.gather(*refs) + + def submit(self, request: "GenerationRequest") -> "GenerationResult": + """ + Low-level API to the executor. Return a "future" GenerationResult + which can be waited. Forwards the request to the workers through RPC. + """ + if request.id is None: + request.set_id(self._get_next_client_id()) + logprob_params = self._get_logprob_params(request) + + with nvtx_range_debug("rpc_submit"): + self.rpc_client.submit(request).remote(need_response=False) + + result = GenerationResult( + request, + background_error_handler=self._handle_background_error, + executor=self, + disaggregated_params=request.disaggregated_params, + logprob_params=logprob_params) + self._results[request.id] = result + + return result + + def start(self): + pass + + def setup_engine_remote(self): + return self.collective_rpc("setup_engine", non_block=False) + + async def setup_engine_remote_async(self): + """Async version of setup_engine_remote for use after async worker initialization.""" + if not self.workers or len(self.workers) == 0: + raise RuntimeError( + "Workers must be initialized before calling setup_engine_remote_async" + ) + + # Setup engine on all workers + result = await self.collective_rpc_async("setup_engine") + logger.info("setup_engine_remote_async finished") + + # Now that engine is set up, start the mainloop for fetching responses + if hasattr(self, '_mainloop_started') and not self._mainloop_started: + logger.info("Starting mainloop after engine setup") + self.setup_mainloop(tasks=[self._fetch_responses_loop_async], + thread_name="ray_executor_main_loop") + self._mainloop_started = True + + return result + + def report_device_ids(self) -> list[str]: + gpu_ids = self.call_all_ray_workers("report_device_id", + leader_only=False, + async_call=False) + return sorted(gpu_ids) + + def abort_request(self, request_id: int) -> None: + self.call_all_ray_workers("abort_request", + leader_only=True, + async_call=False, + request_id=request_id) + + def abort_all_requests(self) -> None: + """Abort all active generation requests.""" + for result in list(self._results.values()): + result.abort() + + def shutdown(self): + if hasattr(self, '_shutdown_event') and self._shutdown_event.is_set(): + return + if hasattr(self, '_shutdown_event'): + self._shutdown_event.set() + + logger_debug(f"Shutting down RayExecutor", color="yellow") + + if hasattr(self, 'main_loop') and self.main_loop and hasattr( + self, 'main_loop_task_obj') and self.main_loop_task_obj: + logger_debug("Cancelling main loop task.", color="yellow") + try: + self.main_loop.call_soon_threadsafe( + self.main_loop_task_obj.cancel) + except Exception as e: + logger_debug(f"Error cancelling main loop task: {e}", + color="yellow") + + if hasattr(self, 'main_loop_thread'): + self.main_loop_thread.join() + + # Then, shutdown the workers + if hasattr(self, 'workers') and self.workers is not None: + try: + shutdown_refs = [ + worker.shutdown.remote() for worker in self.workers + ] + # Add timeout to prevent indefinite hanging + ray.get(shutdown_refs, timeout=30.0) + except ray.exceptions.GetTimeoutError: + logger.warning( + "Timeout waiting for workers to shutdown after 30 seconds") + except Exception as e: + logger.warning(f"Error shutting down: {e}") + + # The engines are already stopped by the shutdown RPC above, so + # kill the actor processes explicitly instead of relying on + # handle garbage collection. ray.kill() only *initiates* an + # asynchronous kill; _wait_for_cluster_resource_release() below + # blocks until Ray has reclaimed the workers' resources. + for worker in self.workers: + try: + ray.kill(worker, no_restart=True) + except Exception as e: + logger.warning(f"Error killing worker: {e}") + + if hasattr(self, 'rpc_client') and self.rpc_client is not None: + try: + self.rpc_client.close() + except Exception as e: + logger_debug(f"Suppressed error during RPC client close: {e}") + + self.workers = None + if hasattr(self, + "placement_group") and self.placement_group is not None: + # Only remove placement group if Ray is still initialized + # to avoid triggering auto_init_ray() during program exit + if ray.is_initialized(): + ray.util.remove_placement_group(self.placement_group) + self.placement_group = None + self.bundle_indices = None + + # ray.kill() and remove_placement_group() above are asynchronous. + # Block until Ray has reclaimed the workers' resources so their GPU + # cleanup has completed before shutdown() returns. + self._wait_for_cluster_resource_release(timeout=30.0) + + if self.has_start_local_cluser and ray.is_initialized(): + logger.debug("Shutting down Ray cluster") + ray.shutdown() + + def _wait_for_cluster_resource_release(self, timeout: float = 30.0) -> None: + """Block until Ray returns the workers' resources to the cluster. + + ray.kill() and remove_placement_group() only initiate an + asynchronous teardown; Ray reclaims an actor's logical resources + after the raylet has reaped the worker process, by which point the + CUDA driver has already destroyed its context (GPU memory and IPC + mappings). Waiting here therefore guarantees that a subsequent LLM + instance will not race against the dying workers, which can + otherwise fail spuriously (e.g. cudaErrorMapBufferObjectFailed when + opening CUDA IPC handles). Full availability can only be expected on + a cluster dedicated to this executor, so the wait is skipped when + attached to an external cluster. Best-effort: logs a warning on + timeout instead of raising, since shutdown must not fail. + """ + if not self.has_start_local_cluser or not ray.is_initialized(): + return + deadline = time.monotonic() + timeout + busy = {} + while time.monotonic() < deadline: + try: + cluster = ray.cluster_resources() + available = ray.available_resources() + except Exception as e: + logger.debug(f"Could not query Ray resources: {e}") + return + busy = { + key: cluster[key] - available.get(key, 0.0) + for key in ("GPU", "CPU") if key in cluster and cluster[key] - + available.get(key, 0.0) > 1e-6 + } + if not busy: + return + time.sleep(0.1) + logger.warning( + f"Timed out after {timeout}s waiting for Ray to reclaim cluster " + f"resources; still in use: {busy}.") + + def _get_worker_ready_futures(self): + return [worker.__ray_ready__.remote() for worker in self.workers] + + def _get_placement_group( + self, + tp_size: int, + worker_kwargs: Dict = None) -> Tuple[Any, List[int]]: + """ + Obtain placement group(s) and bundle indices for workers. + + Priorities: + 1. `ray_placement_config` in `llm_args`. + 2. `TRTLLM_RAY_BUNDLE_INDICES` environment variable (uses current placement group). + 3. Default creation: A PACK placement group where each bundle has `tp_size` GPUs. + - When `tp_size` <= GPUs per node, keep one TP group per node. + - When `tp_size` > GPUs per node, allow a TP group to span nodes. + - rank 0 is forced onto the driver node. + + Returns: + Tuple[Union[PlacementGroup, List[PlacementGroup]], List[int]]: + - placement_group(s): A single `PlacementGroup` (shared by all workers) or a list of `PlacementGroup` (one per worker). + - bundle_indices: A list of bundle indices. + If `placement_group(s)` is a single object, `bundle_indices[i]` maps worker `i` to that bundle in the group. + If `placement_group(s)` is a list, `bundle_indices[i]` maps worker `i` to that bundle in `placement_groups[i]`. + """ + llm_args = worker_kwargs.get("llm_args") if worker_kwargs else None + + placement_config = getattr(llm_args, 'ray_placement_config', + None) if llm_args else None + + def _get_from_placement_config(placement_config): + total_workers = sum( + len(indices) + for indices in placement_config.placement_bundle_indices) + if total_workers != self.world_size: + raise ValueError( + f"Total bundle indices ({total_workers}) must equal world_size ({self.world_size})" + ) + + logger.info( + f"Creating {self.world_size} workers with external placement groups" + ) + + flat_pgs = [] + flat_indices = [] + for pg, indices in zip(placement_config.placement_groups, + placement_config.placement_bundle_indices): + for idx in indices: + flat_pgs.append(pg) + flat_indices.append(idx) + + return flat_pgs, flat_indices + + def _get_from_env(bundle_indices): + pg = get_current_placement_group() + if pg is not None: + bundle_indices = list(map(int, bundle_indices.split(","))) + assert len(bundle_indices) == self.world_size, ( + f"Need {self.world_size} bundle indices for world_size, got {bundle_indices=}" + ) + assert len(set(bundle_indices)) == len(bundle_indices), ( + f"TRTLLM_RAY_BUNDLE_INDICES cannot have duplicate values, but got {bundle_indices=}." + ) + assert max(bundle_indices) < len(pg.bundle_specs), ( + f"{bundle_indices=} out of range for PG with {len(pg.bundle_specs)} bundles" + ) + return pg, bundle_indices + else: + raise ValueError("No global placement group is found.") + + def _get_default(tp_size): + head_tag = f"node:{self.master_address}" + nodes = ray.nodes() + gpus_per_node = int(nodes[0]["Resources"].get( + "GPU", 0)) # assume symmetric across nodes + + bundle_cpu = bundle_gpu = min(tp_size, gpus_per_node) + + bundles, bundle_indices = [], [] + current = 0 + for rank in range(self.world_size): + if current == 0: + bundle = {"GPU": bundle_gpu, "CPU": bundle_cpu} + if len(bundles) == 0: + bundle[ + head_tag] = 0.01 # to force placement on head node + bundles.append(bundle) + + bundle_indices.append(len(bundles) - 1) + current = (current + 1) % bundle_gpu + + strategy = "PACK" + logger.debug( + f"[Strategy={strategy}] Bundles: {bundles} for tp_size: {tp_size} and world_size: {self.world_size}" + ) + pg = placement_group(bundles, strategy=strategy) + + return pg, bundle_indices + + if self.world_size % tp_size != 0: + raise ValueError( + f"world_size {self.world_size} must be a multiple of tp_size {tp_size}" + ) + + # path 0 + if placement_config and placement_config.placement_groups is not None: + return _get_from_placement_config(placement_config) + # path 1 + if bundle_indices := os.getenv("TRTLLM_RAY_BUNDLE_INDICES", None): + return _get_from_env(bundle_indices) + # path 2 + return _get_default(tp_size) + + @property + def enable_postprocess_parallel(self) -> bool: + ret = super().enable_postprocess_parallel + assert ret == False, "Postprocess parallel is not supported in RayExecutor" + return ret diff --git a/tensorrt_llm/executor/ray/gpu_worker.py b/tensorrt_llm/executor/ray/gpu_worker.py new file mode 100644 index 000000000000..cd5165fbba0a --- /dev/null +++ b/tensorrt_llm/executor/ray/gpu_worker.py @@ -0,0 +1,378 @@ +# 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. +import gc +import importlib +import os +import tempfile +from functools import wraps +from pathlib import Path +from queue import Queue +from typing import Any, List, Optional, Type + +import ray +import torch + +from tensorrt_llm._torch.utils import get_device_uuid +from tensorrt_llm._torch.virtual_memory import (materialize_with_tag, + release_with_tag) +from tensorrt_llm.executor.ray.utils import control_action_decorator + +from ... import TorchLlmArgs +from ...bindings import executor as tllm +from ...llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType +from ...llmapi.tokenizer import TokenizerBase +from ...llmapi.utils import configure_cpu_affinity +from ...sampling_params import BatchedLogitsProcessor +from ..base_worker import BaseWorker +from ..postproc_worker import PostprocWorkerConfig +from ..request import GenerationRequest +from ..result import GenerationResult +from ..rpc_worker_mixin import RpcWorkerMixin + +__all__ = [ + "RayGPUWorker", + "RayWorkerWrapper", +] + + +def resolve_obj_by_qualname(qualname: str) -> Any: + """Resolve an object by its fully qualified name.""" + module_name, obj_name = qualname.rsplit(".", 1) + module = importlib.import_module(module_name) + return getattr(module, obj_name) + + +@ray.remote +class RayWorkerWrapper: + + def __init__(self, worker_cls, worker_kwargs, world_size, rank): + self.master_address = os.environ["MASTER_ADDR"] + self.world_size = world_size + self.rank = rank + # Ray can't pickle TensorRT logger + global logger + from tensorrt_llm.logger import logger + + # Expect to see global counts w/ RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1, + # unless CUDA_VISIBLE_DEVICES is set. + logger.debug( + f"CUDA device count visible to Ray: {torch.cuda.device_count()}") + + # Physical gpu id + self.gpu = int(ray.get_gpu_ids()[0]) + self.local_gpu = self.physical_to_local_id(self.gpu) + + # Per-worker DeepGemm JIT cache to avoid rename race across co-located workers + os.environ["DG_JIT_CACHE_DIR"] = os.path.join( + tempfile.gettempdir(), f"deep_gemm_rank{rank}_gpu{self.gpu}") + + torch.cuda.set_device(self.local_gpu) + + self.worker_cls = RayWorkerWrapper._inject_worker_extension( + worker_cls, worker_kwargs.pop("ray_worker_extension_cls", None)) + self.worker_kwargs = worker_kwargs + + def _create_tcp_store(self, + port: Optional[int] = None + ) -> torch.distributed.TCPStore: + # port=0 means let the OS pick an available port (only valid for master) + # For non-master, port must be specified to connect to master's port + actual_port = port if port is not None else 0 + return torch.distributed.TCPStore(host_name=self.master_address, + port=actual_port, + world_size=self.world_size, + is_master=(self.rank == 0), + wait_for_workers=False) + + def setup_tcp_store(self): + if self.rank != 0: + raise RuntimeError("Only the master worker can setup TCP store") + self.store = self._create_tcp_store() + return self.store.port + + def setup_distributed_env_and_worker(self, port: int): + if self.rank != 0: + self.store = self._create_tcp_store(port) + + torch.distributed.init_process_group(backend="cuda:nccl,cpu:gloo", + store=self.store, + world_size=self.world_size, + rank=self.rank) + assert torch.distributed.get_world_size( + ) == self.world_size, "Process group world size must match the expected world size" + logger.info( + f"[Rank {self.rank}] Finished PG init. Global GPU ID: {self.gpu}, local GPU ID: {self.local_gpu}" + ) + + self.worker = self.worker_cls(device_id=self.local_gpu, + **self.worker_kwargs) + self._has_setup_distributed_env_and_worker = True + + @property + def has_setup_distributed_env_and_worker(self) -> bool: + return getattr(self, '_has_setup_distributed_env_and_worker', False) + + def ensure_distributed_setup(func): + + @wraps(func) + def wrapper(self, *args, **kwargs): + if not self.has_setup_distributed_env_and_worker: + raise RuntimeError( + "Have not setup distributed environment and worker yet") + return func(self, *args, **kwargs) + + return wrapper + + @ensure_distributed_setup + def submit(self, request: GenerationRequest) -> GenerationResult: + return self.worker.submit(request) + + @ensure_distributed_setup + def enqueue_request(self, + request: GenerationRequest, + result_wait_queue: Queue | None = None) -> int: + return self.worker.enqueue_request(request, result_wait_queue) + + @ensure_distributed_setup + def abort_request(self, request_id: int) -> None: + self.worker.abort_request(request_id) + + @ensure_distributed_setup + def report_device_id(self) -> str: + local_id = self.physical_to_local_id(self.gpu) + return get_device_uuid(local_id) + + @ensure_distributed_setup + def call_worker_method(self, method_name: str, *args, **kwargs): + """Generic method to call any method on the underlying worker.""" + if hasattr(self.worker, method_name): + method = getattr(self.worker, method_name) + if callable(method): + return method(*args, **kwargs) + else: + raise AttributeError( + f"'{method_name}' is not a callable method of RayGPUWorker." + ) + else: + raise AttributeError( + f"The RayGPUWorker has no method called '{method_name}'.") + + def shutdown(self): + if hasattr(self, 'worker'): + self.worker.shutdown() + + def __repr__(self) -> str: + """Customizes the actor's prefix in the Ray logs. + + This makes it easier to identify which worker is producing specific log messages. + Refer to https://github.com/NVIDIA-NeMo/RL/blob/faad02113c3c502437ccb339cb848796334aedd9/nemo_rl/models/policy/dtensor_policy_worker_v2.py#L95 + """ + if torch.distributed.is_initialized(): + return f"{self.__class__.__qualname__}[rank={torch.distributed.get_rank()}]" + else: + return f"{self.__class__.__qualname__}" + + @staticmethod + def physical_to_local_id(phys_id: int) -> int: + visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + if not visible_devices: + return phys_id + id_mapping = list(map(int, visible_devices.split(","))) + return id_mapping.index(phys_id) + + @staticmethod + def _inject_worker_extension( + worker_class: Type[BaseWorker], + extension_cls_name: Optional[str]) -> Type[BaseWorker]: + """Inject worker extension into the worker class if specified.""" + if not extension_cls_name: + return worker_class + + try: + extension_cls = resolve_obj_by_qualname(extension_cls_name) + except (ImportError, AttributeError, ValueError) as e: + raise RuntimeError( + f"Failed to load worker extension '{extension_cls_name}'" + ) from e + + # Check for conflicts + for attr in dir(extension_cls): + if attr.startswith("__"): + continue + if hasattr(worker_class, attr): + raise ValueError( + f"Worker class {worker_class.__name__} already defines '{attr}', " + f"which conflicts with extension {extension_cls.__name__}.") + + derived_name = f"{worker_class.__name__}With{extension_cls.__name__}" + ExtendedWorker = type(derived_name, (worker_class, extension_cls), + {'__module__': worker_class.__module__}) + return ExtendedWorker + + +class RayGPUWorker(RpcWorkerMixin, BaseWorker): + + def __init__( + self, + device_id: int, + engine: Path, + executor_config: Optional[tllm.ExecutorConfig] = None, + batched_logits_processor: Optional[BatchedLogitsProcessor] = None, + postproc_worker_config: Optional[PostprocWorkerConfig] = None, + is_llm_executor: Optional[bool] = None, + hf_model_dir: Optional[Path] = None, + tokenizer: Optional[TokenizerBase] = None, + llm_args: Optional[BaseLlmArgs] = None, + rpc_addr: Optional[str] = None, + hmac_key: bytes = b"", + ) -> None: + global logger + from tensorrt_llm.logger import logger + + super().__init__( + engine=engine, + executor_config=executor_config, + batched_logits_processor=batched_logits_processor, + postproc_worker_config=postproc_worker_config, + is_llm_executor=is_llm_executor, + hf_model_dir=hf_model_dir, + tokenizer=tokenizer, + llm_args=llm_args, + ) + + self.device_id = device_id + self.global_rank = torch.distributed.get_rank() + if self.global_rank > 1: + logger.set_rank(self.global_rank) + + if rpc_addr is None: + raise RuntimeError( + "RPC mode enabled but no rpc_addr provided to RayGPUWorker") + self.init_rpc_worker(self.global_rank, rpc_addr, hmac_key) + self.start_rpc_server() + + def setup_engine(self): + if torch.distributed.is_initialized( + ) and torch.distributed.get_world_size() > 1: + torch.distributed.barrier() + super().setup_engine() + + def enqueue_request(self, + request: GenerationRequest, + result_wait_queue: Queue | None = None) -> int: + return self._enqueue_request(request, result_wait_queue) + + @control_action_decorator + def sleep(self, sleep_tags: List[str]): + assert isinstance(self.llm_args, + TorchLlmArgs), "sleep() only available for TorchLLM" + + if self.llm_args.sleep_config is None: + raise ValueError( + "Sleep feature is not enabled, please set sleep_config in the LLM arguments." + ) + try: + tags = [ExecutorMemoryType(tag) for tag in sleep_tags] + logger.info(f"Sleep: {tags}") + torch.cuda.synchronize() + release_with_tag(*tags) + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + except Exception as e: + logger.error(f"Encountered an error in sleep: {e}") + raise e + + @control_action_decorator + def wakeup(self, wakeup_tags: List[str]): + assert isinstance(self.llm_args, + TorchLlmArgs), "wakeup() only available for TorchLLM" + + if self.llm_args.sleep_config is None: + raise ValueError( + "Sleep feature is not enabled, please set sleep_config in the LLM arguments." + ) + try: + tags = [ExecutorMemoryType(tag) for tag in wakeup_tags] + logger.info(f"Wakeup: {tags}") + torch.cuda.synchronize() + materialize_with_tag(*tags) + torch.cuda.synchronize() + except Exception as e: + logger.error(f"Encountered an error in wakeup") + raise e + + def start(self): + pass + + def shutdown(self): + + if self.doing_shutdown: + return + else: + self.doing_shutdown = True + + logger.debug(f'Worker {self.rank} shutting down...') + + if hasattr(self, 'shutdown_event'): + self.shutdown_event.set() + + if hasattr(self, 'rpc_server') and self.rpc_server is not None: + logger.info(f"[Rank {self.global_rank}] Shutting down RPC server") + try: + self.rpc_server.shutdown() + except Exception as e: + # Suppress errors during RPC server shutdown + # These can occur if the server is already closed or during cleanup + logger.debug( + f"[Rank {self.global_rank}] Suppressed error during RPC server shutdown: {e}" + ) + self.rpc_server = None + + if self.engine is not None: + self.engine.shutdown() + self.engine = None + + assert self._executor_config is None, "An empty executor_config is expected in shutdown when LLM arguments are defined." + if (self.llm_args.backend == "pytorch" + and hasattr(self, "checkpoint_loader") + and self.checkpoint_loader is not None): + self.checkpoint_loader.cleanup() + self.checkpoint_loader = None + + # Check if there are any errors from the threads before shutdown. + self._handle_background_error() + + logger.debug(f"Worker {self.rank} shutdown done.") + + def _get_comm_ranks_device_id(self): + # Make sure C++ executor would use same devices/ranks as py_executor + global_rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + comm_ranks = [None] * world_size + device_ids = [None] * world_size + + torch.distributed.all_gather_object(comm_ranks, global_rank) + torch.distributed.all_gather_object(device_ids, self.device_id) + + configure_cpu_affinity(self.device_id) + + return comm_ranks, device_ids + + def __enter__(self): + return self + + def __del__(self): + self.shutdown() diff --git a/tensorrt_llm/executor/ray/stub.py b/tensorrt_llm/executor/ray/stub.py new file mode 100644 index 000000000000..a18136a6ceb0 --- /dev/null +++ b/tensorrt_llm/executor/ray/stub.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. +from functools import wraps as _wraps + +from tensorrt_llm._utils import mpi_disabled as _mpi_disabled + +# Don't raise error on import - only when Ray functionality is actually used +_RAY_NOT_INSTALLED_MSG = "Ray requested (TLLM_DISABLE_MPI=1), but not installed. Please install Ray." + + +def remote(*args, **kwargs): + + def decorator(func): + # Returns a function that always raises. + # Decorated class depends on ray, but ray is not installed. + @_wraps(func) + def stub_checker(*_, **__): + raise RuntimeError( + f'Ray not installed, so the remote function / actor "{func.__name__}" is not available.' + ) + + return stub_checker + + if len(args) == 1 and len(kwargs) == 0 and callable(args[0]): + return decorator(args[0]) + + return decorator + + +def __getattr__(name): + msg = f'Ray not installed, so "ray.{name}" is unavailable.' + if _mpi_disabled(): + msg = _RAY_NOT_INSTALLED_MSG + raise RuntimeError(msg) diff --git a/tensorrt_llm/executor/ray/utils.py b/tensorrt_llm/executor/ray/utils.py new file mode 100644 index 000000000000..8df771ece793 --- /dev/null +++ b/tensorrt_llm/executor/ray/utils.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. +import functools +from contextlib import contextmanager +from typing import Callable, Optional + +try: + import ray +except ImportError: + import tensorrt_llm.executor.ray.stub as ray + + +@contextmanager +def unwrap_ray_errors(): + try: + yield + except ray.exceptions.RayTaskError as e: + raise e.as_instanceof_cause() from e + + +def control_action_decorator(func: Optional[Callable] = None, + *, + drain: bool = True) -> Callable: + """Wrap a method in the ``control_action`` context manager. + + Supports both bare and parameterized forms:: + + @control_action_decorator # drain=True (default) + def shutdown(self): ... + + @control_action_decorator(drain=False) # non-draining variant + def update_weights_via_ipc_zmq(self): ... + """ + + def decorator(f: Callable) -> Callable: + + @functools.wraps(f) + def wrapper(self, *args, **kwargs): + with self.engine.control_action(drain=drain): + return f(self, *args, **kwargs) + + return wrapper + + if func is None: + return decorator + return decorator(func) diff --git a/tensorrt_llm/executor/ray_executor.py b/tensorrt_llm/executor/ray_executor.py index accd9efc1632..39fe9ae19e9b 100644 --- a/tensorrt_llm/executor/ray_executor.py +++ b/tensorrt_llm/executor/ray_executor.py @@ -1,550 +1,37 @@ -import asyncio -import os -import time -from typing import Any, Dict, List, Optional, Tuple - -try: - import ray -except ModuleNotFoundError as e: - e.msg = """Cannot import Ray. Please install 'ray' package to use ray orchestrator""" - raise - -from ray.util.placement_group import (get_current_placement_group, - placement_group) - -try: - # Ray >= 2.55.0 - from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -except ImportError: - # Older Ray re-exported it from ray.util.placement_group - from ray.util.placement_group import PlacementGroupSchedulingStrategy - -from tensorrt_llm._ray_utils import unwrap_ray_errors -from tensorrt_llm._utils import nvtx_range_debug -from tensorrt_llm.logger import logger - -from ..llmapi.utils import logger_debug -from .executor import GenerationExecutor -from .postproc_worker import PostprocWorkerConfig -from .ray_gpu_worker import RayGPUWorker, RayWorkerWrapper -from .request import GenerationRequest -from .result import GenerationResult -from .rpc_proxy_mixin import RpcExecutorMixin -from .utils import has_event_loop +# 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. +"""Compatibility shim for ``tensorrt_llm.executor.ray_executor``. + +Will be removed once all usages are migrated to +``tensorrt_llm.executor.ray.executor``. + +DO NOT ADD ANYTHING TO THIS FILE. +""" + +import warnings + +from tensorrt_llm.executor.ray.executor import RayExecutor # noqa: F401 + +warnings.warn( + "tensorrt_llm.executor.ray_executor has moved to " + "tensorrt_llm.executor.ray.executor and will be removed in a future " + "release.", + FutureWarning, + stacklevel=2, +) __all__ = [ "RayExecutor", ] - - -class RayExecutor(RpcExecutorMixin, GenerationExecutor): - - def __init__(self, - worker_kwargs: Dict, - model_world_size: int, - postproc_worker_config: PostprocWorkerConfig, - is_llm_executor: bool, - tp_size=1): - os.environ['RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES'] = '1' - os.environ["RAY_DEDUP_LOGS"] = "0" # for debug - - super().__init__(model_world_size, postproc_worker_config, - is_llm_executor) - - self.has_start_local_cluser = False - runtime_env = { - "env_vars": { - "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1" - } - } - - ray_init_args = { - "include_dashboard": False, - "namespace": "trtllm", - "ignore_reinit_error": True, - "runtime_env": runtime_env - } - - try: - if os.environ.get("TLLM_RAY_FORCE_LOCAL_CLUSTER", "0") != "1": - try: - ray.init(address="auto", **ray_init_args) - logger.info(f"Attached to an existing Ray cluster.") - except ConnectionError: - logger.info(f"Ray cluster not found, starting a new one.") - - if not ray.is_initialized(): - ray.init(**ray_init_args) - self.has_start_local_cluser = True - else: - ray.init(address="local", **ray_init_args) - self.has_start_local_cluser = True - - self.world_size = model_world_size - self.tp_size = tp_size - self.master_address = ray.util.get_node_ip_address() - - self.worker_kwargs = dict( - **worker_kwargs, - postproc_worker_config=postproc_worker_config, - is_llm_executor=is_llm_executor) - - self.init_rpc_executor() - # Inject the generated HMAC key into worker_kwargs for workers - self.worker_kwargs['hmac_key'] = self.hmac_key - self.worker_kwargs['rpc_addr'] = self.rpc_addr - - placement_config = getattr(self.worker_kwargs['llm_args'], - 'ray_placement_config', None) - defer_workers_init = placement_config.defer_workers_init if placement_config else False - - if defer_workers_init: - self.workers = [ - ] # Placeholder, will be initialized in setup_async - self._mainloop_started = False # DO NOT start mainloop until after setup_engine_remote_async is called - else: - if not has_event_loop(): - self.init_workers_sync() - self.setup_engine_remote() - self.setup_mainloop(tasks=[self._fetch_responses_loop_async], - thread_name="ray_executor_main_loop") - - except Exception as e: - self.shutdown() - logger.error(f"Failed to initialize RayExecutor: {e}") - raise e - - def create_workers(self, worker_cls, worker_kwargs): - llm_args = worker_kwargs.get("llm_args") - placement_config = getattr(llm_args, 'ray_placement_config', - None) if llm_args else None - ray_worker_nsight_options = getattr( - llm_args, 'ray_worker_nsight_options', None) if llm_args else None - - # When set to be a fraction, it allows Ray to schedule - # multiple actors on a single GPU for colocate use cases. - num_gpus = float(os.getenv("TRTLLM_RAY_PER_WORKER_GPUS", "1.0")) - if placement_config and placement_config.per_worker_gpu_share is not None: - num_gpus = placement_config.per_worker_gpu_share - - logger.debug(f"{num_gpus=} for each worker.") - - runtime_env = ray.runtime_env.RuntimeEnv() - # Exclude node-local env vars. e.g., The raylet that spawns each worker sets - # RAY_RAYLET_PID to its own PID at exec time. - _NODE_LOCAL_VARS = { - "RAY_RAYLET_PID", - "RAY_NODE_IP_ADDRESS", - } - - runtime_env["env_vars"] = { - k: v - for k, v in os.environ.items() if k not in _NODE_LOCAL_VARS - } - runtime_env["env_vars"].update({ - "TLLM_DISABLE_MPI": "1", - "MASTER_ADDR": self.master_address, # head-IP for NCCL/Gloo - }) - if ray_worker_nsight_options: - runtime_env["nsight"] = ray_worker_nsight_options - - placement_groups, self.bundle_indices = self._get_placement_group( - tp_size=self.tp_size, worker_kwargs=worker_kwargs) - - if isinstance(placement_groups, list): - self.placement_group = None - else: - self.placement_group = placement_groups - - self.workers = [] - for rank in range(self.world_size): - pg = placement_groups[rank] if isinstance( - placement_groups, list) else placement_groups - worker = RayWorkerWrapper.options( - num_gpus=num_gpus, - runtime_env=runtime_env, - scheduling_strategy=PlacementGroupSchedulingStrategy( - placement_group=pg, - placement_group_bundle_index=self.bundle_indices[rank], - )).remote(worker_cls, worker_kwargs, self.world_size, rank) - self.workers.append(worker) - - def init_workers_sync(self): - self.create_workers(RayGPUWorker, self.worker_kwargs) - try: - ray.get(self._get_worker_ready_futures()) - except ray.exceptions.ActorDiedError as e: - raise RuntimeError("RayGPUWorker died during initialization") from e - port = self.call_all_ray_workers("setup_tcp_store", - leader_only=True, - async_call=False)[0] - self.call_all_ray_workers("setup_distributed_env_and_worker", - leader_only=False, - async_call=False, - port=port) - - async def init_workers_async(self): - self.create_workers(RayGPUWorker, self.worker_kwargs) - try: - await asyncio.gather(*self._get_worker_ready_futures()) - except ray.exceptions.ActorDiedError as e: - raise RuntimeError("RayGPUWorker died during initialization") from e - port = (await asyncio.gather(*self.call_all_ray_workers( - "setup_tcp_store", leader_only=True, async_call=True)))[0] - await asyncio.gather( - *self.call_all_ray_workers("setup_distributed_env_and_worker", - leader_only=False, - async_call=True, - port=port)) - - @unwrap_ray_errors() - def call_all_ray_workers(self, func: str, leader_only: bool, - async_call: bool, *args, **kwargs): - workers = (self.workers[0], ) if leader_only else self.workers - if async_call: - return [ - getattr(worker, func).remote(*args, **kwargs) - for worker in workers - ] - else: - return ray.get([ - getattr(worker, func).remote(*args, **kwargs) - for worker in workers - ]) - - @unwrap_ray_errors() - def collective_rpc( - self, - method: str, - args: tuple = (), - kwargs: Optional[dict] = None, - non_block: bool = False, - unique_reply_rank: Optional[int] = None, - target_ranks: int | list[int] | None = None) -> list[Any]: - if target_ranks is None: - target_ranks = unique_reply_rank - workers = (self.workers if target_ranks is None else - [self.workers[rank] for rank in target_ranks] if isinstance( - target_ranks, list) else [self.workers[target_ranks]]) - kwargs = kwargs or {} - - refs = [] - for w in workers: - try: - refs.append(getattr(w, method).remote(*args, **kwargs)) - except AttributeError: - # Here worker is the RayWorkerWrapper. - # For extended worker methods, we need to use call_worker_method since - # Ray actor doesn't work with __getattr__ delegation. - refs.append(w.call_worker_method.remote(method, *args, - **kwargs)) - return refs if non_block else ray.get(refs) - - @unwrap_ray_errors() - async def collective_rpc_async( - self, - method: str, - args: tuple = (), - kwargs: Optional[dict] = None, - unique_reply_rank: Optional[int] = None, - target_ranks: int | list[int] | None = None) -> list[Any]: - refs = self.collective_rpc(method, - args, - kwargs, - non_block=True, - unique_reply_rank=unique_reply_rank, - target_ranks=target_ranks) - return await asyncio.gather(*refs) - - def submit(self, request: "GenerationRequest") -> "GenerationResult": - """ - Low-level API to the executor. Return a "future" GenerationResult - which can be waited. Forwards the request to the workers through RPC. - """ - if request.id is None: - request.set_id(self._get_next_client_id()) - logprob_params = self._get_logprob_params(request) - - with nvtx_range_debug("rpc_submit"): - self.rpc_client.submit(request).remote(need_response=False) - - result = GenerationResult( - request, - background_error_handler=self._handle_background_error, - executor=self, - disaggregated_params=request.disaggregated_params, - logprob_params=logprob_params) - self._results[request.id] = result - - return result - - def start(self): - pass - - def setup_engine_remote(self): - return self.collective_rpc("setup_engine", non_block=False) - - async def setup_engine_remote_async(self): - """Async version of setup_engine_remote for use after async worker initialization.""" - if not self.workers or len(self.workers) == 0: - raise RuntimeError( - "Workers must be initialized before calling setup_engine_remote_async" - ) - - # Setup engine on all workers - result = await self.collective_rpc_async("setup_engine") - logger.info("setup_engine_remote_async finished") - - # Now that engine is set up, start the mainloop for fetching responses - if hasattr(self, '_mainloop_started') and not self._mainloop_started: - logger.info("Starting mainloop after engine setup") - self.setup_mainloop(tasks=[self._fetch_responses_loop_async], - thread_name="ray_executor_main_loop") - self._mainloop_started = True - - return result - - def report_device_ids(self) -> list[str]: - gpu_ids = self.call_all_ray_workers("report_device_id", - leader_only=False, - async_call=False) - return sorted(gpu_ids) - - def abort_request(self, request_id: int) -> None: - self.call_all_ray_workers("abort_request", - leader_only=True, - async_call=False, - request_id=request_id) - - def abort_all_requests(self) -> None: - """Abort all active generation requests.""" - for result in list(self._results.values()): - result.abort() - - def shutdown(self): - if hasattr(self, '_shutdown_event') and self._shutdown_event.is_set(): - return - if hasattr(self, '_shutdown_event'): - self._shutdown_event.set() - - logger_debug(f"Shutting down RayExecutor", color="yellow") - - if hasattr(self, 'main_loop') and self.main_loop and hasattr( - self, 'main_loop_task_obj') and self.main_loop_task_obj: - logger_debug("Cancelling main loop task.", color="yellow") - try: - self.main_loop.call_soon_threadsafe( - self.main_loop_task_obj.cancel) - except Exception as e: - logger_debug(f"Error cancelling main loop task: {e}", - color="yellow") - - if hasattr(self, 'main_loop_thread'): - self.main_loop_thread.join() - - # Then, shutdown the workers - if hasattr(self, 'workers') and self.workers is not None: - try: - shutdown_refs = [ - worker.shutdown.remote() for worker in self.workers - ] - # Add timeout to prevent indefinite hanging - ray.get(shutdown_refs, timeout=30.0) - except ray.exceptions.GetTimeoutError: - logger.warning( - "Timeout waiting for workers to shutdown after 30 seconds") - except Exception as e: - logger.warning(f"Error shutting down: {e}") - - # The engines are already stopped by the shutdown RPC above, so - # kill the actor processes explicitly instead of relying on - # handle garbage collection. ray.kill() only *initiates* an - # asynchronous kill; _wait_for_cluster_resource_release() below - # blocks until Ray has reclaimed the workers' resources. - for worker in self.workers: - try: - ray.kill(worker, no_restart=True) - except Exception as e: - logger.warning(f"Error killing worker: {e}") - - if hasattr(self, 'rpc_client') and self.rpc_client is not None: - try: - self.rpc_client.close() - except Exception as e: - logger_debug(f"Suppressed error during RPC client close: {e}") - - self.workers = None - if hasattr(self, - "placement_group") and self.placement_group is not None: - # Only remove placement group if Ray is still initialized - # to avoid triggering auto_init_ray() during program exit - if ray.is_initialized(): - ray.util.remove_placement_group(self.placement_group) - self.placement_group = None - self.bundle_indices = None - - # ray.kill() and remove_placement_group() above are asynchronous. - # Block until Ray has reclaimed the workers' resources so their GPU - # cleanup has completed before shutdown() returns. - self._wait_for_cluster_resource_release(timeout=30.0) - - if self.has_start_local_cluser and ray.is_initialized(): - logger.debug("Shutting down Ray cluster") - ray.shutdown() - - def _wait_for_cluster_resource_release(self, timeout: float = 30.0) -> None: - """Block until Ray returns the workers' resources to the cluster. - - ray.kill() and remove_placement_group() only initiate an - asynchronous teardown; Ray reclaims an actor's logical resources - after the raylet has reaped the worker process, by which point the - CUDA driver has already destroyed its context (GPU memory and IPC - mappings). Waiting here therefore guarantees that a subsequent LLM - instance will not race against the dying workers, which can - otherwise fail spuriously (e.g. cudaErrorMapBufferObjectFailed when - opening CUDA IPC handles). Full availability can only be expected on - a cluster dedicated to this executor, so the wait is skipped when - attached to an external cluster. Best-effort: logs a warning on - timeout instead of raising, since shutdown must not fail. - """ - if not self.has_start_local_cluser or not ray.is_initialized(): - return - deadline = time.monotonic() + timeout - busy = {} - while time.monotonic() < deadline: - try: - cluster = ray.cluster_resources() - available = ray.available_resources() - except Exception as e: - logger.debug(f"Could not query Ray resources: {e}") - return - busy = { - key: cluster[key] - available.get(key, 0.0) - for key in ("GPU", "CPU") if key in cluster and cluster[key] - - available.get(key, 0.0) > 1e-6 - } - if not busy: - return - time.sleep(0.1) - logger.warning( - f"Timed out after {timeout}s waiting for Ray to reclaim cluster " - f"resources; still in use: {busy}.") - - def _get_worker_ready_futures(self): - return [worker.__ray_ready__.remote() for worker in self.workers] - - def _get_placement_group( - self, - tp_size: int, - worker_kwargs: Dict = None) -> Tuple[Any, List[int]]: - """ - Obtain placement group(s) and bundle indices for workers. - - Priorities: - 1. `ray_placement_config` in `llm_args`. - 2. `TRTLLM_RAY_BUNDLE_INDICES` environment variable (uses current placement group). - 3. Default creation: A PACK placement group where each bundle has `tp_size` GPUs. - - When `tp_size` <= GPUs per node, keep one TP group per node. - - When `tp_size` > GPUs per node, allow a TP group to span nodes. - - rank 0 is forced onto the driver node. - - Returns: - Tuple[Union[PlacementGroup, List[PlacementGroup]], List[int]]: - - placement_group(s): A single `PlacementGroup` (shared by all workers) or a list of `PlacementGroup` (one per worker). - - bundle_indices: A list of bundle indices. - If `placement_group(s)` is a single object, `bundle_indices[i]` maps worker `i` to that bundle in the group. - If `placement_group(s)` is a list, `bundle_indices[i]` maps worker `i` to that bundle in `placement_groups[i]`. - """ - llm_args = worker_kwargs.get("llm_args") if worker_kwargs else None - - placement_config = getattr(llm_args, 'ray_placement_config', - None) if llm_args else None - - def _get_from_placement_config(placement_config): - total_workers = sum( - len(indices) - for indices in placement_config.placement_bundle_indices) - if total_workers != self.world_size: - raise ValueError( - f"Total bundle indices ({total_workers}) must equal world_size ({self.world_size})" - ) - - logger.info( - f"Creating {self.world_size} workers with external placement groups" - ) - - flat_pgs = [] - flat_indices = [] - for pg, indices in zip(placement_config.placement_groups, - placement_config.placement_bundle_indices): - for idx in indices: - flat_pgs.append(pg) - flat_indices.append(idx) - - return flat_pgs, flat_indices - - def _get_from_env(bundle_indices): - pg = get_current_placement_group() - if pg is not None: - bundle_indices = list(map(int, bundle_indices.split(","))) - assert len(bundle_indices) == self.world_size, ( - f"Need {self.world_size} bundle indices for world_size, got {bundle_indices=}" - ) - assert len(set(bundle_indices)) == len(bundle_indices), ( - f"TRTLLM_RAY_BUNDLE_INDICES cannot have duplicate values, but got {bundle_indices=}." - ) - assert max(bundle_indices) < len(pg.bundle_specs), ( - f"{bundle_indices=} out of range for PG with {len(pg.bundle_specs)} bundles" - ) - return pg, bundle_indices - else: - raise ValueError("No global placement group is found.") - - def _get_default(tp_size): - head_tag = f"node:{self.master_address}" - nodes = ray.nodes() - gpus_per_node = int(nodes[0]["Resources"].get( - "GPU", 0)) # assume symmetric across nodes - - bundle_cpu = bundle_gpu = min(tp_size, gpus_per_node) - - bundles, bundle_indices = [], [] - current = 0 - for rank in range(self.world_size): - if current == 0: - bundle = {"GPU": bundle_gpu, "CPU": bundle_cpu} - if len(bundles) == 0: - bundle[ - head_tag] = 0.01 # to force placement on head node - bundles.append(bundle) - - bundle_indices.append(len(bundles) - 1) - current = (current + 1) % bundle_gpu - - strategy = "PACK" - logger.debug( - f"[Strategy={strategy}] Bundles: {bundles} for tp_size: {tp_size} and world_size: {self.world_size}" - ) - pg = placement_group(bundles, strategy=strategy) - - return pg, bundle_indices - - if self.world_size % tp_size != 0: - raise ValueError( - f"world_size {self.world_size} must be a multiple of tp_size {tp_size}" - ) - - # path 0 - if placement_config and placement_config.placement_groups is not None: - return _get_from_placement_config(placement_config) - # path 1 - if bundle_indices := os.getenv("TRTLLM_RAY_BUNDLE_INDICES", None): - return _get_from_env(bundle_indices) - # path 2 - return _get_default(tp_size) - - @property - def enable_postprocess_parallel(self) -> bool: - ret = super().enable_postprocess_parallel - assert ret == False, "Postprocess parallel is not supported in RayExecutor" - return ret diff --git a/tensorrt_llm/executor/ray_gpu_worker.py b/tensorrt_llm/executor/ray_gpu_worker.py index 7a312aa003f4..0f1b855121bf 100644 --- a/tensorrt_llm/executor/ray_gpu_worker.py +++ b/tensorrt_llm/executor/ray_gpu_worker.py @@ -1,364 +1,38 @@ -import gc -import importlib -import os -import tempfile -from functools import wraps -from pathlib import Path -from queue import Queue -from typing import Any, List, Optional, Type - -import ray -import torch - -from tensorrt_llm._ray_utils import control_action_decorator -from tensorrt_llm._torch.utils import get_device_uuid -from tensorrt_llm._torch.virtual_memory import (materialize_with_tag, - release_with_tag) - -from .. import TorchLlmArgs -from ..bindings import executor as tllm -from ..llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType -from ..llmapi.tokenizer import TokenizerBase -from ..llmapi.utils import configure_cpu_affinity -from ..sampling_params import BatchedLogitsProcessor -from .base_worker import BaseWorker -from .postproc_worker import PostprocWorkerConfig -from .request import GenerationRequest -from .result import GenerationResult -from .rpc_worker_mixin import RpcWorkerMixin +# 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. +"""Compatibility shim for ``tensorrt_llm.executor.ray_gpu_worker``. + +Will be removed once all usages are migrated to +``tensorrt_llm.executor.ray.gpu_worker``. + +DO NOT ADD ANYTHING TO THIS FILE. +""" + +import warnings + +from tensorrt_llm.executor.ray.gpu_worker import RayGPUWorker, RayWorkerWrapper # noqa: F401 + +warnings.warn( + "tensorrt_llm.executor.ray_gpu_worker has moved to " + "tensorrt_llm.executor.ray.gpu_worker and will be removed in a future " + "release.", + FutureWarning, + stacklevel=2, +) __all__ = [ "RayGPUWorker", "RayWorkerWrapper", ] - - -def resolve_obj_by_qualname(qualname: str) -> Any: - """Resolve an object by its fully qualified name.""" - module_name, obj_name = qualname.rsplit(".", 1) - module = importlib.import_module(module_name) - return getattr(module, obj_name) - - -@ray.remote -class RayWorkerWrapper: - - def __init__(self, worker_cls, worker_kwargs, world_size, rank): - self.master_address = os.environ["MASTER_ADDR"] - self.world_size = world_size - self.rank = rank - # Ray can't pickle TensorRT logger - global logger - from tensorrt_llm.logger import logger - - # Expect to see global counts w/ RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1, - # unless CUDA_VISIBLE_DEVICES is set. - logger.debug( - f"CUDA device count visible to Ray: {torch.cuda.device_count()}") - - # Physical gpu id - self.gpu = int(ray.get_gpu_ids()[0]) - self.local_gpu = self.physical_to_local_id(self.gpu) - - # Per-worker DeepGemm JIT cache to avoid rename race across co-located workers - os.environ["DG_JIT_CACHE_DIR"] = os.path.join( - tempfile.gettempdir(), f"deep_gemm_rank{rank}_gpu{self.gpu}") - - torch.cuda.set_device(self.local_gpu) - - self.worker_cls = RayWorkerWrapper._inject_worker_extension( - worker_cls, worker_kwargs.pop("ray_worker_extension_cls", None)) - self.worker_kwargs = worker_kwargs - - def _create_tcp_store(self, - port: Optional[int] = None - ) -> torch.distributed.TCPStore: - # port=0 means let the OS pick an available port (only valid for master) - # For non-master, port must be specified to connect to master's port - actual_port = port if port is not None else 0 - return torch.distributed.TCPStore(host_name=self.master_address, - port=actual_port, - world_size=self.world_size, - is_master=(self.rank == 0), - wait_for_workers=False) - - def setup_tcp_store(self): - if self.rank != 0: - raise RuntimeError("Only the master worker can setup TCP store") - self.store = self._create_tcp_store() - return self.store.port - - def setup_distributed_env_and_worker(self, port: int): - if self.rank != 0: - self.store = self._create_tcp_store(port) - - torch.distributed.init_process_group(backend="cuda:nccl,cpu:gloo", - store=self.store, - world_size=self.world_size, - rank=self.rank) - assert torch.distributed.get_world_size( - ) == self.world_size, "Process group world size must match the expected world size" - logger.info( - f"[Rank {self.rank}] Finished PG init. Global GPU ID: {self.gpu}, local GPU ID: {self.local_gpu}" - ) - - self.worker = self.worker_cls(device_id=self.local_gpu, - **self.worker_kwargs) - self._has_setup_distributed_env_and_worker = True - - @property - def has_setup_distributed_env_and_worker(self) -> bool: - return getattr(self, '_has_setup_distributed_env_and_worker', False) - - def ensure_distributed_setup(func): - - @wraps(func) - def wrapper(self, *args, **kwargs): - if not self.has_setup_distributed_env_and_worker: - raise RuntimeError( - "Have not setup distributed environment and worker yet") - return func(self, *args, **kwargs) - - return wrapper - - @ensure_distributed_setup - def submit(self, request: GenerationRequest) -> GenerationResult: - return self.worker.submit(request) - - @ensure_distributed_setup - def enqueue_request(self, - request: GenerationRequest, - result_wait_queue: Queue | None = None) -> int: - return self.worker.enqueue_request(request, result_wait_queue) - - @ensure_distributed_setup - def abort_request(self, request_id: int) -> None: - self.worker.abort_request(request_id) - - @ensure_distributed_setup - def report_device_id(self) -> str: - local_id = self.physical_to_local_id(self.gpu) - return get_device_uuid(local_id) - - @ensure_distributed_setup - def call_worker_method(self, method_name: str, *args, **kwargs): - """Generic method to call any method on the underlying worker.""" - if hasattr(self.worker, method_name): - method = getattr(self.worker, method_name) - if callable(method): - return method(*args, **kwargs) - else: - raise AttributeError( - f"'{method_name}' is not a callable method of RayGPUWorker." - ) - else: - raise AttributeError( - f"The RayGPUWorker has no method called '{method_name}'.") - - def shutdown(self): - if hasattr(self, 'worker'): - self.worker.shutdown() - - def __repr__(self) -> str: - """Customizes the actor's prefix in the Ray logs. - - This makes it easier to identify which worker is producing specific log messages. - Refer to https://github.com/NVIDIA-NeMo/RL/blob/faad02113c3c502437ccb339cb848796334aedd9/nemo_rl/models/policy/dtensor_policy_worker_v2.py#L95 - """ - if torch.distributed.is_initialized(): - return f"{self.__class__.__qualname__}[rank={torch.distributed.get_rank()}]" - else: - return f"{self.__class__.__qualname__}" - - @staticmethod - def physical_to_local_id(phys_id: int) -> int: - visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") - if not visible_devices: - return phys_id - id_mapping = list(map(int, visible_devices.split(","))) - return id_mapping.index(phys_id) - - @staticmethod - def _inject_worker_extension( - worker_class: Type[BaseWorker], - extension_cls_name: Optional[str]) -> Type[BaseWorker]: - """Inject worker extension into the worker class if specified.""" - if not extension_cls_name: - return worker_class - - try: - extension_cls = resolve_obj_by_qualname(extension_cls_name) - except (ImportError, AttributeError, ValueError) as e: - raise RuntimeError( - f"Failed to load worker extension '{extension_cls_name}'" - ) from e - - # Check for conflicts - for attr in dir(extension_cls): - if attr.startswith("__"): - continue - if hasattr(worker_class, attr): - raise ValueError( - f"Worker class {worker_class.__name__} already defines '{attr}', " - f"which conflicts with extension {extension_cls.__name__}.") - - derived_name = f"{worker_class.__name__}With{extension_cls.__name__}" - ExtendedWorker = type(derived_name, (worker_class, extension_cls), - {'__module__': worker_class.__module__}) - return ExtendedWorker - - -class RayGPUWorker(RpcWorkerMixin, BaseWorker): - - def __init__( - self, - device_id: int, - engine: Path, - executor_config: Optional[tllm.ExecutorConfig] = None, - batched_logits_processor: Optional[BatchedLogitsProcessor] = None, - postproc_worker_config: Optional[PostprocWorkerConfig] = None, - is_llm_executor: Optional[bool] = None, - hf_model_dir: Optional[Path] = None, - tokenizer: Optional[TokenizerBase] = None, - llm_args: Optional[BaseLlmArgs] = None, - rpc_addr: Optional[str] = None, - hmac_key: bytes = b"", - ) -> None: - global logger - from tensorrt_llm.logger import logger - - super().__init__( - engine=engine, - executor_config=executor_config, - batched_logits_processor=batched_logits_processor, - postproc_worker_config=postproc_worker_config, - is_llm_executor=is_llm_executor, - hf_model_dir=hf_model_dir, - tokenizer=tokenizer, - llm_args=llm_args, - ) - - self.device_id = device_id - self.global_rank = torch.distributed.get_rank() - if self.global_rank > 1: - logger.set_rank(self.global_rank) - - if rpc_addr is None: - raise RuntimeError( - "RPC mode enabled but no rpc_addr provided to RayGPUWorker") - self.init_rpc_worker(self.global_rank, rpc_addr, hmac_key) - self.start_rpc_server() - - def setup_engine(self): - if torch.distributed.is_initialized( - ) and torch.distributed.get_world_size() > 1: - torch.distributed.barrier() - super().setup_engine() - - def enqueue_request(self, - request: GenerationRequest, - result_wait_queue: Queue | None = None) -> int: - return self._enqueue_request(request, result_wait_queue) - - @control_action_decorator - def sleep(self, sleep_tags: List[str]): - assert isinstance(self.llm_args, - TorchLlmArgs), "sleep() only available for TorchLLM" - - if self.llm_args.sleep_config is None: - raise ValueError( - "Sleep feature is not enabled, please set sleep_config in the LLM arguments." - ) - try: - tags = [ExecutorMemoryType(tag) for tag in sleep_tags] - logger.info(f"Sleep: {tags}") - torch.cuda.synchronize() - release_with_tag(*tags) - torch.cuda.synchronize() - gc.collect() - torch.cuda.empty_cache() - except Exception as e: - logger.error(f"Encountered an error in sleep: {e}") - raise e - - @control_action_decorator - def wakeup(self, wakeup_tags: List[str]): - assert isinstance(self.llm_args, - TorchLlmArgs), "wakeup() only available for TorchLLM" - - if self.llm_args.sleep_config is None: - raise ValueError( - "Sleep feature is not enabled, please set sleep_config in the LLM arguments." - ) - try: - tags = [ExecutorMemoryType(tag) for tag in wakeup_tags] - logger.info(f"Wakeup: {tags}") - torch.cuda.synchronize() - materialize_with_tag(*tags) - torch.cuda.synchronize() - except Exception as e: - logger.error(f"Encountered an error in wakeup") - raise e - - def start(self): - pass - - def shutdown(self): - - if self.doing_shutdown: - return - else: - self.doing_shutdown = True - - logger.debug(f'Worker {self.rank} shutting down...') - - if hasattr(self, 'shutdown_event'): - self.shutdown_event.set() - - if hasattr(self, 'rpc_server') and self.rpc_server is not None: - logger.info(f"[Rank {self.global_rank}] Shutting down RPC server") - try: - self.rpc_server.shutdown() - except Exception as e: - # Suppress errors during RPC server shutdown - # These can occur if the server is already closed or during cleanup - logger.debug( - f"[Rank {self.global_rank}] Suppressed error during RPC server shutdown: {e}" - ) - self.rpc_server = None - - if self.engine is not None: - self.engine.shutdown() - self.engine = None - - assert self._executor_config is None, "An empty executor_config is expected in shutdown when LLM arguments are defined." - if (self.llm_args.backend == "pytorch" - and hasattr(self, "checkpoint_loader") - and self.checkpoint_loader is not None): - self.checkpoint_loader.cleanup() - self.checkpoint_loader = None - - # Check if there are any errors from the threads before shutdown. - self._handle_background_error() - - logger.debug(f"Worker {self.rank} shutdown done.") - - def _get_comm_ranks_device_id(self): - # Make sure C++ executor would use same devices/ranks as py_executor - global_rank = torch.distributed.get_rank() - world_size = torch.distributed.get_world_size() - comm_ranks = [None] * world_size - device_ids = [None] * world_size - - torch.distributed.all_gather_object(comm_ranks, global_rank) - torch.distributed.all_gather_object(device_ids, self.device_id) - - configure_cpu_affinity(self.device_id) - - return comm_ranks, device_ids - - def __enter__(self): - return self - - def __del__(self): - self.shutdown() diff --git a/tensorrt_llm/llmapi/rlhf_utils.py b/tensorrt_llm/llmapi/rlhf_utils.py index 3b39bb5039cf..57ac5b7731f8 100644 --- a/tensorrt_llm/llmapi/rlhf_utils.py +++ b/tensorrt_llm/llmapi/rlhf_utils.py @@ -8,9 +8,9 @@ import torch -from tensorrt_llm._ray_utils import control_action_decorator from tensorrt_llm._torch.modules.fused_moe.moe_load_balancer import MoeLoadBalancer from tensorrt_llm._torch.utils import get_device_uuid +from tensorrt_llm.executor.ray.utils import control_action_decorator from tensorrt_llm.llmapi import serialization from tensorrt_llm.logger import logger diff --git a/tensorrt_llm/ray_stub.py b/tensorrt_llm/ray_stub.py index 34d3b4e97cd4..cefe3c0cea4e 100644 --- a/tensorrt_llm/ray_stub.py +++ b/tensorrt_llm/ray_stub.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -12,35 +12,35 @@ # 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. -from functools import wraps as _wraps - -from tensorrt_llm._utils import mpi_disabled as _mpi_disabled - -# Don't raise error on import - only when Ray functionality is actually used -_RAY_NOT_INSTALLED_MSG = "Ray requested (TLLM_DISABLE_MPI=1), but not installed. Please install Ray." - - -def remote(*args, **kwargs): - - def decorator(func): - # Returns a function that always raises. - # Decorated class depends on ray, but ray is not installed. - @_wraps(func) - def stub_checker(*_, **__): - raise RuntimeError( - f'Ray not installed, so the remote function / actor "{func.__name__}" is not available.' - ) - - return stub_checker - - if len(args) == 1 and len(kwargs) == 0 and callable(args[0]): - return decorator(args[0]) - - return decorator - - -def __getattr__(name): - msg = f'Ray not installed, so "ray.{name}" is unavailable.' - if _mpi_disabled(): - msg = _RAY_NOT_INSTALLED_MSG - raise RuntimeError(msg) +"""Compatibility shim for ``tensorrt_llm.ray_stub``. + +Will be removed once all usages are migrated to +``tensorrt_llm.executor.ray.stub``. + +DO NOT ADD ANYTHING TO THIS FILE. +""" + +import warnings + +# Bound by attribute rather than with ``from ... import remote``: the target +# module answers every unknown name from a module-level ``__getattr__`` that +# raises ``RuntimeError``, and ``from X import Y`` first probes +# ``hasattr(X, "__path__")``, which only swallows ``AttributeError`` -- so that +# spelling raises while merely importing this file. ``__getattr__`` is +# forwarded as well: raising for every other name is what the target module is +# for, and re-exporting only ``remote`` would answer ``AttributeError`` here. +from tensorrt_llm.executor.ray import stub as _stub + +remote = _stub.remote +__getattr__ = _stub.__getattr__ + +warnings.warn( + "tensorrt_llm.ray_stub has moved to tensorrt_llm.executor.ray.stub " + "and will be removed in a future release.", + FutureWarning, + stacklevel=2, +) + +__all__ = [ + "remote", +] diff --git a/tests/integration/defs/examples/test_ray.py b/tests/integration/defs/examples/test_ray.py index 44743f030d48..b56f3291db49 100644 --- a/tests/integration/defs/examples/test_ray.py +++ b/tests/integration/defs/examples/test_ray.py @@ -4,7 +4,7 @@ try: import ray except ImportError: - import tensorrt_llm.ray_stub as ray + import tensorrt_llm.executor.ray.stub as ray import pytest from defs.common import venv_check_call, wait_for_server diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 599698d4ffd7..f7031bd39666 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -35,6 +35,7 @@ l0_cpu: - unittest/executor/test_event_loop_error_broadcast.py - unittest/executor/test_stats_serializer.py - unittest/executor/test_spec_dec_perf_metrics.py + - unittest/executor/test_ray_stub.py - unittest/inputs - unittest/llmapi/apps/test_chat_utils.py - unittest/llmapi/apps/test_harmony_channel_validation.py diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 88130605c153..8c17a48dde98 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -255,6 +255,7 @@ l0_h100: - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logprobs[False-TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logprobs[True-TinyLlama-1.1B-Chat-v1.0] - unittest/_torch/executor/test_overlap_scheduler.py + - unittest/executor/test_shim_ray.py - unittest/_torch/ray_orchestrator/single_gpu/test_llm_sleep.py - unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py -m "part0" - unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py -m "part1" diff --git a/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_inflight_weight_update.py b/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_inflight_weight_update.py index 6ece0fdbb1a9..0577588aaeaa 100644 --- a/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_inflight_weight_update.py +++ b/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_inflight_weight_update.py @@ -45,7 +45,7 @@ from utils.util import skip_pre_hopper from tensorrt_llm import AsyncLLM -from tensorrt_llm._ray_utils import control_action_decorator +from tensorrt_llm.executor.ray.utils import control_action_decorator from tensorrt_llm.llmapi import KvCacheConfig, SamplingParams from tensorrt_llm.llmapi.rlhf_utils import WorkerExtension diff --git a/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py b/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py index 4805c850e7e6..5a454e4e0b40 100644 --- a/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py +++ b/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py @@ -34,12 +34,12 @@ from utils.util import skip_pre_blackwell, skip_pre_hopper from tensorrt_llm import LLM -from tensorrt_llm._ray_utils import control_action_decorator from tensorrt_llm._torch.auto_deploy.custom_ops.quantization.torch_quant import ( _dequantize_nvfp4, _quantize_nvfp4, ) from tensorrt_llm._torch.utils import get_device_uuid +from tensorrt_llm.executor.ray.utils import control_action_decorator from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, MoeConfig, SamplingParams from tensorrt_llm.llmapi.rlhf_utils import WorkerExtension diff --git a/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py b/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py index b43d8f42986d..e22fcde614b4 100644 --- a/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py +++ b/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_ops.py @@ -8,7 +8,7 @@ try: import ray except ModuleNotFoundError: - from tensorrt_llm import ray_stub as ray + from tensorrt_llm.executor.ray import stub as ray from tensorrt_llm._torch.distributed.communicator import TorchDist from tensorrt_llm.functional import AllReduceFusionOp, AllReduceStrategy @@ -25,7 +25,7 @@ def __init__(self, rank, world_size): assert len(ray.get_gpu_ids()) == 1 self.gpu = int(ray.get_gpu_ids()[0]) - from tensorrt_llm.executor.ray_gpu_worker import RayWorkerWrapper + from tensorrt_llm.executor.ray.gpu_worker import RayWorkerWrapper local_gpu = RayWorkerWrapper.physical_to_local_id(self.gpu) torch.cuda.set_device(local_gpu) @@ -273,7 +273,7 @@ def __init__(self, rank, world_size, tp_size, cp_size): assert len(ray.get_gpu_ids()) == 1 self.gpu = int(ray.get_gpu_ids()[0]) - from tensorrt_llm.executor.ray_gpu_worker import RayWorkerWrapper + from tensorrt_llm.executor.ray.gpu_worker import RayWorkerWrapper local_gpu = RayWorkerWrapper.physical_to_local_id(self.gpu) torch.cuda.set_device(local_gpu) diff --git a/tests/unittest/conftest.py b/tests/unittest/conftest.py index 4d3d24efb824..02b99ebbc848 100644 --- a/tests/unittest/conftest.py +++ b/tests/unittest/conftest.py @@ -25,7 +25,7 @@ try: import ray except ModuleNotFoundError: - from tensorrt_llm import ray_stub as ray + from tensorrt_llm.executor.ray import stub as ray import _pytest.outcomes import pytest diff --git a/tests/unittest/executor/test_ray_stub.py b/tests/unittest/executor/test_ray_stub.py new file mode 100644 index 000000000000..0d428e84d017 --- /dev/null +++ b/tests/unittest/executor/test_ray_stub.py @@ -0,0 +1,90 @@ +# 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 stand-in for ``ray`` used when Ray is not installed. + +It must be inert at import and decoration time and fail only when Ray +functionality is actually used, so that a default install can still import +``tensorrt_llm``. +""" + +import importlib.util + +import pytest + +from tensorrt_llm.executor.ray import stub + +# This file's home is the CPU-Generic stage, because that is where Ray is absent +# and the stub is the thing actually exercised (l0_cpu.yml). Those stages run +# `pytest -m cpu_only` (jenkins/L0_Test.groovy:1476), and their conftest ignores +# any test file whose text lacks the literal string "pytest.mark.cpu_only" +# (tests/unittest/conftest.py:239). Without this marker all six tests are +# deselected, pytest exits 5 (no tests collected), and the test_unittests_v2 +# wrapper reports that as a failure rather than as an empty run. +pytestmark = pytest.mark.cpu_only + +_RAY_INSTALLED = importlib.util.find_spec("ray") is not None + + +def test_import_is_inert() -> None: + """Importing the stub must not raise; only *using* Ray may.""" + assert stub.remote is not None + + +def test_bare_decorator_defers_the_failure() -> None: + """``@ray.remote`` must decorate cleanly and fail only when called.""" + + @stub.remote + def train_step(x: int) -> int: + return x + + assert train_step.__name__ == "train_step" + + with pytest.raises(RuntimeError, match="train_step"): + train_step(1) + + +def test_called_decorator_defers_the_failure() -> None: + """``@ray.remote(...)`` -- the parameterised form -- behaves the same.""" + + @stub.remote(num_gpus=1) + class Worker: + def run(self) -> None: + return None + + with pytest.raises(RuntimeError, match="Worker"): + Worker() + + +def test_unknown_attribute_raises_and_names_itself() -> None: + """Any other ``ray.`` must raise, and say which name was wanted.""" + with pytest.raises(RuntimeError, match=r"ray\.init"): + getattr(stub, "init") + + +def test_unknown_attribute_says_to_install_ray_when_ray_was_requested( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """With ``TLLM_DISABLE_MPI=1`` the user asked for Ray, so say so.""" + monkeypatch.setenv("TLLM_DISABLE_MPI", "1") + with pytest.raises(RuntimeError, match="Please install Ray"): + getattr(stub, "init") + + +@pytest.mark.skipif(_RAY_INSTALLED, reason="Ray is installed, so the fallback is not taken here") +def test_distributed_layer_falls_back_to_the_stub() -> None: + """Without Ray, the distributed layer must import and resolve to the stub.""" + from tensorrt_llm._torch.distributed import communicator + + assert communicator.ray.__name__ == "tensorrt_llm.executor.ray.stub" diff --git a/tests/unittest/executor/test_shim_ray.py b/tests/unittest/executor/test_shim_ray.py new file mode 100644 index 000000000000..d23c6d2288ed --- /dev/null +++ b/tests/unittest/executor/test_shim_ray.py @@ -0,0 +1,109 @@ +# 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 compatibility modules kept at the pre-move Ray import paths. + +Each one must forward every name in its own ``__all__`` to the module it +replaces, as the same object. The names are read from the module rather than +restated here, so changing the export set changes what is verified. +""" + +import importlib +import importlib.util +import sys +import warnings +from types import ModuleType + +import pytest + +# Two of the modules below import Ray unconditionally, so the whole file needs +# it -- the test list schedules this where Ray is installed, and this keeps the +# file honest anywhere else. +pytestmark = pytest.mark.skipif( + importlib.util.find_spec("ray") is None, reason="the modules under test import Ray" +) + +# The forwarding module -> the module it forwards to. +FORWARDS: dict[str, str] = { + "tensorrt_llm._ray_utils": "tensorrt_llm.executor.ray.utils", + "tensorrt_llm.ray_stub": "tensorrt_llm.executor.ray.stub", + "tensorrt_llm.executor.ray_executor": "tensorrt_llm.executor.ray.executor", + "tensorrt_llm.executor.ray_gpu_worker": "tensorrt_llm.executor.ray.gpu_worker", +} + + +def _import_pair(forwarding_path: str) -> tuple[ModuleType, ModuleType, str]: + """Import a forwarding module and the module it forwards to. + + Returns: + The forwarding module, the module it forwards to, and the latter's + dotted name. + """ + target_path = FORWARDS[forwarding_path] + return ( + importlib.import_module(forwarding_path), + importlib.import_module(target_path), + target_path, + ) + + +@pytest.mark.parametrize("forwarding_path", sorted(FORWARDS)) +def test_published_names_are_the_same_objects(forwarding_path: str) -> None: + """Identity, not equality. + + A re-implementation would compare equal and still break ``isinstance`` and + unpickling for callers that kept the old import path. + """ + forwarding, target, target_path = _import_pair(forwarding_path) + + published = getattr(forwarding, "__all__", None) + assert published, f"{forwarding_path} publishes nothing, so it forwards nothing" + + for name in published: + assert hasattr(target, name), f"{target_path} has no {name!r} to forward to" + assert getattr(forwarding, name) is getattr(target, name), name + + +@pytest.mark.parametrize("forwarding_path", sorted(FORWARDS)) +def test_objects_report_the_module_that_defines_them(forwarding_path: str) -> None: + """Forwarding must not rewrite ``__module__``. + + Objects have to keep pointing at where they are defined, or pickles written + now would record a path that is going away. + """ + forwarding, target, target_path = _import_pair(forwarding_path) + + for name in forwarding.__all__: + module = getattr(getattr(target, name), "__module__", None) + if module is None: # not every object carries one + continue + assert module == target_path, f"{name}.__module__ is {module!r}" + + +@pytest.mark.parametrize("forwarding_path", sorted(FORWARDS)) +def test_import_warns_once_and_names_the_new_path(forwarding_path: str) -> None: + """Importing the old path has to say so, and say where to go instead.""" + target_path = FORWARDS[forwarding_path] + + # A module body runs once per interpreter, so drop it before re-importing. + sys.modules.pop(forwarding_path, None) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + importlib.import_module(forwarding_path) + + # Deliberately not pinned to a warning category: what matters is that one + # warning is raised and that it points at the replacement. + about_this_module = [w for w in caught if forwarding_path in str(w.message)] + assert len(about_this_module) == 1, [str(w.message) for w in caught] + assert target_path in str(about_this_module[0].message)