wehub-resource-sync
1936 lines
76 KiB
Python
1936 lines
76 KiB
Python
import copy
|
|
import json
|
|
import logging
|
|
import os
|
|
import signal
|
|
import socket
|
|
import sys
|
|
import threading
|
|
import time
|
|
import uuid
|
|
import warnings
|
|
from threading import Event
|
|
from typing import Dict, Optional, Tuple, Type
|
|
|
|
import requests
|
|
import yaml
|
|
from packaging.version import Version
|
|
|
|
import ray
|
|
import ray._private.services
|
|
from .databricks_hook import DefaultDatabricksRayOnSparkStartHook
|
|
from .start_hook_base import RayOnSparkStartHook
|
|
from .utils import (
|
|
_get_cpu_cores,
|
|
_get_local_ray_node_slots,
|
|
_get_num_physical_gpus,
|
|
_wait_service_up,
|
|
calc_mem_ray_head_node,
|
|
exec_cmd,
|
|
gen_cmd_exec_failure_msg,
|
|
get_avail_mem_per_ray_worker_node,
|
|
get_configured_spark_executor_memory_bytes,
|
|
get_max_num_concurrent_tasks,
|
|
get_random_unused_port,
|
|
get_spark_application_driver_host,
|
|
get_spark_session,
|
|
get_spark_task_assigned_physical_gpus,
|
|
is_in_databricks_runtime,
|
|
is_port_in_use,
|
|
)
|
|
from ray._common.network_utils import (
|
|
build_address,
|
|
get_all_interfaces_ip,
|
|
parse_address,
|
|
)
|
|
from ray._common.utils import load_class
|
|
from ray.autoscaler._private.spark.node_provider import HEAD_NODE_ID
|
|
from ray.util.annotations import DeveloperAPI, PublicAPI
|
|
|
|
_logger = logging.getLogger("ray.util.spark")
|
|
_logger.setLevel(logging.INFO)
|
|
|
|
|
|
RAY_ON_SPARK_START_HOOK = "RAY_ON_SPARK_START_HOOK"
|
|
|
|
MAX_NUM_WORKER_NODES = -1
|
|
|
|
RAY_ON_SPARK_COLLECT_LOG_TO_PATH = "RAY_ON_SPARK_COLLECT_LOG_TO_PATH"
|
|
RAY_ON_SPARK_START_RAY_PARENT_PID = "RAY_ON_SPARK_START_RAY_PARENT_PID"
|
|
|
|
|
|
def _check_system_environment():
|
|
if os.name != "posix":
|
|
raise RuntimeError("Ray on spark only supports running on POSIX system.")
|
|
|
|
spark_dependency_error = "ray.util.spark module requires pyspark >= 3.3"
|
|
try:
|
|
import pyspark
|
|
|
|
if Version(pyspark.__version__).release < (3, 3, 0):
|
|
raise RuntimeError(spark_dependency_error)
|
|
except ImportError:
|
|
raise RuntimeError(spark_dependency_error)
|
|
|
|
|
|
class RayClusterOnSpark:
|
|
"""
|
|
This class is the type of instance returned by the `_setup_ray_cluster` interface.
|
|
Its main functionality is to:
|
|
Connect to, disconnect from, and shutdown the Ray cluster running on Apache Spark.
|
|
Serve as a Python context manager for the `RayClusterOnSpark` instance.
|
|
|
|
Args
|
|
address: The url for the ray head node (defined as the hostname and unused
|
|
port on Spark driver node)
|
|
head_proc: Ray head process
|
|
spark_job_group_id: The Spark job id for a submitted ray job
|
|
num_workers_node: The number of workers in the ray cluster.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
address,
|
|
head_proc,
|
|
min_worker_nodes,
|
|
max_worker_nodes,
|
|
temp_dir,
|
|
cluster_unique_id,
|
|
start_hook,
|
|
ray_dashboard_port,
|
|
spark_job_server,
|
|
global_cluster_lock_fd,
|
|
ray_client_server_port,
|
|
):
|
|
self.address = address
|
|
self.head_proc = head_proc
|
|
self.min_worker_nodes = min_worker_nodes
|
|
self.max_worker_nodes = max_worker_nodes
|
|
self.temp_dir = temp_dir
|
|
self.cluster_unique_id = cluster_unique_id
|
|
self.start_hook = start_hook
|
|
self.ray_dashboard_port = ray_dashboard_port
|
|
self.spark_job_server = spark_job_server
|
|
self.global_cluster_lock_fd = global_cluster_lock_fd
|
|
self.ray_client_server_port = ray_client_server_port
|
|
|
|
self.is_shutdown = False
|
|
self.spark_job_is_canceled = False
|
|
self.background_job_exception = None
|
|
|
|
# Ray client context returns by `ray.init`
|
|
self.ray_ctx = None
|
|
|
|
def wait_until_ready(self):
|
|
import ray
|
|
|
|
if self.is_shutdown:
|
|
raise RuntimeError(
|
|
"The ray cluster has been shut down or it failed to start."
|
|
)
|
|
|
|
try:
|
|
ray.init(address=self.address)
|
|
|
|
if self.ray_dashboard_port is not None and _wait_service_up(
|
|
parse_address(self.address)[0],
|
|
self.ray_dashboard_port,
|
|
_RAY_DASHBOARD_STARTUP_TIMEOUT,
|
|
):
|
|
self.start_hook.on_ray_dashboard_created(self.ray_dashboard_port)
|
|
else:
|
|
try:
|
|
__import__("ray.dashboard.optional_deps")
|
|
except ModuleNotFoundError as e:
|
|
_logger.warning(
|
|
"Dependencies to launch the optional dashboard API "
|
|
"server cannot be found. They can be installed with "
|
|
f"pip install ray[default], root cause: ({repr(e)})"
|
|
)
|
|
|
|
last_alive_worker_count = 0
|
|
last_progress_move_time = time.time()
|
|
while True:
|
|
time.sleep(_RAY_CLUSTER_STARTUP_PROGRESS_CHECKING_INTERVAL)
|
|
|
|
# Inside the waiting ready loop,
|
|
# checking `self.background_job_exception`, if it is not None,
|
|
# it means the background spark job has failed,
|
|
# in this case, raise error directly.
|
|
if self.background_job_exception is not None:
|
|
raise RuntimeError(
|
|
"Ray workers failed to start."
|
|
) from self.background_job_exception
|
|
|
|
cur_alive_worker_count = (
|
|
len([node for node in ray.nodes() if node["Alive"]]) - 1
|
|
) # Minus 1 means excluding the head node.
|
|
|
|
if cur_alive_worker_count >= self.min_worker_nodes:
|
|
_logger.info(
|
|
f"Started {cur_alive_worker_count} Ray worker nodes, "
|
|
f"meet the minimum number of Ray worker nodes required."
|
|
)
|
|
return
|
|
|
|
if cur_alive_worker_count > last_alive_worker_count:
|
|
last_alive_worker_count = cur_alive_worker_count
|
|
last_progress_move_time = time.time()
|
|
_logger.info(
|
|
"Ray worker nodes are starting. Progress: "
|
|
f"({cur_alive_worker_count} / {self.max_worker_nodes})"
|
|
)
|
|
else:
|
|
if (
|
|
time.time() - last_progress_move_time
|
|
> _RAY_CONNECT_CLUSTER_POLL_PROGRESS_TIMEOUT
|
|
):
|
|
if cur_alive_worker_count == 0:
|
|
(
|
|
job_server_host,
|
|
job_server_port,
|
|
) = self.spark_job_server.server_address[:2]
|
|
response = requests.post(
|
|
url=(
|
|
f"http://{build_address(job_server_host, job_server_port)}"
|
|
"/query_last_worker_err"
|
|
),
|
|
json={"spark_job_group_id": None},
|
|
)
|
|
response.raise_for_status()
|
|
|
|
decoded_resp = response.content.decode("utf-8")
|
|
json_res = json.loads(decoded_resp)
|
|
last_worker_err = json_res["last_worker_err"]
|
|
|
|
if last_worker_err:
|
|
raise RuntimeError(
|
|
"Starting Ray worker node failed, error:\n"
|
|
f"{last_worker_err}"
|
|
)
|
|
else:
|
|
raise RuntimeError(
|
|
"Current spark cluster has no resources to launch "
|
|
"Ray worker nodes."
|
|
)
|
|
_logger.warning(
|
|
"Timeout in waiting for minimal ray workers to start. "
|
|
"Started / Total requested: "
|
|
f"({cur_alive_worker_count} / {self.min_worker_nodes}). "
|
|
"Current spark cluster does not have sufficient resources "
|
|
"to launch requested minimal number of Ray worker nodes."
|
|
)
|
|
return
|
|
finally:
|
|
ray.shutdown()
|
|
|
|
def connect(self):
|
|
if ray.is_initialized():
|
|
raise RuntimeError("Already connected to Ray cluster.")
|
|
self.ray_ctx = ray.init(address=self.address)
|
|
|
|
def disconnect(self):
|
|
ray.shutdown()
|
|
self.ray_ctx = None
|
|
|
|
def shutdown(self):
|
|
"""
|
|
Shutdown the ray cluster created by the `setup_ray_cluster` API.
|
|
"""
|
|
import fcntl
|
|
|
|
if not self.is_shutdown:
|
|
try:
|
|
self.disconnect()
|
|
except Exception:
|
|
pass
|
|
os.environ.pop("RAY_ADDRESS", None)
|
|
|
|
if self.global_cluster_lock_fd is not None:
|
|
# release global mode cluster lock.
|
|
fcntl.flock(self.global_cluster_lock_fd, fcntl.LOCK_UN)
|
|
|
|
self.spark_job_server.shutdown()
|
|
try:
|
|
self.head_proc.terminate()
|
|
except Exception as e:
|
|
# swallow exception.
|
|
_logger.warning(
|
|
"An Error occurred during shutdown of ray head node: " f"{repr(e)}"
|
|
)
|
|
self.is_shutdown = True
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
self.shutdown()
|
|
|
|
|
|
def _convert_ray_node_option(key, value):
|
|
converted_key = f"--{key.replace('_', '-')}"
|
|
if key in ["system_config", "resources", "labels"]:
|
|
return f"{converted_key}={json.dumps(value)}"
|
|
if value is None:
|
|
return converted_key
|
|
return f"{converted_key}={str(value)}"
|
|
|
|
|
|
def _convert_ray_node_options(options):
|
|
return [_convert_ray_node_option(k, v) for k, v in options.items()]
|
|
|
|
|
|
_RAY_HEAD_STARTUP_TIMEOUT = 20
|
|
_RAY_DASHBOARD_STARTUP_TIMEOUT = 60
|
|
_BACKGROUND_JOB_STARTUP_WAIT = int(
|
|
os.environ.get("RAY_ON_SPARK_BACKGROUND_JOB_STARTUP_WAIT", "30")
|
|
)
|
|
_RAY_CLUSTER_STARTUP_PROGRESS_CHECKING_INTERVAL = 3
|
|
_RAY_WORKER_NODE_STARTUP_INTERVAL = int(
|
|
os.environ.get("RAY_ON_SPARK_RAY_WORKER_NODE_STARTUP_INTERVAL", "10")
|
|
)
|
|
_RAY_CONNECT_CLUSTER_POLL_PROGRESS_TIMEOUT = 120
|
|
|
|
|
|
def _preallocate_ray_worker_port_range():
|
|
"""
|
|
If we start multiple ray workers on a machine concurrently, some ray worker
|
|
processes might fail due to ray port conflicts, this is because race condition
|
|
on getting free port and opening the free port.
|
|
To address the issue, this function use an exclusive file lock to delay the
|
|
worker processes to ensure that port acquisition does not create a resource
|
|
contention issue due to a race condition.
|
|
|
|
After acquiring lock, it will allocate port range for worker ports
|
|
(for ray node config --min-worker-port and --max-worker-port).
|
|
Because on a spark cluster, multiple ray cluster might be created, so on one spark
|
|
worker machine, there might be multiple ray worker nodes running, these worker
|
|
nodes might belong to different ray cluster, and we must ensure these ray nodes on
|
|
the same machine using non-overlapping worker port range, to achieve this, in this
|
|
function, it creates a file `/tmp/ray_on_spark_worker_port_allocation.txt` file,
|
|
the file format is composed of multiple lines, each line contains 2 number: `pid`
|
|
and `port_range_slot_index`, each port range slot allocates 1000 ports, and
|
|
corresponding port range is:
|
|
- range_begin (inclusive): 20000 + port_range_slot_index * 1000
|
|
- range_end (exclusive): range_begin + 1000
|
|
In this function, it first scans `/tmp/ray_on_spark_worker_port_allocation.txt`
|
|
file, removing lines that containing dead process pid, then find the first unused
|
|
port_range_slot_index, then regenerate this file, and return the allocated port
|
|
range.
|
|
|
|
Returns: Allocated port range for current worker ports
|
|
"""
|
|
import fcntl
|
|
|
|
import psutil
|
|
|
|
def acquire_lock(file_path):
|
|
mode = os.O_RDWR | os.O_CREAT | os.O_TRUNC
|
|
try:
|
|
fd = os.open(file_path, mode)
|
|
# The lock file must be readable / writable to all users.
|
|
os.chmod(file_path, 0o0777)
|
|
# Allow for retrying getting a file lock a maximum number of seconds
|
|
max_lock_iter = 600
|
|
for _ in range(max_lock_iter):
|
|
try:
|
|
fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
except BlockingIOError:
|
|
# Lock is used by other processes, continue loop to wait for lock
|
|
# available
|
|
pass
|
|
else:
|
|
# Acquire lock successfully.
|
|
return fd
|
|
time.sleep(10)
|
|
raise TimeoutError(f"Acquiring lock on file {file_path} timeout.")
|
|
except Exception:
|
|
os.close(fd)
|
|
|
|
lock_file_path = "/tmp/ray_on_spark_worker_startup_barrier_lock.lock"
|
|
try:
|
|
lock_fd = acquire_lock(lock_file_path)
|
|
except TimeoutError:
|
|
# If timeout happens, the file lock might be hold by another process and that
|
|
# process does not release the lock in time by some unexpected reason.
|
|
# In this case, remove the existing lock file and create the file again, and
|
|
# then acquire file lock on the new file.
|
|
try:
|
|
os.remove(lock_file_path)
|
|
except Exception:
|
|
pass
|
|
lock_fd = acquire_lock(lock_file_path)
|
|
|
|
def release_lock():
|
|
fcntl.flock(lock_fd, fcntl.LOCK_UN)
|
|
os.close(lock_fd)
|
|
|
|
try:
|
|
port_alloc_file = "/tmp/ray_on_spark_worker_port_allocation.txt"
|
|
|
|
# NB: reading / writing `port_alloc_file` is protected by exclusive lock
|
|
# on file `lock_file_path`
|
|
if os.path.exists(port_alloc_file):
|
|
with open(port_alloc_file, mode="r") as fp:
|
|
port_alloc_data = fp.read()
|
|
port_alloc_table = [
|
|
line.split(" ") for line in port_alloc_data.strip().split("\n")
|
|
]
|
|
port_alloc_table = [
|
|
(int(pid_str), int(slot_index_str))
|
|
for pid_str, slot_index_str in port_alloc_table
|
|
]
|
|
else:
|
|
port_alloc_table = []
|
|
with open(port_alloc_file, mode="w"):
|
|
pass
|
|
# The port range allocation file must be readable / writable to all users.
|
|
os.chmod(port_alloc_file, 0o0777)
|
|
|
|
port_alloc_map = {
|
|
pid: slot_index
|
|
for pid, slot_index in port_alloc_table
|
|
if psutil.pid_exists(pid) # remove slot used by dead process
|
|
}
|
|
|
|
allocated_slot_set = set(port_alloc_map.values())
|
|
|
|
if len(allocated_slot_set) == 0:
|
|
new_slot_index = 0
|
|
else:
|
|
new_slot_index = max(allocated_slot_set) + 1
|
|
for index in range(new_slot_index):
|
|
if index not in allocated_slot_set:
|
|
new_slot_index = index
|
|
break
|
|
|
|
port_alloc_map[os.getpid()] = new_slot_index
|
|
|
|
with open(port_alloc_file, mode="w") as fp:
|
|
for pid, slot_index in port_alloc_map.items():
|
|
fp.write(f"{pid} {slot_index}\n")
|
|
|
|
worker_port_range_begin = 20000 + new_slot_index * 1000
|
|
worker_port_range_end = worker_port_range_begin + 1000
|
|
|
|
if worker_port_range_end > 65536:
|
|
raise RuntimeError(
|
|
"Too many ray worker nodes are running on this machine, cannot "
|
|
"allocate worker port range for new ray worker node."
|
|
)
|
|
except Exception:
|
|
release_lock()
|
|
raise
|
|
|
|
def hold_lock():
|
|
time.sleep(_RAY_WORKER_NODE_STARTUP_INTERVAL)
|
|
release_lock()
|
|
|
|
threading.Thread(target=hold_lock, args=()).start()
|
|
|
|
return worker_port_range_begin, worker_port_range_end
|
|
|
|
|
|
def _append_default_spilling_dir_config(head_node_options, object_spilling_dir):
|
|
if "system_config" not in head_node_options:
|
|
head_node_options["system_config"] = {}
|
|
sys_conf = head_node_options["system_config"]
|
|
if "object_spilling_config" not in sys_conf:
|
|
sys_conf["object_spilling_config"] = json.dumps(
|
|
{
|
|
"type": "filesystem",
|
|
"params": {
|
|
"directory_path": object_spilling_dir,
|
|
},
|
|
}
|
|
)
|
|
return head_node_options
|
|
|
|
|
|
def _append_resources_config(node_options, resources):
|
|
if "resources" not in node_options:
|
|
node_options["resources"] = {}
|
|
|
|
node_options["resources"].update(resources)
|
|
return node_options
|
|
|
|
|
|
def _get_default_ray_tmp_dir():
|
|
return os.path.join(os.environ.get("RAY_TMPDIR", "/tmp"), "ray")
|
|
|
|
|
|
def _create_hook_entry(is_global):
|
|
if RAY_ON_SPARK_START_HOOK in os.environ:
|
|
return load_class(os.environ[RAY_ON_SPARK_START_HOOK])()
|
|
elif is_in_databricks_runtime():
|
|
return DefaultDatabricksRayOnSparkStartHook(is_global)
|
|
else:
|
|
return RayOnSparkStartHook(is_global)
|
|
|
|
|
|
def _setup_ray_cluster(
|
|
*,
|
|
max_worker_nodes: int,
|
|
min_worker_nodes: int,
|
|
num_cpus_worker_node: int,
|
|
num_cpus_head_node: int,
|
|
num_gpus_worker_node: int,
|
|
num_gpus_head_node: int,
|
|
using_stage_scheduling: bool,
|
|
heap_memory_worker_node: int,
|
|
heap_memory_head_node: int,
|
|
object_store_memory_worker_node: int,
|
|
object_store_memory_head_node: int,
|
|
head_node_options: Dict,
|
|
worker_node_options: Dict,
|
|
ray_temp_root_dir: str,
|
|
collect_log_to_path: str,
|
|
autoscale_upscaling_speed: float,
|
|
autoscale_idle_timeout_minutes: float,
|
|
is_global: bool,
|
|
) -> Type[RayClusterOnSpark]:
|
|
"""
|
|
The public API `ray.util.spark.setup_ray_cluster` does some argument
|
|
validation and then pass validated arguments to this interface.
|
|
and it returns a `RayClusterOnSpark` instance.
|
|
|
|
The returned instance can be used to connect to, disconnect from and shutdown the
|
|
ray cluster. This instance can also be used as a context manager (used by
|
|
encapsulating operations within `with _setup_ray_cluster(...):`). Upon entering the
|
|
managed scope, the ray cluster is initiated and connected to. When exiting the
|
|
scope, the ray cluster is disconnected and shut down.
|
|
|
|
Note: This function interface is stable and can be used for
|
|
instrumentation logging patching.
|
|
"""
|
|
import fcntl
|
|
|
|
start_hook = _create_hook_entry(is_global)
|
|
spark = get_spark_session()
|
|
|
|
ray_head_ip = socket.gethostbyname(get_spark_application_driver_host(spark))
|
|
ray_head_port = get_random_unused_port(ray_head_ip, min_port=9000, max_port=10000)
|
|
port_exclude_list = [ray_head_port]
|
|
|
|
# Make a copy for head_node_options to avoid changing original dict in user code.
|
|
head_node_options = head_node_options.copy()
|
|
include_dashboard = head_node_options.pop("include_dashboard", None)
|
|
ray_dashboard_port = head_node_options.pop("dashboard_port", None)
|
|
|
|
if is_global:
|
|
ray_client_server_port = 10001
|
|
else:
|
|
ray_client_server_port = get_random_unused_port(
|
|
ray_head_ip,
|
|
min_port=9000,
|
|
max_port=10000,
|
|
exclude_list=port_exclude_list,
|
|
)
|
|
|
|
port_exclude_list.append(ray_client_server_port)
|
|
|
|
spark_job_server_port = get_random_unused_port(
|
|
ray_head_ip,
|
|
min_port=9000,
|
|
max_port=10000,
|
|
exclude_list=port_exclude_list,
|
|
)
|
|
port_exclude_list.append(spark_job_server_port)
|
|
|
|
if include_dashboard is None or include_dashboard is True:
|
|
if ray_dashboard_port is None:
|
|
ray_dashboard_port = get_random_unused_port(
|
|
ray_head_ip,
|
|
min_port=9000,
|
|
max_port=10000,
|
|
exclude_list=port_exclude_list,
|
|
)
|
|
port_exclude_list.append(ray_dashboard_port)
|
|
ray_dashboard_agent_port = get_random_unused_port(
|
|
ray_head_ip,
|
|
min_port=9000,
|
|
max_port=10000,
|
|
exclude_list=port_exclude_list,
|
|
)
|
|
port_exclude_list.append(ray_dashboard_agent_port)
|
|
|
|
dashboard_options = [
|
|
f"--dashboard-host={get_all_interfaces_ip()}",
|
|
f"--dashboard-port={ray_dashboard_port}",
|
|
f"--dashboard-agent-listen-port={ray_dashboard_agent_port}",
|
|
]
|
|
# If include_dashboard is None, we don't set `--include-dashboard` option,
|
|
# in this case Ray will decide whether dashboard can be started
|
|
# (e.g. checking any missing dependencies).
|
|
if include_dashboard is True:
|
|
dashboard_options += ["--include-dashboard=true"]
|
|
else:
|
|
dashboard_options = [
|
|
"--include-dashboard=false",
|
|
]
|
|
|
|
_logger.info(
|
|
f"Ray head hostname: {ray_head_ip}, port: {ray_head_port}, "
|
|
f"ray client server port: {ray_client_server_port}."
|
|
)
|
|
|
|
cluster_unique_id = uuid.uuid4().hex[:8]
|
|
|
|
if is_global:
|
|
# global mode enabled
|
|
# for global mode, Ray always uses default temp dir
|
|
# so that local Ray client can discover it without specifying
|
|
# head node address.
|
|
if ray_temp_root_dir is not None:
|
|
raise ValueError(
|
|
"Ray on spark global mode cluster does not allow you to set "
|
|
"'ray_temp_root_dir' argument."
|
|
)
|
|
|
|
# We only allow user to launch one active Ray on spark global cluster
|
|
# at a time. So acquiring a global file lock before setting up a new
|
|
# Ray on spark global cluster.
|
|
global_cluster_lock_fd = os.open(
|
|
"/tmp/ray_on_spark_global_cluster.lock", os.O_RDWR | os.O_CREAT | os.O_TRUNC
|
|
)
|
|
|
|
try:
|
|
# acquiring exclusive lock to ensure copy logs and removing dir safely.
|
|
fcntl.flock(global_cluster_lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
except BlockingIOError:
|
|
# acquiring global lock failed.
|
|
raise ValueError(
|
|
"Acquiring global lock failed for setting up new global mode Ray on "
|
|
"spark cluster. If there is an active global mode Ray on spark "
|
|
"cluster, please shut down it before you create a new one."
|
|
)
|
|
|
|
ray_temp_dir = None
|
|
ray_default_tmp_dir = _get_default_ray_tmp_dir()
|
|
os.makedirs(ray_default_tmp_dir, exist_ok=True)
|
|
object_spilling_dir = os.path.join(ray_default_tmp_dir, "spill")
|
|
else:
|
|
global_cluster_lock_fd = None
|
|
if ray_temp_root_dir is None:
|
|
ray_temp_root_dir = start_hook.get_default_temp_root_dir()
|
|
ray_temp_dir = os.path.join(
|
|
ray_temp_root_dir, f"ray-{ray_head_port}-{cluster_unique_id}"
|
|
)
|
|
os.makedirs(ray_temp_dir, exist_ok=True)
|
|
object_spilling_dir = os.path.join(ray_temp_dir, "spill")
|
|
|
|
os.makedirs(object_spilling_dir, exist_ok=True)
|
|
|
|
head_node_options = _append_default_spilling_dir_config(
|
|
head_node_options, object_spilling_dir
|
|
)
|
|
|
|
from ray.autoscaler._private.spark.spark_job_server import (
|
|
_start_spark_job_server,
|
|
)
|
|
|
|
ray_node_custom_env = start_hook.custom_environment_variables()
|
|
spark_job_server = _start_spark_job_server(
|
|
ray_head_ip, spark_job_server_port, spark, ray_node_custom_env
|
|
)
|
|
autoscaling_cluster = AutoscalingCluster(
|
|
head_resources={
|
|
"CPU": num_cpus_head_node,
|
|
"GPU": num_gpus_head_node,
|
|
"memory": heap_memory_head_node,
|
|
"object_store_memory": object_store_memory_head_node,
|
|
},
|
|
worker_node_types={
|
|
"ray.worker": {
|
|
"resources": {
|
|
"CPU": num_cpus_worker_node,
|
|
"GPU": num_gpus_worker_node,
|
|
"memory": heap_memory_worker_node,
|
|
"object_store_memory": object_store_memory_worker_node,
|
|
},
|
|
"node_config": {},
|
|
"min_workers": min_worker_nodes,
|
|
"max_workers": max_worker_nodes,
|
|
},
|
|
},
|
|
extra_provider_config={
|
|
"ray_head_ip": ray_head_ip,
|
|
"ray_head_port": ray_head_port,
|
|
"cluster_unique_id": cluster_unique_id,
|
|
"using_stage_scheduling": using_stage_scheduling,
|
|
"ray_temp_dir": ray_temp_dir,
|
|
"worker_node_options": worker_node_options,
|
|
"collect_log_to_path": collect_log_to_path,
|
|
"spark_job_server_port": spark_job_server_port,
|
|
},
|
|
upscaling_speed=autoscale_upscaling_speed,
|
|
idle_timeout_minutes=autoscale_idle_timeout_minutes,
|
|
)
|
|
ray_head_proc, tail_output_deque = autoscaling_cluster.start(
|
|
ray_head_ip,
|
|
ray_head_port,
|
|
ray_client_server_port,
|
|
ray_temp_dir,
|
|
dashboard_options,
|
|
head_node_options,
|
|
collect_log_to_path,
|
|
ray_node_custom_env,
|
|
)
|
|
ray_head_node_cmd = autoscaling_cluster.ray_head_node_cmd
|
|
|
|
# wait ray head node spin up.
|
|
time.sleep(_RAY_HEAD_STARTUP_TIMEOUT)
|
|
|
|
if not is_port_in_use(ray_head_ip, ray_head_port):
|
|
if ray_head_proc.poll() is None:
|
|
# Ray head GCS service is down. Kill ray head node.
|
|
ray_head_proc.terminate()
|
|
# wait killing complete.
|
|
time.sleep(0.5)
|
|
|
|
cmd_exec_failure_msg = gen_cmd_exec_failure_msg(
|
|
ray_head_node_cmd, ray_head_proc.returncode, tail_output_deque
|
|
)
|
|
raise RuntimeError("Start Ray head node failed!\n" + cmd_exec_failure_msg)
|
|
|
|
_logger.info("Ray head node started.")
|
|
|
|
cluster_address = build_address(ray_head_ip, ray_head_port)
|
|
# Set RAY_ADDRESS environment variable to the cluster address.
|
|
os.environ["RAY_ADDRESS"] = cluster_address
|
|
|
|
ray_cluster_handler = RayClusterOnSpark(
|
|
address=cluster_address,
|
|
head_proc=ray_head_proc,
|
|
min_worker_nodes=min_worker_nodes,
|
|
max_worker_nodes=max_worker_nodes,
|
|
temp_dir=ray_temp_dir,
|
|
cluster_unique_id=cluster_unique_id,
|
|
start_hook=start_hook,
|
|
ray_dashboard_port=ray_dashboard_port,
|
|
spark_job_server=spark_job_server,
|
|
global_cluster_lock_fd=global_cluster_lock_fd,
|
|
ray_client_server_port=ray_client_server_port,
|
|
)
|
|
|
|
start_hook.on_cluster_created(ray_cluster_handler)
|
|
|
|
return ray_cluster_handler
|
|
|
|
|
|
_active_ray_cluster = None
|
|
_active_ray_cluster_rwlock = threading.RLock()
|
|
|
|
|
|
def _create_resource_profile(num_cpus_per_node, num_gpus_per_node):
|
|
from pyspark.resource.profile import ResourceProfileBuilder
|
|
from pyspark.resource.requests import TaskResourceRequests
|
|
|
|
task_res_req = TaskResourceRequests().cpus(num_cpus_per_node)
|
|
if num_gpus_per_node > 0:
|
|
task_res_req = task_res_req.resource("gpu", num_gpus_per_node)
|
|
return ResourceProfileBuilder().require(task_res_req).build
|
|
|
|
|
|
# A dict storing blocked key to replacement argument you should use.
|
|
_head_node_option_block_keys = {
|
|
"temp_dir": "ray_temp_root_dir",
|
|
"block": None,
|
|
"head": None,
|
|
"node_ip_address": None,
|
|
"port": None,
|
|
"num_cpus": None,
|
|
"num_gpus": None,
|
|
"dashboard_host": None,
|
|
"dashboard_agent_listen_port": None,
|
|
}
|
|
|
|
_worker_node_option_block_keys = {
|
|
"temp_dir": "ray_temp_root_dir",
|
|
"block": None,
|
|
"head": None,
|
|
"address": None,
|
|
"num_cpus": "num_cpus_worker_node",
|
|
"num_gpus": "num_gpus_worker_node",
|
|
"memory": None,
|
|
"object_store_memory": "object_store_memory_worker_node",
|
|
"dashboard_agent_listen_port": None,
|
|
"min_worker_port": None,
|
|
"max_worker_port": None,
|
|
}
|
|
|
|
|
|
def _verify_node_options(node_options, block_keys, node_type):
|
|
for key in node_options:
|
|
if key.startswith("--") or "-" in key:
|
|
raise ValueError(
|
|
"For a ray node option like '--foo-bar', you should convert it to "
|
|
"following format 'foo_bar' in 'head_node_options' / "
|
|
"'worker_node_options' arguments."
|
|
)
|
|
|
|
if key in block_keys:
|
|
common_err_msg = (
|
|
f"Setting the option '{key}' for {node_type} nodes is not allowed."
|
|
)
|
|
replacement_arg = block_keys[key]
|
|
if replacement_arg:
|
|
raise ValueError(
|
|
f"{common_err_msg} You should set the '{replacement_arg}' option "
|
|
"instead."
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"{common_err_msg} This option is controlled by Ray on Spark."
|
|
)
|
|
|
|
|
|
def _setup_ray_cluster_internal(
|
|
max_worker_nodes: int,
|
|
min_worker_nodes: Optional[int],
|
|
num_cpus_worker_node: Optional[int],
|
|
num_cpus_head_node: Optional[int],
|
|
num_gpus_worker_node: Optional[int],
|
|
num_gpus_head_node: Optional[int],
|
|
heap_memory_worker_node: Optional[int],
|
|
heap_memory_head_node: Optional[int],
|
|
object_store_memory_worker_node: Optional[int],
|
|
object_store_memory_head_node: Optional[int],
|
|
head_node_options: Optional[Dict],
|
|
worker_node_options: Optional[Dict],
|
|
ray_temp_root_dir: Optional[str],
|
|
strict_mode: bool,
|
|
collect_log_to_path: Optional[str],
|
|
autoscale_upscaling_speed: Optional[float],
|
|
autoscale_idle_timeout_minutes: Optional[float],
|
|
is_global: bool,
|
|
**kwargs,
|
|
) -> Tuple[str, str]:
|
|
global _active_ray_cluster
|
|
|
|
_check_system_environment()
|
|
_install_sigterm_signal()
|
|
|
|
head_node_options = head_node_options or {}
|
|
worker_node_options = worker_node_options or {}
|
|
|
|
_verify_node_options(
|
|
head_node_options,
|
|
_head_node_option_block_keys,
|
|
"Ray head node on spark",
|
|
)
|
|
_verify_node_options(
|
|
worker_node_options,
|
|
_worker_node_option_block_keys,
|
|
"Ray worker node on spark",
|
|
)
|
|
|
|
if _active_ray_cluster is not None:
|
|
raise RuntimeError(
|
|
"Current active ray cluster on spark haven't shut down. Please call "
|
|
"`ray.util.spark.shutdown_ray_cluster()` before initiating a new Ray "
|
|
"cluster on spark."
|
|
)
|
|
|
|
if ray.is_initialized():
|
|
raise RuntimeError(
|
|
"Current python process already initialized Ray, Please shut down it "
|
|
"by `ray.shutdown()` before initiating a Ray cluster on spark."
|
|
)
|
|
|
|
spark = get_spark_session()
|
|
|
|
spark_master = spark.sparkContext.master
|
|
|
|
is_spark_local_mode = spark_master == "local" or spark_master.startswith("local[")
|
|
|
|
if not (
|
|
spark_master.startswith("spark://")
|
|
or spark_master.startswith("local-cluster[")
|
|
or spark_master == "yarn"
|
|
or is_spark_local_mode
|
|
):
|
|
raise RuntimeError(
|
|
"Ray on Spark only supports spark cluster in standalone mode, "
|
|
"local-cluster mode, spark on yarn mode or spark local mode."
|
|
)
|
|
|
|
if is_spark_local_mode:
|
|
support_stage_scheduling = False
|
|
elif (
|
|
is_in_databricks_runtime()
|
|
and Version(os.environ["DATABRICKS_RUNTIME_VERSION"]).major >= 12
|
|
):
|
|
support_stage_scheduling = True
|
|
else:
|
|
import pyspark
|
|
|
|
if Version(pyspark.__version__).release >= (3, 4, 0):
|
|
support_stage_scheduling = True
|
|
else:
|
|
support_stage_scheduling = False
|
|
|
|
if "num_cpus_per_node" in kwargs:
|
|
if num_cpus_worker_node is not None:
|
|
raise ValueError(
|
|
"'num_cpus_per_node' and 'num_cpus_worker_node' arguments are "
|
|
"equivalent. Only set 'num_cpus_worker_node'."
|
|
)
|
|
num_cpus_worker_node = kwargs["num_cpus_per_node"]
|
|
warnings.warn(
|
|
"'num_cpus_per_node' argument is deprecated, please use "
|
|
"'num_cpus_worker_node' argument instead.",
|
|
DeprecationWarning,
|
|
)
|
|
|
|
if "num_gpus_per_node" in kwargs:
|
|
if num_gpus_worker_node is not None:
|
|
raise ValueError(
|
|
"'num_gpus_per_node' and 'num_gpus_worker_node' arguments are "
|
|
"equivalent. Only set 'num_gpus_worker_node'."
|
|
)
|
|
num_gpus_worker_node = kwargs["num_gpus_per_node"]
|
|
warnings.warn(
|
|
"'num_gpus_per_node' argument is deprecated, please use "
|
|
"'num_gpus_worker_node' argument instead.",
|
|
DeprecationWarning,
|
|
)
|
|
|
|
if "object_store_memory_per_node" in kwargs:
|
|
if object_store_memory_worker_node is not None:
|
|
raise ValueError(
|
|
"'object_store_memory_per_node' and 'object_store_memory_worker_node' "
|
|
"arguments are equivalent. Only set "
|
|
"'object_store_memory_worker_node'."
|
|
)
|
|
object_store_memory_worker_node = kwargs["object_store_memory_per_node"]
|
|
warnings.warn(
|
|
"'object_store_memory_per_node' argument is deprecated, please use "
|
|
"'object_store_memory_worker_node' argument instead.",
|
|
DeprecationWarning,
|
|
)
|
|
|
|
# Environment configurations within the Spark Session that dictate how many cpus
|
|
# and gpus to use for each submitted spark task.
|
|
num_spark_task_cpus = int(spark.sparkContext.getConf().get("spark.task.cpus", "1"))
|
|
|
|
if num_cpus_worker_node is not None and num_cpus_worker_node <= 0:
|
|
raise ValueError("Argument `num_cpus_worker_node` value must be > 0.")
|
|
|
|
# note: spark.task.resource.gpu.amount config might be fractional value like 0.5
|
|
default_num_spark_task_gpus = float(
|
|
spark.sparkContext.getConf().get("spark.task.resource.gpu.amount", "0")
|
|
)
|
|
rounded_num_spark_task_gpus = int(default_num_spark_task_gpus)
|
|
if default_num_spark_task_gpus > 0:
|
|
warn_msg = (
|
|
"You configured 'spark.task.resource.gpu.amount' to "
|
|
f"{default_num_spark_task_gpus},"
|
|
"we recommend setting this value to 0 so that Spark jobs do not "
|
|
"reserve GPU resources, preventing Ray-on-Spark workloads from having the "
|
|
"maximum number of GPUs available."
|
|
)
|
|
|
|
if is_in_databricks_runtime():
|
|
from ray.util.spark.databricks_hook import (
|
|
get_databricks_display_html_function,
|
|
)
|
|
|
|
get_databricks_display_html_function()(
|
|
f"<b style='color:red;'>{warn_msg}</b>"
|
|
)
|
|
else:
|
|
_logger.warning(warn_msg)
|
|
|
|
if num_gpus_worker_node is not None and num_gpus_worker_node < 0:
|
|
raise ValueError("Argument `num_gpus_worker_node` value must be >= 0.")
|
|
|
|
def _get_spark_worker_resources(_):
|
|
from ray.util.spark.utils import (
|
|
_get_cpu_cores,
|
|
_get_num_physical_gpus,
|
|
_get_spark_worker_total_physical_memory,
|
|
)
|
|
|
|
num_cpus_spark_worker = _get_cpu_cores()
|
|
num_gpus_spark_worker = _get_num_physical_gpus()
|
|
total_mem_bytes = _get_spark_worker_total_physical_memory()
|
|
|
|
return (
|
|
num_cpus_spark_worker,
|
|
num_gpus_spark_worker,
|
|
total_mem_bytes,
|
|
)
|
|
|
|
(num_cpus_spark_worker, num_gpus_spark_worker, spark_worker_mem_bytes,) = (
|
|
spark.sparkContext.parallelize([1], 1)
|
|
.map(_get_spark_worker_resources)
|
|
.collect()[0]
|
|
)
|
|
|
|
if num_cpus_worker_node is not None and num_gpus_worker_node is not None:
|
|
if support_stage_scheduling:
|
|
using_stage_scheduling = True
|
|
res_profile = _create_resource_profile(
|
|
num_cpus_worker_node, num_gpus_worker_node
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
"Current spark version does not support stage scheduling, so that "
|
|
"you cannot set the argument `num_cpus_worker_node` and "
|
|
"`num_gpus_worker_node` values. Without setting the 2 arguments, "
|
|
"per-Ray worker node will be assigned with number of "
|
|
f"'spark.task.cpus' (equals to {num_spark_task_cpus}) cpu cores "
|
|
"and rounded down number of 'spark.task.resource.gpu.amount' "
|
|
f"(equals to {rounded_num_spark_task_gpus}) GPUs. To enable spark "
|
|
f"stage scheduling, you need to upgrade spark to 3.4 version or use "
|
|
"Databricks Runtime 12.x, and you cannot use spark local mode."
|
|
)
|
|
elif num_cpus_worker_node is None and num_gpus_worker_node is None:
|
|
if support_stage_scheduling:
|
|
# Make one Ray worker node using maximum CPU / GPU resources
|
|
# of the whole spark worker node, this is the optimal
|
|
# configuration.
|
|
num_cpus_worker_node = num_cpus_spark_worker
|
|
num_gpus_worker_node = num_gpus_spark_worker
|
|
using_stage_scheduling = True
|
|
res_profile = _create_resource_profile(
|
|
num_cpus_worker_node, num_gpus_worker_node
|
|
)
|
|
else:
|
|
using_stage_scheduling = False
|
|
res_profile = None
|
|
|
|
num_cpus_worker_node = num_spark_task_cpus
|
|
num_gpus_worker_node = rounded_num_spark_task_gpus
|
|
else:
|
|
raise ValueError(
|
|
"'num_cpus_worker_node' and 'num_gpus_worker_node' arguments must be"
|
|
"set together or unset together."
|
|
)
|
|
|
|
(
|
|
ray_worker_node_heap_mem_bytes,
|
|
ray_worker_node_object_store_mem_bytes,
|
|
) = get_avail_mem_per_ray_worker_node(
|
|
spark,
|
|
heap_memory_worker_node,
|
|
object_store_memory_worker_node,
|
|
num_cpus_worker_node,
|
|
num_gpus_worker_node,
|
|
)
|
|
|
|
spark_worker_ray_node_slots = _get_local_ray_node_slots(
|
|
num_cpus_spark_worker,
|
|
num_gpus_spark_worker,
|
|
num_cpus_worker_node,
|
|
num_gpus_worker_node,
|
|
)
|
|
|
|
spark_executor_memory_bytes = get_configured_spark_executor_memory_bytes(spark)
|
|
spark_worker_required_memory_bytes = (
|
|
spark_executor_memory_bytes
|
|
+ spark_worker_ray_node_slots
|
|
* (ray_worker_node_heap_mem_bytes + ray_worker_node_object_store_mem_bytes)
|
|
)
|
|
if spark_worker_required_memory_bytes > 0.8 * spark_worker_mem_bytes:
|
|
warn_msg = (
|
|
"In each spark worker node, we recommend making the sum of "
|
|
"'spark_executor_memory + num_Ray_worker_nodes_per_spark_worker * "
|
|
"(memory_worker_node + object_store_memory_worker_node)' to be less than "
|
|
"'spark_worker_physical_memory * 0.8', otherwise it might lead to "
|
|
"spark worker physical memory exhaustion and Ray task OOM errors."
|
|
)
|
|
|
|
if is_in_databricks_runtime():
|
|
from ray.util.spark.databricks_hook import (
|
|
get_databricks_display_html_function,
|
|
)
|
|
|
|
get_databricks_display_html_function()(
|
|
f"<b style='background-color:Cyan;'>{warn_msg}<br></b>"
|
|
)
|
|
else:
|
|
_logger.warning(warn_msg)
|
|
|
|
if "num_worker_nodes" in kwargs:
|
|
raise ValueError(
|
|
"'num_worker_nodes' argument is removed, please set "
|
|
"'max_worker_nodes' and 'min_worker_nodes' argument instead."
|
|
)
|
|
|
|
if max_worker_nodes == MAX_NUM_WORKER_NODES:
|
|
if min_worker_nodes is not None:
|
|
raise ValueError(
|
|
"If you set 'max_worker_nodes' to 'MAX_NUM_WORKER_NODES', autoscaling "
|
|
"is not supported, so that you cannot set 'min_worker_nodes' argument "
|
|
"and 'min_worker_nodes' is automatically set to be equal to "
|
|
"'max_worker_nodes'."
|
|
)
|
|
|
|
# max_worker_nodes=MAX_NUM_WORKER_NODES represents using all available
|
|
# spark task slots
|
|
max_worker_nodes = get_max_num_concurrent_tasks(spark.sparkContext, res_profile)
|
|
min_worker_nodes = max_worker_nodes
|
|
elif max_worker_nodes <= 0:
|
|
raise ValueError(
|
|
"The value of 'max_worker_nodes' argument must be either a positive "
|
|
"integer or 'ray.util.spark.MAX_NUM_WORKER_NODES'."
|
|
)
|
|
|
|
if "autoscale" in kwargs:
|
|
raise ValueError(
|
|
"'autoscale' argument is removed. You can set 'min_worker_nodes' argument "
|
|
"to be less than 'max_worker_nodes' to make autoscaling enabled."
|
|
)
|
|
|
|
if min_worker_nodes is None:
|
|
min_worker_nodes = max_worker_nodes
|
|
elif not (0 <= min_worker_nodes <= max_worker_nodes):
|
|
raise ValueError(
|
|
"The value of 'max_worker_nodes' argument must be an integer >= 0 "
|
|
"and <= 'max_worker_nodes'"
|
|
)
|
|
|
|
insufficient_resources = []
|
|
|
|
if num_cpus_worker_node < 4:
|
|
insufficient_resources.append(
|
|
"The provided CPU resources for each ray worker are inadequate to start "
|
|
"a ray cluster. Based on the total cpu resources available and the "
|
|
"configured task sizing, each ray worker node would start with "
|
|
f"{num_cpus_worker_node} CPU cores. This is less than the recommended "
|
|
"value of `4` CPUs per worker. On spark version >= 3.4 or Databricks "
|
|
"Runtime 12.x, you can set the argument `num_cpus_worker_node` to "
|
|
"a value >= 4 to address it, otherwise you need to increase the spark "
|
|
"application configuration 'spark.task.cpus' to a minimum of `4` to "
|
|
"address it."
|
|
)
|
|
|
|
if ray_worker_node_heap_mem_bytes < 10 * 1024 * 1024 * 1024:
|
|
insufficient_resources.append(
|
|
"The provided memory resources for each ray worker node are inadequate. "
|
|
"Based on the total memory available on the spark cluster and the "
|
|
"configured task sizing, each ray worker would start with "
|
|
f"{ray_worker_node_heap_mem_bytes} bytes heap memory. This is less than "
|
|
"the recommended value of 10GB. The ray worker node heap memory size is "
|
|
"calculated by "
|
|
"(SPARK_WORKER_PHYSICAL_MEMORY / num_local_spark_task_slots * 0.8) - "
|
|
"object_store_memory_worker_node. To increase the heap space available, "
|
|
"increase the memory in the spark cluster by using instance types with "
|
|
"larger memory, or increase number of CPU/GPU per Ray worker node "
|
|
"(so it leads to less Ray worker node slots per spark worker node), "
|
|
"or apply a lower `object_store_memory_worker_node`."
|
|
)
|
|
if insufficient_resources:
|
|
if strict_mode:
|
|
raise ValueError(
|
|
"You are creating ray cluster on spark with strict mode (it can be "
|
|
"disabled by setting argument 'strict_mode=False' when calling API "
|
|
"'setup_ray_cluster'), strict mode requires the spark cluster config "
|
|
"satisfying following criterion: "
|
|
"\n".join(insufficient_resources)
|
|
)
|
|
else:
|
|
_logger.warning("\n".join(insufficient_resources))
|
|
|
|
if num_cpus_head_node is None:
|
|
if is_global:
|
|
num_cpus_head_node = _get_cpu_cores()
|
|
else:
|
|
num_cpus_head_node = 0
|
|
else:
|
|
if num_cpus_head_node < 0:
|
|
raise ValueError(
|
|
"Argument `num_cpus_head_node` value must be >= 0. "
|
|
f"Current value is {num_cpus_head_node}."
|
|
)
|
|
|
|
if num_gpus_head_node is None:
|
|
if is_global:
|
|
try:
|
|
num_gpus_head_node = _get_num_physical_gpus()
|
|
except Exception:
|
|
num_gpus_head_node = 0
|
|
else:
|
|
num_gpus_head_node = 0
|
|
else:
|
|
if num_gpus_head_node < 0:
|
|
raise ValueError(
|
|
"Argument `num_gpus_head_node` value must be >= 0."
|
|
f"Current value is {num_gpus_head_node}."
|
|
)
|
|
|
|
if (
|
|
num_cpus_head_node == 0
|
|
and num_gpus_head_node == 0
|
|
and object_store_memory_head_node is None
|
|
):
|
|
# Because tasks that require CPU or GPU resources are not scheduled to Ray
|
|
# head node, and user does not set `object_store_memory_head_node` explicitly,
|
|
# limit the heap memory and object store memory allocation to the
|
|
# head node, in order to save spark driver memory.
|
|
heap_memory_head_node = 1024 * 1024 * 1024
|
|
object_store_memory_head_node = 1024 * 1024 * 1024
|
|
else:
|
|
heap_memory_head_node, object_store_memory_head_node = calc_mem_ray_head_node(
|
|
heap_memory_head_node, object_store_memory_head_node
|
|
)
|
|
|
|
with _active_ray_cluster_rwlock:
|
|
cluster = _setup_ray_cluster(
|
|
max_worker_nodes=max_worker_nodes,
|
|
min_worker_nodes=min_worker_nodes,
|
|
num_cpus_worker_node=num_cpus_worker_node,
|
|
num_cpus_head_node=num_cpus_head_node,
|
|
num_gpus_worker_node=num_gpus_worker_node,
|
|
num_gpus_head_node=num_gpus_head_node,
|
|
using_stage_scheduling=using_stage_scheduling,
|
|
heap_memory_worker_node=ray_worker_node_heap_mem_bytes,
|
|
heap_memory_head_node=heap_memory_head_node,
|
|
object_store_memory_worker_node=ray_worker_node_object_store_mem_bytes,
|
|
object_store_memory_head_node=object_store_memory_head_node,
|
|
head_node_options=head_node_options,
|
|
worker_node_options=worker_node_options,
|
|
ray_temp_root_dir=ray_temp_root_dir,
|
|
collect_log_to_path=collect_log_to_path,
|
|
autoscale_upscaling_speed=autoscale_upscaling_speed,
|
|
autoscale_idle_timeout_minutes=autoscale_idle_timeout_minutes,
|
|
is_global=is_global,
|
|
)
|
|
# set global _active_ray_cluster to be the
|
|
# started cluster.
|
|
_active_ray_cluster = cluster
|
|
|
|
try:
|
|
cluster.wait_until_ready() # NB: this line might raise error.
|
|
except Exception as e:
|
|
try:
|
|
shutdown_ray_cluster()
|
|
except Exception:
|
|
pass
|
|
raise RuntimeError("Launch Ray-on-Spark cluster failed") from e
|
|
|
|
head_ip = parse_address(cluster.address)[0]
|
|
remote_connection_address = (
|
|
f"ray://{build_address(head_ip, cluster.ray_client_server_port)}"
|
|
)
|
|
return cluster.address, remote_connection_address
|
|
|
|
|
|
@PublicAPI
|
|
def setup_ray_cluster(
|
|
*,
|
|
max_worker_nodes: int,
|
|
min_worker_nodes: Optional[int] = None,
|
|
num_cpus_worker_node: Optional[int] = None,
|
|
num_cpus_head_node: Optional[int] = None,
|
|
num_gpus_worker_node: Optional[int] = None,
|
|
num_gpus_head_node: Optional[int] = None,
|
|
memory_worker_node: Optional[int] = None,
|
|
memory_head_node: Optional[int] = None,
|
|
object_store_memory_worker_node: Optional[int] = None,
|
|
object_store_memory_head_node: Optional[int] = None,
|
|
head_node_options: Optional[Dict] = None,
|
|
worker_node_options: Optional[Dict] = None,
|
|
ray_temp_root_dir: Optional[str] = None,
|
|
strict_mode: bool = False,
|
|
collect_log_to_path: Optional[str] = None,
|
|
autoscale_upscaling_speed: Optional[float] = 1.0,
|
|
autoscale_idle_timeout_minutes: Optional[float] = 1.0,
|
|
**kwargs,
|
|
) -> Tuple[str, str]:
|
|
"""
|
|
Set up a ray cluster on the spark cluster by starting a ray head node in the
|
|
spark application's driver side node.
|
|
After creating the head node, a background spark job is created that
|
|
generates an instance of `RayClusterOnSpark` that contains configuration for the
|
|
ray cluster that will run on the Spark cluster's worker nodes.
|
|
After a ray cluster is set up, "RAY_ADDRESS" environment variable is set to
|
|
the cluster address, so you can call `ray.init()` without specifying ray cluster
|
|
address to connect to the cluster. To shut down the cluster you can call
|
|
`ray.util.spark.shutdown_ray_cluster()`.
|
|
Note: If the active ray cluster haven't shut down, you cannot create a new ray
|
|
cluster.
|
|
|
|
Args:
|
|
max_worker_nodes: This argument represents maximum ray worker nodes to start
|
|
for the ray cluster. you can
|
|
specify the `max_worker_nodes` as `ray.util.spark.MAX_NUM_WORKER_NODES`
|
|
represents a ray cluster
|
|
configuration that will use all available resources configured for the
|
|
spark application.
|
|
To create a spark application that is intended to exclusively run a
|
|
shared ray cluster in non-scaling, it is recommended to set this argument
|
|
to `ray.util.spark.MAX_NUM_WORKER_NODES`.
|
|
min_worker_nodes: Minimal number of worker nodes (default `None`),
|
|
if "max_worker_nodes" value is equal to "min_worker_nodes" argument,
|
|
or "min_worker_nodes" argument value is None, then autoscaling is disabled
|
|
and Ray cluster is launched with fixed number "max_worker_nodes" of
|
|
Ray worker nodes, otherwise autoscaling is enabled.
|
|
num_cpus_worker_node: Number of cpus available to per-ray worker node, if not
|
|
provided, if spark stage scheduling is supported, 'num_cpus_head_node'
|
|
value equals to number of cpu cores per spark worker node, otherwise
|
|
it uses spark application configuration 'spark.task.cpus' instead.
|
|
**Limitation** Only spark version >= 3.4 or Databricks Runtime 12.x
|
|
supports setting this argument.
|
|
num_cpus_head_node: Number of cpus available to Ray head node, if not provide,
|
|
if it is global mode Ray cluster, use number of cpu cores in spark driver
|
|
node, otherwise use 0 instead.
|
|
use 0 instead. Number 0 means tasks requiring CPU resources are not
|
|
scheduled to Ray head node.
|
|
num_gpus_worker_node: Number of gpus available to per-ray worker node, if not
|
|
provided, if spark stage scheduling is supported, 'num_gpus_worker_node'
|
|
value equals to number of GPUs per spark worker node, otherwise
|
|
it uses rounded down value of spark application configuration
|
|
'spark.task.resource.gpu.amount' instead.
|
|
This argument is only available on spark cluster that is configured with
|
|
'gpu' resources.
|
|
**Limitation** Only spark version >= 3.4 or Databricks Runtime 12.x
|
|
supports setting this argument.
|
|
num_gpus_head_node: Number of gpus available to Ray head node, if not provide,
|
|
if it is global mode Ray cluster, use number of GPUs in spark driver node,
|
|
otherwise use 0 instead.
|
|
This argument is only available on spark cluster which spark driver node
|
|
has GPUs.
|
|
memory_worker_node: Optional[int]:
|
|
Heap memory configured for Ray worker node. This is basically setting
|
|
`--memory` option when starting Ray node by `ray start` command.
|
|
memory_head_node: Optional[int]:
|
|
Heap memory configured for Ray head node. This is basically setting
|
|
`--memory` option when starting Ray node by `ray start` command.
|
|
object_store_memory_worker_node: Object store memory available to per-ray worker
|
|
node, but it is capped by
|
|
"dev_shm_available_size * 0.8 / num_tasks_per_spark_worker".
|
|
The default value equals to
|
|
"0.3 * spark_worker_physical_memory * 0.8 / num_tasks_per_spark_worker".
|
|
object_store_memory_head_node: Object store memory available to Ray head
|
|
node, but it is capped by "dev_shm_available_size * 0.8".
|
|
The default value equals to
|
|
"0.3 * spark_driver_physical_memory * 0.8".
|
|
head_node_options: A dict representing Ray head node extra options, these
|
|
options will be passed to `ray start` script. Note you need to convert
|
|
`ray start` options key from `--foo-bar` format to `foo_bar` format.
|
|
For flag options (e.g. '--disable-usage-stats'), you should set the value
|
|
to None in the option dict, like `{"disable_usage_stats": None}`.
|
|
Note: Short name options (e.g. '-v') are not supported.
|
|
worker_node_options: A dict representing Ray worker node extra options,
|
|
these options will be passed to `ray start` script. Note you need to
|
|
convert `ray start` options key from `--foo-bar` format to `foo_bar`
|
|
format.
|
|
For flag options (e.g. '--disable-usage-stats'), you should set the value
|
|
to None in the option dict, like `{"disable_usage_stats": None}`.
|
|
Note: Short name options (e.g. '-v') are not supported.
|
|
ray_temp_root_dir: A local disk path to store the ray temporary data. The
|
|
created cluster will create a subdirectory
|
|
"ray-{head_port}-{random_suffix}" beneath this path.
|
|
strict_mode: Boolean flag to fast-fail initialization of the ray cluster if
|
|
the available spark cluster does not have sufficient resources to fulfill
|
|
the resource allocation for memory, cpu and gpu. When set to true, if the
|
|
requested resources are not available for recommended minimum recommended
|
|
functionality, an exception will be raised that details the inadequate
|
|
spark cluster configuration settings. If overridden as `False`,
|
|
a warning is raised.
|
|
collect_log_to_path: If specified, after ray head / worker nodes terminated,
|
|
collect their logs to the specified path. On Databricks Runtime, we
|
|
recommend you to specify a local path starts with '/dbfs/', because the
|
|
path mounts with a centralized storage device and stored data is persisted
|
|
after Databricks spark cluster terminated.
|
|
autoscale_upscaling_speed: If autoscale enabled, it represents the number of
|
|
nodes allowed to be pending as a multiple of the current number of nodes.
|
|
The higher the value, the more aggressive upscaling will be. For example,
|
|
if this is set to 1.0, the cluster can grow in size by at most 100% at any
|
|
time, so if the cluster currently has 20 nodes, at most 20 pending launches
|
|
are allowed. The minimum number of pending launches is 5 regardless of
|
|
this setting.
|
|
Default value is 1.0, minimum value is 1.0
|
|
autoscale_idle_timeout_minutes: If autoscale enabled, it represents the number
|
|
of minutes that need to pass before an idle worker node is removed by the
|
|
autoscaler. The smaller the value, the more aggressive downscaling will be.
|
|
Worker nodes are considered idle when they hold no active tasks, actors,
|
|
or referenced objects (either in-memory or spilled to disk). This parameter
|
|
does not affect the head node.
|
|
Default value is 1.0, minimum value is 0
|
|
**kwargs: Additional keyword arguments forwarded to
|
|
``_setup_ray_cluster_internal`` for advanced/experimental options.
|
|
|
|
Returns:
|
|
returns a tuple of (address, remote_connection_address)
|
|
"address" is in format of "<ray_head_node_ip>:<port>"
|
|
"remote_connection_address" is in format of
|
|
"ray://<ray_head_node_ip>:<ray-client-server-port>",
|
|
if your client runs on a machine that also hosts a Ray cluster node locally,
|
|
you can connect to the Ray cluster via ``ray.init(address)``,
|
|
otherwise you can connect to the Ray cluster via
|
|
``ray.init(remote_connection_address)``.
|
|
"""
|
|
|
|
return _setup_ray_cluster_internal(
|
|
max_worker_nodes=max_worker_nodes,
|
|
min_worker_nodes=min_worker_nodes,
|
|
num_cpus_worker_node=num_cpus_worker_node,
|
|
num_cpus_head_node=num_cpus_head_node,
|
|
num_gpus_worker_node=num_gpus_worker_node,
|
|
num_gpus_head_node=num_gpus_head_node,
|
|
heap_memory_worker_node=memory_worker_node,
|
|
heap_memory_head_node=memory_head_node,
|
|
object_store_memory_worker_node=object_store_memory_worker_node,
|
|
object_store_memory_head_node=object_store_memory_head_node,
|
|
head_node_options=head_node_options,
|
|
worker_node_options=worker_node_options,
|
|
ray_temp_root_dir=ray_temp_root_dir,
|
|
strict_mode=strict_mode,
|
|
collect_log_to_path=collect_log_to_path,
|
|
autoscale_upscaling_speed=autoscale_upscaling_speed,
|
|
autoscale_idle_timeout_minutes=autoscale_idle_timeout_minutes,
|
|
is_global=False,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
@PublicAPI
|
|
def setup_global_ray_cluster(
|
|
*,
|
|
max_worker_nodes: int,
|
|
is_blocking: bool = True,
|
|
min_worker_nodes: Optional[int] = None,
|
|
num_cpus_worker_node: Optional[int] = None,
|
|
num_cpus_head_node: Optional[int] = None,
|
|
num_gpus_worker_node: Optional[int] = None,
|
|
num_gpus_head_node: Optional[int] = None,
|
|
memory_worker_node: Optional[int] = None,
|
|
memory_head_node: Optional[int] = None,
|
|
object_store_memory_worker_node: Optional[int] = None,
|
|
object_store_memory_head_node: Optional[int] = None,
|
|
head_node_options: Optional[Dict] = None,
|
|
worker_node_options: Optional[Dict] = None,
|
|
strict_mode: bool = False,
|
|
collect_log_to_path: Optional[str] = None,
|
|
autoscale_upscaling_speed: Optional[float] = 1.0,
|
|
autoscale_idle_timeout_minutes: Optional[float] = 1.0,
|
|
):
|
|
"""
|
|
Set up a global mode cluster.
|
|
The global Ray on spark cluster means:
|
|
- You can only create one active global Ray on spark cluster at a time.
|
|
On databricks cluster, the global Ray cluster can be used by all users,
|
|
- as contrast, non-global Ray cluster can only be used by current notebook
|
|
user.
|
|
- It is up persistently without automatic shutdown.
|
|
- On databricks notebook, you can connect to the global cluster by calling
|
|
``ray.init()`` without specifying its address, it will discover the
|
|
global cluster automatically if it is up.
|
|
|
|
For global mode, the ``ray_temp_root_dir`` argument is not supported.
|
|
Global model Ray cluster always use the default Ray temporary directory
|
|
path.
|
|
|
|
All arguments are the same with ``setup_ray_cluster`` API except that:
|
|
- the ``ray_temp_root_dir`` argument is not supported.
|
|
Global model Ray cluster always use the default Ray temporary directory
|
|
path.
|
|
- A new argument "is_blocking" (default ``True``) is added.
|
|
If "is_blocking" is True,
|
|
then keep the call blocking until it is interrupted.
|
|
once the call is interrupted, the global Ray on spark cluster is shut down and
|
|
`setup_global_ray_cluster` call terminates.
|
|
If "is_blocking" is False,
|
|
once Ray cluster setup completes, return immediately.
|
|
"""
|
|
|
|
cluster_address = _setup_ray_cluster_internal(
|
|
max_worker_nodes=max_worker_nodes,
|
|
min_worker_nodes=min_worker_nodes,
|
|
num_cpus_worker_node=num_cpus_worker_node,
|
|
num_cpus_head_node=num_cpus_head_node,
|
|
num_gpus_worker_node=num_gpus_worker_node,
|
|
num_gpus_head_node=num_gpus_head_node,
|
|
heap_memory_worker_node=memory_worker_node,
|
|
heap_memory_head_node=memory_head_node,
|
|
object_store_memory_worker_node=object_store_memory_worker_node,
|
|
object_store_memory_head_node=object_store_memory_head_node,
|
|
head_node_options=head_node_options,
|
|
worker_node_options=worker_node_options,
|
|
ray_temp_root_dir=None,
|
|
strict_mode=strict_mode,
|
|
collect_log_to_path=collect_log_to_path,
|
|
autoscale_upscaling_speed=autoscale_upscaling_speed,
|
|
autoscale_idle_timeout_minutes=autoscale_idle_timeout_minutes,
|
|
is_global=True,
|
|
)
|
|
|
|
if not is_blocking:
|
|
return cluster_address
|
|
|
|
global _global_ray_cluster_cancel_event
|
|
try:
|
|
_global_ray_cluster_cancel_event = Event()
|
|
# serve forever until user cancel the command.
|
|
_global_ray_cluster_cancel_event.wait()
|
|
finally:
|
|
_global_ray_cluster_cancel_event = None
|
|
# once the program is interrupted,
|
|
# or the corresponding databricks notebook command is interrupted
|
|
# shut down the Ray cluster.
|
|
shutdown_ray_cluster()
|
|
|
|
|
|
def _start_ray_worker_nodes(
|
|
*,
|
|
spark_job_server,
|
|
spark_job_group_id,
|
|
spark_job_group_desc,
|
|
num_worker_nodes,
|
|
using_stage_scheduling,
|
|
ray_head_ip,
|
|
ray_head_port,
|
|
ray_temp_dir,
|
|
num_cpus_per_node,
|
|
num_gpus_per_node,
|
|
heap_memory_per_node,
|
|
object_store_memory_per_node,
|
|
worker_node_options,
|
|
collect_log_to_path,
|
|
node_id,
|
|
):
|
|
# NB:
|
|
# In order to start ray worker nodes on spark cluster worker machines,
|
|
# We launch a background spark job:
|
|
# 1. Each spark task launches one ray worker node. This design ensures all ray
|
|
# worker nodes have the same shape (same cpus / gpus / memory configuration).
|
|
# If ray worker nodes have a non-uniform shape, the Ray cluster setup will
|
|
# be non-deterministic and could create issues with node sizing.
|
|
# 2. A ray worker node is started via the `ray start` CLI. In each spark task,
|
|
# a child process is started and will execute a `ray start ...` command in
|
|
# blocking mode.
|
|
# 3. Each task will acquire a file lock for 10s to ensure that the ray worker
|
|
# init will acquire a port connection to the ray head node that does not
|
|
# contend with other worker processes on the same Spark worker node.
|
|
# 4. When the ray cluster is shutdown, killing ray worker nodes is implemented by
|
|
# `sparkContext.cancelJobGroup` to cancel the background spark job, sending a
|
|
# SIGKILL signal to all spark tasks. Once the spark tasks are killed,
|
|
# `ray_start_node` process detects parent died event then it kills ray
|
|
# worker node.
|
|
spark = spark_job_server.spark
|
|
spark_job_server_port = spark_job_server.server_address[1]
|
|
ray_node_custom_env = spark_job_server.ray_node_custom_env
|
|
|
|
def ray_cluster_job_mapper(_):
|
|
from pyspark.taskcontext import TaskContext
|
|
|
|
_worker_logger = logging.getLogger("ray.util.spark.worker")
|
|
|
|
context = TaskContext.get()
|
|
|
|
(
|
|
worker_port_range_begin,
|
|
worker_port_range_end,
|
|
) = _preallocate_ray_worker_port_range()
|
|
|
|
# 10001 is used as ray client server port of global mode ray cluster.
|
|
ray_worker_node_dashboard_agent_port = get_random_unused_port(
|
|
ray_head_ip, min_port=10002, max_port=20000
|
|
)
|
|
ray_worker_node_cmd = [
|
|
sys.executable,
|
|
"-m",
|
|
"ray.util.spark.start_ray_node",
|
|
f"--num-cpus={num_cpus_per_node}",
|
|
"--block",
|
|
f"--address={build_address(ray_head_ip, ray_head_port)}",
|
|
f"--memory={heap_memory_per_node}",
|
|
f"--object-store-memory={object_store_memory_per_node}",
|
|
f"--min-worker-port={worker_port_range_begin}",
|
|
f"--max-worker-port={worker_port_range_end - 1}",
|
|
f"--dashboard-agent-listen-port={ray_worker_node_dashboard_agent_port}",
|
|
*_convert_ray_node_options(worker_node_options),
|
|
]
|
|
if ray_temp_dir is not None:
|
|
ray_worker_node_cmd.append(f"--temp-dir={ray_temp_dir}")
|
|
|
|
ray_worker_node_extra_envs = {
|
|
RAY_ON_SPARK_COLLECT_LOG_TO_PATH: collect_log_to_path or "",
|
|
RAY_ON_SPARK_START_RAY_PARENT_PID: str(os.getpid()),
|
|
"RAY_ENABLE_WINDOWS_OR_OSX_CLUSTER": "1",
|
|
**ray_node_custom_env,
|
|
}
|
|
|
|
if num_gpus_per_node > 0:
|
|
task_resources = context.resources()
|
|
|
|
if "gpu" not in task_resources:
|
|
raise RuntimeError(
|
|
"Couldn't get the gpu id, Please check the GPU resource "
|
|
"configuration"
|
|
)
|
|
gpu_addr_list = [
|
|
int(addr.strip()) for addr in task_resources["gpu"].addresses
|
|
]
|
|
|
|
available_physical_gpus = get_spark_task_assigned_physical_gpus(
|
|
gpu_addr_list
|
|
)
|
|
ray_worker_node_cmd.append(
|
|
f"--num-gpus={len(available_physical_gpus)}",
|
|
)
|
|
ray_worker_node_extra_envs["CUDA_VISIBLE_DEVICES"] = ",".join(
|
|
[str(gpu_id) for gpu_id in available_physical_gpus]
|
|
)
|
|
|
|
_worker_logger.info(
|
|
f"Start Ray worker, command: {' '.join(ray_worker_node_cmd)}"
|
|
)
|
|
|
|
try:
|
|
is_task_reschedule_failure = False
|
|
# Check node id availability
|
|
response = requests.post(
|
|
url=(
|
|
f"http://{build_address(ray_head_ip, spark_job_server_port)}"
|
|
"/check_node_id_availability"
|
|
),
|
|
json={
|
|
"node_id": node_id,
|
|
"spark_job_group_id": spark_job_group_id,
|
|
},
|
|
)
|
|
if not response.json()["available"]:
|
|
# The case happens when a Ray node is down unexpected
|
|
# caused by spark worker node down and spark tries to
|
|
# reschedule the spark task, so it triggers node
|
|
# creation with duplicated node id.
|
|
# in this case, finish the spark task immediately
|
|
# so spark won't try to reschedule this task
|
|
# and Ray autoscaler will trigger a new node creation
|
|
# with new node id, and a new spark job will be created
|
|
# for holding it.
|
|
is_task_reschedule_failure = True
|
|
raise RuntimeError(
|
|
"Starting Ray worker node twice with the same node id "
|
|
"is not allowed."
|
|
)
|
|
|
|
# Notify job server the task has been launched.
|
|
requests.post(
|
|
url=(
|
|
f"http://{build_address(ray_head_ip, spark_job_server_port)}"
|
|
"/notify_task_launched"
|
|
),
|
|
json={
|
|
"spark_job_group_id": spark_job_group_id,
|
|
},
|
|
)
|
|
|
|
# Note:
|
|
# When a pyspark job cancelled, the UDF python worker process are killed by
|
|
# signal "SIGKILL", then `start_ray_node` process will detect the parent
|
|
# died event (see `ray.util.spark.start_ray_node.check_parent_alive`) and
|
|
# then kill ray worker node process and execute cleanup routine.
|
|
exec_cmd(
|
|
ray_worker_node_cmd,
|
|
synchronous=True,
|
|
extra_env=ray_worker_node_extra_envs,
|
|
)
|
|
except Exception as e:
|
|
# In the following 2 cases, exception is raised:
|
|
# (1)
|
|
# Starting Ray worker node fails, the `e` will contain detail
|
|
# subprocess stdout/stderr output.
|
|
# (2)
|
|
# In autoscaling mode, when Ray worker node is down, autoscaler will
|
|
# try to start new Ray worker node if necessary,
|
|
# and it creates a new spark job to launch Ray worker node process,
|
|
# note the old spark job will reschedule the failed spark task
|
|
# and raise error of "Starting Ray worker node twice with the same
|
|
# node id is not allowed".
|
|
#
|
|
# For either case (1) or case (2),
|
|
# to avoid Spark triggers more spark task retries, we swallow
|
|
# exception here to make spark the task exit normally.
|
|
err_msg = f"Ray worker node process exit, reason: {e}."
|
|
_logger.warning(err_msg)
|
|
|
|
yield err_msg, is_task_reschedule_failure
|
|
|
|
spark.sparkContext.setJobGroup(
|
|
spark_job_group_id,
|
|
spark_job_group_desc,
|
|
)
|
|
|
|
# Starting a normal spark job (not barrier spark job) to run ray worker
|
|
# nodes, the design purpose is:
|
|
# 1. Using normal spark job, spark tasks can automatically retry
|
|
# individually, we don't need to write additional retry logic, But, in
|
|
# barrier mode, if one spark task fails, it will cause all other spark
|
|
# tasks killed.
|
|
# 2. Using normal spark job, we can support failover when a spark worker
|
|
# physical machine crashes. (spark will try to re-schedule the spark task
|
|
# to other spark worker nodes)
|
|
# 3. Using barrier mode job, if the cluster resources does not satisfy
|
|
# "idle spark task slots >= argument num_spark_task", then the barrier
|
|
# job gets stuck and waits until enough idle task slots available, this
|
|
# behavior is not user-friendly, on a shared spark cluster, user is hard
|
|
# to estimate how many idle tasks available at a time, But, if using normal
|
|
# spark job, it can launch job with less spark tasks (i.e. user will see a
|
|
# ray cluster setup with less worker number initially), and when more task
|
|
# slots become available, it continues to launch tasks on new available
|
|
# slots, and user can see the ray cluster worker number increases when more
|
|
# slots available.
|
|
job_rdd = spark.sparkContext.parallelize(
|
|
list(range(num_worker_nodes)), num_worker_nodes
|
|
)
|
|
|
|
if using_stage_scheduling:
|
|
resource_profile = _create_resource_profile(
|
|
num_cpus_per_node,
|
|
num_gpus_per_node,
|
|
)
|
|
job_rdd = job_rdd.withResources(resource_profile)
|
|
|
|
hook_entry = _create_hook_entry(is_global=(ray_temp_dir is None))
|
|
hook_entry.on_spark_job_created(spark_job_group_id)
|
|
|
|
err_msg, is_task_reschedule_failure = job_rdd.mapPartitions(
|
|
ray_cluster_job_mapper
|
|
).collect()[0]
|
|
if not is_task_reschedule_failure:
|
|
spark_job_server.last_worker_error = err_msg
|
|
return err_msg
|
|
|
|
return None
|
|
|
|
|
|
@PublicAPI
|
|
def shutdown_ray_cluster() -> None:
|
|
"""
|
|
Shut down the active ray cluster.
|
|
"""
|
|
global _active_ray_cluster
|
|
|
|
with _active_ray_cluster_rwlock:
|
|
if _active_ray_cluster is None:
|
|
raise RuntimeError("No active ray cluster to shut down.")
|
|
|
|
_active_ray_cluster.shutdown()
|
|
_active_ray_cluster = None
|
|
|
|
|
|
_global_ray_cluster_cancel_event = None
|
|
|
|
|
|
@DeveloperAPI
|
|
class AutoscalingCluster:
|
|
"""Create a ray on spark autoscaling cluster."""
|
|
|
|
def __init__(
|
|
self,
|
|
head_resources: dict,
|
|
worker_node_types: dict,
|
|
extra_provider_config: dict,
|
|
upscaling_speed: float,
|
|
idle_timeout_minutes: float,
|
|
):
|
|
"""Create the cluster.
|
|
|
|
Args:
|
|
head_resources: resources of the head node, including CPU.
|
|
worker_node_types: autoscaler node types config for worker nodes.
|
|
extra_provider_config: extra fields merged into the autoscaler
|
|
``provider`` config.
|
|
upscaling_speed: maximum allowed in-flight upscaling as a
|
|
multiple of the current cluster size.
|
|
idle_timeout_minutes: minutes an idle worker node must remain
|
|
idle before the autoscaler removes it.
|
|
"""
|
|
self._head_resources = head_resources.copy()
|
|
self._head_resources["NODE_ID_AS_RESOURCE"] = HEAD_NODE_ID
|
|
self._config = self._generate_config(
|
|
head_resources,
|
|
worker_node_types,
|
|
extra_provider_config,
|
|
upscaling_speed,
|
|
idle_timeout_minutes,
|
|
)
|
|
|
|
def _generate_config(
|
|
self,
|
|
head_resources,
|
|
worker_node_types,
|
|
extra_provider_config,
|
|
upscaling_speed,
|
|
idle_timeout_minutes,
|
|
):
|
|
base_config = yaml.safe_load(
|
|
open(
|
|
os.path.join(
|
|
os.path.dirname(ray.__file__),
|
|
"autoscaler/spark/defaults.yaml",
|
|
)
|
|
)
|
|
)
|
|
custom_config = copy.deepcopy(base_config)
|
|
custom_config["available_node_types"] = worker_node_types
|
|
custom_config["available_node_types"]["ray.head.default"] = {
|
|
"resources": head_resources,
|
|
"node_config": {},
|
|
"max_workers": 0,
|
|
}
|
|
|
|
custom_config["max_workers"] = sum(
|
|
v["max_workers"] for _, v in worker_node_types.items()
|
|
)
|
|
|
|
custom_config["provider"].update(extra_provider_config)
|
|
|
|
custom_config["upscaling_speed"] = upscaling_speed
|
|
custom_config["idle_timeout_minutes"] = idle_timeout_minutes
|
|
|
|
return custom_config
|
|
|
|
def start(
|
|
self,
|
|
ray_head_ip,
|
|
ray_head_port,
|
|
ray_client_server_port,
|
|
ray_temp_dir,
|
|
dashboard_options,
|
|
head_node_options,
|
|
collect_log_to_path,
|
|
ray_node_custom_env,
|
|
):
|
|
"""Start the cluster.
|
|
|
|
After this call returns, you can connect to the cluster with
|
|
ray.init("auto").
|
|
"""
|
|
from ray.util.spark.cluster_init import (
|
|
RAY_ON_SPARK_COLLECT_LOG_TO_PATH,
|
|
_append_resources_config,
|
|
_convert_ray_node_options,
|
|
)
|
|
|
|
if ray_temp_dir is not None:
|
|
autoscale_config = os.path.join(ray_temp_dir, "autoscaling_config.json")
|
|
else:
|
|
autoscale_config = os.path.join(
|
|
_get_default_ray_tmp_dir(), "autoscaling_config.json"
|
|
)
|
|
with open(autoscale_config, "w") as f:
|
|
f.write(json.dumps(self._config))
|
|
|
|
(
|
|
worker_port_range_begin,
|
|
worker_port_range_end,
|
|
) = _preallocate_ray_worker_port_range()
|
|
|
|
ray_head_node_cmd = [
|
|
sys.executable,
|
|
"-m",
|
|
"ray.util.spark.start_ray_node",
|
|
"--block",
|
|
"--head",
|
|
f"--node-ip-address={ray_head_ip}",
|
|
f"--port={ray_head_port}",
|
|
f"--ray-client-server-port={ray_client_server_port}",
|
|
f"--autoscaling-config={autoscale_config}",
|
|
f"--min-worker-port={worker_port_range_begin}",
|
|
f"--max-worker-port={worker_port_range_end - 1}",
|
|
*dashboard_options,
|
|
]
|
|
|
|
if ray_temp_dir is not None:
|
|
ray_head_node_cmd.append(f"--temp-dir={ray_temp_dir}")
|
|
|
|
if "CPU" in self._head_resources:
|
|
ray_head_node_cmd.append(
|
|
"--num-cpus={}".format(self._head_resources.pop("CPU"))
|
|
)
|
|
if "GPU" in self._head_resources:
|
|
ray_head_node_cmd.append(
|
|
"--num-gpus={}".format(self._head_resources.pop("GPU"))
|
|
)
|
|
if "memory" in self._head_resources:
|
|
ray_head_node_cmd.append(
|
|
"--memory={}".format(self._head_resources.pop("memory"))
|
|
)
|
|
if "object_store_memory" in self._head_resources:
|
|
ray_head_node_cmd.append(
|
|
"--object-store-memory={}".format(
|
|
self._head_resources.pop("object_store_memory")
|
|
)
|
|
)
|
|
|
|
head_node_options = _append_resources_config(
|
|
head_node_options, self._head_resources
|
|
)
|
|
ray_head_node_cmd.extend(_convert_ray_node_options(head_node_options))
|
|
|
|
extra_env = {
|
|
"AUTOSCALER_UPDATE_INTERVAL_S": "1",
|
|
RAY_ON_SPARK_COLLECT_LOG_TO_PATH: collect_log_to_path or "",
|
|
RAY_ON_SPARK_START_RAY_PARENT_PID: str(os.getpid()),
|
|
**ray_node_custom_env,
|
|
}
|
|
|
|
self.ray_head_node_cmd = ray_head_node_cmd
|
|
|
|
return _start_ray_head_node(
|
|
ray_head_node_cmd, synchronous=False, extra_env=extra_env
|
|
)
|
|
|
|
|
|
def _start_ray_head_node(ray_head_node_cmd, synchronous, extra_env):
|
|
def preexec_function():
|
|
# Make `start_ray_node` script and Ray node process run
|
|
# in a separate group,
|
|
# otherwise Ray node will be in the same group of parent process,
|
|
# if parent process is a Jupyter notebook kernel, when user
|
|
# clicks interrupt cell button, SIGINT signal is sent, then Ray node will
|
|
# receive SIGINT signal, and it causes Ray node process dies.
|
|
# `start_ray_node` script should also run in a separate group
|
|
# because on Databricks Runtime, because if Databricks notebook
|
|
# is detached, if the children processes don't exit within 1s,
|
|
# they will receive SIGKILL, this behavior makes start_ray_node
|
|
# doesn't have enough time to complete cleanup work like removing
|
|
# temp directory and collecting logs.
|
|
os.setpgrp()
|
|
|
|
return exec_cmd(
|
|
ray_head_node_cmd,
|
|
synchronous=synchronous,
|
|
extra_env=extra_env,
|
|
preexec_fn=preexec_function,
|
|
)
|
|
|
|
|
|
_sigterm_signal_installed = False
|
|
|
|
|
|
def _install_sigterm_signal():
|
|
global _sigterm_signal_installed
|
|
|
|
if _sigterm_signal_installed:
|
|
return
|
|
|
|
try:
|
|
_origin_sigterm_handler = signal.getsignal(signal.SIGTERM)
|
|
|
|
def _sigterm_handler(signum, frame):
|
|
try:
|
|
shutdown_ray_cluster()
|
|
except Exception:
|
|
# swallow exception to continue executing the following code in the
|
|
# handler
|
|
pass
|
|
signal.signal(
|
|
signal.SIGTERM, _origin_sigterm_handler
|
|
) # Reset to original signal
|
|
os.kill(
|
|
os.getpid(), signal.SIGTERM
|
|
) # Re-raise the signal to trigger original behavior
|
|
|
|
signal.signal(signal.SIGTERM, _sigterm_handler)
|
|
_sigterm_signal_installed = True
|
|
except Exception:
|
|
_logger.warning("Install Ray-on-Spark SIGTERM handler failed.")
|