update placement api

hchings · hchings · commit 6e9b96704599 · 2025-11-25T05:23:42.000-08:00
diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py
@@ -37,7 +37,6 @@
 from .request import GenerationRequest, LoRARequest, PromptAdapterRequest
 from .result import GenerationResult, IterationResult
 from .utils import IntraProcessQueue, ProcessPoolExecutorSession, RequestError
-from ray.util.placement_group import PlacementGroup
 
 if TYPE_CHECKING:
     from .proxy import GenerationExecutorProxy
@@ -368,8 +367,6 @@ def _create_ray_executor(
         postproc_worker_config: PostprocWorkerConfig,
         is_llm_executor: bool,
         tp_size: int,
-        placement_share: float = 1.0,
-        placement_where: list[tuple[PlacementGroup, list[int]]] = None,
     ):
         logger.warning(f"Orchestrator is creating Ray executor")
         from .ray_executor import RayExecutor
@@ -378,9 +375,7 @@ def _create_ray_executor(
                            model_world_size=model_world_size,
                            postproc_worker_config=postproc_worker_config,
                            is_llm_executor=is_llm_executor,
-                           tp_size=tp_size,
-                           placement_share=placement_share,
-                           placement_where=placement_where)
+                           tp_size=tp_size)
 
     @staticmethod
     def _create_rpc_executor(
@@ -444,8 +439,6 @@ def create(
         hf_model_dir: Optional[Path] = None,
         tokenizer: Optional[TokenizerBase] = None,
         llm_args: Optional[BaseLlmArgs] = None,
-        placement_share: float = 1.0,
-        placement_where: list[tuple[PlacementGroup, list[int]]] = None,
         **args,
     ) -> Union["GenerationExecutorProxy", "GenerationExecutorWorker"]:
         if world_size == 0:
@@ -485,9 +478,7 @@ def create(
                 model_world_size,
                 postproc_worker_config,
                 is_llm_executor=is_llm_executor,
-                tp_size=args.get("tp_size", 1),
-                placement_share=placement_share,
-                placement_where=placement_where)
+                tp_size=args.get("tp_size", 1))
         elif orchestrator_type is not None and orchestrator_type != "rpc":
             raise ValueError(
                 f"Unsupported orchestrator_type: {orchestrator_type}")
diff --git a/tensorrt_llm/executor/ray_executor.py b/tensorrt_llm/executor/ray_executor.py
@@ -8,8 +8,7 @@
     e.msg = """Cannot import Ray. Please install 'ray' package to use ray orchestrator"""
     raise
 
-from ray.util.placement_group import (PlacementGroup,
-                                      PlacementGroupSchedulingStrategy,
+from ray.util.placement_group import (PlacementGroupSchedulingStrategy,
                                       get_current_placement_group,
                                       placement_group)
 
@@ -38,18 +37,13 @@ def __init__(self,
                  model_world_size: int,
                  postproc_worker_config: PostprocWorkerConfig,
                  is_llm_executor: bool,
-                 tp_size=1,
-                 placement_share: float = 1.0,
-                 placement_where: list[tuple[PlacementGroup, list[int]]] = None):
+                 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.placement_share = placement_share
-        self.placement_where = placement_where
-
         self.has_start_local_cluser = False
         runtime_env = {
             "env_vars": {
@@ -125,9 +119,13 @@ def __init__(self,
             raise e
 
     def create_workers(self, worker_cls, worker_kwargs):
+        llm_args = worker_kwargs.get("llm_args")
+
         # 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"))
+        num_gpus = (llm_args.per_worker_gpu_share if llm_args
+                    and llm_args.per_worker_gpu_share is not None else float(
+                        os.getenv("TRTLLM_RAY_PER_WORKER_GPUS", "1.0")))
         logger.debug(f"{num_gpus=} for each worker.")
 
         runtime_env = ray.runtime_env.RuntimeEnv()
@@ -138,21 +136,26 @@ def create_workers(self, worker_cls, worker_kwargs):
             "MASTER_PORT": str(self.master_port)
         })
 
-        rank = 0
-        self.world_size = sum(len(bundle_indices) for _, bundle_indices in self.placement_where)
+        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 pg, bundle_indices in self.placement_where:
-            for bundle_index in bundle_indices:
-                self.workers.append(
-                    RayWorkerWrapper.options(
-                        num_gpus=self.placement_share,
-                        runtime_env=runtime_env,  # per-actor env
-                        scheduling_strategy=PlacementGroupSchedulingStrategy(
-                            placement_group=pg,
-                            placement_group_bundle_index=bundle_index,
-                        )).remote(worker_cls, worker_kwargs, self.world_size, rank)
-                )
-                rank += 1
+        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)
@@ -336,15 +339,48 @@ def shutdown(self):
     def _get_worker_ready_futures(self):
         return [worker.__ray_ready__.remote() for worker in self.workers]
 
-    def _get_placement_group(self,
-                             tp_size: int) -> Tuple[PlacementGroup, List[int]]:
+    def _get_placement_group(
+            self,
+            tp_size: int,
+            worker_kwargs: Dict = None) -> Tuple[Any, List[int]]:
         """
         Either use the existing placement group from driver script (e.g., in the case of RL FW integration),
         or create a default 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 span nodes.
          - rank 0 must be put on the driver node
+
+        Returns:
+            Tuple of (placement_group(s), bundle_indices)
+            - placement_group(s) can be a single PlacementGroup or a List[PlacementGroup]
+            - bundle_indices is always a List[int]
         """
+        llm_args = worker_kwargs.get("llm_args") if worker_kwargs else None
+
+        if llm_args and hasattr(
+                llm_args,
+                'placement_groups') and llm_args.placement_groups is not None:
+            total_workers = sum(
+                len(indices) for indices in llm_args.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(llm_args.placement_groups,
+                                   llm_args.placement_bundle_indices):
+                for idx in indices:
+                    flat_pgs.append(pg)
+                    flat_indices.append(idx)
+
+            return flat_pgs, flat_indices
+
         bundle_indices = os.getenv("TRTLLM_RAY_BUNDLE_INDICES", None)
 
         if bundle_indices:
diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py
@@ -126,17 +126,12 @@ def __init__(self,
                  dtype: str = "auto",
                  revision: Optional[str] = None,
                  tokenizer_revision: Optional[str] = None,
-                 placement_share: float = 1.0,
-                 placement_where: list[tuple[PlacementGroup, list[int]]] = None,
                  **kwargs: Any) -> None:
 
         self._executor_cls = kwargs.pop("executor_cls", GenerationExecutor)
         self._orchestrator_type = kwargs.get("orchestrator_type", None)
         self._llm_id = None
 
-        self.placement_share = placement_share
-        self.placement_where = placement_where
-
         log_level = logger.level
         logger.set_level("info")  # force display the backend
 
@@ -814,14 +809,12 @@ def __init__(self,
                  dtype: str = "auto",
                  revision: Optional[str] = None,
                  tokenizer_revision: Optional[str] = None,
-                 placement_share: float = 1.0,
-                 placement_where: list[tuple[PlacementGroup, list[int]]] = None,
                  **kwargs: Any) -> None:
         # TODO: deprecate backend in LLM kwargs
 
         super().__init__(model, tokenizer, tokenizer_mode, skip_tokenizer_init,
                          trust_remote_code, tensor_parallel_size, dtype,
-                         revision, tokenizer_revision, placement_share, placement_where, **kwargs)
+                         revision, tokenizer_revision, **kwargs)
 
     @property
     def workspace(self) -> Path:
@@ -979,9 +972,7 @@ def _build_model(self):
                 num_postprocess_workers=self.args.num_postprocess_workers,
                 postprocess_tokenizer_dir=self.args.postprocess_tokenizer_dir,
             ),
-            is_llm_executor=True,
-            placement_share=self.placement_share,
-            placement_where=self.placement_where)
+            is_llm_executor=True)
 
 
 @append_docstring(TORCH_LLM_DOCSTRING)
@@ -1002,8 +993,6 @@ def __init__(self,
                  dtype: str = "auto",
                  revision: Optional[str] = None,
                  tokenizer_revision: Optional[str] = None,
-                 placement_share: float = 1.0,
-                 placement_where: list[tuple[PlacementGroup, list[int]]] = None,
                  **kwargs: Any) -> None:
 
         # TODO: deprecate backend in LLM kwargs
@@ -1022,8 +1011,6 @@ def __init__(self,
                          revision,
                          tokenizer_revision,
                          backend=backend,
-                         placement_share=placement_share,
-                         placement_where=placement_where,
                          **kwargs)
 
     @set_api_status("prototype")
@@ -1091,9 +1078,7 @@ def _build_model(self):
             is_llm_executor=True,
             hf_model_dir=self._hf_model_dir,
             tokenizer=self.tokenizer,
-            llm_args=self.args,
-            placement_share=self.placement_share,
-            placement_where=self.placement_where)
+            llm_args=self.args)
 
     def _validate_args_for_torch_backend(self, kwargs: dict) -> None:
         """Validate that users don't pass TrtLlmArgs-specific arguments when using PyTorch backend.
@@ -1129,12 +1114,10 @@ def __init__(self,
                  dtype: str = "auto",
                  revision: Optional[str] = None,
                  tokenizer_revision: Optional[str] = None,
-                 placement_share: float = 1.0,
-                 placement_where: list[tuple[PlacementGroup, list[int]]] = None,
                  **kwargs: Any) -> None:
         super().__init__(model, tokenizer, tokenizer_mode, skip_tokenizer_init,
                          trust_remote_code, tensor_parallel_size, dtype,
-                         revision, tokenizer_revision, placement_share, placement_where, **kwargs)
+                         revision, tokenizer_revision, **kwargs)
 
 
 # sphinx will ignore the LLM's docstring if it is not explicitly set
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
@@ -8,8 +8,9 @@
 from dataclasses import dataclass
 from enum import Enum, EnumMeta
 from pathlib import Path
-from typing import (Any, ClassVar, Dict, List, Literal, Optional, Set, Tuple,
-                    Type, TypeAlias, TypeVar, Union, get_args, get_origin)
+from typing import (TYPE_CHECKING, Any, ClassVar, Dict, List, Literal, Optional,
+                    Set, Tuple, Type, TypeAlias, TypeVar, Union, get_args,
+                    get_origin)
 
 import torch
 import yaml
@@ -19,6 +20,11 @@
 from strenum import StrEnum
 from transformers import PreTrainedTokenizerBase
 
+try:
+    from ray.util.placement_group import PlacementGroup
+except ImportError:
+    PlacementGroup = None
+
 from tensorrt_llm.lora_helper import (LoraConfig,
                                       get_default_trtllm_modules_to_hf_modules)
 
@@ -2695,6 +2701,26 @@ class TorchLlmArgs(BaseLlmArgs):
         "Allows users to extend the functions of the RayGPUWorker class.",
         status="prototype")
 
+    # Ray placement group config. Namings TBD.
+    placement_groups: Optional[List[Any]] = Field(
+        default=None,
+        description="List of Ray placement groups, one per node. "
+        "Each element must be a ray.util.placement_group.PlacementGroup instance.",
+        exclude_from_json=True,
+        status="prototype")
+
+    placement_bundle_indices: Optional[List[List[int]]] = Field(
+        default=None,
+        description="List of bundle indices for each placement group. "
+        "Outer list corresponds to placement_groups, inner list contains bundle indices for that group. ",
+        status="prototype")
+
+    per_worker_gpu_share: Optional[float] = Field(
+        default=None,
+        description="GPU fraction per worker for colocation scenarios. "
+        "Example: 0.1 means 10 actors can share one GPU. Defaults to 1.0 (one actor per GPU).",
+        status="prototype")
+
     enable_sleep: bool = Field(
         default=False,
         description=
@@ -3000,6 +3026,44 @@ def validate_ray_worker_extension_cls(self) -> 'TorchLlmArgs':
             )
         return self
 
+    @model_validator(mode='after')
+    def validate_ray_placement_config(self) -> 'TorchLlmArgs':
+        has_pgs = self.placement_groups is not None
+        has_indices = self.placement_bundle_indices is not None
+
+        if (has_pgs or has_indices) and self.orchestrator_type != "ray":
+            raise ValueError(
+                "placement_groups is only supported with orchestrator_type='ray'"
+            )
+
+        if has_pgs != has_indices:
+            raise ValueError(
+                "placement_groups and placement_bundle_indices must be provided together"
+            )
+
+        if has_pgs:
+            if len(self.placement_groups) != len(self.placement_bundle_indices):
+                raise ValueError(
+                    f"placement_groups length ({len(self.placement_groups)}) must equal "
+                    f"placement_bundle_indices length ({len(self.placement_bundle_indices)})"
+                )
+
+        if self.per_worker_gpu_share is not None:
+            if not (0 < self.per_worker_gpu_share <= 1.0):
+                raise ValueError(
+                    f"per_worker_gpu_share must be between 0 and 1.0, "
+                    f"got {self.per_worker_gpu_share}")
+
+        if has_pgs:
+            if PlacementGroup is not None:
+                for i, pg in enumerate(self.placement_groups):
+                    if not isinstance(pg, PlacementGroup):
+                        raise TypeError(
+                            f"placement_groups[{i}] must be a Ray PlacementGroup, "
+                            f"got {type(pg).__name__}")
+
+        return self
+
     def get_executor_config(
         self,
         _hf_model_dir: Optional[Path] = None,
diff --git a/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_executor.py b/tests/unittest/_torch/ray_orchestrator/multi_gpu/test_executor.py
