chore: import upstream snapshot with attribution

python/ray/tests/spark/discover_2_gpu.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+# This script is used in spark GPU cluster config for discovering available GPU.
+echo "{\"name\":\"gpu\",\"addresses\":[\"0\",\"1\"]}"

python/ray/tests/spark/discover_4_gpu.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+# This script is used in spark GPU cluster config for discovering available GPU.
+echo "{\"name\":\"gpu\",\"addresses\":[\"0\",\"1\",\"2\",\"3\"]}"

python/ray/tests/spark/test_GPU.py

@@ -0,0 +1,228 @@
+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__]))

python/ray/tests/spark/test_basic.py

@@ -0,0 +1,497 @@
+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__]))

python/ray/tests/spark/test_databricks_hook.py

@@ -0,0 +1,87 @@
+import os
+import sys
+import time
+
+import pytest
+from pyspark.sql import SparkSession
+
+import ray
+import ray.util.spark.databricks_hook
+from ray._common.test_utils import wait_for_condition
+from ray.util.spark import setup_ray_cluster
+
+pytestmark = pytest.mark.skipif(
+    not sys.platform.startswith("linux"),
+    reason="Ray on spark only supports running on Linux.",
+)
+
+
+class MockDbApiEntry:
+    def __init__(self):
+        self.created_time = time.time()
+        self.registered_job_groups = []
+
+    def getIdleTimeMillisSinceLastNotebookExecution(self):
+        return (time.time() - self.created_time) * 1000
+
+    def registerBackgroundSparkJobGroup(self, job_group_id):
+        self.registered_job_groups.append(job_group_id)
+
+
+class TestDatabricksHook:
+    @classmethod
+    def setup_class(cls):
+        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()
+        )
+
+    @classmethod
+    def teardown_class(cls):
+        time.sleep(10)  # Wait all background spark job canceled.
+        cls.spark.stop()
+        os.environ.pop("SPARK_WORKER_CORES")
+
+    def test_hook(self, monkeypatch):
+        monkeypatch.setattr(
+            "ray.util.spark.databricks_hook._DATABRICKS_DEFAULT_TMP_ROOT_DIR", "/tmp"
+        )
+        monkeypatch.setenv("DATABRICKS_RUNTIME_VERSION", "12.2")
+        monkeypatch.setenv("DATABRICKS_RAY_ON_SPARK_AUTOSHUTDOWN_MINUTES", "0.5")
+        db_api_entry = MockDbApiEntry()
+        monkeypatch.setattr(
+            "ray.util.spark.databricks_hook.get_db_entry_point", lambda: db_api_entry
+        )
+        monkeypatch.setattr(
+            "ray.util.spark.databricks_hook.get_databricks_display_html_function",
+            lambda: lambda x: print(x),
+        )
+        try:
+            setup_ray_cluster(
+                max_worker_nodes=2,
+                num_cpus_worker_node=1,
+                num_gpus_worker_node=0,
+                head_node_options={"include_dashboard": False},
+            )
+            cluster = ray.util.spark.cluster_init._active_ray_cluster
+            assert not cluster.is_shutdown
+            wait_for_condition(
+                lambda: cluster.is_shutdown,
+                timeout=45,
+                retry_interval_ms=10000,
+            )
+            assert cluster.is_shutdown
+            assert ray.util.spark.cluster_init._active_ray_cluster is None
+        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()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/spark/test_multicores_per_task.py

@@ -0,0 +1,53 @@
+import os
+import sys
+
+import pytest
+from pyspark.sql import SparkSession
+
+from ray.tests.spark.test_basic import _setup_ray_on_spark_envs
+from ray.tests.spark.test_GPU import RayOnSparkGPUClusterTestBase
+
+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 TestMultiCoresPerTaskCluster(RayOnSparkGPUClusterTestBase):
+    @classmethod
+    def setup_class(cls):
+        cls.num_total_cpus = 4
+        cls.num_total_gpus = 4
+        cls.num_cpus_per_spark_task = 2
+        cls.num_gpus_per_spark_task = 2
+        cls.max_spark_tasks = 2
+        gpu_discovery_script_path = os.path.join(
+            os.path.dirname(os.path.abspath(__file__)), "discover_4_gpu.sh"
+        )
+        os.environ["SPARK_WORKER_CORES"] = "4"
+        cls.spark = (
+            SparkSession.builder.master("local-cluster[1, 4, 1024]")
+            .config("spark.task.cpus", "2")
+            .config("spark.task.resource.gpu.amount", "2")
+            .config("spark.executor.cores", "4")
+            .config("spark.worker.resource.gpu.amount", "4")
+            .config("spark.executor.resource.gpu.amount", "4")
+            .config("spark.task.maxFailures", "1")
+            .config(
+                "spark.worker.resource.gpu.discoveryScript", gpu_discovery_script_path
+            )
+            .config("spark.executorEnv.RAY_ON_SPARK_WORKER_CPU_CORES", "4")
+            .config("spark.executorEnv.RAY_ON_SPARK_WORKER_GPU_NUM", "4")
+            .getOrCreate()
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/tests/spark/test_utils.py

@@ -0,0 +1,194 @@
+import os
+import re
+import sys
+from unittest.mock import patch
+
+import pytest
+
+from ray.util.spark.cluster_init import (
+    _append_default_spilling_dir_config,
+    _convert_ray_node_options,
+    _verify_node_options,
+)
+from ray.util.spark.utils import (
+    _calc_mem_per_ray_worker_node,
+    _get_avail_mem_per_ray_worker_node,
+    get_spark_task_assigned_physical_gpus,
+)
+
+pytestmark = pytest.mark.skipif(
+    not sys.platform.startswith("linux"),
+    reason="Ray on spark only supports running on Linux.",
+)
+
+
+def test_get_spark_task_assigned_physical_gpus():
+    with patch.dict(os.environ, {}, clear=True):
+        assert get_spark_task_assigned_physical_gpus([2, 5]) == [2, 5]
+
+    with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3,6"}, clear=True):
+        assert get_spark_task_assigned_physical_gpus([0, 1]) == [2, 3]
+        assert get_spark_task_assigned_physical_gpus([0, 2]) == [2, 6]
+
+
+@patch("ray._private.ray_constants.OBJECT_STORE_MINIMUM_MEMORY_BYTES", 1)
+def test_calc_mem_per_ray_worker_node(monkeypatch):
+    assert _calc_mem_per_ray_worker_node(4, 1000000, 400000, None, 100000) == (
+        120000,
+        80000,
+        None,
+    )
+    assert _calc_mem_per_ray_worker_node(4, 1000000, 400000, None, 70000) == (
+        130000,
+        70000,
+        None,
+    )
+    assert _calc_mem_per_ray_worker_node(4, 1000000, 400000, None, None) == (
+        140000,
+        60000,
+        None,
+    )
+    assert _calc_mem_per_ray_worker_node(4, 1000000, 200000, None, None) == (
+        160000,
+        40000,
+        None,
+    )
+    assert _calc_mem_per_ray_worker_node(4, 1000000, 400000, 150000, 70000) == (
+        150000,
+        70000,
+        None,
+    )
+    monkeypatch.setenv("RAY_OBJECT_STORE_ALLOW_SLOW_STORAGE", "1")
+    assert _calc_mem_per_ray_worker_node(4, 1000000, 400000, None, 100000) == (
+        100000,
+        100000,
+        None,
+    )
+
+
+@patch("ray._private.ray_constants.OBJECT_STORE_MINIMUM_MEMORY_BYTES", 1)
+def test_get_avail_mem_per_ray_worker_node(monkeypatch):
+    monkeypatch.setenv("RAY_ON_SPARK_WORKER_CPU_CORES", "4")
+    monkeypatch.setenv("RAY_ON_SPARK_WORKER_GPU_NUM", "8")
+    monkeypatch.setenv("RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES", "1000000")
+    monkeypatch.setenv("RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES", "500000")
+
+    assert _get_avail_mem_per_ray_worker_node(
+        num_cpus_per_node=1,
+        num_gpus_per_node=2,
+        heap_memory_per_node=None,
+        object_store_memory_per_node=None,
+    ) == (140000, 60000, None, None)
+
+    assert _get_avail_mem_per_ray_worker_node(
+        num_cpus_per_node=1,
+        num_gpus_per_node=2,
+        heap_memory_per_node=None,
+        object_store_memory_per_node=80000,
+    ) == (120000, 80000, None, None)
+
+    assert _get_avail_mem_per_ray_worker_node(
+        num_cpus_per_node=1,
+        num_gpus_per_node=2,
+        heap_memory_per_node=None,
+        object_store_memory_per_node=120000,
+    ) == (100000, 100000, None, None)
+
+    assert _get_avail_mem_per_ray_worker_node(
+        num_cpus_per_node=2,
+        num_gpus_per_node=2,
+        heap_memory_per_node=None,
+        object_store_memory_per_node=None,
+    ) == (280000, 120000, None, None)
+
+    assert _get_avail_mem_per_ray_worker_node(
+        num_cpus_per_node=1,
+        num_gpus_per_node=4,
+        heap_memory_per_node=None,
+        object_store_memory_per_node=None,
+    ) == (280000, 120000, None, None)
+
+    assert _get_avail_mem_per_ray_worker_node(
+        num_cpus_per_node=1,
+        num_gpus_per_node=2,
+        heap_memory_per_node=150000,
+        object_store_memory_per_node=70000,
+    ) == (150000, 70000, None, None)
+
+
+def test_convert_ray_node_options():
+    assert _convert_ray_node_options(
+        {
+            "cluster_name": "aBc",
+            "disable_usage_stats": None,
+            "include_dashboard": False,
+        }
+    ) == ["--cluster-name=aBc", "--disable-usage-stats", "--include-dashboard=False"]
+
+
+def test_verify_node_options():
+    _verify_node_options(
+        node_options={"permitted": "127.0.0.1"},
+        block_keys={"not_permitted": None},
+        node_type="head",
+    )
+
+    with pytest.raises(
+        ValueError,
+        match=re.compile(
+            "Setting the option 'node_ip_address' for head nodes is not allowed.*"
+            "This option is controlled by Ray on Spark"
+        ),
+    ):
+        _verify_node_options(
+            node_options={"node_ip_address": "127.0.0.1"},
+            block_keys={"node_ip_address": None},
+            node_type="head",
+        )
+
+    with pytest.raises(
+        ValueError,
+        match=re.compile(
+            "Setting the option 'not_permitted' for worker nodes is not allowed.*"
+            "You should set the 'permitted' option instead",
+        ),
+    ):
+        _verify_node_options(
+            node_options={"not_permitted": "bad"},
+            block_keys={"not_permitted": "permitted"},
+            node_type="worker",
+        )
+
+
+def test_append_default_spilling_dir_config():
+    assert _append_default_spilling_dir_config({}, "/xx/yy") == {
+        "system_config": {
+            "object_spilling_config": '{"type": "filesystem", "params": {"directory_path": "/xx/yy"}}'  # noqa: E501
+        }
+    }
+    assert _append_default_spilling_dir_config(
+        {"system_config": {"a": 3}}, "/xx/yy"
+    ) == {
+        "system_config": {
+            "a": 3,
+            "object_spilling_config": '{"type": "filesystem", "params": {"directory_path": "/xx/yy"}}',  # noqa: E501
+        }
+    }
+    assert _append_default_spilling_dir_config(
+        {
+            "system_config": {
+                "a": 4,
+                "object_spilling_config": '{"type": "filesystem", "params": {"directory_path": "/aa/bb"}}',  # noqa: E501
+            },
+        },
+        "/xx/yy",
+    ) == {
+        "system_config": {
+            "a": 4,
+            "object_spilling_config": '{"type": "filesystem", "params": {"directory_path": "/aa/bb"}}',  # noqa: E501
+        }
+    }
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))