chore: import upstream snapshot with attribution

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__]))