wehub-resource-sync
498 lines
17 KiB
Python
498 lines
17 KiB
Python
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__]))
|