Method _get_placement_group

tensorrt_llm/executor/ray_executor.py:348–446

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

(
            self,
            tp_size: int,
            worker_kwargs: Dict = None)

        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]]:
        """
        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

        placement_config = getattr(llm_args, 'ray_placement_config',
                                   None) if llm_args else None
        if placement_config and placement_config.placement_groups is not None:
            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

        bundle_indices = os.getenv("TRTLLM_RAY_BUNDLE_INDICES", None)

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