[TRTLLM-14833][chore] Introduce executor/ray/ for Ray executor integration

.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 |

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)

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)
 

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

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",

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

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",

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",
+]

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):

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,

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.
+"""