[Core] Fix incorrect gpu ids if placement group bundle index specified

python/ray/_private/worker.py

@@ -991,6 +991,13 @@ def get_gpu_ids() -> Union[List[int], List[str]]:
     """
     worker = global_worker
     worker.check_connected()
+    # respect placement_group_bundle_index if specified.
+    gpu_ids = worker.get_accelerator_ids_for_accelerator_resource(
+        ray_constants.GPU, f"^{ray_constants.GPU}_group_[0-9]+_[0-9A-Za-z]+$"
+    )
+    if gpu_ids:
+        return gpu_ids
+
     return worker.get_accelerator_ids_for_accelerator_resource(
         ray_constants.GPU, f"^{ray_constants.GPU}_group_[0-9A-Za-z]+$"
     )
@@ -2640,9 +2647,9 @@ def get(
                 port=None,
                 patch_stdstreams=False,
                 quiet=None,
-                breakpoint_uuid=debugger_breakpoint.decode()
-                if debugger_breakpoint
-                else None,
+                breakpoint_uuid=(
+                    debugger_breakpoint.decode() if debugger_breakpoint else None
+                ),
                 debugger_external=worker.ray_debugger_external,
             )
             rdb.set_trace(frame=frame)

python/ray/runtime_context.py

@@ -420,10 +420,16 @@ def get_accelerator_ids(self) -> Dict[str, List[str]]:
         for (
             accelerator_resource_name
         ) in ray._private.accelerators.get_all_accelerator_resource_names():
+            # respect placement_group_bundle_index if specified.
             accelerator_ids = worker.get_accelerator_ids_for_accelerator_resource(
                 accelerator_resource_name,
-                f"^{accelerator_resource_name}_group_[0-9A-Za-z]+$",
+                f"^{accelerator_resource_name}_group_[0-9]+_[0-9A-Za-z]+$",
             )
+            if not accelerator_ids:
+                accelerator_ids = worker.get_accelerator_ids_for_accelerator_resource(
+                    accelerator_resource_name,
+                    f"^{accelerator_resource_name}_group_[0-9A-Za-z]+$",
+                )
             ids_dict[accelerator_resource_name] = [str(id) for id in accelerator_ids]
         return ids_dict
 

python/ray/tests/test_actor_resources.py

@@ -7,6 +7,7 @@
 
 import ray
 import ray.cluster_utils
+from ray.util.placement_group import placement_group, PlacementGroupSchedulingStrategy
 
 
 def test_actor_deletion_with_gpus(shutdown_only):
@@ -678,6 +679,49 @@ def get_cuda_visible_devices(self):
     assert ray.get(actor.get_cuda_visible_devices.remote()) == "0,1"
 
 
+def test_actor_cuda_visible_devices_placement_group_bundle_index(shutdown_only):
+    ray.init(num_cpus=4, num_gpus=2)
+
+    @ray.remote
+    class Actor:
+        def __init__(self) -> None:
+            self.gpu_ids = ray.get_gpu_ids()
+
+        def get_cuda_visible_devices(self):
+            gpu_ids = os.environ["CUDA_VISIBLE_DEVICES"]
+            assert int(gpu_ids) == self.gpu_ids[0]
+            return gpu_ids
+
+    bundles = [{"GPU": 1, "CPU": 2}] * 2
+    pg = placement_group(bundles)
+    pg.ready()
+
+    m1 = [
+        Actor.options(
+            num_gpus=0.1,
+            scheduling_strategy=PlacementGroupSchedulingStrategy(
+                placement_group=pg, placement_group_bundle_index=i
+            ),
+        ).remote()
+        for i in range(2)
+    ]
+
+    m2 = [
+        Actor.options(
+            num_gpus=0.1,
+            scheduling_strategy=PlacementGroupSchedulingStrategy(
+                placement_group=pg, placement_group_bundle_index=i
+            ),
+        ).remote()
+        for i in range(2)
+    ]
+
+    assert ray.get(m1[0].get_cuda_visible_devices.remote()) == "0"
+    assert ray.get(m1[1].get_cuda_visible_devices.remote()) == "1"
+    assert ray.get(m2[0].get_cuda_visible_devices.remote()) == "0"
+    assert ray.get(m2[1].get_cuda_visible_devices.remote()) == "1"
+
+
 if __name__ == "__main__":
     import pytest