chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,528 @@
|
||||
import collections
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
|
||||
from ray._common.network_utils import is_ipv6
|
||||
|
||||
_logger = logging.getLogger("ray.util.spark.utils")
|
||||
|
||||
|
||||
def is_in_databricks_runtime():
|
||||
return "DATABRICKS_RUNTIME_VERSION" in os.environ
|
||||
|
||||
|
||||
def gen_cmd_exec_failure_msg(cmd, return_code, tail_output_deque):
|
||||
cmd_str = " ".join(cmd)
|
||||
tail_output = "".join(tail_output_deque)
|
||||
return (
|
||||
f"Command {cmd_str} failed with return code {return_code}, tail output are "
|
||||
f"included below.\n{tail_output}\n"
|
||||
)
|
||||
|
||||
|
||||
def get_configured_spark_executor_memory_bytes(spark):
|
||||
value_str = spark.conf.get("spark.executor.memory", "1g").lower()
|
||||
value_num = int(value_str[:-1])
|
||||
value_unit = value_str[-1]
|
||||
unit_map = {
|
||||
"k": 1024,
|
||||
"m": 1024 * 1024,
|
||||
"g": 1024 * 1024 * 1024,
|
||||
"t": 1024 * 1024 * 1024 * 1024,
|
||||
}
|
||||
return value_num * unit_map[value_unit]
|
||||
|
||||
|
||||
def exec_cmd(
|
||||
cmd,
|
||||
*,
|
||||
extra_env=None,
|
||||
synchronous=True,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
A convenience wrapper of `subprocess.Popen` for running a command from a Python
|
||||
script.
|
||||
If `synchronous` is True, wait until the process terminated and if subprocess
|
||||
return code is not 0, raise error containing last 100 lines output.
|
||||
If `synchronous` is False, return an `Popen` instance and a deque instance holding
|
||||
tail outputs.
|
||||
The subprocess stdout / stderr output will be streamly redirected to current
|
||||
process stdout.
|
||||
"""
|
||||
illegal_kwargs = set(kwargs.keys()).intersection({"text", "stdout", "stderr"})
|
||||
if illegal_kwargs:
|
||||
raise ValueError(f"`kwargs` cannot contain {list(illegal_kwargs)}")
|
||||
|
||||
env = kwargs.pop("env", None)
|
||||
if extra_env is not None and env is not None:
|
||||
raise ValueError("`extra_env` and `env` cannot be used at the same time")
|
||||
|
||||
env = env if extra_env is None else {**os.environ, **extra_env}
|
||||
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
env=env,
|
||||
text=True,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
tail_output_deque = collections.deque(maxlen=100)
|
||||
|
||||
def redirect_log_thread_fn():
|
||||
for line in process.stdout:
|
||||
# collect tail logs by `tail_output_deque`
|
||||
tail_output_deque.append(line)
|
||||
|
||||
# redirect to stdout.
|
||||
sys.stdout.write(line)
|
||||
|
||||
threading.Thread(target=redirect_log_thread_fn, args=()).start()
|
||||
|
||||
if not synchronous:
|
||||
return process, tail_output_deque
|
||||
|
||||
return_code = process.wait()
|
||||
if return_code != 0:
|
||||
raise RuntimeError(
|
||||
gen_cmd_exec_failure_msg(cmd, return_code, tail_output_deque)
|
||||
)
|
||||
|
||||
|
||||
def is_port_in_use(host, port):
|
||||
import socket
|
||||
from contextlib import closing
|
||||
|
||||
with closing(
|
||||
socket.socket(
|
||||
socket.AF_INET6 if is_ipv6(host) else socket.AF_INET, socket.SOCK_STREAM
|
||||
)
|
||||
) as sock:
|
||||
return sock.connect_ex((host, port)) == 0
|
||||
|
||||
|
||||
def _wait_service_up(host, port, timeout):
|
||||
beg_time = time.time()
|
||||
|
||||
while time.time() - beg_time < timeout:
|
||||
if is_port_in_use(host, port):
|
||||
return True
|
||||
time.sleep(1)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def get_random_unused_port(
|
||||
host, min_port=1024, max_port=65535, max_retries=100, exclude_list=None
|
||||
):
|
||||
"""
|
||||
Get random unused port.
|
||||
"""
|
||||
# Use true random generator
|
||||
rng = random.SystemRandom()
|
||||
|
||||
exclude_list = exclude_list or []
|
||||
for _ in range(max_retries):
|
||||
port = rng.randint(min_port, max_port)
|
||||
if port in exclude_list:
|
||||
continue
|
||||
if not is_port_in_use(host, port):
|
||||
return port
|
||||
raise RuntimeError(
|
||||
f"Get available port between range {min_port} and {max_port} failed."
|
||||
)
|
||||
|
||||
|
||||
def get_spark_session():
|
||||
from pyspark.sql import SparkSession
|
||||
|
||||
spark_session = SparkSession.getActiveSession()
|
||||
if spark_session is None:
|
||||
raise RuntimeError(
|
||||
"Spark session haven't been initiated yet. Please use "
|
||||
"`SparkSession.builder` to create a spark session and connect to a spark "
|
||||
"cluster."
|
||||
)
|
||||
return spark_session
|
||||
|
||||
|
||||
def get_spark_application_driver_host(spark):
|
||||
return spark.conf.get("spark.driver.host")
|
||||
|
||||
|
||||
def get_max_num_concurrent_tasks(spark_context, resource_profile):
|
||||
"""Gets the current max number of concurrent tasks."""
|
||||
# pylint: disable=protected-access=
|
||||
ssc = spark_context._jsc.sc()
|
||||
if resource_profile is not None:
|
||||
|
||||
def dummpy_mapper(_):
|
||||
pass
|
||||
|
||||
# Runs a dummy spark job to register the `res_profile`
|
||||
spark_context.parallelize([1], 1).withResources(resource_profile).map(
|
||||
dummpy_mapper
|
||||
).collect()
|
||||
|
||||
return ssc.maxNumConcurrentTasks(resource_profile._java_resource_profile)
|
||||
else:
|
||||
return ssc.maxNumConcurrentTasks(
|
||||
ssc.resourceProfileManager().defaultResourceProfile()
|
||||
)
|
||||
|
||||
|
||||
def _get_spark_worker_total_physical_memory():
|
||||
import psutil
|
||||
|
||||
if RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES in os.environ:
|
||||
return int(os.environ[RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES])
|
||||
return psutil.virtual_memory().total
|
||||
|
||||
|
||||
def _get_spark_worker_total_shared_memory():
|
||||
import shutil
|
||||
|
||||
if RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES in os.environ:
|
||||
return int(os.environ[RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES])
|
||||
|
||||
return shutil.disk_usage("/dev/shm").total
|
||||
|
||||
|
||||
# The maximum proportion for Ray worker node object store memory size
|
||||
_RAY_ON_SPARK_MAX_OBJECT_STORE_MEMORY_PROPORTION = 0.8
|
||||
|
||||
# The buffer offset for calculating Ray node memory.
|
||||
_RAY_ON_SPARK_NODE_MEMORY_BUFFER_OFFSET = 0.8
|
||||
|
||||
|
||||
def calc_mem_ray_head_node(configured_heap_memory_bytes, configured_object_store_bytes):
|
||||
import shutil
|
||||
|
||||
import psutil
|
||||
|
||||
if RAY_ON_SPARK_DRIVER_PHYSICAL_MEMORY_BYTES in os.environ:
|
||||
available_physical_mem = int(
|
||||
os.environ[RAY_ON_SPARK_DRIVER_PHYSICAL_MEMORY_BYTES]
|
||||
)
|
||||
else:
|
||||
available_physical_mem = psutil.virtual_memory().total
|
||||
|
||||
available_physical_mem = (
|
||||
available_physical_mem * _RAY_ON_SPARK_NODE_MEMORY_BUFFER_OFFSET
|
||||
)
|
||||
|
||||
if RAY_ON_SPARK_DRIVER_SHARED_MEMORY_BYTES in os.environ:
|
||||
available_shared_mem = int(os.environ[RAY_ON_SPARK_DRIVER_SHARED_MEMORY_BYTES])
|
||||
else:
|
||||
available_shared_mem = shutil.disk_usage("/dev/shm").total
|
||||
|
||||
available_shared_mem = (
|
||||
available_shared_mem * _RAY_ON_SPARK_NODE_MEMORY_BUFFER_OFFSET
|
||||
)
|
||||
|
||||
heap_mem_bytes, object_store_bytes, warning_msg = _calc_mem_per_ray_node(
|
||||
available_physical_mem,
|
||||
available_shared_mem,
|
||||
configured_heap_memory_bytes,
|
||||
configured_object_store_bytes,
|
||||
)
|
||||
|
||||
if warning_msg is not None:
|
||||
_logger.warning(warning_msg)
|
||||
|
||||
return heap_mem_bytes, object_store_bytes
|
||||
|
||||
|
||||
def _calc_mem_per_ray_worker_node(
|
||||
num_task_slots,
|
||||
physical_mem_bytes,
|
||||
shared_mem_bytes,
|
||||
configured_heap_memory_bytes,
|
||||
configured_object_store_bytes,
|
||||
):
|
||||
available_physical_mem_per_node = int(
|
||||
physical_mem_bytes / num_task_slots * _RAY_ON_SPARK_NODE_MEMORY_BUFFER_OFFSET
|
||||
)
|
||||
available_shared_mem_per_node = int(
|
||||
shared_mem_bytes / num_task_slots * _RAY_ON_SPARK_NODE_MEMORY_BUFFER_OFFSET
|
||||
)
|
||||
return _calc_mem_per_ray_node(
|
||||
available_physical_mem_per_node,
|
||||
available_shared_mem_per_node,
|
||||
configured_heap_memory_bytes,
|
||||
configured_object_store_bytes,
|
||||
)
|
||||
|
||||
|
||||
def _calc_mem_per_ray_node(
|
||||
available_physical_mem_per_node,
|
||||
available_shared_mem_per_node,
|
||||
configured_heap_memory_bytes,
|
||||
configured_object_store_bytes,
|
||||
):
|
||||
from ray._private.ray_constants import (
|
||||
DEFAULT_OBJECT_STORE_MEMORY_PROPORTION,
|
||||
OBJECT_STORE_MINIMUM_MEMORY_BYTES,
|
||||
)
|
||||
|
||||
warning_msg = None
|
||||
|
||||
object_store_bytes = configured_object_store_bytes or (
|
||||
available_physical_mem_per_node * DEFAULT_OBJECT_STORE_MEMORY_PROPORTION
|
||||
)
|
||||
|
||||
# If allow Ray using slow storage oas object store,
|
||||
# we don't need to cap object store size by /dev/shm capacity
|
||||
if not os.environ.get("RAY_OBJECT_STORE_ALLOW_SLOW_STORAGE"):
|
||||
if object_store_bytes > available_shared_mem_per_node:
|
||||
object_store_bytes = available_shared_mem_per_node
|
||||
|
||||
object_store_bytes_upper_bound = (
|
||||
available_physical_mem_per_node
|
||||
* _RAY_ON_SPARK_MAX_OBJECT_STORE_MEMORY_PROPORTION
|
||||
)
|
||||
|
||||
if object_store_bytes > object_store_bytes_upper_bound:
|
||||
object_store_bytes = object_store_bytes_upper_bound
|
||||
warning_msg = (
|
||||
"Your configured `object_store_memory_per_node` value "
|
||||
"is too high and it is capped by 80% of per-Ray node "
|
||||
"allocated memory."
|
||||
)
|
||||
|
||||
if object_store_bytes < OBJECT_STORE_MINIMUM_MEMORY_BYTES:
|
||||
if object_store_bytes == available_shared_mem_per_node:
|
||||
warning_msg = (
|
||||
"Your operating system is configured with too small /dev/shm "
|
||||
"size, so `object_store_memory_worker_node` value is configured "
|
||||
f"to minimal size ({OBJECT_STORE_MINIMUM_MEMORY_BYTES} bytes),"
|
||||
f"Please increase system /dev/shm size."
|
||||
)
|
||||
else:
|
||||
warning_msg = (
|
||||
"You configured too small Ray node object store memory size, "
|
||||
"so `object_store_memory_worker_node` value is configured "
|
||||
f"to minimal size ({OBJECT_STORE_MINIMUM_MEMORY_BYTES} bytes),"
|
||||
"Please increase 'object_store_memory_worker_node' argument value."
|
||||
)
|
||||
|
||||
object_store_bytes = OBJECT_STORE_MINIMUM_MEMORY_BYTES
|
||||
|
||||
object_store_bytes = int(object_store_bytes)
|
||||
|
||||
if configured_heap_memory_bytes is None:
|
||||
heap_mem_bytes = int(available_physical_mem_per_node - object_store_bytes)
|
||||
else:
|
||||
heap_mem_bytes = int(configured_heap_memory_bytes)
|
||||
|
||||
return heap_mem_bytes, object_store_bytes, warning_msg
|
||||
|
||||
|
||||
# User can manually set these environment variables
|
||||
# if ray on spark code accessing corresponding information failed.
|
||||
# Note these environment variables must be set in spark executor side,
|
||||
# you should set them via setting spark config of
|
||||
# `spark.executorEnv.[EnvironmentVariableName]`
|
||||
RAY_ON_SPARK_WORKER_CPU_CORES = "RAY_ON_SPARK_WORKER_CPU_CORES"
|
||||
RAY_ON_SPARK_WORKER_GPU_NUM = "RAY_ON_SPARK_WORKER_GPU_NUM"
|
||||
RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES = "RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES"
|
||||
RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES = "RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES"
|
||||
|
||||
# User can manually set these environment variables on spark driver node
|
||||
# if ray on spark code accessing corresponding information failed.
|
||||
RAY_ON_SPARK_DRIVER_PHYSICAL_MEMORY_BYTES = "RAY_ON_SPARK_DRIVER_PHYSICAL_MEMORY_BYTES"
|
||||
RAY_ON_SPARK_DRIVER_SHARED_MEMORY_BYTES = "RAY_ON_SPARK_DRIVER_SHARED_MEMORY_BYTES"
|
||||
|
||||
|
||||
def _get_cpu_cores():
|
||||
import multiprocessing
|
||||
|
||||
if RAY_ON_SPARK_WORKER_CPU_CORES in os.environ:
|
||||
# In some cases, spark standalone cluster might configure virtual cpu cores
|
||||
# for spark worker that different with number of physical cpu cores,
|
||||
# but we cannot easily get the virtual cpu cores configured for spark
|
||||
# worker, as a workaround, we provide an environmental variable config
|
||||
# `RAY_ON_SPARK_WORKER_CPU_CORES` for user.
|
||||
return int(os.environ[RAY_ON_SPARK_WORKER_CPU_CORES])
|
||||
|
||||
return multiprocessing.cpu_count()
|
||||
|
||||
|
||||
def _get_num_physical_gpus():
|
||||
if RAY_ON_SPARK_WORKER_GPU_NUM in os.environ:
|
||||
# In some cases, spark standalone cluster might configure part of physical
|
||||
# GPUs for spark worker,
|
||||
# but we cannot easily get related configuration,
|
||||
# as a workaround, we provide an environmental variable config
|
||||
# `RAY_ON_SPARK_WORKER_CPU_CORES` for user.
|
||||
return int(os.environ[RAY_ON_SPARK_WORKER_GPU_NUM])
|
||||
|
||||
if shutil.which("nvidia-smi") is None:
|
||||
# GPU driver is not installed.
|
||||
return 0
|
||||
try:
|
||||
completed_proc = subprocess.run(
|
||||
"nvidia-smi --query-gpu=name --format=csv,noheader",
|
||||
shell=True,
|
||||
check=True,
|
||||
text=True,
|
||||
capture_output=True,
|
||||
)
|
||||
return len(completed_proc.stdout.strip().split("\n"))
|
||||
except Exception as e:
|
||||
_logger.info(
|
||||
"'nvidia-smi --query-gpu=name --format=csv,noheader' command execution "
|
||||
f"failed, error: {repr(e)}"
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def _get_local_ray_node_slots(
|
||||
num_cpus,
|
||||
num_gpus,
|
||||
num_cpus_per_node,
|
||||
num_gpus_per_node,
|
||||
):
|
||||
if num_cpus_per_node > num_cpus:
|
||||
raise ValueError(
|
||||
"cpu number per Ray worker node should be <= spark worker node CPU cores, "
|
||||
f"you set cpu number per Ray worker node to {num_cpus_per_node} but "
|
||||
f"spark worker node CPU core number is {num_cpus}."
|
||||
)
|
||||
num_ray_node_slots = num_cpus // num_cpus_per_node
|
||||
|
||||
if num_gpus_per_node > 0:
|
||||
if num_gpus_per_node > num_gpus:
|
||||
raise ValueError(
|
||||
"gpu number per Ray worker node should be <= spark worker node "
|
||||
"GPU number, you set GPU devices number per Ray worker node to "
|
||||
f"{num_gpus_per_node} but spark worker node GPU devices number "
|
||||
f"is {num_gpus}."
|
||||
)
|
||||
if num_ray_node_slots > num_gpus // num_gpus_per_node:
|
||||
num_ray_node_slots = num_gpus // num_gpus_per_node
|
||||
|
||||
return num_ray_node_slots
|
||||
|
||||
|
||||
def _get_avail_mem_per_ray_worker_node(
|
||||
num_cpus_per_node,
|
||||
num_gpus_per_node,
|
||||
heap_memory_per_node,
|
||||
object_store_memory_per_node,
|
||||
):
|
||||
"""
|
||||
Returns tuple of (
|
||||
ray_worker_node_heap_mem_bytes,
|
||||
ray_worker_node_object_store_bytes,
|
||||
error_message, # always None
|
||||
warning_message,
|
||||
)
|
||||
"""
|
||||
num_cpus = _get_cpu_cores()
|
||||
if num_gpus_per_node > 0:
|
||||
num_gpus = _get_num_physical_gpus()
|
||||
else:
|
||||
num_gpus = 0
|
||||
|
||||
num_ray_node_slots = _get_local_ray_node_slots(
|
||||
num_cpus, num_gpus, num_cpus_per_node, num_gpus_per_node
|
||||
)
|
||||
|
||||
physical_mem_bytes = _get_spark_worker_total_physical_memory()
|
||||
shared_mem_bytes = _get_spark_worker_total_shared_memory()
|
||||
|
||||
(
|
||||
ray_worker_node_heap_mem_bytes,
|
||||
ray_worker_node_object_store_bytes,
|
||||
warning_msg,
|
||||
) = _calc_mem_per_ray_worker_node(
|
||||
num_ray_node_slots,
|
||||
physical_mem_bytes,
|
||||
shared_mem_bytes,
|
||||
heap_memory_per_node,
|
||||
object_store_memory_per_node,
|
||||
)
|
||||
return (
|
||||
ray_worker_node_heap_mem_bytes,
|
||||
ray_worker_node_object_store_bytes,
|
||||
None,
|
||||
warning_msg,
|
||||
)
|
||||
|
||||
|
||||
def get_avail_mem_per_ray_worker_node(
|
||||
spark,
|
||||
heap_memory_per_node,
|
||||
object_store_memory_per_node,
|
||||
num_cpus_per_node,
|
||||
num_gpus_per_node,
|
||||
):
|
||||
"""
|
||||
Return the available heap memory and object store memory for each ray worker,
|
||||
and error / warning message if it has.
|
||||
Return value is a tuple of
|
||||
(ray_worker_node_heap_mem_bytes, ray_worker_node_object_store_bytes,
|
||||
error_message, warning_message)
|
||||
NB: We have one ray node per spark task.
|
||||
"""
|
||||
|
||||
def mapper(_):
|
||||
try:
|
||||
return _get_avail_mem_per_ray_worker_node(
|
||||
num_cpus_per_node,
|
||||
num_gpus_per_node,
|
||||
heap_memory_per_node,
|
||||
object_store_memory_per_node,
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
trace_msg = "\n".join(traceback.format_tb(e.__traceback__))
|
||||
return -1, -1, repr(e) + trace_msg, None
|
||||
|
||||
# Running memory inference routine on spark executor side since the spark worker
|
||||
# nodes may have a different machine configuration compared to the spark driver
|
||||
# node.
|
||||
(
|
||||
inferred_ray_worker_node_heap_mem_bytes,
|
||||
inferred_ray_worker_node_object_store_bytes,
|
||||
err,
|
||||
warning_msg,
|
||||
) = (
|
||||
spark.sparkContext.parallelize([1], 1).map(mapper).collect()[0]
|
||||
)
|
||||
|
||||
if err is not None:
|
||||
raise RuntimeError(
|
||||
f"Inferring ray worker node available memory failed, error: {err}. "
|
||||
"You can bypass this error by setting following spark configs: "
|
||||
"spark.executorEnv.RAY_ON_SPARK_WORKER_CPU_CORES, "
|
||||
"spark.executorEnv.RAY_ON_SPARK_WORKER_GPU_NUM, "
|
||||
"spark.executorEnv.RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES, "
|
||||
"spark.executorEnv.RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES."
|
||||
)
|
||||
if warning_msg is not None:
|
||||
_logger.warning(warning_msg)
|
||||
return (
|
||||
inferred_ray_worker_node_heap_mem_bytes,
|
||||
inferred_ray_worker_node_object_store_bytes,
|
||||
)
|
||||
|
||||
|
||||
def get_spark_task_assigned_physical_gpus(gpu_addr_list):
|
||||
if "CUDA_VISIBLE_DEVICES" in os.environ:
|
||||
visible_cuda_dev_list = [
|
||||
int(dev.strip()) for dev in os.environ["CUDA_VISIBLE_DEVICES"].split(",")
|
||||
]
|
||||
return [visible_cuda_dev_list[addr] for addr in gpu_addr_list]
|
||||
else:
|
||||
return gpu_addr_list
|
||||
Reference in New Issue
Block a user