code: def test_gpu_allocation(self):
      for (
          max_worker_nodes,
          num_cpus_worker_node,
          num_gpus_worker_node,
      ) in [
          (
              self.max_spark_tasks // 2,
              self.num_cpus_per_spark_task,
              self.num_gpus_per_spark_task,
          ),
          (
              self.max_spark_tasks,
              self.num_cpus_per_spark_task,
              self.num_gpus_per_spark_task,
          ),
          (
              self.max_spark_tasks // 2,
              self.num_cpus_per_spark_task * 2,
              self.num_gpus_per_spark_task * 2,
          ),
          (
              self.max_spark_tasks // 2,
              self.num_cpus_per_spark_task,
              self.num_gpus_per_spark_task * 2,
          ),
      ]:
          with _setup_ray_cluster(
              max_worker_nodes=max_worker_nodes,
              num_cpus_worker_node=num_cpus_worker_node,
              num_gpus_worker_node=num_gpus_worker_node,
              head_node_options={"include_dashboard": False},
          ):
              ray.init()
              worker_res_list = self.get_ray_worker_resources_list()
              assert len(worker_res_list) == max_worker_nodes
  
              num_ray_task_slots = self.max_spark_tasks // (
                  num_gpus_worker_node // self.num_gpus_per_spark_task
              )
              (
                  mem_worker_node,
                  object_store_mem_worker_node,
                  _,
              ) = _calc_mem_per_ray_worker_node(
                  num_task_slots=num_ray_task_slots,
                  physical_mem_bytes=_RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES,
                  shared_mem_bytes=_RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES,
                  configured_object_store_bytes=None,
              )
  
              for worker_res in worker_res_list:
                  assert worker_res["CPU"] == num_cpus_worker_node
                  assert worker_res["GPU"] == num_gpus_worker_node
                  assert worker_res["memory"] == mem_worker_node
                  assert (
                      worker_res["object_store_memory"]
                      == object_store_mem_worker_node
                  )
  
              @ray.remote(
                  num_cpus=num_cpus_worker_node,
                  num_gpus=num_gpus_worker_node,
              )
              def f(_):
                  # Add a sleep to avoid the task finishing too fast,
                  # so that it can make all ray tasks concurrently running in all idle
                  # task slots.
                  time.sleep(5)
                  return [
                      int(gpu_id)
                      for gpu_id in os.environ[
                          "CUDA_VISIBLE_DEVICES"
                      ].split(",")
                  ]
  
              futures = [f.remote(i) for i in range(max_worker_nodes)]
              results = ray.get(futures)
              merged_results = functools.reduce(
                  lambda x, y: x + y, results
              )
              # Test all ray tasks are assigned with different GPUs.
              assert sorted(merged_results) == list(
                  range(num_gpus_worker_node * max_worker_nodes)
              )
  
imports: import sys
  import pytest
  import os
  import time
  import functools
  from abc import ABC
  from pyspark.sql import SparkSession
  from ray.tests.spark.test_basic import RayOnSparkCPUClusterTestBase, _setup_ray_cluster, _setup_ray_on_spark_envs, _RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES, _RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES
  from ray.util.spark.utils import _calc_mem_per_ray_worker_node
  from ray._private.test_utils import wait_for_condition
  import ray
label_desc: Use image node_id 1 for calling the RayOnSparkGPUClusterTestBase obj's underlying member method code with example usage: obj.test_gpu_allocation() and returns: unknown