From 1d7f1330db36f1a2df4c8736c91f59177eaefc63 Mon Sep 17 00:00:00 2001 From: Lori Ren Date: Tue, 11 Aug 2026 08:52:52 +0000 Subject: [PATCH 1/2] [TRTLLM-14833][chore] Introduce executor/ray/ for Ray executor integration (#17487) Signed-off-by: Lori Ren --- .pre-commit-config.yaml | 16 ++++---- docs/source/features/ray-orchestrator.md | 2 +- examples/ray_orchestrator/README.md | 2 +- legacy-files.txt | 8 ++-- pyproject.toml | 8 ++-- ruff-legacy-baseline.json | 2 +- ruff-legacy.toml | 8 ++-- .../_torch/distributed/communicator.py | 2 +- tensorrt_llm/executor/executor.py | 2 +- tensorrt_llm/executor/ray/__init__.py | 22 ++++++++++ .../{ray_executor.py => ray/executor.py} | 32 ++++++++++----- .../{ray_gpu_worker.py => ray/gpu_worker.py} | 40 +++++++++++++------ .../{ray_stub.py => executor/ray/stub.py} | 2 +- .../{_ray_utils.py => executor/ray/utils.py} | 4 +- tensorrt_llm/llmapi/rlhf_utils.py | 2 +- tests/integration/defs/examples/test_ray.py | 2 +- .../multi_gpu/test_inflight_weight_update.py | 2 +- .../test_llm_update_weights_multi_gpu.py | 2 +- .../ray_orchestrator/multi_gpu/test_ops.py | 6 +-- tests/unittest/conftest.py | 2 +- 20 files changed, 108 insertions(+), 58 deletions(-) create mode 100644 tensorrt_llm/executor/ray/__init__.py rename tensorrt_llm/executor/{ray_executor.py => ray/executor.py} (95%) rename tensorrt_llm/executor/{ray_gpu_worker.py => ray/gpu_worker.py} (91%) rename tensorrt_llm/{ray_stub.py => executor/ray/stub.py} (96%) rename tensorrt_llm/{_ray_utils.py => executor/ray/utils.py} (93%) 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/_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 similarity index 95% rename from tensorrt_llm/executor/ray_executor.py rename to tensorrt_llm/executor/ray/executor.py index accd9efc1632..6ff1abf666d8 100644 --- a/tensorrt_llm/executor/ray_executor.py +++ b/tensorrt_llm/executor/ray/executor.py @@ -1,3 +1,17 @@ +# 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 @@ -19,18 +33,18 @@ # 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.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 .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 +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", diff --git a/tensorrt_llm/executor/ray_gpu_worker.py b/tensorrt_llm/executor/ray/gpu_worker.py similarity index 91% rename from tensorrt_llm/executor/ray_gpu_worker.py rename to tensorrt_llm/executor/ray/gpu_worker.py index 7a312aa003f4..cd5165fbba0a 100644 --- a/tensorrt_llm/executor/ray_gpu_worker.py +++ b/tensorrt_llm/executor/ray/gpu_worker.py @@ -1,3 +1,17 @@ +# 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 @@ -10,22 +24,22 @@ 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 +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", diff --git a/tensorrt_llm/ray_stub.py b/tensorrt_llm/executor/ray/stub.py similarity index 96% rename from tensorrt_llm/ray_stub.py rename to tensorrt_llm/executor/ray/stub.py index 34d3b4e97cd4..a18136a6ceb0 100644 --- a/tensorrt_llm/ray_stub.py +++ b/tensorrt_llm/executor/ray/stub.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); diff --git a/tensorrt_llm/_ray_utils.py b/tensorrt_llm/executor/ray/utils.py similarity index 93% rename from tensorrt_llm/_ray_utils.py rename to tensorrt_llm/executor/ray/utils.py index 489dad9cc8f2..8df771ece793 100644 --- a/tensorrt_llm/_ray_utils.py +++ b/tensorrt_llm/executor/ray/utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); @@ -19,7 +19,7 @@ try: import ray except ImportError: - import tensorrt_llm.ray_stub as ray + import tensorrt_llm.executor.ray.stub as ray @contextmanager 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/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/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 From 4dc7f4465c9459b637ad4ba43726d5dc221cc09f Mon Sep 17 00:00:00 2001 From: Lori Ren Date: Thu, 13 Aug 2026 07:27:31 +0000 Subject: [PATCH 2/2] [TRTLLM-14833][chore] Add compatibility shims for the relocated Ray modules (#17487) Signed-off-by: Lori Ren --- tensorrt_llm/_ray_utils.py | 40 +++++++ tensorrt_llm/executor/ray_executor.py | 37 ++++++ tensorrt_llm/executor/ray_gpu_worker.py | 38 ++++++ tensorrt_llm/ray_stub.py | 46 ++++++++ .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../test_lists/test-db/l0_h100.yml | 1 + tests/unittest/executor/test_ray_stub.py | 90 +++++++++++++++ tests/unittest/executor/test_shim_ray.py | 109 ++++++++++++++++++ 8 files changed, 362 insertions(+) create mode 100644 tensorrt_llm/_ray_utils.py create mode 100644 tensorrt_llm/executor/ray_executor.py create mode 100644 tensorrt_llm/executor/ray_gpu_worker.py create mode 100644 tensorrt_llm/ray_stub.py create mode 100644 tests/unittest/executor/test_ray_stub.py create mode 100644 tests/unittest/executor/test_shim_ray.py diff --git a/tensorrt_llm/_ray_utils.py b/tensorrt_llm/_ray_utils.py new file mode 100644 index 000000000000..433489bfdcf7 --- /dev/null +++ b/tensorrt_llm/_ray_utils.py @@ -0,0 +1,40 @@ +# 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._ray_utils``. + +Will be removed once all usages are migrated to +``tensorrt_llm.executor.ray.utils``. + +DO NOT ADD ANYTHING TO THIS FILE. +""" + +import warnings + +from tensorrt_llm.executor.ray.utils import ( # noqa: F401 + control_action_decorator, + unwrap_ray_errors, +) + +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, +) + +__all__ = [ + "control_action_decorator", + "unwrap_ray_errors", +] diff --git a/tensorrt_llm/executor/ray_executor.py b/tensorrt_llm/executor/ray_executor.py new file mode 100644 index 000000000000..39fe9ae19e9b --- /dev/null +++ b/tensorrt_llm/executor/ray_executor.py @@ -0,0 +1,37 @@ +# 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", +] diff --git a/tensorrt_llm/executor/ray_gpu_worker.py b/tensorrt_llm/executor/ray_gpu_worker.py new file mode 100644 index 000000000000..0f1b855121bf --- /dev/null +++ b/tensorrt_llm/executor/ray_gpu_worker.py @@ -0,0 +1,38 @@ +# 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", +] diff --git a/tensorrt_llm/ray_stub.py b/tensorrt_llm/ray_stub.py new file mode 100644 index 000000000000..cefe3c0cea4e --- /dev/null +++ b/tensorrt_llm/ray_stub.py @@ -0,0 +1,46 @@ +# 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.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/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/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)