code: def test_gpu_ids(shutdown_only):
      num_gpus = 3
      ray.init(num_cpus=num_gpus, num_gpus=num_gpus)
  
      def get_gpu_ids(num_gpus_per_worker):
          gpu_ids = ray.get_gpu_ids()
          assert len(gpu_ids) == num_gpus_per_worker
          neuron_core_ids = (
              ray.get_runtime_context().get_accelerator_ids()[
                  "neuron_cores"
              ]
          )
          gpu_ids_from_runtime_context = (
              ray.get_runtime_context().get_accelerator_ids()["GPU"]
          )
          assert len(gpu_ids) == len(gpu_ids_from_runtime_context)
          assert len(neuron_core_ids) == 0
          assert os.environ["CUDA_VISIBLE_DEVICES"] == ",".join(
              [str(i) for i in gpu_ids]  # noqa
          )
          for gpu_id in gpu_ids:
              assert gpu_id in range(num_gpus)
          return gpu_ids
  
      f0 = ray.remote(num_gpus=0)(lambda: get_gpu_ids(0))
      f1 = ray.remote(num_gpus=1)(lambda: get_gpu_ids(1))
      f2 = ray.remote(num_gpus=2)(lambda: get_gpu_ids(2))
  
      # Wait for all workers to start up.
      @ray.remote
      def f():
          time.sleep(0.2)
          return os.getpid()
  
      start_time = time.time()
      while True:
          num_workers_started = len(
              set(ray.get([f.remote() for _ in range(num_gpus)]))
          )
          if num_workers_started == num_gpus:
              break
          if time.time() > start_time + 10:
              raise RayTestTimeoutException(
                  "Timed out while waiting for workers to start up."
              )
  
      list_of_ids = ray.get([f0.remote() for _ in range(10)])
      assert list_of_ids == 10 * [[]]
      ray.get([f1.remote() for _ in range(10)])
      ray.get([f2.remote() for _ in range(10)])
  
      # Test that actors have CUDA_VISIBLE_DEVICES set properly.
  
      @ray.remote
      class Actor0:
          def __init__(self):
              gpu_ids = ray.get_gpu_ids()
              assert len(gpu_ids) == 0
              assert os.environ["CUDA_VISIBLE_DEVICES"] == ",".join(
                  [str(i) for i in gpu_ids]  # noqa
              )
              # Set self.x to make sure that we got here.
              self.x = 1
  
          def test(self):
              gpu_ids = ray.get_gpu_ids()
              assert len(gpu_ids) == 0
              assert os.environ["CUDA_VISIBLE_DEVICES"] == ",".join(
                  [str(i) for i in gpu_ids]
              )
              return self.x
  
      @ray.remote(num_gpus=1)
      class Actor1:
          def __init__(self):
              gpu_ids = ray.get_gpu_ids()
              assert len(gpu_ids) == 1
              assert os.environ["CUDA_VISIBLE_DEVICES"] == ",".join(
                  [str(i) for i in gpu_ids]
              )
              # Set self.x to make sure that we got here.
              self.x = 1
  
          def test(self):
              gpu_ids = ray.get_gpu_ids()
              assert len(gpu_ids) == 1
              assert os.environ["CUDA_VISIBLE_DEVICES"] == ",".join(
                  [str(i) for i in gpu_ids]
              )
              return self.x
  
      a0 = Actor0.remote()
      ray.get(a0.test.remote())
  
      a1 = Actor1.remote()
      ray.get(a1.test.remote())
  
imports: import logging
  import os
  import sys
  import time
  import numpy
  import pytest
  import ray
  import ray.cluster_utils
  from ray._private.test_utils import RayTestTimeoutException, wait_for_condition
  from ray.util.placement_group import placement_group
  from ray.util.accelerators import AWS_NEURON_CORE
  from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
label_desc: Use image node_id 1 for calling a global function with example usage: test_gpu_ids(shutdown_only) and returns: gpu_ids, os, self, self