import functools import os import sys import time from abc import ABC import pytest from pyspark.sql import SparkSession import ray from ray._common.test_utils import wait_for_condition from ray.tests.spark.test_basic import ( _RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES, _RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES, RayOnSparkCPUClusterTestBase, _setup_ray_cluster, _setup_ray_on_spark_envs, ) from ray.util.spark.utils import _calc_mem_per_ray_worker_node pytestmark = [ pytest.mark.skipif( os.name != "posix", reason="Ray on spark only supports running on POSIX system.", ), pytest.mark.timeout(1500), ] def setup_module(): _setup_ray_on_spark_envs() class RayOnSparkGPUClusterTestBase(RayOnSparkCPUClusterTestBase, ABC): num_total_gpus = None num_gpus_per_spark_task = None 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_heap_memory_bytes=None, 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) ) def test_gpu_autoscaling(self): for max_worker_nodes, num_cpus_worker_node, num_gpus_worker_node in [ ( 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, ), ]: 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_heap_memory_bytes=None, configured_object_store_bytes=None, ) 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}, min_worker_nodes=0, autoscale_idle_timeout_minutes=0.1, ): ray.init() worker_res_list = self.get_ray_worker_resources_list() assert len(worker_res_list) == 0 @ray.remote( num_cpus=num_cpus_worker_node, num_gpus=num_gpus_worker_node ) def f(x): import time time.sleep(5) return x * x # Test scale up futures = [f.remote(i) for i in range(8)] results = ray.get(futures) assert results == [i * i for i in range(8)] worker_res_list = self.get_ray_worker_resources_list() assert len(worker_res_list) == max_worker_nodes and all( worker_res_list[i]["CPU"] == num_cpus_worker_node and worker_res_list[i]["GPU"] == num_gpus_worker_node and worker_res_list[i]["memory"] == mem_worker_node and worker_res_list[i]["object_store_memory"] == object_store_mem_worker_node for i in range(max_worker_nodes) ) # Test scale down wait_for_condition( lambda: len(self.get_ray_worker_resources_list()) == 0, timeout=60, retry_interval_ms=1000, ) def test_default_resource_allocation(self): with _setup_ray_cluster( max_worker_nodes=1, head_node_options={"include_dashboard": False}, ): ray.init() worker_res_list = self.get_ray_worker_resources_list() assert worker_res_list[0]["CPU"] == self.num_total_gpus assert worker_res_list[0]["GPU"] == self.num_total_cpus class TestBasicSparkGPUCluster(RayOnSparkGPUClusterTestBase): @classmethod def setup_class(cls): cls.num_total_cpus = 2 cls.num_total_gpus = 2 cls.num_cpus_per_spark_task = 1 cls.num_gpus_per_spark_task = 1 cls.max_spark_tasks = 2 gpu_discovery_script_path = os.path.join( os.path.dirname(os.path.abspath(__file__)), "discover_2_gpu.sh" ) os.environ["SPARK_WORKER_CORES"] = "4" cls.spark = ( SparkSession.builder.master("local-cluster[1, 2, 1024]") .config("spark.task.cpus", "1") .config("spark.task.resource.gpu.amount", "1") .config("spark.executor.cores", "2") .config("spark.worker.resource.gpu.amount", "2") .config("spark.executor.resource.gpu.amount", "2") .config("spark.task.maxFailures", "1") .config( "spark.worker.resource.gpu.discoveryScript", gpu_discovery_script_path ) .config("spark.executorEnv.RAY_ON_SPARK_WORKER_CPU_CORES", "2") .config("spark.executorEnv.RAY_ON_SPARK_WORKER_GPU_NUM", "2") .getOrCreate() ) if __name__ == "__main__": sys.exit(pytest.main(["-sv", __file__]))