import logging import os import re import shutil import socket import sys import tempfile import threading import time from abc import ABC from contextlib import contextmanager from unittest import mock import pytest from pyspark.sql import SparkSession import ray import ray.util.spark.cluster_init from ray._common.test_utils import wait_for_condition from ray.util.spark import ( MAX_NUM_WORKER_NODES, setup_global_ray_cluster, setup_ray_cluster, shutdown_ray_cluster, ) 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), ] _RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES = 2000000000 _RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES = 10000000000 def _setup_ray_on_spark_envs(): os.environ["RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES"] = str( _RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES ) os.environ["RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES"] = str( _RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES ) def setup_module(): _setup_ray_on_spark_envs() @contextmanager def _setup_ray_cluster(*args, **kwds): setup_ray_cluster(*args, **kwds) try: yield ray.util.spark.cluster_init._active_ray_cluster finally: shutdown_ray_cluster() _logger = logging.getLogger(__name__) class RayOnSparkCPUClusterTestBase(ABC): spark = None num_total_cpus = None num_total_gpus = None num_cpus_per_spark_task = None num_gpus_per_spark_task = None max_spark_tasks = None @classmethod def teardown_class(cls): time.sleep(10) # Wait all background spark job canceled. os.environ.pop("SPARK_WORKER_CORES", None) cls.spark.stop() @staticmethod def get_ray_worker_resources_list(): wr_list = [] for node in ray.nodes(): # exclude dead node and head node (with 0 CPU resource) if node["Alive"] and node["Resources"].get("CPU", 0) > 0: wr_list.append(node["Resources"]) return wr_list def test_cpu_allocation(self): for max_worker_nodes, num_cpus_worker_node, max_worker_nodes_arg in [ ( self.max_spark_tasks // 2, self.num_cpus_per_spark_task, self.max_spark_tasks // 2, ), (self.max_spark_tasks, self.num_cpus_per_spark_task, MAX_NUM_WORKER_NODES), ( self.max_spark_tasks // 2, self.num_cpus_per_spark_task * 2, MAX_NUM_WORKER_NODES, ), ( self.max_spark_tasks // 2, self.num_cpus_per_spark_task * 2, self.max_spark_tasks // 2 + 1, ), # Test case: requesting resources exceeding all cluster resources ]: num_ray_task_slots = self.max_spark_tasks // ( num_cpus_worker_node // self.num_cpus_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_arg, num_cpus_worker_node=num_cpus_worker_node, num_gpus_worker_node=0, 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 for worker_res in worker_res_list: assert ( worker_res["CPU"] == num_cpus_worker_node and worker_res["memory"] == mem_worker_node and worker_res["object_store_memory"] == object_store_mem_worker_node ) def test_public_api(self): try: ray_temp_root_dir = tempfile.mkdtemp(dir="/tmp") collect_log_to_path = tempfile.mkdtemp(dir="/tmp") # Test the case that `collect_log_to_path` directory does not exist. shutil.rmtree(collect_log_to_path, ignore_errors=True) setup_ray_cluster( max_worker_nodes=MAX_NUM_WORKER_NODES, num_cpus_worker_node=1, num_gpus_worker_node=0, collect_log_to_path=collect_log_to_path, ray_temp_root_dir=ray_temp_root_dir, head_node_options={"include_dashboard": True}, ) assert ( os.environ["RAY_ADDRESS"] == ray.util.spark.cluster_init._active_ray_cluster.address ) ray.init() @ray.remote def f(x): return x * x futures = [f.remote(i) for i in range(32)] results = ray.get(futures) assert results == [i * i for i in range(32)] shutdown_ray_cluster() assert "RAY_ADDRESS" not in os.environ time.sleep(7) # assert temp dir is removed. assert len(os.listdir(ray_temp_root_dir)) == 1 and os.listdir( ray_temp_root_dir )[0].endswith(".lock") # assert logs are copied to specified path listed_items = os.listdir(collect_log_to_path) assert len(listed_items) == 1 and listed_items[0].startswith("ray-") log_dest_dir = os.path.join( collect_log_to_path, listed_items[0], socket.gethostname() ) assert os.path.exists(log_dest_dir) and len(os.listdir(log_dest_dir)) > 0 finally: if ray.util.spark.cluster_init._active_ray_cluster is not None: # if the test raised error and does not destroy cluster, # destroy it here. ray.util.spark.cluster_init._active_ray_cluster.shutdown() time.sleep(5) shutil.rmtree(ray_temp_root_dir, ignore_errors=True) shutil.rmtree(collect_log_to_path, ignore_errors=True) def test_autoscaling(self): for max_worker_nodes, num_cpus_worker_node, min_worker_nodes in [ (self.max_spark_tasks, self.num_cpus_per_spark_task, 0), (self.max_spark_tasks // 2, self.num_cpus_per_spark_task * 2, 0), (self.max_spark_tasks, self.num_cpus_per_spark_task, 1), ]: num_ray_task_slots = self.max_spark_tasks // ( num_cpus_worker_node // self.num_cpus_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, min_worker_nodes=min_worker_nodes, num_cpus_worker_node=num_cpus_worker_node, num_gpus_worker_node=0, head_node_options={"include_dashboard": False}, autoscale_idle_timeout_minutes=0.1, ): ray.init() worker_res_list = self.get_ray_worker_resources_list() assert len(worker_res_list) == min_worker_nodes @ray.remote(num_cpus=num_cpus_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]["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()) == min_worker_nodes, timeout=60, retry_interval_ms=1000, ) if min_worker_nodes > 0: # Test scaling down keeps nodes number >= min_worker_nodes time.sleep(30) assert len(self.get_ray_worker_resources_list()) == min_worker_nodes class TestBasicSparkCluster(RayOnSparkCPUClusterTestBase): @classmethod def setup_class(cls): cls.num_total_cpus = 2 cls.num_total_gpus = 0 cls.num_cpus_per_spark_task = 1 cls.num_gpus_per_spark_task = 0 cls.max_spark_tasks = 2 os.environ["SPARK_WORKER_CORES"] = "2" cls.spark = ( SparkSession.builder.master("local-cluster[1, 2, 1024]") .config("spark.task.cpus", "1") .config("spark.task.maxFailures", "1") .config("spark.executorEnv.RAY_ON_SPARK_WORKER_CPU_CORES", "2") .getOrCreate() ) class TestSparkLocalCluster: @classmethod def setup_class(cls): cls.spark = ( SparkSession.builder.master("local[2]") .config("spark.task.cpus", "1") .config("spark.task.maxFailures", "1") .getOrCreate() ) @classmethod def teardown_class(cls): time.sleep(10) # Wait all background spark job canceled. cls.spark.stop() def test_basic(self): local_addr, remote_addr = setup_ray_cluster( max_worker_nodes=2, head_node_options={"include_dashboard": False}, collect_log_to_path="/tmp/ray_log_collect", ) for cluster_addr in [local_addr, remote_addr]: ray.init(address=cluster_addr) @ray.remote def f(x): return x * x futures = [f.remote(i) for i in range(32)] results = ray.get(futures) assert results == [i * i for i in range(32)] ray.shutdown() shutdown_ray_cluster() def test_use_driver_resources(self): setup_ray_cluster( max_worker_nodes=1, num_cpus_head_node=3, num_gpus_head_node=2, object_store_memory_head_node=256 * 1024 * 1024, head_node_options={"include_dashboard": False}, min_worker_nodes=0, ) ray.init() head_resources_list = [] for node in ray.nodes(): if node["Alive"] and node["Resources"].get("CPU", 0) == 3: head_resources_list.append(node["Resources"]) assert len(head_resources_list) == 1 head_resources = head_resources_list[0] assert head_resources.get("GPU", 0) == 2 shutdown_ray_cluster() def test_setup_global_ray_cluster(self): shutil.rmtree("/tmp/ray", ignore_errors=True) assert ray.util.spark.cluster_init._global_ray_cluster_cancel_event is None def start_serve_thread(): def serve(): try: with mock.patch( "ray.util.spark.cluster_init.get_spark_session", return_value=self.spark, ): setup_global_ray_cluster( max_worker_nodes=1, min_worker_nodes=0, ) except BaseException: # For debugging testing failure. import traceback traceback.print_exc() raise threading.Thread(target=serve, daemon=True).start() start_serve_thread() wait_for_condition( ( lambda: ray.util.spark.cluster_init._global_ray_cluster_cancel_event is not None ), timeout=120, retry_interval_ms=10000, ) # assert it uses default temp directory assert os.path.exists("/tmp/ray") # assert we can connect to it on client server port 10001 assert ( ray.util.spark.cluster_init._active_ray_cluster.ray_client_server_port == 10001 ) with mock.patch("ray.util.spark.cluster_init._active_ray_cluster", None): # assert we cannot create another global mode cluster at a time with pytest.raises( ValueError, match=re.compile( "Acquiring global lock failed for setting up new global mode " "Ray on spark cluster" ), ): setup_global_ray_cluster( max_worker_nodes=1, min_worker_nodes=0, ) # shut down the cluster ray.util.spark.cluster_init._global_ray_cluster_cancel_event.set() # assert temp directory is deleted wait_for_condition( lambda: not os.path.exists("/tmp/ray"), timeout=60, retry_interval_ms=10000, ) def test_autoscaling_config_generation(self): from ray.util.spark.cluster_init import AutoscalingCluster autoscaling_cluster = AutoscalingCluster( head_resources={ "CPU": 3, "GPU": 4, "memory": 10000000, "object_store_memory": 20000000, }, worker_node_types={ "ray.worker": { "resources": { "CPU": 5, "GPU": 6, "memory": 30000000, "object_store_memory": 40000000, }, "node_config": {}, "min_workers": 0, "max_workers": 100, }, }, extra_provider_config={ "extra_aa": "abc", "extra_bb": 789, }, upscaling_speed=2.0, idle_timeout_minutes=3.0, ) config = autoscaling_cluster._config assert config["max_workers"] == 100 assert config["available_node_types"]["ray.head.default"] == { "resources": { "CPU": 3, "GPU": 4, "memory": 10000000, "object_store_memory": 20000000, }, "node_config": {}, "max_workers": 0, } assert config["available_node_types"]["ray.worker"] == { "resources": { "CPU": 5, "GPU": 6, "memory": 30000000, "object_store_memory": 40000000, }, "node_config": {}, "min_workers": 0, "max_workers": 100, } assert config["upscaling_speed"] == 2.0 assert config["idle_timeout_minutes"] == 3.0 assert config["provider"]["extra_aa"] == "abc" assert config["provider"]["extra_bb"] == 789 def test_start_ray_node_in_new_process_group(self): from ray.util.spark.cluster_init import _start_ray_head_node proc, _ = _start_ray_head_node( [ sys.executable, "-m", "ray.util.spark.start_ray_node", "--head", "--block", "--port=44335", ], synchronous=False, extra_env={ "RAY_ON_SPARK_COLLECT_LOG_TO_PATH": "", "RAY_ON_SPARK_START_RAY_PARENT_PID": str(os.getpid()), }, ) time.sleep(10) # Assert the created Ray head node process has a different # group id from parent process group id. assert os.getpgid(proc.pid) != os.getpgrp() proc.terminate() if __name__ == "__main__": sys.exit(pytest.main(["-sv", __file__]))