chore: import upstream snapshot with attribution

python/ray/util/spark/__init__.py

@@ -0,0 +1,13 @@
+from ray.util.spark.cluster_init import (
+    MAX_NUM_WORKER_NODES,
+    setup_global_ray_cluster,
+    setup_ray_cluster,
+    shutdown_ray_cluster,
+)
+
+__all__ = [
+    "setup_ray_cluster",
+    "shutdown_ray_cluster",
+    "MAX_NUM_WORKER_NODES",
+    "setup_global_ray_cluster",
+]

python/ray/util/spark/databricks_hook.py

@@ -0,0 +1,221 @@
+import logging
+import os
+import threading
+import time
+
+from .start_hook_base import RayOnSparkStartHook
+from .utils import get_spark_session
+
+_logger = logging.getLogger(__name__)
+
+DATABRICKS_HOST = "DATABRICKS_HOST"
+DATABRICKS_TOKEN = "DATABRICKS_TOKEN"
+DATABRICKS_CLIENT_ID = "DATABRICKS_CLIENT_ID"
+DATABRICKS_CLIENT_SECRET = "DATABRICKS_CLIENT_SECRET"
+
+
+def verify_databricks_auth_env():
+    return (DATABRICKS_HOST in os.environ and DATABRICKS_TOKEN in os.environ) or (
+        DATABRICKS_HOST in os.environ
+        and DATABRICKS_CLIENT_ID in os.environ
+        and DATABRICKS_CLIENT_SECRET in os.environ
+    )
+
+
+def get_databricks_function(func_name):
+    import IPython
+
+    ip_shell = IPython.get_ipython()
+    if ip_shell is None:
+        raise RuntimeError("No IPython environment.")
+    return ip_shell.ns_table["user_global"][func_name]
+
+
+def get_databricks_display_html_function():
+    return get_databricks_function("displayHTML")
+
+
+def get_db_entry_point():
+    """
+    Return databricks entry_point instance, it is for calling some
+    internal API in databricks runtime
+    """
+    from dbruntime import UserNamespaceInitializer
+
+    user_namespace_initializer = UserNamespaceInitializer.getOrCreate()
+    return user_namespace_initializer.get_spark_entry_point()
+
+
+def display_databricks_driver_proxy_url(spark_context, port, title):
+    """
+    This helper function create a proxy URL for databricks driver webapp forwarding.
+    In databricks runtime, user does not have permission to directly access web
+    service binding on driver machine port, but user can visit it by a proxy URL with
+    following format: "/driver-proxy/o/{orgId}/{clusterId}/{port}/".
+    """
+    driverLocal = spark_context._jvm.com.databricks.backend.daemon.driver.DriverLocal
+    commandContextTags = driverLocal.commandContext().get().toStringMap().apply("tags")
+    orgId = commandContextTags.apply("orgId")
+    clusterId = commandContextTags.apply("clusterId")
+
+    proxy_link = f"/driver-proxy/o/{orgId}/{clusterId}/{port}/"
+    proxy_url = f"https://dbc-dp-{orgId}.cloud.databricks.com{proxy_link}"
+
+    print("To monitor and debug Ray from Databricks, view the dashboard at ")
+    print(f" {proxy_url}")
+
+    get_databricks_display_html_function()(
+        f"""
+      <div style="margin-top: 16px;margin-bottom: 16px">
+          <a href="{proxy_link}">
+              Open {title} in a new tab
+          </a>
+      </div>
+    """
+    )
+
+
+DATABRICKS_AUTO_SHUTDOWN_POLL_INTERVAL_SECONDS = 3
+DATABRICKS_RAY_ON_SPARK_AUTOSHUTDOWN_MINUTES = (
+    "DATABRICKS_RAY_ON_SPARK_AUTOSHUTDOWN_MINUTES"
+)
+
+
+_DATABRICKS_DEFAULT_TMP_ROOT_DIR = "/local_disk0/tmp"
+
+
+class DefaultDatabricksRayOnSparkStartHook(RayOnSparkStartHook):
+    def get_default_temp_root_dir(self):
+        return _DATABRICKS_DEFAULT_TMP_ROOT_DIR
+
+    def on_ray_dashboard_created(self, port):
+        display_databricks_driver_proxy_url(
+            get_spark_session().sparkContext, port, "Ray Cluster Dashboard"
+        )
+
+    def on_cluster_created(self, ray_cluster_handler):
+        db_api_entry = get_db_entry_point()
+
+        if self.is_global:
+            # Disable auto shutdown if
+            # 1) autoscaling enabled
+            #  because in autoscaling mode, background spark job will be killed
+            #  automatically when ray cluster is idle.
+            # 2) global mode cluster
+            #  Because global mode cluster is designed to keep running until
+            #  user request to shut down it, and global mode cluster is shared
+            #  by other users, the code here cannot track usage from other users
+            #  so that we don't know whether it is safe to shut down the global
+            #  cluster automatically.
+            auto_shutdown_minutes = 0
+        else:
+            auto_shutdown_minutes = float(
+                os.environ.get(DATABRICKS_RAY_ON_SPARK_AUTOSHUTDOWN_MINUTES, "30")
+            )
+        if auto_shutdown_minutes == 0:
+            _logger.info(
+                "The Ray cluster will keep running until you manually detach the "
+                "Databricks notebook or call "
+                "`ray.util.spark.shutdown_ray_cluster()`."
+            )
+            return
+        if auto_shutdown_minutes < 0:
+            raise ValueError(
+                "You must set "
+                f"'{DATABRICKS_RAY_ON_SPARK_AUTOSHUTDOWN_MINUTES}' "
+                "to a value >= 0."
+            )
+
+        try:
+            db_api_entry.getIdleTimeMillisSinceLastNotebookExecution()
+        except Exception:
+            _logger.warning(
+                "Failed to retrieve idle time since last notebook execution, "
+                "so that we cannot automatically shut down Ray cluster when "
+                "Databricks notebook is inactive for the specified minutes. "
+                "You need to manually detach Databricks notebook "
+                "or call `ray.util.spark.shutdown_ray_cluster()` to shut down "
+                "Ray cluster on spark."
+            )
+            return
+
+        _logger.info(
+            "The Ray cluster will be shut down automatically if you don't run "
+            "commands on the Databricks notebook for "
+            f"{auto_shutdown_minutes} minutes. You can change the "
+            "auto-shutdown minutes by setting "
+            f"'{DATABRICKS_RAY_ON_SPARK_AUTOSHUTDOWN_MINUTES}' environment "
+            "variable, setting it to 0 means that the Ray cluster keeps running "
+            "until you manually call `ray.util.spark.shutdown_ray_cluster()` or "
+            "detach Databricks notebook."
+        )
+
+        def auto_shutdown_watcher():
+            auto_shutdown_millis = auto_shutdown_minutes * 60 * 1000
+            while True:
+                if ray_cluster_handler.is_shutdown:
+                    # The cluster is shut down. The watcher thread exits.
+                    return
+
+                idle_time = db_api_entry.getIdleTimeMillisSinceLastNotebookExecution()
+
+                if idle_time > auto_shutdown_millis:
+                    from ray.util.spark import cluster_init
+
+                    with cluster_init._active_ray_cluster_rwlock:
+                        if ray_cluster_handler is cluster_init._active_ray_cluster:
+                            cluster_init.shutdown_ray_cluster()
+                    return
+
+                time.sleep(DATABRICKS_AUTO_SHUTDOWN_POLL_INTERVAL_SECONDS)
+
+        threading.Thread(target=auto_shutdown_watcher, daemon=True).start()
+
+    def on_spark_job_created(self, job_group_id):
+        db_api_entry = get_db_entry_point()
+        db_api_entry.registerBackgroundSparkJobGroup(job_group_id)
+
+    def custom_environment_variables(self):
+        conf = {
+            **super().custom_environment_variables(),
+            # Hardcode `GLOO_SOCKET_IFNAME` to `eth0` for Databricks runtime.
+            # Torch on DBR does not reliably detect the correct interface to use,
+            # and ends up selecting the loopback interface, breaking cross-node
+            # commnication.
+            "GLOO_SOCKET_IFNAME": "eth0",
+            # 'DISABLE_MLFLOW_INTEGRATION' is the environmental variable to disable
+            # huggingface transformers MLflow integration,
+            # it doesn't work well in Databricks runtime,
+            # So disable it by default.
+            "DISABLE_MLFLOW_INTEGRATION": "TRUE",
+        }
+
+        if verify_databricks_auth_env():
+            conf[DATABRICKS_HOST] = os.environ[DATABRICKS_HOST]
+            if DATABRICKS_TOKEN in os.environ:
+                # PAT auth
+                conf[DATABRICKS_TOKEN] = os.environ[DATABRICKS_TOKEN]
+            else:
+                # OAuth
+                conf[DATABRICKS_CLIENT_ID] = os.environ[DATABRICKS_CLIENT_ID]
+                conf[DATABRICKS_CLIENT_SECRET] = os.environ[DATABRICKS_CLIENT_SECRET]
+        else:
+            warn_msg = (
+                "MLflow support is not correctly configured within Ray tasks."
+                "To enable MLflow integration, "
+                "you need to set environmental variables DATABRICKS_HOST + "
+                "DATABRICKS_TOKEN, or set environmental variables "
+                "DATABRICKS_HOST + DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET "
+                "before calling `ray.util.spark.setup_ray_cluster`, these variables "
+                "are used to set up authentication with Databricks MLflow "
+                "service. For details, you can refer to Databricks documentation at "
+                "<a href='https://docs.databricks.com/en/dev-tools/auth/pat.html'>"
+                "Databricks PAT auth</a> or "
+                "<a href='https://docs.databricks.com/en/dev-tools/auth/"
+                "oauth-m2m.html'>Databricks OAuth</a>."
+            )
+            get_databricks_display_html_function()(
+                f"<b style='color:red;'>{warn_msg}<br></b>"
+            )
+
+        return conf

python/ray/util/spark/start_hook_base.py

@@ -0,0 +1,18 @@
+class RayOnSparkStartHook:
+    def __init__(self, is_global):
+        self.is_global = is_global
+
+    def get_default_temp_root_dir(self):
+        return "/tmp"
+
+    def on_ray_dashboard_created(self, port):
+        pass
+
+    def on_cluster_created(self, ray_cluster_handler):
+        pass
+
+    def on_spark_job_created(self, job_group_id):
+        pass
+
+    def custom_environment_variables(self):
+        return {}

python/ray/util/spark/start_ray_node.py

@@ -0,0 +1,211 @@
+import fcntl
+import logging
+import os.path
+import shutil
+import signal
+import socket
+import subprocess
+import sys
+import threading
+import time
+
+from ray._private.ray_process_reaper import SIGTERM_GRACE_PERIOD_SECONDS
+from ray.util.spark.cluster_init import (
+    RAY_ON_SPARK_COLLECT_LOG_TO_PATH,
+    RAY_ON_SPARK_START_RAY_PARENT_PID,
+)
+
+# Spark on ray implementation does not directly invoke `ray start ...` script to create
+# ray node subprocess, instead, it creates a subprocess to run this
+# `ray.util.spark.start_ray_node` module, and in this module it invokes `ray start ...`
+# script to start ray node, the purpose of `start_ray_node` module is to set up a
+# exit handler for cleaning ray temp directory when ray node exits.
+# When spark driver python process dies, or spark python worker dies, because
+# `start_ray_node` starts a daemon thread of `check_parent_alive`, it will detect
+# parent process died event and then trigger cleanup work.
+
+
+_logger = logging.getLogger(__name__)
+
+
+if __name__ == "__main__":
+    arg_list = sys.argv[1:]
+
+    collect_log_to_path = os.environ[RAY_ON_SPARK_COLLECT_LOG_TO_PATH]
+
+    temp_dir_arg_prefix = "--temp-dir="
+    temp_dir = None
+
+    for arg in arg_list:
+        if arg.startswith(temp_dir_arg_prefix):
+            temp_dir = arg[len(temp_dir_arg_prefix) :]
+
+    if temp_dir is not None:
+        temp_dir = os.path.normpath(temp_dir)
+    else:
+        # This case is for global mode Ray on spark cluster
+        from ray.util.spark.cluster_init import _get_default_ray_tmp_dir
+
+        temp_dir = _get_default_ray_tmp_dir()
+
+    # Multiple Ray nodes might be launched in the same machine,
+    # so set `exist_ok` to True
+    os.makedirs(temp_dir, exist_ok=True)
+
+    ray_cli_cmd = "ray"
+    lock_file = temp_dir + ".lock"
+
+    lock_fd = os.open(lock_file, os.O_RDWR | os.O_CREAT | os.O_TRUNC)
+
+    # Mutilple ray nodes might start on the same machine, and they are using the
+    # same temp directory, adding a shared lock representing current ray node is
+    # using the temp directory.
+    fcntl.flock(lock_fd, fcntl.LOCK_SH)
+
+    process = subprocess.Popen(
+        # 'ray start ...' command uses python that is set by
+        # Shebang #! ..., the Shebang line is hardcoded in ray script,
+        # it can't be changed to other python executable path.
+        # to enforce using current python executable,
+        # turn the subprocess command to
+        # '`sys.executable` `which ray` start ...'
+        [sys.executable, shutil.which(ray_cli_cmd), "start", *arg_list],
+        text=True,
+    )
+
+    exit_handler_executed = False
+    sigterm_handler_executed = False
+    ON_EXIT_HANDLER_WAIT_TIME = 3
+
+    def on_exit_handler():
+        global exit_handler_executed
+
+        if exit_handler_executed:
+            # wait for exit_handler execution completed in other threads.
+            time.sleep(ON_EXIT_HANDLER_WAIT_TIME)
+            return
+
+        exit_handler_executed = True
+
+        try:
+            # Wait for a while to ensure the children processes of the ray node all
+            # exited.
+            time.sleep(SIGTERM_GRACE_PERIOD_SECONDS + 0.5)
+
+            if process.poll() is None:
+                # "ray start ..." command process is still alive. Force to kill it.
+                process.kill()
+
+            # Release the shared lock, representing current ray node does not use the
+            # temp dir.
+            fcntl.flock(lock_fd, fcntl.LOCK_UN)
+
+            try:
+                # acquiring exclusive lock to ensure copy logs and removing dir safely.
+                fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                lock_acquired = True
+            except BlockingIOError:
+                # The file has active shared lock or exclusive lock, representing there
+                # are other ray nodes running, or other node running cleanup temp-dir
+                # routine. skip cleaning temp-dir, and skip copy logs to destination
+                # directory as well.
+                lock_acquired = False
+
+            if lock_acquired:
+                # This is the final terminated ray node on current spark worker,
+                # start copy logs (including all local ray nodes logs) to destination.
+                if collect_log_to_path:
+                    try:
+                        log_dir_prefix = os.path.basename(temp_dir)
+                        if log_dir_prefix == "ray":
+                            # global mode cluster case, append a timestamp to it to
+                            # avoid name conflict with last Ray global cluster log dir.
+                            log_dir_prefix = (
+                                log_dir_prefix + f"-global-{int(time.time())}"
+                            )
+                        base_dir = os.path.join(
+                            collect_log_to_path, log_dir_prefix + "-logs"
+                        )
+                        # Note: multiple Ray node launcher process might
+                        # execute this line code, so we set exist_ok=True here.
+                        os.makedirs(base_dir, exist_ok=True)
+                        copy_log_dest_path = os.path.join(
+                            base_dir,
+                            socket.gethostname(),
+                        )
+                        ray_session_dir = os.readlink(
+                            os.path.join(temp_dir, "session_latest")
+                        )
+                        shutil.copytree(
+                            os.path.join(ray_session_dir, "logs"),
+                            copy_log_dest_path,
+                        )
+                    except Exception as e:
+                        _logger.warning(
+                            "Collect logs to destination directory failed, "
+                            f"error: {repr(e)}."
+                        )
+
+                # Start cleaning the temp-dir,
+                shutil.rmtree(temp_dir, ignore_errors=True)
+        except Exception:
+            # swallow any exception.
+            pass
+        finally:
+            fcntl.flock(lock_fd, fcntl.LOCK_UN)
+            os.close(lock_fd)
+
+    def check_parent_alive() -> None:
+        orig_parent_pid = int(os.environ[RAY_ON_SPARK_START_RAY_PARENT_PID])
+        while True:
+            time.sleep(0.5)
+            if os.getppid() != orig_parent_pid:
+                # Note raising SIGTERM signal in a background thread
+                # doesn't work
+                sigterm_handler()
+                break
+
+    threading.Thread(target=check_parent_alive, daemon=True).start()
+
+    try:
+
+        def sighup_handler(*args):
+            pass
+
+        # When spark application is terminated, this process will receive
+        # SIGHUP (comes from pyspark application termination).
+        # Ignore the SIGHUP signal, because in this case,
+        # `check_parent_alive` will capture parent process died event
+        # and execute killing node and cleanup routine
+        # but if we enable default SIGHUP handler, it will kill
+        # the process immediately and it causes `check_parent_alive`
+        # have no time to exeucte cleanup routine.
+        signal.signal(signal.SIGHUP, sighup_handler)
+
+        def sigterm_handler(*args):
+            global sigterm_handler_executed
+            if not sigterm_handler_executed:
+                sigterm_handler_executed = True
+                process.terminate()
+                on_exit_handler()
+            else:
+                # wait for exit_handler execution completed in other threads.
+                time.sleep(ON_EXIT_HANDLER_WAIT_TIME)
+            # Sigterm exit code is 143.
+            os._exit(143)
+
+        signal.signal(signal.SIGTERM, sigterm_handler)
+        while True:
+            try:
+                ret_code = process.wait()
+                break
+            except KeyboardInterrupt:
+                # Jupyter notebook interrupt button triggers SIGINT signal and
+                # `start_ray_node` (subprocess) will receive SIGINT signal and it
+                # causes KeyboardInterrupt exception being raised.
+                pass
+        on_exit_handler()
+        sys.exit(ret_code)
+    except Exception:
+        on_exit_handler()
+        raise

python/ray/util/spark/utils.py

@@ -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