[Placement Groups] Support CUDA_VISIBLE_DEVICES (#10053)

python/ray/tests/test_placement_group.py

@@ -1,9 +1,11 @@
 import pytest
+import os
+import sys
+
 try:
     import pytest_timeout
 except ImportError:
     pytest_timeout = None
-import sys
 
 import ray
 import ray.test_utils
@@ -172,5 +174,23 @@ def test_placement_group_hang(ray_start_cluster):
     assert "CPU_group_" in list(resources.keys())[0], resources
 
 
+def test_cuda_visible_devices(ray_start_cluster):
+    @ray.remote(num_gpus=1)
+    def f():
+        return os.environ["CUDA_VISIBLE_DEVICES"]
+
+    cluster = ray_start_cluster
+    num_nodes = 1
+    for _ in range(num_nodes):
+        cluster.add_node(num_gpus=1)
+    ray.init(address=cluster.address)
+
+    g1 = ray.experimental.placement_group([{"CPU": 1, "GPU": 1}])
+    o1 = f.options(placement_group_id=g1).remote()
+
+    devices = ray.get(o1)
+    assert devices == "0", devices
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))