chore: import upstream snapshot with attribution

python/ray/_private/runtime_env/BUILD.bazel

@@ -0,0 +1,42 @@
+# TODO(hjiang): All existing pythons are not using bazel as build system, which leads to missing BUILD file and targets.
+# Revisit if we decide to support bazel build in the future.
+
+load("@rules_python//python:defs.bzl", "py_library")
+
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "validation",
+    srcs = ["validation.py"],
+)
+
+py_library(
+    name = "utils",
+    srcs = ["utils.py"],
+)
+
+py_library(
+    name = "virtualenv_utils",
+    srcs = ["virtualenv_utils.py"],
+    deps = [
+        ":utils",
+    ],
+)
+
+py_library(
+    name = "dependency_utils",
+    srcs = ["dependency_utils.py"],
+    deps = [
+        ":utils",
+    ],
+)
+
+py_library(
+    name = "uv",
+    srcs = ["uv.py"],
+    deps = [
+        ":dependency_utils",
+        ":utils",
+        ":virtualenv_utils",
+    ],
+)

python/ray/_private/runtime_env/__init__.py

@@ -0,0 +1,3 @@
+# List of files to exclude from the Ray directory when using runtime_env for
+# Ray development. These are not necessary in the Ray workers.
+RAY_WORKER_DEV_EXCLUDES = ["raylet", "gcs_server", "cpp/", "tests/", "core/src"]

python/ray/_private/runtime_env/_clonevirtualenv.py

@@ -0,0 +1,334 @@
+#!/usr/bin/env python
+
+from __future__ import with_statement
+
+import logging
+import optparse
+import os
+import os.path
+import re
+import shutil
+import subprocess
+import sys
+import itertools
+
+__version__ = "0.5.7"
+
+
+logger = logging.getLogger()
+
+
+env_bin_dir = "bin"
+if sys.platform == "win32":
+    env_bin_dir = "Scripts"
+    _WIN32 = True
+else:
+    _WIN32 = False
+
+
+class UserError(Exception):
+    pass
+
+
+def _dirmatch(path, matchwith):
+    """Check if path is within matchwith's tree.
+    >>> _dirmatch('/home/foo/bar', '/home/foo/bar')
+    True
+    >>> _dirmatch('/home/foo/bar/', '/home/foo/bar')
+    True
+    >>> _dirmatch('/home/foo/bar/etc', '/home/foo/bar')
+    True
+    >>> _dirmatch('/home/foo/bar2', '/home/foo/bar')
+    False
+    >>> _dirmatch('/home/foo/bar2/etc', '/home/foo/bar')
+    False
+    """
+    matchlen = len(matchwith)
+    if path.startswith(matchwith) and path[matchlen : matchlen + 1] in [os.sep, ""]:
+        return True
+    return False
+
+
+def _virtualenv_sys(venv_path):
+    """obtain version and path info from a virtualenv."""
+    executable = os.path.join(venv_path, env_bin_dir, "python")
+    if _WIN32:
+        env = os.environ.copy()
+    else:
+        env = {}
+    # Must use "executable" as the first argument rather than as the
+    # keyword argument "executable" to get correct value from sys.path
+    p = subprocess.Popen(
+        [
+            executable,
+            "-c",
+            "import sys;"
+            'print ("%d.%d" % (sys.version_info.major, sys.version_info.minor));'
+            'print ("\\n".join(sys.path));',
+        ],
+        env=env,
+        stdout=subprocess.PIPE,
+    )
+    stdout, err = p.communicate()
+    assert not p.returncode and stdout
+    lines = stdout.decode("utf-8").splitlines()
+    return lines[0], list(filter(bool, lines[1:]))
+
+
+def clone_virtualenv(src_dir, dst_dir):
+    if not os.path.exists(src_dir):
+        raise UserError("src dir %r does not exist" % src_dir)
+    if os.path.exists(dst_dir):
+        raise UserError("dest dir %r exists" % dst_dir)
+    # sys_path = _virtualenv_syspath(src_dir)
+    logger.info("cloning virtualenv '%s' => '%s'..." % (src_dir, dst_dir))
+    shutil.copytree(
+        src_dir, dst_dir, symlinks=True, ignore=shutil.ignore_patterns("*.pyc")
+    )
+    version, sys_path = _virtualenv_sys(dst_dir)
+    logger.info("fixing scripts in bin...")
+    fixup_scripts(src_dir, dst_dir, version)
+
+    has_old = lambda s: any(i for i in s if _dirmatch(i, src_dir))  # noqa: E731
+
+    if has_old(sys_path):
+        # only need to fix stuff in sys.path if we have old
+        # paths in the sys.path of new python env. right?
+        logger.info("fixing paths in sys.path...")
+        fixup_syspath_items(sys_path, src_dir, dst_dir)
+    v_sys = _virtualenv_sys(dst_dir)
+    remaining = has_old(v_sys[1])
+    assert not remaining, v_sys
+    fix_symlink_if_necessary(src_dir, dst_dir)
+
+
+def fix_symlink_if_necessary(src_dir, dst_dir):
+    # sometimes the source virtual environment has symlinks that point to itself
+    # one example is $OLD_VIRTUAL_ENV/local/lib points to $OLD_VIRTUAL_ENV/lib
+    # this function makes sure
+    # $NEW_VIRTUAL_ENV/local/lib will point to $NEW_VIRTUAL_ENV/lib
+    # usually this goes unnoticed unless one tries to upgrade a package though pip,
+    # so this bug is hard to find.
+    logger.info("scanning for internal symlinks that point to the original virtual env")
+    for dirpath, dirnames, filenames in os.walk(dst_dir):
+        for a_file in itertools.chain(filenames, dirnames):
+            full_file_path = os.path.join(dirpath, a_file)
+            if os.path.islink(full_file_path):
+                target = os.path.realpath(full_file_path)
+                if target.startswith(src_dir):
+                    new_target = target.replace(src_dir, dst_dir)
+                    logger.debug("fixing symlink in %s" % (full_file_path,))
+                    os.remove(full_file_path)
+                    os.symlink(new_target, full_file_path)
+
+
+def fixup_scripts(old_dir, new_dir, version, rewrite_env_python=False):
+    bin_dir = os.path.join(new_dir, env_bin_dir)
+    root, dirs, files = next(os.walk(bin_dir))
+    pybinre = re.compile(r"pythonw?([0-9]+(\.[0-9]+(\.[0-9]+)?)?)?$")
+    for file_ in files:
+        filename = os.path.join(root, file_)
+        if file_ in ["python", "python%s" % version, "activate_this.py"]:
+            continue
+        elif file_.startswith("python") and pybinre.match(file_):
+            # ignore other possible python binaries
+            continue
+        elif file_.endswith(".pyc"):
+            # ignore compiled files
+            continue
+        elif file_ == "activate" or file_.startswith("activate."):
+            fixup_activate(os.path.join(root, file_), old_dir, new_dir)
+        elif os.path.islink(filename):
+            fixup_link(filename, old_dir, new_dir)
+        elif os.path.isfile(filename):
+            fixup_script_(
+                root,
+                file_,
+                old_dir,
+                new_dir,
+                version,
+                rewrite_env_python=rewrite_env_python,
+            )
+
+
+def fixup_script_(root, file_, old_dir, new_dir, version, rewrite_env_python=False):
+    old_shebang = "#!%s/bin/python" % os.path.normcase(os.path.abspath(old_dir))
+    new_shebang = "#!%s/bin/python" % os.path.normcase(os.path.abspath(new_dir))
+    env_shebang = "#!/usr/bin/env python"
+
+    filename = os.path.join(root, file_)
+    with open(filename, "rb") as f:
+        if f.read(2) != b"#!":
+            # no shebang
+            return
+        f.seek(0)
+        lines = f.readlines()
+
+    if not lines:
+        # warn: empty script
+        return
+
+    def rewrite_shebang(version=None):
+        logger.debug("fixing %s" % filename)
+        shebang = new_shebang
+        if version:
+            shebang = shebang + version
+        shebang = (shebang + "\n").encode("utf-8")
+        with open(filename, "wb") as f:
+            f.write(shebang)
+            f.writelines(lines[1:])
+
+    try:
+        bang = lines[0].decode("utf-8").strip()
+    except UnicodeDecodeError:
+        # binary file
+        return
+
+    # This takes care of the scheme in which shebang is of type
+    # '#!/venv/bin/python3' while the version of system python
+    # is of type 3.x e.g. 3.5.
+    short_version = bang[len(old_shebang) :]
+
+    if not bang.startswith("#!"):
+        return
+    elif bang == old_shebang:
+        rewrite_shebang()
+    elif bang.startswith(old_shebang) and bang[len(old_shebang) :] == version:
+        rewrite_shebang(version)
+    elif (
+        bang.startswith(old_shebang)
+        and short_version
+        and bang[len(old_shebang) :] == short_version
+    ):
+        rewrite_shebang(short_version)
+    elif rewrite_env_python and bang.startswith(env_shebang):
+        if bang == env_shebang:
+            rewrite_shebang()
+        elif bang[len(env_shebang) :] == version:
+            rewrite_shebang(version)
+    else:
+        # can't do anything
+        return
+
+
+def fixup_activate(filename, old_dir, new_dir):
+    logger.debug("fixing %s" % filename)
+    with open(filename, "rb") as f:
+        data = f.read().decode("utf-8")
+
+    data = data.replace(old_dir, new_dir)
+    with open(filename, "wb") as f:
+        f.write(data.encode("utf-8"))
+
+
+def fixup_link(filename, old_dir, new_dir, target=None):
+    logger.debug("fixing %s" % filename)
+    if target is None:
+        target = os.readlink(filename)
+
+    origdir = os.path.dirname(os.path.abspath(filename)).replace(new_dir, old_dir)
+    if not os.path.isabs(target):
+        target = os.path.abspath(os.path.join(origdir, target))
+        rellink = True
+    else:
+        rellink = False
+
+    if _dirmatch(target, old_dir):
+        if rellink:
+            # keep relative links, but don't keep original in case it
+            # traversed up out of, then back into the venv.
+            # so, recreate a relative link from absolute.
+            target = target[len(origdir) :].lstrip(os.sep)
+        else:
+            target = target.replace(old_dir, new_dir, 1)
+
+    # else: links outside the venv, replaced with absolute path to target.
+    _replace_symlink(filename, target)
+
+
+def _replace_symlink(filename, newtarget):
+    tmpfn = "%s.new" % filename
+    os.symlink(newtarget, tmpfn)
+    os.rename(tmpfn, filename)
+
+
+def fixup_syspath_items(syspath, old_dir, new_dir):
+    for path in syspath:
+        if not os.path.isdir(path):
+            continue
+        path = os.path.normcase(os.path.abspath(path))
+        if _dirmatch(path, old_dir):
+            path = path.replace(old_dir, new_dir, 1)
+            if not os.path.exists(path):
+                continue
+        elif not _dirmatch(path, new_dir):
+            continue
+        root, dirs, files = next(os.walk(path))
+        for file_ in files:
+            filename = os.path.join(root, file_)
+            if filename.endswith(".pth"):
+                fixup_pth_file(filename, old_dir, new_dir)
+            elif filename.endswith(".egg-link"):
+                fixup_egglink_file(filename, old_dir, new_dir)
+
+
+def fixup_pth_file(filename, old_dir, new_dir):
+    logger.debug("fixup_pth_file %s" % filename)
+
+    with open(filename, "r") as f:
+        lines = f.readlines()
+
+    has_change = False
+
+    for num, line in enumerate(lines):
+        line = (line.decode("utf-8") if hasattr(line, "decode") else line).strip()
+
+        if not line or line.startswith("#") or line.startswith("import "):
+            continue
+        elif _dirmatch(line, old_dir):
+            lines[num] = line.replace(old_dir, new_dir, 1)
+            has_change = True
+
+    if has_change:
+        with open(filename, "w") as f:
+            payload = os.linesep.join([line.strip() for line in lines]) + os.linesep
+            f.write(payload)
+
+
+def fixup_egglink_file(filename, old_dir, new_dir):
+    logger.debug("fixing %s" % filename)
+    with open(filename, "rb") as f:
+        link = f.read().decode("utf-8").strip()
+    if _dirmatch(link, old_dir):
+        link = link.replace(old_dir, new_dir, 1)
+        with open(filename, "wb") as f:
+            link = (link + "\n").encode("utf-8")
+            f.write(link)
+
+
+def main():
+    parser = optparse.OptionParser(
+        "usage: %prog [options] /path/to/existing/venv /path/to/cloned/venv"
+    )
+    parser.add_option(
+        "-v", action="count", dest="verbose", default=False, help="verbosity"
+    )
+    options, args = parser.parse_args()
+    try:
+        old_dir, new_dir = args
+    except ValueError:
+        print("virtualenv-clone %s" % (__version__,))
+        parser.error("not enough arguments given.")
+    old_dir = os.path.realpath(old_dir)
+    new_dir = os.path.realpath(new_dir)
+    loglevel = (logging.WARNING, logging.INFO, logging.DEBUG)[min(2, options.verbose)]
+    logging.basicConfig(level=loglevel, format="%(message)s")
+    try:
+        clone_virtualenv(old_dir, new_dir)
+    except UserError:
+        e = sys.exc_info()[1]
+        parser.error(str(e))
+
+
+if __name__ == "__main__":
+    main()

python/ray/_private/runtime_env/agent/main.py

@@ -0,0 +1,266 @@
+import argparse
+import logging
+import os
+import socket
+import sys
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray._common.utils import (
+    get_or_create_event_loop,
+)
+from ray._private import logging_utils
+from ray._private.authentication.http_token_authentication import (
+    get_token_auth_middleware,
+)
+from ray._private.process_watcher import create_check_raylet_task
+from ray._raylet import RUNTIME_ENV_AGENT_PORT_NAME, GcsClient, persist_port
+from ray.core.generated import (
+    runtime_env_agent_pb2,
+)
+
+
+def import_libs():
+    my_dir = os.path.abspath(os.path.dirname(__file__))
+    sys.path.insert(0, os.path.join(my_dir, "thirdparty_files"))  # for aiohttp
+    sys.path.insert(0, my_dir)  # for runtime_env_agent and runtime_env_consts
+
+
+import_libs()
+
+import aiohttp  # noqa: E402
+import runtime_env_consts  # noqa: E402
+from aiohttp import web  # noqa: E402
+from runtime_env_agent import RuntimeEnvAgent  # noqa: E402
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Runtime env agent.")
+    parser.add_argument(
+        "--node-id",
+        required=True,
+        type=str,
+        help="the unique ID of this node.",
+    )
+    parser.add_argument(
+        "--node-ip-address",
+        required=True,
+        type=str,
+        help="the IP address of this node.",
+    )
+    parser.add_argument(
+        "--runtime-env-agent-port",
+        required=True,
+        type=int,
+        default=None,
+        help="The port on which the runtime env agent will receive HTTP requests.",
+    )
+    parser.add_argument(
+        "--session-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="The path of this ray session directory.",
+    )
+
+    parser.add_argument(
+        "--gcs-address", required=True, type=str, help="The address (ip:port) of GCS."
+    )
+    parser.add_argument(
+        "--cluster-id-hex", required=True, type=str, help="The cluster id in hex."
+    )
+    parser.add_argument(
+        "--runtime-env-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="Specify the path of the resource directory used by runtime_env.",
+    )
+
+    parser.add_argument(
+        "--logging-level",
+        required=False,
+        type=lambda s: logging.getLevelName(s.upper()),
+        default=ray_constants.LOGGER_LEVEL,
+        choices=ray_constants.LOGGER_LEVEL_CHOICES,
+        help=ray_constants.LOGGER_LEVEL_HELP,
+    )
+    parser.add_argument(
+        "--logging-format",
+        required=False,
+        type=str,
+        default=ray_constants.LOGGER_FORMAT,
+        help=ray_constants.LOGGER_FORMAT_HELP,
+    )
+    parser.add_argument(
+        "--logging-filename",
+        required=False,
+        type=str,
+        default=runtime_env_consts.RUNTIME_ENV_AGENT_LOG_FILENAME,
+        help="Specify the name of log file, "
+        'log to stdout if set empty, default is "{}".'.format(
+            runtime_env_consts.RUNTIME_ENV_AGENT_LOG_FILENAME
+        ),
+    )
+    parser.add_argument(
+        "--logging-rotate-bytes",
+        required=True,
+        type=int,
+        help="Specify the max bytes for rotating log file",
+    )
+    parser.add_argument(
+        "--logging-rotate-backup-count",
+        required=True,
+        type=int,
+        help="Specify the backup count of rotated log file",
+    )
+    parser.add_argument(
+        "--log-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="Specify the path of log directory.",
+    )
+    parser.add_argument(
+        "--temp-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="Specify the path of the temporary directory use by Ray process.",
+    )
+    parser.add_argument(
+        "--stdout-filepath",
+        required=False,
+        type=str,
+        default="",
+        help="The filepath to dump runtime env agent stdout.",
+    )
+    parser.add_argument(
+        "--stderr-filepath",
+        required=False,
+        type=str,
+        default="",
+        help="The filepath to dump runtime env agent stderr.",
+    )
+
+    args = parser.parse_args()
+
+    # Disable log rotation for windows platform.
+    logging_rotation_bytes = args.logging_rotate_bytes if sys.platform != "win32" else 0
+    logging_rotation_backup_count = (
+        args.logging_rotate_backup_count if sys.platform != "win32" else 1
+    )
+
+    logging_params = dict(
+        logging_level=args.logging_level,
+        logging_format=args.logging_format,
+        log_dir=args.log_dir,
+        filename=args.logging_filename,
+        max_bytes=logging_rotation_bytes,
+        backup_count=logging_rotation_backup_count,
+    )
+
+    # Setup stdout/stderr redirect files if redirection enabled.
+    logging_utils.redirect_stdout_stderr_if_needed(
+        args.stdout_filepath,
+        args.stderr_filepath,
+        logging_rotation_bytes,
+        logging_rotation_backup_count,
+    )
+
+    gcs_client = GcsClient(address=args.gcs_address, cluster_id=args.cluster_id_hex)
+    agent = RuntimeEnvAgent(
+        runtime_env_dir=args.runtime_env_dir,
+        logging_params=logging_params,
+        gcs_client=gcs_client,
+        temp_dir=args.temp_dir,
+        address=args.node_ip_address,
+        runtime_env_agent_port=args.runtime_env_agent_port,
+    )
+
+    ray._raylet.setproctitle(ray_constants.AGENT_PROCESS_TYPE_RUNTIME_ENV_AGENT)
+
+    # POST /get_or_create_runtime_env
+    # body is serialzied protobuf GetOrCreateRuntimeEnvRequest
+    # reply is serialzied protobuf GetOrCreateRuntimeEnvReply
+    async def get_or_create_runtime_env(request: web.Request) -> web.Response:
+        data = await request.read()
+        request = runtime_env_agent_pb2.GetOrCreateRuntimeEnvRequest()
+        request.ParseFromString(data)
+        reply = await agent.GetOrCreateRuntimeEnv(request)
+        return web.Response(
+            body=reply.SerializeToString(), content_type="application/octet-stream"
+        )
+
+    # POST /delete_runtime_env_if_possible
+    # body is serialzied protobuf DeleteRuntimeEnvIfPossibleRequest
+    # reply is serialzied protobuf DeleteRuntimeEnvIfPossibleReply
+    async def delete_runtime_env_if_possible(request: web.Request) -> web.Response:
+        data = await request.read()
+        request = runtime_env_agent_pb2.DeleteRuntimeEnvIfPossibleRequest()
+        request.ParseFromString(data)
+        reply = await agent.DeleteRuntimeEnvIfPossible(request)
+        return web.Response(
+            body=reply.SerializeToString(), content_type="application/octet-stream"
+        )
+
+    # POST /get_runtime_envs_info
+    # body is serialzied protobuf GetRuntimeEnvsInfoRequest
+    # reply is serialzied protobuf GetRuntimeEnvsInfoReply
+    async def get_runtime_envs_info(request: web.Request) -> web.Response:
+        data = await request.read()
+        request = runtime_env_agent_pb2.GetRuntimeEnvsInfoRequest()
+        request.ParseFromString(data)
+        reply = await agent.GetRuntimeEnvsInfo(request)
+        return web.Response(
+            body=reply.SerializeToString(), content_type="application/octet-stream"
+        )
+
+    app = web.Application(middlewares=[get_token_auth_middleware(aiohttp)])
+
+    app.router.add_post("/get_or_create_runtime_env", get_or_create_runtime_env)
+    app.router.add_post(
+        "/delete_runtime_env_if_possible", delete_runtime_env_if_possible
+    )
+    app.router.add_post("/get_runtime_envs_info", get_runtime_envs_info)
+
+    loop = get_or_create_event_loop()
+    check_raylet_task = None
+    if sys.platform not in ["win32", "cygwin"]:
+
+        def parent_dead_callback(msg):
+            agent._logger.info(
+                "Raylet is dead! Exiting Runtime Env Agent. "
+                f"addr: {args.node_ip_address}, "
+                f"port: {args.runtime_env_agent_port}\n"
+                f"{msg}"
+            )
+
+        # No need to await this task.
+        check_raylet_task = create_check_raylet_task(
+            args.log_dir, gcs_client, parent_dead_callback, loop
+        )
+
+    port = args.runtime_env_agent_port or 0
+    infos = socket.getaddrinfo(args.node_ip_address, port, type=socket.SOCK_STREAM)
+    family, socktype, proto, _, sockaddr = infos[0]
+    sock = socket.socket(family, socktype, proto)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    sock.bind(sockaddr)
+
+    bound_port = sock.getsockname()[1]
+    persist_port(
+        args.session_dir,
+        args.node_id,
+        RUNTIME_ENV_AGENT_PORT_NAME,
+        bound_port,
+    )
+
+    try:
+        web.run_app(app, sock=sock, loop=loop)
+    except SystemExit as e:
+        agent._logger.info(f"SystemExit! {e}")
+        # We have to poke the task exception, or there's an error message
+        # "task exception was never retrieved".
+        if check_raylet_task is not None:
+            check_raylet_task.exception()
+        sys.exit(e.code)

python/ray/_private/runtime_env/agent/runtime_env_agent.py

@@ -0,0 +1,620 @@
+import asyncio
+import logging
+import os
+import time
+import traceback
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Set, Tuple
+
+import ray
+import ray._private.runtime_env.agent.runtime_env_consts as runtime_env_consts
+from ray._common.utils import get_or_create_event_loop
+from ray._private.ray_constants import (
+    DEFAULT_RUNTIME_ENV_TIMEOUT_SECONDS,
+)
+from ray._private.ray_logging import setup_component_logger
+from ray._private.runtime_env.conda import CondaPlugin
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.default_impl import get_image_uri_plugin_cls
+from ray._private.runtime_env.image_uri import ContainerPlugin
+from ray._private.runtime_env.java_jars import JavaJarsPlugin
+from ray._private.runtime_env.nsight import NsightPlugin
+from ray._private.runtime_env.pip import PipPlugin
+from ray._private.runtime_env.plugin import (
+    RuntimeEnvPlugin,
+    RuntimeEnvPluginManager,
+    create_for_plugin_if_needed,
+)
+from ray._private.runtime_env.py_executable import PyExecutablePlugin
+from ray._private.runtime_env.py_modules import PyModulesPlugin
+from ray._private.runtime_env.rocprof_sys import RocProfSysPlugin
+from ray._private.runtime_env.uv import UvPlugin
+from ray._private.runtime_env.working_dir import WorkingDirPlugin
+from ray._raylet import GcsClient
+from ray.core.generated import runtime_env_agent_pb2
+from ray.core.generated.runtime_env_common_pb2 import (
+    RuntimeEnvState as ProtoRuntimeEnvState,
+)
+from ray.runtime_env import RuntimeEnv, RuntimeEnvConfig
+
+default_logger = logging.getLogger(__name__)
+
+# TODO(edoakes): this is used for unit tests. We should replace it with a
+# better pluggability mechanism once available.
+SLEEP_FOR_TESTING_S = os.environ.get("RAY_RUNTIME_ENV_SLEEP_FOR_TESTING_S")
+
+
+@dataclass
+class CreatedEnvResult:
+    # Whether or not the env was installed correctly.
+    success: bool
+    # If success is True, will be a serialized RuntimeEnvContext
+    # If success is False, will be an error message.
+    result: str
+    # The time to create a runtime env in ms.
+    creation_time_ms: int
+
+
+# e.g., "working_dir"
+UriType = str
+
+
+class ReferenceTable:
+    """
+    The URI reference table which is used for GC.
+    When the reference count is decreased to zero,
+    the URI should be removed from this table and
+    added to cache if needed.
+    """
+
+    def __init__(
+        self,
+        uris_parser: Callable[[RuntimeEnv], Tuple[str, UriType]],
+        unused_uris_callback: Callable[[List[Tuple[str, UriType]]], None],
+        unused_runtime_env_callback: Callable[[str], None],
+    ):
+        # Runtime Environment reference table. The key is serialized runtime env and
+        # the value is reference count.
+        self._runtime_env_reference: Dict[str, int] = defaultdict(int)
+        # URI reference table. The key is URI parsed from runtime env and the value
+        # is reference count.
+        self._uri_reference: Dict[str, int] = defaultdict(int)
+        self._uris_parser = uris_parser
+        self._unused_uris_callback = unused_uris_callback
+        self._unused_runtime_env_callback = unused_runtime_env_callback
+        # send the `DeleteRuntimeEnvIfPossible` RPC when the client exits. The URI won't
+        # be leaked now because the reference count will be reset to zero when the job
+        # finished.
+        self._reference_exclude_sources: Set[str] = {
+            "client_server",
+        }
+
+    def _increase_reference_for_uris(self, uris):
+        default_logger.debug(f"Increase reference for uris {uris}.")
+        for uri, _ in uris:
+            self._uri_reference[uri] += 1
+
+    def _decrease_reference_for_uris(self, uris):
+        default_logger.debug(f"Decrease reference for uris {uris}.")
+        unused_uris = list()
+        for uri, uri_type in uris:
+            if self._uri_reference[uri] > 0:
+                self._uri_reference[uri] -= 1
+                if self._uri_reference[uri] == 0:
+                    unused_uris.append((uri, uri_type))
+                    del self._uri_reference[uri]
+            else:
+                default_logger.warning(f"URI {uri} does not exist.")
+        if unused_uris:
+            default_logger.info(f"Unused uris {unused_uris}.")
+            self._unused_uris_callback(unused_uris)
+        return unused_uris
+
+    def _increase_reference_for_runtime_env(self, serialized_env: str):
+        default_logger.debug(f"Increase reference for runtime env {serialized_env}.")
+        self._runtime_env_reference[serialized_env] += 1
+
+    def _decrease_reference_for_runtime_env(self, serialized_env: str):
+        """Decrease reference count for the given [serialized_env]. Throw exception if we cannot decrement reference."""
+        default_logger.debug(f"Decrease reference for runtime env {serialized_env}.")
+        unused = False
+        if self._runtime_env_reference[serialized_env] > 0:
+            self._runtime_env_reference[serialized_env] -= 1
+            if self._runtime_env_reference[serialized_env] == 0:
+                unused = True
+                del self._runtime_env_reference[serialized_env]
+        else:
+            default_logger.warning(f"Runtime env {serialized_env} does not exist.")
+            raise ValueError(
+                f"{serialized_env} cannot decrement reference since the reference count is 0"
+            )
+        if unused:
+            default_logger.info(f"Unused runtime env {serialized_env}.")
+            self._unused_runtime_env_callback(serialized_env)
+
+    def increase_reference(
+        self, runtime_env: RuntimeEnv, serialized_env: str, source_process: str
+    ) -> None:
+        if source_process in self._reference_exclude_sources:
+            return
+        self._increase_reference_for_runtime_env(serialized_env)
+        uris = self._uris_parser(runtime_env)
+        self._increase_reference_for_uris(uris)
+
+    def decrease_reference(
+        self, runtime_env: RuntimeEnv, serialized_env: str, source_process: str
+    ) -> None:
+        """Decrease reference count for runtime env and uri. Throw exception if decrement reference count fails."""
+        if source_process in self._reference_exclude_sources:
+            return
+        self._decrease_reference_for_runtime_env(serialized_env)
+        uris = self._uris_parser(runtime_env)
+        self._decrease_reference_for_uris(uris)
+
+    @property
+    def runtime_env_refs(self) -> Dict[str, int]:
+        """Return the runtime_env -> ref count mapping.
+
+        Returns:
+            The mapping of serialized runtime env -> ref count.
+        """
+        return self._runtime_env_reference
+
+
+class RuntimeEnvAgent:
+    """An RPC server to create and delete runtime envs.
+
+    Attributes:
+        dashboard_agent: The DashboardAgent object contains global config.
+    """
+
+    def __init__(
+        self,
+        runtime_env_dir: str,
+        logging_params: dict,
+        gcs_client: GcsClient,
+        temp_dir: str,
+        address: str,
+        runtime_env_agent_port: int,
+    ):
+        """Initialize the runtime env agent.
+
+        Args:
+            runtime_env_dir: Directory used to store runtime env resources.
+            logging_params: Keyword arguments forwarded to
+                :func:`setup_component_logger` to configure the agent logger.
+            gcs_client: GCS client used to fetch package data.
+            temp_dir: Temporary directory used by plugins (e.g. container plugin).
+            address: IP address that the agent is listening on, used for logging.
+            runtime_env_agent_port: Port that the agent is listening on, used for
+                logging.
+        """
+        super().__init__()
+
+        self._logger = default_logger
+        self._logging_params = logging_params
+        self._logger = setup_component_logger(
+            logger_name=default_logger.name, **self._logging_params
+        )
+        # Don't propagate logs to the root logger, because these logs
+        # might contain sensitive information. Instead, these logs should
+        # be confined to the runtime env agent log file `self.LOG_FILENAME`.
+        self._logger.propagate = False
+
+        self._logger.info("Starting runtime env agent at pid %s", os.getpid())
+        self._logger.info(f"Parent raylet pid is {os.environ.get('RAY_RAYLET_PID')}")
+
+        self._runtime_env_dir = runtime_env_dir
+        self._per_job_logger_cache = dict()
+        # Cache the results of creating envs to avoid repeatedly calling into
+        # conda and other slow calls.
+        self._env_cache: Dict[str, CreatedEnvResult] = dict()
+        # Maps a serialized runtime env to a lock that is used
+        # to prevent multiple concurrent installs of the same env.
+        self._env_locks: Dict[str, asyncio.Lock] = dict()
+        self._gcs_client = gcs_client
+
+        self._pip_plugin = PipPlugin(self._runtime_env_dir)
+        self._uv_plugin = UvPlugin(self._runtime_env_dir)
+        self._conda_plugin = CondaPlugin(self._runtime_env_dir)
+        self._py_modules_plugin = PyModulesPlugin(
+            self._runtime_env_dir, self._gcs_client
+        )
+        self._py_executable_plugin = PyExecutablePlugin()
+        self._java_jars_plugin = JavaJarsPlugin(self._runtime_env_dir, self._gcs_client)
+        self._working_dir_plugin = WorkingDirPlugin(
+            self._runtime_env_dir, self._gcs_client
+        )
+        self._container_plugin = ContainerPlugin(temp_dir)
+        # TODO(jonathan-anyscale): change the plugin to ProfilerPlugin
+        # and unify with nsight and other profilers.
+        self._nsight_plugin = NsightPlugin(self._runtime_env_dir)
+        self._rocprof_sys_plugin = RocProfSysPlugin(self._runtime_env_dir)
+        self._image_uri_plugin = get_image_uri_plugin_cls()(temp_dir)
+
+        # TODO(architkulkarni): "base plugins" and third-party plugins should all go
+        # through the same code path.  We should never need to refer to
+        # self._xxx_plugin, we should just iterate through self._plugins.
+        self._base_plugins: List[RuntimeEnvPlugin] = [
+            self._working_dir_plugin,
+            self._uv_plugin,
+            self._pip_plugin,
+            self._conda_plugin,
+            self._py_modules_plugin,
+            self._py_executable_plugin,
+            self._java_jars_plugin,
+            self._container_plugin,
+            self._nsight_plugin,
+            self._rocprof_sys_plugin,
+            self._image_uri_plugin,
+        ]
+        self._plugin_manager = RuntimeEnvPluginManager()
+        for plugin in self._base_plugins:
+            self._plugin_manager.add_plugin(plugin)
+
+        self._reference_table = ReferenceTable(
+            self.uris_parser,
+            self.unused_uris_processor,
+            self.unused_runtime_env_processor,
+        )
+
+        self._logger.info(
+            "Listening to address %s, port %d", address, runtime_env_agent_port
+        )
+
+        try:
+            self._node_ip = ray.util.get_node_ip_address()
+            self._node_prefix = f"[Node {self._node_ip}] "
+        except Exception as e:
+            self._logger.warning(f"Failed to get node IP address, using fallback: {e}")
+            self._node_prefix = "[Node unknown] "
+
+    def uris_parser(self, runtime_env: RuntimeEnv):
+        result = list()
+        for name, plugin_setup_context in self._plugin_manager.plugins.items():
+            plugin = plugin_setup_context.class_instance
+            uris = plugin.get_uris(runtime_env)
+            for uri in uris:
+                result.append((uri, UriType(name)))
+        return result
+
+    def unused_uris_processor(self, unused_uris: List[Tuple[str, UriType]]) -> None:
+        for uri, uri_type in unused_uris:
+            self._plugin_manager.plugins[str(uri_type)].uri_cache.mark_unused(uri)
+
+    def unused_runtime_env_processor(self, unused_runtime_env: str) -> None:
+        def delete_runtime_env():
+            del self._env_cache[unused_runtime_env]
+            self._logger.info(
+                "Runtime env %s removed from env-level cache.", unused_runtime_env
+            )
+
+        if unused_runtime_env in self._env_cache:
+            if not self._env_cache[unused_runtime_env].success:
+                loop = get_or_create_event_loop()
+                # Cache the bad runtime env result by ttl seconds.
+                loop.call_later(
+                    runtime_env_consts.BAD_RUNTIME_ENV_CACHE_TTL_SECONDS,
+                    delete_runtime_env,
+                )
+            else:
+                delete_runtime_env()
+
+    def get_or_create_logger(self, job_id: bytes, log_files: List[str]):
+        job_id = job_id.decode()
+        if job_id not in self._per_job_logger_cache:
+            params = self._logging_params.copy()
+            params["filename"] = [f"runtime_env_setup-{job_id}.log", *log_files]
+            params["logger_name"] = f"runtime_env_{job_id}"
+            params["propagate"] = False
+            per_job_logger = setup_component_logger(**params)
+            self._per_job_logger_cache[job_id] = per_job_logger
+        return self._per_job_logger_cache[job_id]
+
+    async def GetOrCreateRuntimeEnv(self, request):
+        self._logger.debug(
+            f"Got request from {request.source_process} to increase "
+            "reference for runtime env: "
+            f"{request.serialized_runtime_env}."
+        )
+
+        async def _setup_runtime_env(
+            runtime_env: RuntimeEnv,
+            runtime_env_config: RuntimeEnvConfig,
+        ):
+            log_files = runtime_env_config.get("log_files", [])
+            # Use a separate logger for each job.
+            per_job_logger = self.get_or_create_logger(request.job_id, log_files)
+            context = RuntimeEnvContext(env_vars=runtime_env.env_vars())
+
+            # Warn about unrecognized fields in the runtime env.
+            for name, _ in runtime_env.plugins():
+                if name not in self._plugin_manager.plugins:
+                    per_job_logger.warning(
+                        f"runtime_env field {name} is not recognized by "
+                        "Ray and will be ignored.  In the future, unrecognized "
+                        "fields in the runtime_env will raise an exception."
+                    )
+
+            # Creates each runtime env URI by their priority. `working_dir` is special
+            # because it needs to be created before other plugins. All other plugins are
+            # created in the priority order (smaller priority value -> earlier to
+            # create), with a special environment variable being set to the working dir.
+            # ${RAY_RUNTIME_ENV_CREATE_WORKING_DIR}
+
+            # First create working dir...
+            working_dir_ctx = self._plugin_manager.plugins[WorkingDirPlugin.name]
+            await create_for_plugin_if_needed(
+                runtime_env,
+                working_dir_ctx.class_instance,
+                working_dir_ctx.uri_cache,
+                context,
+                per_job_logger,
+            )
+
+            # Then within the working dir, create the other plugins.
+            working_dir_uri_or_none = runtime_env.working_dir_uri()
+            with self._working_dir_plugin.with_working_dir_env(working_dir_uri_or_none):
+                """Run setup for each plugin unless it has already been cached."""
+                for (
+                    plugin_setup_context
+                ) in self._plugin_manager.sorted_plugin_setup_contexts():
+                    plugin = plugin_setup_context.class_instance
+                    if plugin.name != WorkingDirPlugin.name:
+                        uri_cache = plugin_setup_context.uri_cache
+                        await create_for_plugin_if_needed(
+                            runtime_env, plugin, uri_cache, context, per_job_logger
+                        )
+            return context
+
+        async def _create_runtime_env_with_retry(
+            runtime_env: RuntimeEnv,
+            setup_timeout_seconds: int,
+            runtime_env_config: RuntimeEnvConfig,
+        ) -> Tuple[bool, str, str]:
+            """Create runtime env with retry times. This function won't raise exceptions.
+
+            Args:
+                runtime_env: The instance of RuntimeEnv class.
+                setup_timeout_seconds: The timeout of runtime environment creation for
+                    each attempt.
+                runtime_env_config: The configuration for the runtime environment.
+
+            Returns:
+                Tuple[bool, str, str]: A tuple containing:
+                    - result (bool): Whether the creation was successful
+                    - runtime_env_context (str): The serialized context if successful, None otherwise
+                    - error_message (str): Error message if failed, None otherwise
+            """
+            self._logger.info(
+                f"Creating runtime env: {serialized_env} with timeout "
+                f"{setup_timeout_seconds} seconds."
+            )
+            num_retries = runtime_env_consts.RUNTIME_ENV_RETRY_TIMES
+            error_message = None
+            serialized_context = None
+            for i in range(num_retries):
+                # Only sleep when retrying.
+                if i != 0:
+                    await asyncio.sleep(
+                        runtime_env_consts.RUNTIME_ENV_RETRY_INTERVAL_MS / 1000
+                    )
+
+                try:
+                    runtime_env_setup_task = _setup_runtime_env(
+                        runtime_env, runtime_env_config
+                    )
+                    runtime_env_context = await asyncio.wait_for(
+                        runtime_env_setup_task, timeout=setup_timeout_seconds
+                    )
+                    serialized_context = runtime_env_context.serialize()
+                    error_message = None
+                    break
+                except Exception as e:
+                    err_msg = f"Failed to create runtime env {serialized_env}."
+                    self._logger.exception(err_msg)
+                    error_message = "".join(
+                        traceback.format_exception(type(e), e, e.__traceback__)
+                    )
+                    if isinstance(e, asyncio.TimeoutError):
+                        hint = (
+                            f"Failed to install runtime_env within the "
+                            f"timeout of {setup_timeout_seconds} seconds. Consider "
+                            "increasing the timeout in the runtime_env config. "
+                            "For example: \n"
+                            '    runtime_env={"config": {"setup_timeout_seconds":'
+                            " 1800}, ...}\n"
+                            "If not provided, the default timeout is "
+                            f"{DEFAULT_RUNTIME_ENV_TIMEOUT_SECONDS} seconds. "
+                        )
+                        error_message = hint + error_message
+
+            if error_message:
+                self._logger.error(
+                    "runtime_env creation failed %d times, giving up.",
+                    num_retries,
+                )
+                return False, None, error_message
+            else:
+                self._logger.info(
+                    "Successfully created runtime env: %s, context: %s",
+                    serialized_env,
+                    serialized_context,
+                )
+                return True, serialized_context, None
+
+        try:
+            serialized_env = request.serialized_runtime_env
+            runtime_env = RuntimeEnv.deserialize(serialized_env)
+        except Exception as e:
+            self._logger.exception(
+                "[Increase] Failed to parse runtime env: " f"{serialized_env}"
+            )
+
+            error_message = "".join(
+                traceback.format_exception(type(e), e, e.__traceback__)
+            )
+
+            return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply(
+                status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED,
+                error_message=f"{self._node_prefix}{error_message}",
+            )
+
+        # Increase reference
+        self._reference_table.increase_reference(
+            runtime_env, serialized_env, request.source_process
+        )
+
+        if serialized_env not in self._env_locks:
+            # async lock to prevent the same env being concurrently installed
+            self._env_locks[serialized_env] = asyncio.Lock()
+
+        async with self._env_locks[serialized_env]:
+            if serialized_env in self._env_cache:
+                serialized_context = self._env_cache[serialized_env]
+                result = self._env_cache[serialized_env]
+                if result.success:
+                    context = result.result
+                    self._logger.info(
+                        "Runtime env already created "
+                        f"successfully. Env: {serialized_env}, "
+                        f"context: {context}"
+                    )
+                    return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply(
+                        status=runtime_env_agent_pb2.AGENT_RPC_STATUS_OK,
+                        serialized_runtime_env_context=context,
+                    )
+                else:
+                    error_message = result.result
+                    self._logger.info(
+                        "Runtime env already failed. "
+                        f"Env: {serialized_env}, "
+                        f"err: {error_message}"
+                    )
+                    # Recover the reference.
+                    self._reference_table.decrease_reference(
+                        runtime_env, serialized_env, request.source_process
+                    )
+                    return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply(
+                        status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED,
+                        error_message=f"{self._node_prefix}{error_message}",
+                    )
+
+            if SLEEP_FOR_TESTING_S:
+                self._logger.info(f"Sleeping for {SLEEP_FOR_TESTING_S}s.")
+                time.sleep(int(SLEEP_FOR_TESTING_S))
+
+            runtime_env_config = RuntimeEnvConfig.from_proto(request.runtime_env_config)
+
+            # accroding to the document of `asyncio.wait_for`,
+            # None means disable timeout logic
+            setup_timeout_seconds = (
+                None
+                if runtime_env_config["setup_timeout_seconds"] == -1
+                else runtime_env_config["setup_timeout_seconds"]
+            )
+
+            start = time.perf_counter()
+            (
+                successful,
+                serialized_context,
+                error_message,
+            ) = await _create_runtime_env_with_retry(
+                runtime_env,
+                setup_timeout_seconds,
+                runtime_env_config,
+            )
+            creation_time_ms = int(round((time.perf_counter() - start) * 1000, 0))
+            if not successful:
+                # Recover the reference.
+                self._reference_table.decrease_reference(
+                    runtime_env, serialized_env, request.source_process
+                )
+            # Add the result to env cache.
+            self._env_cache[serialized_env] = CreatedEnvResult(
+                successful,
+                serialized_context if successful else error_message,
+                creation_time_ms,
+            )
+            # Reply the RPC
+            return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply(
+                status=runtime_env_agent_pb2.AGENT_RPC_STATUS_OK
+                if successful
+                else runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED,
+                serialized_runtime_env_context=serialized_context,
+                error_message=f"{self._node_prefix}{error_message}"
+                if not successful
+                else "",
+            )
+
+    async def DeleteRuntimeEnvIfPossible(self, request):
+        self._logger.info(
+            f"Got request from {request.source_process} to decrease "
+            "reference for runtime env: "
+            f"{request.serialized_runtime_env}."
+        )
+
+        try:
+            runtime_env = RuntimeEnv.deserialize(request.serialized_runtime_env)
+        except Exception as e:
+            self._logger.exception(
+                "[Decrease] Failed to parse runtime env: "
+                f"{request.serialized_runtime_env}"
+            )
+
+            error_message = "".join(
+                traceback.format_exception(type(e), e, e.__traceback__)
+            )
+
+            return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply(
+                status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED,
+                error_message=f"{self._node_prefix}{error_message}",
+            )
+
+        try:
+            self._reference_table.decrease_reference(
+                runtime_env, request.serialized_runtime_env, request.source_process
+            )
+        except Exception as e:
+            return runtime_env_agent_pb2.DeleteRuntimeEnvIfPossibleReply(
+                status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED,
+                error_message=f"{self._node_prefix}Failed to decrement reference for runtime env for {str(e)}",
+            )
+
+        return runtime_env_agent_pb2.DeleteRuntimeEnvIfPossibleReply(
+            status=runtime_env_agent_pb2.AGENT_RPC_STATUS_OK
+        )
+
+    async def GetRuntimeEnvsInfo(self, request):
+        """Return the runtime env information of the node."""
+        # TODO(sang): Currently, it only includes runtime_env information.
+        # We should include the URI information which includes,
+        # URIs
+        # Caller
+        # Ref counts
+        # Cache information
+        # Metrics (creation time & success)
+        # Deleted URIs
+        limit = request.limit if request.HasField("limit") else -1
+        runtime_env_states = defaultdict(ProtoRuntimeEnvState)
+        runtime_env_refs = self._reference_table.runtime_env_refs
+        for runtime_env, ref_cnt in runtime_env_refs.items():
+            runtime_env_states[runtime_env].runtime_env = runtime_env
+            runtime_env_states[runtime_env].ref_cnt = ref_cnt
+        for runtime_env, result in self._env_cache.items():
+            runtime_env_states[runtime_env].runtime_env = runtime_env
+            runtime_env_states[runtime_env].success = result.success
+            if not result.success:
+                runtime_env_states[runtime_env].error = result.result
+            runtime_env_states[runtime_env].creation_time_ms = result.creation_time_ms
+
+        reply = runtime_env_agent_pb2.GetRuntimeEnvsInfoReply()
+        count = 0
+        for runtime_env_state in runtime_env_states.values():
+            if limit != -1 and count >= limit:
+                break
+            count += 1
+            reply.runtime_env_states.append(runtime_env_state)
+        reply.total = len(runtime_env_states)
+        return reply

python/ray/_private/runtime_env/agent/runtime_env_consts.py

@@ -0,0 +1,20 @@
+import ray._private.ray_constants as ray_constants
+
+RUNTIME_ENV_RETRY_TIMES = ray_constants.env_integer("RUNTIME_ENV_RETRY_TIMES", 3)
+
+RUNTIME_ENV_RETRY_INTERVAL_MS = ray_constants.env_integer(
+    "RUNTIME_ENV_RETRY_INTERVAL_MS", 1000
+)
+
+# Cache TTL for bad runtime env. After this time, delete the cache and retry to create
+# runtime env if needed.
+BAD_RUNTIME_ENV_CACHE_TTL_SECONDS = ray_constants.env_integer(
+    "BAD_RUNTIME_ENV_CACHE_TTL_SECONDS", 60 * 10
+)
+
+RUNTIME_ENV_LOG_FILENAME = "runtime_env.log"
+RUNTIME_ENV_AGENT_PORT_PREFIX = "RUNTIME_ENV_AGENT_PORT_PREFIX:"
+RUNTIME_ENV_AGENT_LOG_FILENAME = "runtime_env_agent.log"
+RUNTIME_ENV_AGENT_CHECK_PARENT_INTERVAL_S_ENV_NAME = (
+    "RAY_RUNTIME_ENV_AGENT_CHECK_PARENT_INTERVAL_S"  # noqa
+)

python/ray/_private/runtime_env/conda.py

@@ -0,0 +1,400 @@
+import hashlib
+import json
+import logging
+import os
+import runpy
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import yaml
+from filelock import FileLock
+
+import ray
+from ray._common.utils import (
+    get_or_create_event_loop,
+    try_to_create_directory,
+)
+from ray._private.runtime_env.conda_utils import (
+    create_conda_env_if_needed,
+    delete_conda_env,
+    get_conda_activate_commands,
+    get_conda_envs,
+    get_conda_info_json,
+)
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.packaging import Protocol, parse_uri
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.runtime_env.validation import parse_and_validate_conda
+from ray._private.utils import (
+    get_directory_size_bytes,
+    get_master_wheel_url,
+    get_release_wheel_url,
+    get_wheel_filename,
+)
+
+default_logger = logging.getLogger(__name__)
+
+_WIN32 = os.name == "nt"
+
+
+def _resolve_current_ray_path() -> str:
+    # When ray is built from source with pip install -e,
+    # ray.__file__ returns .../python/ray/__init__.py and this function returns
+    # ".../python".
+    # When ray is installed from a prebuilt binary, ray.__file__ returns
+    # .../site-packages/ray/__init__.py and this function returns
+    # ".../site-packages".
+    return os.path.split(os.path.split(ray.__file__)[0])[0]
+
+
+def _get_ray_setup_spec():
+    """Find the Ray setup_spec from the currently running Ray.
+
+    This function works even when Ray is built from source with pip install -e.
+    """
+    ray_source_python_path = _resolve_current_ray_path()
+    setup_py_path = os.path.join(ray_source_python_path, "setup.py")
+    return runpy.run_path(setup_py_path)["setup_spec"]
+
+
+def _resolve_install_from_source_ray_dependencies():
+    """Find the Ray dependencies when Ray is installed from source."""
+    deps = (
+        _get_ray_setup_spec().install_requires + _get_ray_setup_spec().extras["default"]
+    )
+    # Remove duplicates
+    return list(set(deps))
+
+
+def _inject_ray_to_conda_site(
+    conda_path, logger: Optional[logging.Logger] = default_logger
+):
+    """Write the current Ray site package directory to a new site"""
+    if _WIN32:
+        python_binary = os.path.join(conda_path, "python")
+    else:
+        python_binary = os.path.join(conda_path, "bin/python")
+    site_packages_path = (
+        subprocess.check_output(
+            [
+                python_binary,
+                "-c",
+                "import sysconfig; print(sysconfig.get_paths()['purelib'])",
+            ]
+        )
+        .decode()
+        .strip()
+    )
+
+    ray_path = _resolve_current_ray_path()
+    logger.warning(
+        f"Injecting {ray_path} to environment site-packages {site_packages_path} "
+        "because _inject_current_ray flag is on."
+    )
+
+    maybe_ray_dir = os.path.join(site_packages_path, "ray")
+    if os.path.isdir(maybe_ray_dir):
+        logger.warning(f"Replacing existing ray installation with {ray_path}")
+        shutil.rmtree(maybe_ray_dir)
+
+    # See usage of *.pth file at
+    # https://docs.python.org/3/library/site.html
+    with open(os.path.join(site_packages_path, "ray_shared.pth"), "w") as f:
+        f.write(ray_path)
+
+
+def _current_py_version():
+    return ".".join(map(str, sys.version_info[:3]))  # like 3.6.10
+
+
+def current_ray_pip_specifier(
+    logger: Optional[logging.Logger] = default_logger,
+) -> Optional[str]:
+    """The pip requirement specifier for the running version of Ray.
+
+    Args:
+        logger: Logger used to warn when the running Ray version cannot be
+            detected (e.g. when running a source build).
+
+    Returns:
+        A string which can be passed to `pip install` to install the
+        currently running Ray version, or None if running on a version
+        built from source locally (likely if you are developing Ray).
+
+    Examples:
+        Returns "https://s3-us-west-2.amazonaws.com/ray-wheels/[..].whl"
+            if running a stable release, a nightly or a specific commit
+    """
+    if os.environ.get("RAY_CI_POST_WHEEL_TESTS"):
+        # Running in Buildkite CI after the wheel has been built.
+        # Wheels are at in the ray/.whl directory, but use relative path to
+        # allow for testing locally if needed.
+        return os.path.join(
+            Path(ray.__file__).resolve().parents[2], ".whl", get_wheel_filename()
+        )
+    elif ray.__commit__ == "{{RAY_COMMIT_SHA}}":
+        # Running on a version built from source locally.
+        if os.environ.get("RAY_RUNTIME_ENV_LOCAL_DEV_MODE") != "1":
+            logger.warning(
+                "Current Ray version could not be detected, most likely "
+                "because you have manually built Ray from source.  To use "
+                "runtime_env in this case, set the environment variable "
+                "RAY_RUNTIME_ENV_LOCAL_DEV_MODE=1."
+            )
+        return None
+    elif "dev" in ray.__version__:
+        # Running on a nightly wheel.
+        return get_master_wheel_url()
+    else:
+        return get_release_wheel_url()
+
+
+def inject_dependencies(
+    conda_dict: Dict[Any, Any],
+    py_version: str,
+    pip_dependencies: Optional[List[str]] = None,
+) -> Dict[Any, Any]:
+    """Add Ray, Python and (optionally) extra pip dependencies to a conda dict.
+
+    Args:
+        conda_dict: A dict representing the JSON-serialized conda
+            environment YAML file.  This dict will be modified and returned.
+        py_version: A string representing a Python version to inject
+            into the conda dependencies, e.g. "3.7.7"
+        pip_dependencies: A list of pip dependencies that
+            will be prepended to the list of pip dependencies in
+            the conda dict.  If the conda dict does not already have a "pip"
+            field, one will be created.
+    Returns:
+        The modified dict.  (Note: the input argument conda_dict is modified
+        and returned.)
+    """
+    if pip_dependencies is None:
+        pip_dependencies = []
+    if conda_dict.get("dependencies") is None:
+        conda_dict["dependencies"] = []
+
+    # Inject Python dependency.
+    deps = conda_dict["dependencies"]
+
+    # Add current python dependency.  If the user has already included a
+    # python version dependency, conda will raise a readable error if the two
+    # are incompatible, e.g:
+    #   ResolvePackageNotFound: - python[version='3.5.*,>=3.6']
+    deps.append(f"python={py_version}")
+
+    if "pip" not in deps:
+        deps.append("pip")
+
+    # Insert pip dependencies.
+    found_pip_dict = False
+    for dep in deps:
+        if isinstance(dep, dict) and dep.get("pip") and isinstance(dep["pip"], list):
+            dep["pip"] = pip_dependencies + dep["pip"]
+            found_pip_dict = True
+            break
+    if not found_pip_dict:
+        deps.append({"pip": pip_dependencies})
+
+    return conda_dict
+
+
+def _get_conda_env_hash(conda_dict: Dict) -> str:
+    # Set `sort_keys=True` so that different orderings yield the same hash.
+    serialized_conda_spec = json.dumps(conda_dict, sort_keys=True)
+    hash = hashlib.sha1(serialized_conda_spec.encode("utf-8")).hexdigest()
+    return hash
+
+
+def get_uri(runtime_env: Dict) -> Optional[str]:
+    """Return `"conda://<hashed_dependencies>"`, or None if no GC required."""
+    conda = runtime_env.get("conda")
+    if conda is not None:
+        if isinstance(conda, str):
+            # User-preinstalled conda env.  We don't garbage collect these, so
+            # we don't track them with URIs.
+            uri = None
+        elif isinstance(conda, dict):
+            uri = f"conda://{_get_conda_env_hash(conda_dict=conda)}"
+        else:
+            raise TypeError(
+                "conda field received by RuntimeEnvAgent must be "
+                f"str or dict, not {type(conda).__name__}."
+            )
+    else:
+        uri = None
+    return uri
+
+
+def _get_conda_dict_with_ray_inserted(
+    runtime_env: "RuntimeEnv",  # noqa: F821
+    logger: Optional[logging.Logger] = default_logger,
+) -> Dict[str, Any]:
+    """Returns the conda spec with the Ray and `python` dependency inserted."""
+    conda_dict = json.loads(runtime_env.conda_config())
+    assert conda_dict is not None
+
+    ray_pip = current_ray_pip_specifier(logger=logger)
+    if ray_pip:
+        extra_pip_dependencies = [ray_pip, "ray[default]"]
+    elif runtime_env.get_extension("_inject_current_ray"):
+        extra_pip_dependencies = _resolve_install_from_source_ray_dependencies()
+    else:
+        extra_pip_dependencies = []
+    conda_dict = inject_dependencies(
+        conda_dict, _current_py_version(), extra_pip_dependencies
+    )
+    return conda_dict
+
+
+class CondaPlugin(RuntimeEnvPlugin):
+
+    name = "conda"
+
+    def __init__(self, resources_dir: str):
+        self._resources_dir = os.path.join(resources_dir, "conda")
+        try_to_create_directory(self._resources_dir)
+
+        # It is not safe for multiple processes to install conda envs
+        # concurrently, even if the envs are different, so use a global
+        # lock for all conda installs and deletions.
+        # See https://github.com/ray-project/ray/issues/17086
+        self._installs_and_deletions_file_lock = os.path.join(
+            self._resources_dir, "ray-conda-installs-and-deletions.lock"
+        )
+        # A set of named conda environments (instead of yaml or dict)
+        # that are validated to exist.
+        # NOTE: It has to be only used within the same thread, which
+        # is an event loop.
+        # Also, we don't need to GC this field because it is pretty small.
+        self._validated_named_conda_env = set()
+
+    def _get_path_from_hash(self, hash: str) -> str:
+        """Generate a path from the hash of a conda or pip spec.
+
+        The output path also functions as the name of the conda environment
+        when using the `--prefix` option to `conda create` and `conda remove`.
+
+        Example output:
+            /tmp/ray/session_2021-11-03_16-33-59_356303_41018/runtime_resources
+                /conda/ray-9a7972c3a75f55e976e620484f58410c920db091
+        """
+        return os.path.join(self._resources_dir, hash)
+
+    def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]:  # noqa: F821
+        """Return the conda URI from the RuntimeEnv if it exists, else return []."""
+        conda_uri = runtime_env.conda_uri()
+        if conda_uri:
+            return [conda_uri]
+        return []
+
+    def delete_uri(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ) -> int:
+        """Delete URI and return the number of bytes deleted."""
+        logger.info(f"Got request to delete URI {uri}")
+        protocol, hash = parse_uri(uri)
+        if protocol != Protocol.CONDA:
+            raise ValueError(
+                "CondaPlugin can only delete URIs with protocol "
+                f"conda.  Received protocol {protocol}, URI {uri}"
+            )
+
+        conda_env_path = self._get_path_from_hash(hash)
+        local_dir_size = get_directory_size_bytes(conda_env_path)
+
+        with FileLock(self._installs_and_deletions_file_lock):
+            successful = delete_conda_env(prefix=conda_env_path, logger=logger)
+        if not successful:
+            logger.warning(f"Error when deleting conda env {conda_env_path}. ")
+            return 0
+
+        return local_dir_size
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger = default_logger,
+    ) -> int:
+        if not runtime_env.has_conda():
+            return 0
+
+        def _create():
+            result = parse_and_validate_conda(runtime_env.get("conda"))
+
+            if isinstance(result, str):
+                # The conda env name is given.
+                # In this case, we only verify if the given
+                # conda env exists.
+
+                # If the env is already validated, do nothing.
+                if result in self._validated_named_conda_env:
+                    return 0
+
+                conda_info = get_conda_info_json()
+                envs = get_conda_envs(conda_info)
+
+                # We accept `result` as a conda name or full path.
+                if not any(result == env[0] or result == env[1] for env in envs):
+                    raise ValueError(
+                        f"The given conda environment '{result}' "
+                        f"from the runtime env {runtime_env} doesn't "
+                        "exist from the output of `conda info --json`. "
+                        "You can only specify an env that already exists. "
+                        f"Please make sure to create an env {result} "
+                    )
+                self._validated_named_conda_env.add(result)
+                return 0
+
+            logger.debug(
+                "Setting up conda for runtime_env: " f"{runtime_env.serialize()}"
+            )
+            protocol, hash = parse_uri(uri)
+            conda_env_name = self._get_path_from_hash(hash)
+
+            conda_dict = _get_conda_dict_with_ray_inserted(runtime_env, logger=logger)
+
+            logger.info(f"Setting up conda environment with {runtime_env}")
+            with FileLock(self._installs_and_deletions_file_lock):
+                try:
+                    conda_yaml_file = os.path.join(
+                        self._resources_dir, "environment.yml"
+                    )
+                    with open(conda_yaml_file, "w") as file:
+                        yaml.dump(conda_dict, file)
+                    create_conda_env_if_needed(
+                        conda_yaml_file, prefix=conda_env_name, logger=logger
+                    )
+                finally:
+                    os.remove(conda_yaml_file)
+
+                if runtime_env.get_extension("_inject_current_ray"):
+                    _inject_ray_to_conda_site(conda_path=conda_env_name, logger=logger)
+            logger.info(f"Finished creating conda environment at {conda_env_name}")
+            return get_directory_size_bytes(conda_env_name)
+
+        loop = get_or_create_event_loop()
+        return await loop.run_in_executor(None, _create)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        if not runtime_env.has_conda():
+            return
+
+        if runtime_env.conda_env_name():
+            conda_env_name = runtime_env.conda_env_name()
+        else:
+            protocol, hash = parse_uri(runtime_env.conda_uri())
+            conda_env_name = self._get_path_from_hash(hash)
+        context.py_executable = "python"
+        context.command_prefix += get_conda_activate_commands(conda_env_name)

python/ray/_private/runtime_env/conda_utils.py

@@ -0,0 +1,285 @@
+import hashlib
+import json
+import logging
+import os
+import shutil
+import subprocess
+from typing import List, Optional, Tuple, Union
+
+"""Utilities for conda.  Adapted from https://github.com/mlflow/mlflow."""
+
+# Name of environment variable indicating a path to a conda installation. Ray
+# will default to running "conda" if unset.
+RAY_CONDA_HOME = "RAY_CONDA_HOME"
+
+_WIN32 = os.name == "nt"
+
+
+def get_conda_activate_commands(conda_env_name: str) -> List[str]:
+    """
+    Get a list of commands to run to silently activate the given conda env.
+    """
+    #  Checking for newer conda versions
+    if not _WIN32 and ("CONDA_EXE" in os.environ or RAY_CONDA_HOME in os.environ):
+        conda_path = get_conda_bin_executable("conda")
+        activate_conda_env = [
+            ".",
+            f"{os.path.dirname(conda_path)}/../etc/profile.d/conda.sh",
+            "&&",
+        ]
+        activate_conda_env += ["conda", "activate", conda_env_name]
+
+    else:
+        activate_path = get_conda_bin_executable("activate")
+        if not _WIN32:
+            # Use bash command syntax
+            activate_conda_env = ["source", activate_path, conda_env_name]
+        else:
+            conda_path = get_conda_bin_executable("conda")
+            activate_conda_env = [conda_path, "activate", conda_env_name]
+    return activate_conda_env + ["1>&2", "&&"]
+
+
+def get_conda_bin_executable(executable_name: str) -> str:
+    """
+    Return path to the specified executable, assumed to be discoverable within
+    a conda installation.
+
+    The conda home directory (expected to contain a 'bin' subdirectory on
+    linux) is configurable via the ``RAY_CONDA_HOME`` environment variable. If
+    ``RAY_CONDA_HOME`` is unspecified, try the ``CONDA_EXE`` environment
+    variable set by activating conda. If neither is specified, this method
+    returns `executable_name`.
+    """
+    conda_home = os.environ.get(RAY_CONDA_HOME)
+    if conda_home:
+        if _WIN32:
+            candidate = os.path.join(conda_home, "%s.exe" % executable_name)
+            if os.path.exists(candidate):
+                return candidate
+            candidate = os.path.join(conda_home, "%s.bat" % executable_name)
+            if os.path.exists(candidate):
+                return candidate
+        else:
+            return os.path.join(conda_home, "bin/%s" % executable_name)
+    else:
+        conda_home = "."
+    # Use CONDA_EXE as per https://github.com/conda/conda/issues/7126
+    if "CONDA_EXE" in os.environ:
+        conda_bin_dir = os.path.dirname(os.environ["CONDA_EXE"])
+        if _WIN32:
+            candidate = os.path.join(conda_home, "%s.exe" % executable_name)
+            if os.path.exists(candidate):
+                return candidate
+            candidate = os.path.join(conda_home, "%s.bat" % executable_name)
+            if os.path.exists(candidate):
+                return candidate
+        else:
+            return os.path.join(conda_bin_dir, executable_name)
+    if _WIN32:
+        return executable_name + ".bat"
+    return executable_name
+
+
+def _get_conda_env_name(conda_env_path: str) -> str:
+    conda_env_contents = open(conda_env_path).read()
+    return "ray-%s" % hashlib.sha1(conda_env_contents.encode("utf-8")).hexdigest()
+
+
+def create_conda_env_if_needed(
+    conda_yaml_file: str, prefix: str, logger: Optional[logging.Logger] = None
+) -> None:
+    """
+    Given a conda YAML, creates a conda environment containing the required
+    dependencies if such a conda environment doesn't already exist.
+    Args:
+        conda_yaml_file: The path to a conda `environment.yml` file.
+        prefix: Directory to install the environment into via
+            the `--prefix` option to conda create.  This also becomes the name
+            of the conda env; i.e. it can be passed into `conda activate` and
+            `conda remove`
+        logger: Logger used to surface progress and errors; defaults to the
+            module logger when not provided.
+    """
+    if logger is None:
+        logger = logging.getLogger(__name__)
+
+    conda_path = get_conda_bin_executable("conda")
+    try:
+        exec_cmd([conda_path, "--help"], throw_on_error=False)
+    except (EnvironmentError, FileNotFoundError):
+        raise ValueError(
+            f"Could not find Conda executable at '{conda_path}'. "
+            "Ensure Conda is installed as per the instructions at "
+            "https://conda.io/projects/conda/en/latest/"
+            "user-guide/install/index.html. "
+            "You can also configure Ray to look for a specific "
+            f"Conda executable by setting the {RAY_CONDA_HOME} "
+            "environment variable to the path of the Conda executable."
+        )
+
+    _, stdout, _ = exec_cmd([conda_path, "env", "list", "--json"])
+    envs = json.loads(stdout[stdout.index("{") :])["envs"]
+
+    if prefix in envs:
+        logger.info(f"Conda environment {prefix} already exists.")
+        return
+
+    create_cmd = [
+        conda_path,
+        "env",
+        "create",
+        "--file",
+        conda_yaml_file,
+        "--prefix",
+        prefix,
+    ]
+
+    logger.info(f"Creating conda environment {prefix}")
+    exit_code, output = exec_cmd_stream_to_logger(create_cmd, logger)
+    if exit_code != 0:
+        if os.path.exists(prefix):
+            shutil.rmtree(prefix)
+        raise RuntimeError(
+            f"Failed to install conda environment {prefix}:\nOutput:\n{output}"
+        )
+
+
+def delete_conda_env(prefix: str, logger: Optional[logging.Logger] = None) -> bool:
+    if logger is None:
+        logger = logging.getLogger(__name__)
+
+    logger.info(f"Deleting conda environment {prefix}")
+
+    conda_path = get_conda_bin_executable("conda")
+    delete_cmd = [conda_path, "remove", "-p", prefix, "--all", "-y"]
+    exit_code, output = exec_cmd_stream_to_logger(delete_cmd, logger)
+
+    if exit_code != 0:
+        logger.debug(f"Failed to delete conda environment {prefix}:\n{output}")
+        return False
+
+    return True
+
+
+def get_conda_env_list() -> list:
+    """
+    Get conda env list in full paths.
+    """
+    conda_path = get_conda_bin_executable("conda")
+    try:
+        exec_cmd([conda_path, "--help"], throw_on_error=False)
+    except EnvironmentError:
+        raise ValueError(f"Could not find Conda executable at {conda_path}.")
+    _, stdout, _ = exec_cmd([conda_path, "env", "list", "--json"])
+    envs = json.loads(stdout)["envs"]
+    return envs
+
+
+def get_conda_info_json() -> dict:
+    """
+    Get `conda info --json` output.
+
+    Returns dict of conda info. See [1] for more details. We mostly care about these
+    keys:
+
+    - `conda_prefix`: str The path to the conda installation.
+    - `envs`: List[str] absolute paths to conda environments.
+
+    [1] https://github.com/conda/conda/blob/main/conda/cli/main_info.py
+    """
+    conda_path = get_conda_bin_executable("conda")
+    try:
+        exec_cmd([conda_path, "--help"], throw_on_error=False)
+    except EnvironmentError:
+        raise ValueError(f"Could not find Conda executable at {conda_path}.")
+    _, stdout, _ = exec_cmd([conda_path, "info", "--json"])
+    return json.loads(stdout)
+
+
+def get_conda_envs(conda_info: dict) -> List[Tuple[str, str]]:
+    """
+    Gets the conda environments, as a list of (name, path) tuples.
+    """
+    prefix = conda_info["conda_prefix"]
+    ret = []
+    for env in conda_info["envs"]:
+        if env == prefix:
+            ret.append(("base", env))
+        else:
+            ret.append((os.path.basename(env), env))
+    return ret
+
+
+class ShellCommandException(Exception):
+    pass
+
+
+def exec_cmd(
+    cmd: List[str], throw_on_error: bool = True, logger: Optional[logging.Logger] = None
+) -> Union[int, Tuple[int, str, str]]:
+    """
+    Runs a command as a child process.
+
+    A convenience wrapper for running a command from a Python script.
+
+    Note on the return value: A tuple of the exit code,
+    standard output and standard error is returned.
+
+    Args:
+        cmd: the command to run, as a list of strings
+        throw_on_error: if true, raises an Exception if the exit code of the
+            program is nonzero
+        logger: Unused; retained for API compatibility.
+
+    Returns:
+        A tuple of (exit_code, stdout, stderr) from the child process.
+    """
+    child = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stdin=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        universal_newlines=True,
+    )
+    (stdout, stderr) = child.communicate()
+    exit_code = child.wait()
+    if throw_on_error and exit_code != 0:
+        raise ShellCommandException(
+            "Non-zero exit code: %s\n\nSTDOUT:\n%s\n\nSTDERR:%s"
+            % (exit_code, stdout, stderr)
+        )
+    return exit_code, stdout, stderr
+
+
+def exec_cmd_stream_to_logger(
+    cmd: List[str], logger: logging.Logger, n_lines: int = 50, **kwargs
+) -> Tuple[int, str]:
+    """Runs a command as a child process, streaming output to the logger.
+
+    The last n_lines lines of output are also returned (stdout and stderr).
+    """
+    if "env" in kwargs and _WIN32 and "PATH" not in [x.upper() for x in kwargs.keys]:
+        raise ValueError("On windows, Popen requires 'PATH' in 'env'")
+    child = subprocess.Popen(
+        cmd,
+        universal_newlines=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        **kwargs,
+    )
+    last_n_lines = []
+    with child.stdout:
+        for line in iter(child.stdout.readline, b""):
+            exit_code = child.poll()
+            if exit_code is not None:
+                break
+            line = line.strip()
+            if not line:
+                continue
+            last_n_lines.append(line.strip())
+            last_n_lines = last_n_lines[-n_lines:]
+            logger.info(line.strip())
+
+    exit_code = child.wait()
+    return exit_code, "\n".join(last_n_lines)

python/ray/_private/runtime_env/constants.py

@@ -0,0 +1,28 @@
+# Env var set by job manager to pass runtime env and metadata to subprocess
+RAY_JOB_CONFIG_JSON_ENV_VAR = "RAY_JOB_CONFIG_JSON_ENV_VAR"
+
+# The plugin config which should be loaded when ray cluster starts.
+# It is a json formatted config,
+# e.g. [{"class": "xxx.xxx.xxx_plugin", "priority": 10}].
+RAY_RUNTIME_ENV_PLUGINS_ENV_VAR = "RAY_RUNTIME_ENV_PLUGINS"
+
+# The field name of plugin class in the plugin config.
+RAY_RUNTIME_ENV_CLASS_FIELD_NAME = "class"
+
+# The field name of priority in the plugin config.
+RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME = "priority"
+
+# The default priority of runtime env plugin.
+RAY_RUNTIME_ENV_PLUGIN_DEFAULT_PRIORITY = 10
+
+# The minimum priority of runtime env plugin.
+RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY = 0
+
+# The maximum priority of runtime env plugin.
+RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY = 100
+
+# The schema files or directories of plugins which should be loaded in workers.
+RAY_RUNTIME_ENV_PLUGIN_SCHEMAS_ENV_VAR = "RAY_RUNTIME_ENV_PLUGIN_SCHEMAS"
+
+# The file suffix of runtime env plugin schemas.
+RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX = ".json"

python/ray/_private/runtime_env/context.py

@@ -0,0 +1,108 @@
+import json
+import logging
+import os
+import shlex
+import subprocess
+import sys
+from typing import Dict, List, Optional
+
+from ray._private.services import get_ray_jars_dir
+from ray._private.utils import update_envs
+from ray.core.generated.common_pb2 import Language
+from ray.util.annotations import DeveloperAPI
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class RuntimeEnvContext:
+    """A context used to describe the created runtime env."""
+
+    def __init__(
+        self,
+        command_prefix: List[str] = None,
+        env_vars: Dict[str, str] = None,
+        py_executable: Optional[str] = None,
+        override_worker_entrypoint: Optional[str] = None,
+        java_jars: List[str] = None,
+    ):
+        self.command_prefix = command_prefix or []
+        self.env_vars = env_vars or {}
+        self.py_executable = py_executable or sys.executable
+        self.override_worker_entrypoint: Optional[str] = override_worker_entrypoint
+        self.java_jars = java_jars or []
+
+    def serialize(self) -> str:
+        return json.dumps(self.__dict__)
+
+    @staticmethod
+    def deserialize(json_string):
+        return RuntimeEnvContext(**json.loads(json_string))
+
+    def exec_worker(self, passthrough_args: List[str], language: Language):
+        update_envs(self.env_vars)
+
+        if language == Language.PYTHON and sys.platform == "win32":
+            executable = [self.py_executable]
+        elif language == Language.PYTHON:
+            executable = ["exec", self.py_executable]
+        elif language == Language.JAVA:
+            executable = ["java"]
+            ray_jars = os.path.join(get_ray_jars_dir(), "*")
+
+            local_java_jars = []
+            for java_jar in self.java_jars:
+                local_java_jars.append(f"{java_jar}/*")
+                local_java_jars.append(java_jar)
+
+            class_path_args = ["-cp", ray_jars + ":" + str(":".join(local_java_jars))]
+            passthrough_args = class_path_args + passthrough_args
+        elif sys.platform == "win32":
+            executable = []
+        else:
+            executable = ["exec"]
+
+        # By default, raylet uses the path to default_worker.py on host.
+        # However, the path to default_worker.py inside the container
+        # can be different. We need the user to specify the path to
+        # default_worker.py inside the container.
+        if self.override_worker_entrypoint:
+            logger.debug(
+                f"Changing the worker entrypoint from {passthrough_args[0]} to "
+                f"{self.override_worker_entrypoint}."
+            )
+            passthrough_args[0] = self.override_worker_entrypoint
+
+        if sys.platform == "win32":
+
+            def quote(s):
+                s = s.replace("&", "%26")
+                return s
+
+            passthrough_args = [quote(s) for s in passthrough_args]
+
+            cmd = [*self.command_prefix, *executable, *passthrough_args]
+            logger.debug(f"Exec'ing worker with command: {cmd}")
+            subprocess.Popen(cmd, shell=True).wait()
+        else:
+            # We use shlex to do the necessary shell escape
+            # of special characters in passthrough_args.
+            passthrough_args = [shlex.quote(s) for s in passthrough_args]
+            cmd = [*self.command_prefix, *executable, *passthrough_args]
+            # TODO(SongGuyang): We add this env to command for macOS because it doesn't
+            # work for the C++ process of `os.execvp`. We should find a better way to
+            # fix it.
+            MACOS_LIBRARY_PATH_ENV_NAME = "DYLD_LIBRARY_PATH"
+            if MACOS_LIBRARY_PATH_ENV_NAME in os.environ:
+                cmd.insert(
+                    0,
+                    f"{MACOS_LIBRARY_PATH_ENV_NAME}="
+                    f"{os.environ[MACOS_LIBRARY_PATH_ENV_NAME]}",
+                )
+            logger.debug(f"Exec'ing worker with command: {cmd}")
+            # PyCharm will monkey patch the os.execvp at
+            # .pycharm_helpers/pydev/_pydev_bundle/pydev_monkey.py
+            # The monkey patched os.execvp function has a different
+            # signature. So, we use os.execvp("executable", args=[])
+            # instead of os.execvp(file="executable", args=[])
+            os.execvp("bash", args=["bash", "-c", " ".join(cmd)])

python/ray/_private/runtime_env/default_impl.py

@@ -0,0 +1,5 @@
+from ray._private.runtime_env.image_uri import ImageURIPlugin
+
+
+def get_image_uri_plugin_cls():
+    return ImageURIPlugin

python/ray/_private/runtime_env/dependency_utils.py

@@ -0,0 +1,118 @@
+"""Util functions to manage dependency requirements."""
+
+import logging
+import os
+import tempfile
+from contextlib import asynccontextmanager
+from typing import List, Optional, Tuple
+
+from ray._private.runtime_env import virtualenv_utils
+from ray._private.runtime_env.utils import check_output_cmd
+
+INTERNAL_PIP_FILENAME = "ray_runtime_env_internal_pip_requirements.txt"
+MAX_INTERNAL_PIP_FILENAME_TRIES = 100
+
+
+def gen_requirements_txt(requirements_file: str, pip_packages: List[str]):
+    """Dump [pip_packages] to the given [requirements_file] for later env setup."""
+    with open(requirements_file, "w") as file:
+        for line in pip_packages:
+            file.write(line + "\n")
+
+
+@asynccontextmanager
+async def check_ray(python: str, cwd: str, logger: logging.Logger):
+    """A context manager to check ray is not overwritten.
+
+    Currently, we only check ray version and path. It works for virtualenv,
+        - ray is in Python's site-packages.
+        - ray is overwritten during yield.
+        - ray is in virtualenv's site-packages.
+    """
+
+    async def _get_ray_version_and_path() -> Tuple[str, str]:
+        with tempfile.TemporaryDirectory(
+            prefix="check_ray_version_tempfile"
+        ) as tmp_dir:
+            ray_version_path = os.path.join(tmp_dir, "ray_version.txt")
+            check_ray_cmd = [
+                python,
+                "-c",
+                """
+import ray
+with open(r"{ray_version_path}", "wt") as f:
+    f.write(ray.__version__)
+    f.write(" ")
+    f.write(ray.__path__[0])
+                """.format(
+                    ray_version_path=ray_version_path
+                ),
+            ]
+            if virtualenv_utils._WIN32:
+                env = os.environ.copy()
+            else:
+                env = {}
+            output = await check_output_cmd(
+                check_ray_cmd, logger=logger, cwd=cwd, env=env
+            )
+            logger.info(f"try to write ray version information in: {ray_version_path}")
+            with open(ray_version_path, "rt") as f:
+                output = f.read()
+            # print after import ray may have  endings, so we strip them by *_
+            ray_version, ray_path, *_ = [s.strip() for s in output.split()]
+        return ray_version, ray_path
+
+    version, path = await _get_ray_version_and_path()
+    yield
+    actual_version, actual_path = await _get_ray_version_and_path()
+    if actual_version != version:
+        raise RuntimeError(
+            "Changing the ray version is not allowed: \n"
+            f"  current version: {actual_version}, "
+            f"  expect version: {version}, "
+            f"  current path: {actual_path}, "
+            f"  expect path: {path}, "
+            "Please ensure the dependencies in the runtime_env pip field "
+            "do not install a different version of Ray."
+        )
+    if actual_path != path:
+        logger.info(
+            f"Detected new Ray package with the same version at {actual_path} (vs system {path})."
+        )
+
+
+def get_requirements_file(target_dir: str, pip_list: Optional[List[str]]) -> str:
+    """Returns the path to the requirements file to use for this runtime env.
+
+    If pip_list is not None, we will check if the internal pip filename is in any of
+    the entries of pip_list. If so, we will append numbers to the end of the
+    filename until we find one that doesn't conflict. This prevents infinite
+    recursion if the user specifies the internal pip filename in their pip list.
+
+    Args:
+        target_dir: The directory to store the requirements file in.
+        pip_list: A list of pip requirements specified by the user.
+
+    Returns:
+        The path to the requirements file to use for this runtime env.
+    """
+
+    def filename_in_pip_list(filename: str) -> bool:
+        for pip_entry in pip_list:
+            if filename in pip_entry:
+                return True
+        return False
+
+    filename = INTERNAL_PIP_FILENAME
+    if pip_list is not None:
+        i = 1
+        while filename_in_pip_list(filename) and i < MAX_INTERNAL_PIP_FILENAME_TRIES:
+            filename = f"{INTERNAL_PIP_FILENAME}.{i}"
+            i += 1
+        if i == MAX_INTERNAL_PIP_FILENAME_TRIES:
+            raise RuntimeError(
+                "Could not find a valid filename for the internal "
+                "pip requirements file. Please specify a different "
+                "pip list in your runtime env."
+            )
+    return os.path.join(target_dir, filename)

python/ray/_private/runtime_env/image_uri.py

@@ -0,0 +1,229 @@
+import asyncio
+import logging
+import os
+import tempfile
+from typing import List, Optional
+
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+
+default_logger = logging.getLogger(__name__)
+
+
+async def _create_impl(image_uri: str, logger: logging.Logger):
+    # Pull image if it doesn't exist
+    # Also get path to `default_worker.py` inside the image.
+    with tempfile.TemporaryDirectory() as tmpdir:
+        os.chmod(tmpdir, 0o777)
+        result_file = os.path.join(tmpdir, "worker_path.txt")
+        get_worker_path_script = """
+import ray._private.workers.default_worker as dw
+with open('/shared/worker_path.txt', 'w') as f:
+    f.write(dw.__file__)
+"""
+        cmd = [
+            "podman",
+            "run",
+            "--rm",
+            "-v",
+            f"{tmpdir}:/shared:Z",
+            image_uri,
+            "python",
+            "-c",
+            get_worker_path_script,
+        ]
+
+        logger.info("Pulling image %s", image_uri)
+
+        process = await asyncio.create_subprocess_exec(
+            *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+        )
+
+        stdout, stderr = await process.communicate()
+
+        if process.returncode != 0:
+            raise RuntimeError(
+                f"Podman command failed: cmd={cmd}, returncode={process.returncode}, stdout={stdout.decode()}, stderr={stderr.decode()}"
+            )
+
+        if not os.path.exists(result_file):
+            raise FileNotFoundError(
+                f"Worker path file not created when getting worker path for image {image_uri}"
+            )
+
+        with open(result_file, "r") as f:
+            worker_path = f.read().strip()
+
+        if not worker_path.endswith(".py"):
+            raise ValueError(
+                f"Invalid worker path inferred in image {image_uri}: {worker_path}"
+            )
+
+        logger.info(f"Inferred worker path in image {image_uri}: {worker_path}")
+        return worker_path
+
+
+def _modify_context_impl(
+    image_uri: str,
+    worker_path: str,
+    run_options: Optional[List[str]],
+    context: RuntimeEnvContext,
+    logger: logging.Logger,
+    ray_tmp_dir: str,
+):
+    context.override_worker_entrypoint = worker_path
+
+    container_driver = "podman"
+    container_command = [
+        container_driver,
+        "run",
+        "-v",
+        ray_tmp_dir + ":" + ray_tmp_dir,
+        "--cgroup-manager=cgroupfs",
+        "--network=host",
+        "--pid=host",
+        "--ipc=host",
+        # NOTE(zcin): Mounted volumes in rootless containers are
+        # owned by the user `root`. The user on host (which will
+        # usually be `ray` if this is being run in a ray docker
+        # image) who started the container is mapped using user
+        # namespaces to the user `root` in a rootless container. In
+        # order for the Ray Python worker to access the mounted ray
+        # tmp dir, we need to use keep-id mode which maps the user
+        # as itself (instead of as `root`) into the container.
+        # https://www.redhat.com/sysadmin/rootless-podman-user-namespace-modes
+        "--userns=keep-id",
+    ]
+
+    # Environment variables to set in container
+    env_vars = dict()
+
+    # Propagate all host environment variables that have the prefix "RAY_"
+    # This should include RAY_RAYLET_PID
+    for env_var_name, env_var_value in os.environ.items():
+        if env_var_name.startswith("RAY_"):
+            env_vars[env_var_name] = env_var_value
+
+    # Support for runtime_env['env_vars']
+    env_vars.update(context.env_vars)
+
+    # Set environment variables
+    for env_var_name, env_var_value in env_vars.items():
+        container_command.append("--env")
+        container_command.append(f"{env_var_name}='{env_var_value}'")
+
+    # The RAY_JOB_ID environment variable is needed for the default worker.
+    # It won't be set at the time setup() is called, but it will be set
+    # when worker command is executed, so we use RAY_JOB_ID=$RAY_JOB_ID
+    # for the container start command
+    container_command.append("--env")
+    container_command.append("RAY_JOB_ID=$RAY_JOB_ID")
+
+    if run_options:
+        container_command.extend(run_options)
+    # TODO(chenk008): add resource limit
+    container_command.append("--entrypoint")
+    container_command.append("python")
+    container_command.append(image_uri)
+
+    # Example:
+    # podman run -v /tmp/ray:/tmp/ray
+    # --cgroup-manager=cgroupfs --network=host --pid=host --ipc=host
+    # --userns=keep-id --env RAY_RAYLET_PID=23478 --env RAY_JOB_ID=$RAY_JOB_ID
+    # --entrypoint python rayproject/ray:nightly-py39
+    container_command_str = " ".join(container_command)
+    logger.info(f"Starting worker in container with prefix {container_command_str}")
+
+    context.py_executable = container_command_str
+
+
+class ImageURIPlugin(RuntimeEnvPlugin):
+    """Starts worker in a container of a custom image."""
+
+    name = "image_uri"
+
+    @staticmethod
+    def get_compatible_keys():
+        return {"image_uri", "config", "env_vars"}
+
+    def __init__(self, ray_tmp_dir: str):
+        self._ray_tmp_dir = ray_tmp_dir
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger,
+    ) -> float:
+        if not runtime_env.image_uri():
+            return
+
+        self.worker_path = await _create_impl(runtime_env.image_uri(), logger)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        if not runtime_env.image_uri():
+            return
+
+        _modify_context_impl(
+            runtime_env.image_uri(),
+            self.worker_path,
+            [],
+            context,
+            logger,
+            self._ray_tmp_dir,
+        )
+
+
+class ContainerPlugin(RuntimeEnvPlugin):
+    """Starts worker in container."""
+
+    name = "container"
+
+    def __init__(self, ray_tmp_dir: str):
+        self._ray_tmp_dir = ray_tmp_dir
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger,
+    ) -> float:
+        if not runtime_env.has_py_container() or not runtime_env.py_container_image():
+            return
+
+        self.worker_path = await _create_impl(runtime_env.py_container_image(), logger)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        if not runtime_env.has_py_container() or not runtime_env.py_container_image():
+            return
+
+        if runtime_env.py_container_worker_path():
+            logger.warning(
+                "You are using `container.worker_path`, but the path to "
+                "`default_worker.py` is now automatically detected from the image. "
+                "`container.worker_path` is deprecated and will be removed in future "
+                "versions."
+            )
+
+        _modify_context_impl(
+            runtime_env.py_container_image(),
+            runtime_env.py_container_worker_path() or self.worker_path,
+            runtime_env.py_container_run_options(),
+            context,
+            logger,
+            self._ray_tmp_dir,
+        )

python/ray/_private/runtime_env/java_jars.py

@@ -0,0 +1,104 @@
+import logging
+import os
+from typing import Dict, List, Optional
+
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.packaging import (
+    delete_package,
+    download_and_unpack_package,
+    get_local_dir_from_uri,
+    is_jar_uri,
+)
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.utils import get_directory_size_bytes
+from ray._raylet import GcsClient
+from ray.exceptions import RuntimeEnvSetupError
+
+default_logger = logging.getLogger(__name__)
+
+
+class JavaJarsPlugin(RuntimeEnvPlugin):
+
+    name = "java_jars"
+
+    def __init__(self, resources_dir: str, gcs_client: GcsClient):
+        self._resources_dir = os.path.join(resources_dir, "java_jars_files")
+        self._gcs_client = gcs_client
+        try_to_create_directory(self._resources_dir)
+
+    def _get_local_dir_from_uri(self, uri: str):
+        return get_local_dir_from_uri(uri, self._resources_dir)
+
+    def delete_uri(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ) -> int:
+        """Delete URI and return the number of bytes deleted."""
+        local_dir = get_local_dir_from_uri(uri, self._resources_dir)
+        local_dir_size = get_directory_size_bytes(local_dir)
+
+        deleted = delete_package(uri, self._resources_dir)
+        if not deleted:
+            logger.warning(f"Tried to delete nonexistent URI: {uri}.")
+            return 0
+
+        return local_dir_size
+
+    def get_uris(self, runtime_env: dict) -> List[str]:
+        return runtime_env.java_jars()
+
+    async def _download_jars(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ):
+        """Download a jar URI."""
+        try:
+            jar_file = await download_and_unpack_package(
+                uri, self._resources_dir, self._gcs_client, logger=logger
+            )
+        except Exception as e:
+            raise RuntimeEnvSetupError(
+                "Failed to download jar file: {}".format(e)
+            ) from e
+        module_dir = self._get_local_dir_from_uri(uri)
+        logger.debug(f"Succeeded to download jar file {jar_file} .")
+        return module_dir
+
+    async def create(
+        self,
+        uri: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ) -> int:
+        if not uri:
+            return 0
+        if is_jar_uri(uri):
+            module_dir = await self._download_jars(uri=uri, logger=logger)
+        else:
+            try:
+                module_dir = await download_and_unpack_package(
+                    uri, self._resources_dir, self._gcs_client, logger=logger
+                )
+            except Exception as e:
+                raise RuntimeEnvSetupError(
+                    "Failed to download jar file: {}".format(e)
+                ) from e
+
+        return get_directory_size_bytes(module_dir)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env_dict: Dict,
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        for uri in uris:
+            module_dir = self._get_local_dir_from_uri(uri)
+            if not module_dir.exists():
+                raise ValueError(
+                    f"Local directory {module_dir} for URI {uri} does "
+                    "not exist on the cluster. Something may have gone wrong while "
+                    "downloading, unpacking or installing the java jar files."
+                )
+            context.java_jars.append(str(module_dir))

python/ray/_private/runtime_env/nsight.py

@@ -0,0 +1,149 @@
+import asyncio
+import copy
+import logging
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+from ray._common.utils import (
+    try_to_create_directory,
+)
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray.exceptions import RuntimeEnvSetupError
+
+default_logger = logging.getLogger(__name__)
+
+# Nsight options used when runtime_env={"_nsight": "default"}
+NSIGHT_DEFAULT_CONFIG = {
+    "t": "cuda,cudnn,cublas,nvtx",
+    "o": "'worker_process_%p'",
+    "stop-on-exit": "true",
+}
+
+
+def parse_nsight_config(nsight_config: Dict[str, str]) -> List[str]:
+    """
+    Function to convert dictionary of nsight options into
+    nsight command line
+
+    The function returns:
+    - List[str]: nsys profile cmd line split into list of str
+    """
+    nsight_cmd = ["nsys", "profile"]
+    for option, option_val in nsight_config.items():
+        # option standard based on
+        # https://www.gnu.org/software/libc/manual/html_node/Argument-Syntax.html
+        if len(option) > 1:
+            nsight_cmd.append(f"--{option}={option_val}")
+        else:
+            nsight_cmd += [f"-{option}", option_val]
+    return nsight_cmd
+
+
+class NsightPlugin(RuntimeEnvPlugin):
+    name = "_nsight"
+
+    def __init__(self, resources_dir: str):
+        self.nsight_cmd = []
+
+        # replace this with better way to get logs dir
+        session_dir, runtime_dir = os.path.split(resources_dir)
+        self._nsight_dir = Path(session_dir) / "logs" / "nsight"
+        try_to_create_directory(self._nsight_dir)
+
+    async def _check_nsight_script(
+        self, nsight_config: Dict[str, str]
+    ) -> Tuple[bool, str]:
+        """
+        Function to validate if nsight_config is a valid nsight profile options
+        Args:
+            nsight_config: dictionary mapping nsight option to it's value
+        Returns:
+            a tuple consists of a boolean indicating if the nsight_config
+            is valid option and an error message if the nsight_config is invalid
+        """
+
+        # use empty as nsight report test filename
+        nsight_config_copy = copy.deepcopy(nsight_config)
+        nsight_config_copy["o"] = str(Path(self._nsight_dir) / "empty")
+        nsight_cmd = parse_nsight_config(nsight_config_copy)
+        try:
+            nsight_cmd = nsight_cmd + [sys.executable, "-c", '""']
+            process = await asyncio.create_subprocess_exec(
+                *nsight_cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            stdout, stderr = await process.communicate()
+            error_msg = stderr.strip() if stderr.strip() != "" else stdout.strip()
+
+            # cleanup test.nsys-rep file
+            clean_up_cmd = ["rm", f"{nsight_config_copy['o']}.nsys-rep"]
+            cleanup_process = await asyncio.create_subprocess_exec(
+                *clean_up_cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            _, _ = await cleanup_process.communicate()
+            if process.returncode == 0:
+                return True, None
+            else:
+                return False, error_msg
+        except FileNotFoundError:
+            return False, ("nsight is not installed")
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger = default_logger,
+    ) -> int:
+        nsight_config = runtime_env.nsight()
+        if not nsight_config:
+            return 0
+
+        if nsight_config and sys.platform != "linux":
+            raise RuntimeEnvSetupError(
+                "Nsight CLI is only available in Linux.\n"
+                "More information can be found in "
+                "https://docs.nvidia.com/nsight-compute/NsightComputeCli/index.html"
+            )
+
+        if isinstance(nsight_config, str):
+            if nsight_config == "default":
+                nsight_config = NSIGHT_DEFAULT_CONFIG
+            else:
+                raise RuntimeEnvSetupError(
+                    f"Unsupported nsight config: {nsight_config}. "
+                    "The supported config is 'default' or "
+                    "Dictionary of nsight options"
+                )
+
+        is_valid_nsight_cmd, error_msg = await self._check_nsight_script(nsight_config)
+        if not is_valid_nsight_cmd:
+            logger.warning(error_msg)
+            raise RuntimeEnvSetupError(
+                "nsight profile failed to run with the following "
+                f"error message:\n {error_msg}"
+            )
+        # add set output path to logs dir
+        nsight_config["o"] = str(
+            Path(self._nsight_dir) / nsight_config.get("o", NSIGHT_DEFAULT_CONFIG["o"])
+        )
+
+        self.nsight_cmd = parse_nsight_config(nsight_config)
+        return 0
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        logger.info("Running nsight profiler")
+        context.py_executable = " ".join(self.nsight_cmd) + " python"

python/ray/_private/runtime_env/pip.py

@@ -0,0 +1,422 @@
+import asyncio
+import hashlib
+import json
+import logging
+import os
+import re
+import shutil
+import sys
+from asyncio import create_task, get_running_loop
+from typing import Dict, List, Optional
+
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env import dependency_utils, virtualenv_utils
+from ray._private.runtime_env.packaging import Protocol, parse_uri
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.runtime_env.utils import check_output_cmd
+from ray._private.utils import get_directory_size_bytes
+
+default_logger = logging.getLogger(__name__)
+
+# Matches unresolved environment variable placeholders such as
+# ${RAY_RUNTIME_ENV_CREATE_WORKING_DIR} or $VAR. Such placeholders are
+# expanded by the runtime env agent only after the driver-side hash is
+# computed, so any path containing one is not a real path on the driver
+# and must not be opened during hash computation.
+_ENV_VAR_PATTERN = re.compile(r"\$\{[^}]+\}|\$[A-Za-z_][A-Za-z0-9_]*")
+
+
+def _parse_requirements_file(file_path: str) -> List[str]:
+    packages = []
+    with open(file_path, "r", encoding="utf-8") as f:
+        for line in f:
+            # Strip whitespace and remove inline comments (preceded by a space)
+            line = line.split(" #")[0].strip()
+            if not line or line.startswith("#"):
+                continue
+            packages.append(line)
+    return packages
+
+
+def _get_pip_hash(pip_dict: Dict) -> str:
+    pip_dict_copy = pip_dict.copy()
+    # Using a list as a stack for iterative processing to handle nested requirements.
+    # Each item is a tuple (package_spec, parent_dir), where parent_dir is the directory
+    # of the file that contained this package spec (None for top-level packages)
+    packages_to_process = [
+        (pkg, None) for pkg in reversed(pip_dict_copy.get("packages", []))
+    ]
+    expanded_packages = []
+    # Track visited files using absolute paths to prevent circular references
+    visited_files = set()
+
+    while packages_to_process:
+        pkg, parent_dir = packages_to_process.pop()
+
+        file_path = None
+        if pkg.startswith("-r"):
+            file_path = pkg[2:].lstrip()
+        elif pkg.startswith("--requirement"):
+            file_path = pkg[len("--requirement") :].lstrip()
+            if file_path.startswith("="):
+                file_path = file_path[1:].lstrip()
+        else:
+            expanded_packages.append(pkg)
+            continue
+
+        if file_path is not None:
+            # If the path contains an unresolved environment variable
+            # placeholder (e.g. ${RAY_RUNTIME_ENV_CREATE_WORKING_DIR}),
+            # we cannot open it on the driver. Keep the original spec
+            # string in the hash input so the URI still reflects the user
+            # input, and let the runtime env agent expand and read the
+            # file on the worker side.
+            if _ENV_VAR_PATTERN.search(file_path):
+                expanded_packages.append(pkg)
+                continue
+
+            if parent_dir and not os.path.isabs(file_path):
+                file_path = os.path.join(parent_dir, file_path)
+
+            try:
+                abs_file_path = os.path.abspath(file_path)
+            except Exception:
+                default_logger.warning(f"Invalid path: {file_path}")
+                continue
+
+            if abs_file_path in visited_files:
+                default_logger.warning(
+                    f"Skipping circular reference to {abs_file_path}"
+                )
+                continue
+            visited_files.add(abs_file_path)
+
+            file_dir = os.path.dirname(abs_file_path)
+            packages_from_file = _parse_requirements_file(abs_file_path)
+            packages_to_process.extend(
+                [(p, file_dir) for p in reversed(packages_from_file)]
+            )
+    pip_dict_copy["packages"] = expanded_packages
+    serialized_pip_spec = json.dumps(pip_dict_copy, sort_keys=True)
+    hash_val = hashlib.sha1(serialized_pip_spec.encode("utf-8")).hexdigest()
+    return hash_val
+
+
+def get_uri(runtime_env: Dict) -> Optional[str]:
+    """Return `"pip://<hashed_dependencies>"`, or None if no GC required."""
+    pip = runtime_env.get("pip")
+    if pip is not None:
+        if isinstance(pip, dict):
+            uri = "pip://" + _get_pip_hash(pip_dict=pip)
+        elif isinstance(pip, list):
+            uri = "pip://" + _get_pip_hash(pip_dict=dict(packages=pip))
+        else:
+            raise TypeError(
+                "pip field received by RuntimeEnvAgent must be "
+                f"list or dict, not {type(pip).__name__}."
+            )
+    else:
+        uri = None
+    return uri
+
+
+class PipProcessor:
+    def __init__(
+        self,
+        target_dir: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        try:
+            import virtualenv  # noqa: F401 ensure virtualenv exists.
+        except ImportError:
+            raise RuntimeError(
+                f"Please install virtualenv "
+                f"`{sys.executable} -m pip install virtualenv`"
+                f"to enable pip runtime env."
+            )
+        logger.debug("Setting up pip for runtime_env: %s", runtime_env)
+        self._target_dir = target_dir
+        self._runtime_env = runtime_env
+        self._logger = logger
+
+        self._pip_config = self._runtime_env.pip_config()
+        self._pip_env = os.environ.copy()
+        self._pip_env.update(self._runtime_env.env_vars())
+
+    @classmethod
+    async def _ensure_pip_version(
+        cls,
+        path: str,
+        pip_version: Optional[str],
+        cwd: str,
+        pip_env: Dict,
+        logger: logging.Logger,
+    ):
+        """Run the pip command to reinstall pip to the specified version."""
+        if not pip_version:
+            return
+
+        python = virtualenv_utils.get_virtualenv_python(path)
+        # Ensure pip version.
+        pip_reinstall_cmd = [
+            python,
+            "-m",
+            "pip",
+            "install",
+            "--disable-pip-version-check",
+            f"pip{pip_version}",
+        ]
+        logger.info("Installing pip with version %s", pip_version)
+
+        await check_output_cmd(pip_reinstall_cmd, logger=logger, cwd=cwd, env=pip_env)
+
+    async def _pip_check(
+        self,
+        path: str,
+        pip_check: bool,
+        cwd: str,
+        pip_env: Dict,
+        logger: logging.Logger,
+    ):
+        """Run the pip check command to check python dependency conflicts.
+        If exists conflicts, the exit code of pip check command will be non-zero.
+        """
+        if not pip_check:
+            logger.info("Skip pip check.")
+            return
+        python = virtualenv_utils.get_virtualenv_python(path)
+
+        await check_output_cmd(
+            [python, "-m", "pip", "check", "--disable-pip-version-check"],
+            logger=logger,
+            cwd=cwd,
+            env=pip_env,
+        )
+
+        logger.info("Pip check on %s successfully.", path)
+
+    async def _install_pip_packages(
+        self,
+        path: str,
+        pip_packages: List[str],
+        cwd: str,
+        pip_env: Dict,
+        logger: logging.Logger,
+    ):
+        virtualenv_path = virtualenv_utils.get_virtualenv_path(path)
+        python = virtualenv_utils.get_virtualenv_python(path)
+        # TODO(fyrestone): Support -i, --no-deps, --no-cache-dir, ...
+        pip_requirements_file = dependency_utils.get_requirements_file(
+            path, pip_packages
+        )
+
+        # Avoid blocking the event loop.
+        loop = get_running_loop()
+        await loop.run_in_executor(
+            None,
+            dependency_utils.gen_requirements_txt,
+            pip_requirements_file,
+            pip_packages,
+        )
+
+        # Install all dependencies
+        # The default options for pip install are
+        #
+        # --disable-pip-version-check
+        #   Don't periodically check PyPI to determine whether a new version
+        #   of pip is available for download.
+        #
+        # --no-cache-dir
+        #   Disable the cache, the pip runtime env is a one-time installation,
+        #   and we don't need to handle the pip cache broken.
+        #
+        # Allow users to specify their own options to install packages via `pip`.
+        pip_install_cmd = [
+            python,
+            "-m",
+            "pip",
+            "install",
+            "-r",
+            pip_requirements_file,
+        ]
+
+        pip_opt_list = self._pip_config.get(
+            "pip_install_options", ["--disable-pip-version-check", "--no-cache-dir"]
+        )
+        pip_install_cmd.extend(pip_opt_list)
+
+        logger.info("Installing python requirements to %s", virtualenv_path)
+
+        await check_output_cmd(pip_install_cmd, logger=logger, cwd=cwd, env=pip_env)
+
+    async def _run(self):
+        path = self._target_dir
+        logger = self._logger
+        pip_packages = self._pip_config["packages"]
+        # We create an empty directory for exec cmd so that the cmd will
+        # run more stable. e.g. if cwd has ray, then checking ray will
+        # look up ray in cwd instead of site packages.
+        exec_cwd = os.path.join(path, "exec_cwd")
+        os.makedirs(exec_cwd, exist_ok=True)
+        try:
+            await virtualenv_utils.create_or_get_virtualenv(path, exec_cwd, logger)
+            python = virtualenv_utils.get_virtualenv_python(path)
+            async with dependency_utils.check_ray(python, exec_cwd, logger):
+                # Ensure pip version.
+                await self._ensure_pip_version(
+                    path,
+                    self._pip_config.get("pip_version", None),
+                    exec_cwd,
+                    self._pip_env,
+                    logger,
+                )
+                # Install pip packages.
+                await self._install_pip_packages(
+                    path,
+                    pip_packages,
+                    exec_cwd,
+                    self._pip_env,
+                    logger,
+                )
+                # Check python environment for conflicts.
+                await self._pip_check(
+                    path,
+                    self._pip_config.get("pip_check", False),
+                    exec_cwd,
+                    self._pip_env,
+                    logger,
+                )
+        except Exception:
+            logger.info("Delete incomplete virtualenv: %s", path)
+            shutil.rmtree(path, ignore_errors=True)
+            logger.exception("Failed to install pip packages.")
+            raise
+
+    def __await__(self):
+        return self._run().__await__()
+
+
+class PipPlugin(RuntimeEnvPlugin):
+    name = "pip"
+
+    def __init__(self, resources_dir: str):
+        self._pip_resources_dir = os.path.join(resources_dir, "pip")
+        self._creating_task = {}
+        # Maps a URI to a lock that is used to prevent multiple concurrent
+        # installs of the same virtualenv, see #24513
+        self._create_locks: Dict[str, asyncio.Lock] = {}
+        # Key: created hashes. Value: size of the pip dir.
+        self._created_hash_bytes: Dict[str, int] = {}
+        try_to_create_directory(self._pip_resources_dir)
+
+    def _get_path_from_hash(self, hash_val: str) -> str:
+        """Generate a path from the hash of a pip spec.
+
+        Example output:
+            /tmp/ray/session_2021-11-03_16-33-59_356303_41018/runtime_resources
+                /pip/ray-9a7972c3a75f55e976e620484f58410c920db091
+        """
+        return os.path.join(self._pip_resources_dir, hash_val)
+
+    def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]:  # noqa: F821
+        """Return the pip URI from the RuntimeEnv if it exists, else return []."""
+        pip_uri = runtime_env.pip_uri()
+        if pip_uri:
+            return [pip_uri]
+        return []
+
+    def delete_uri(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ) -> int:
+        """Delete URI and return the number of bytes deleted."""
+        logger.info("Got request to delete pip URI %s", uri)
+        protocol, hash_val = parse_uri(uri)
+        if protocol != Protocol.PIP:
+            raise ValueError(
+                "PipPlugin can only delete URIs with protocol "
+                f"pip. Received protocol {protocol}, URI {uri}"
+            )
+
+        # Cancel running create task.
+        task = self._creating_task.pop(hash_val, None)
+        if task is not None:
+            task.cancel()
+
+        del self._created_hash_bytes[hash_val]
+
+        pip_env_path = self._get_path_from_hash(hash_val)
+        local_dir_size = get_directory_size_bytes(pip_env_path)
+        del self._create_locks[uri]
+        try:
+            shutil.rmtree(pip_env_path)
+        except OSError as e:
+            logger.warning(f"Error when deleting pip env {pip_env_path}: {str(e)}")
+            return 0
+
+        return local_dir_size
+
+    async def create(
+        self,
+        uri: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: "RuntimeEnvContext",  # noqa: F821
+        logger: Optional[logging.Logger] = default_logger,
+    ) -> int:
+        if not runtime_env.has_pip():
+            return 0
+
+        protocol, hash_val = parse_uri(uri)
+        target_dir = self._get_path_from_hash(hash_val)
+
+        async def _create_for_hash():
+            await PipProcessor(
+                target_dir,
+                runtime_env,
+                logger,
+            )
+
+            loop = get_running_loop()
+            return await loop.run_in_executor(
+                None, get_directory_size_bytes, target_dir
+            )
+
+        if uri not in self._create_locks:
+            # async lock to prevent the same virtualenv being concurrently installed
+            self._create_locks[uri] = asyncio.Lock()
+
+        async with self._create_locks[uri]:
+            if hash_val in self._created_hash_bytes:
+                return self._created_hash_bytes[hash_val]
+            self._creating_task[hash_val] = task = create_task(_create_for_hash())
+            task.add_done_callback(lambda _: self._creating_task.pop(hash_val, None))
+            pip_dir_bytes = await task
+            self._created_hash_bytes[hash_val] = pip_dir_bytes
+            return pip_dir_bytes
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: "RuntimeEnvContext",  # noqa: F821
+        logger: logging.Logger = default_logger,
+    ):
+        if not runtime_env.has_pip():
+            return
+        # PipPlugin only uses a single URI.
+        uri = uris[0]
+        # Update py_executable.
+        protocol, hash_val = parse_uri(uri)
+        target_dir = self._get_path_from_hash(hash_val)
+        virtualenv_python = virtualenv_utils.get_virtualenv_python(target_dir)
+
+        if not os.path.exists(virtualenv_python):
+            raise ValueError(
+                f"Local directory {target_dir} for URI {uri} does "
+                "not exist on the cluster. Something may have gone wrong while "
+                "installing the runtime_env `pip` packages."
+            )
+        context.py_executable = virtualenv_python
+        context.command_prefix += virtualenv_utils.get_virtualenv_activate_command(
+            target_dir
+        )

python/ray/_private/runtime_env/plugin.py

@@ -0,0 +1,265 @@
+import json
+import logging
+import os
+from abc import ABC
+from typing import Any, Dict, List, Optional, Type
+
+from ray._common.utils import import_attr
+from ray._private.runtime_env.constants import (
+    RAY_RUNTIME_ENV_CLASS_FIELD_NAME,
+    RAY_RUNTIME_ENV_PLUGIN_DEFAULT_PRIORITY,
+    RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY,
+    RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY,
+    RAY_RUNTIME_ENV_PLUGINS_ENV_VAR,
+    RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME,
+)
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.uri_cache import URICache
+from ray.util.annotations import DeveloperAPI
+
+default_logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class RuntimeEnvPlugin(ABC):
+    """Abstract base class for runtime environment plugins."""
+
+    name: str = None
+    priority: int = RAY_RUNTIME_ENV_PLUGIN_DEFAULT_PRIORITY
+
+    @staticmethod
+    def validate(runtime_env_dict: dict) -> None:
+        """Validate user entry for this plugin.
+
+        The method is invoked upon installation of runtime env.
+
+        Args:
+            runtime_env_dict: The user-supplied runtime environment dict.
+
+        Raises:
+            ValueError: If the validation fails.
+        """
+        pass
+
+    def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]:  # noqa: F821
+        return []
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger,
+    ) -> float:
+        """Create and install the runtime environment.
+
+        Gets called in the runtime env agent at install time. The URI can be
+        used as a caching mechanism.
+
+        Args:
+            uri: A URI uniquely describing this resource.
+            runtime_env: The RuntimeEnv object.
+            context: Auxiliary information supplied by Ray.
+            logger: A logger to log messages during the context modification.
+
+        Returns:
+            float: The disk space taken up by this plugin installation for this
+                environment. e.g. for working_dir, this downloads the files to the
+                local node.
+        """
+        return 0
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger,
+    ) -> None:
+        """Modify context to change worker startup behavior.
+
+        For example, you can use this to prepend "cd <dir>" command to worker
+        startup, or add new environment variables.
+
+        Args:
+            uris: The URIs used by this resource.
+            runtime_env: The RuntimeEnv object.
+            context: Auxiliary information supplied by Ray.
+            logger: A logger to log messages during the context modification.
+        """
+        return
+
+    def delete_uri(self, uri: str, logger: logging.Logger) -> float:
+        """Delete the runtime environment given uri.
+
+        Args:
+            uri: A URI uniquely describing this resource.
+            logger: The logger used to log messages during the deletion.
+
+        Returns:
+            float: The amount of space reclaimed by the deletion.
+        """
+        return 0
+
+
+class PluginSetupContext:
+    def __init__(
+        self,
+        name: str,
+        class_instance: RuntimeEnvPlugin,
+        priority: int,
+        uri_cache: URICache,
+    ):
+        self.name = name
+        self.class_instance = class_instance
+        self.priority = priority
+        self.uri_cache = uri_cache
+
+
+class RuntimeEnvPluginManager:
+    """This manager is used to load plugins in runtime env agent."""
+
+    def __init__(self):
+        self.plugins: Dict[str, PluginSetupContext] = {}
+        plugin_config_str = os.environ.get(RAY_RUNTIME_ENV_PLUGINS_ENV_VAR)
+        if plugin_config_str:
+            plugin_configs = json.loads(plugin_config_str)
+            self.load_plugins(plugin_configs)
+
+    def validate_plugin_class(self, plugin_class: Type[RuntimeEnvPlugin]) -> None:
+        if not issubclass(plugin_class, RuntimeEnvPlugin):
+            raise RuntimeError(
+                f"Invalid runtime env plugin class {plugin_class}. "
+                "The plugin class must inherit "
+                "ray._private.runtime_env.plugin.RuntimeEnvPlugin."
+            )
+        if not plugin_class.name:
+            raise RuntimeError(f"No valid name in runtime env plugin {plugin_class}.")
+        if plugin_class.name in self.plugins:
+            raise RuntimeError(
+                f"The name of runtime env plugin {plugin_class} conflicts "
+                f"with {self.plugins[plugin_class.name]}.",
+            )
+
+    def validate_priority(self, priority: Any) -> None:
+        if (
+            not isinstance(priority, int)
+            or priority < RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY
+            or priority > RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY
+        ):
+            raise RuntimeError(
+                f"Invalid runtime env priority {priority}, "
+                "it should be an integer between "
+                f"{RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY} "
+                f"and {RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY}."
+            )
+
+    def load_plugins(self, plugin_configs: List[Dict]) -> None:
+        """Load runtime env plugins and create URI caches for them."""
+        for plugin_config in plugin_configs:
+            if (
+                not isinstance(plugin_config, dict)
+                or RAY_RUNTIME_ENV_CLASS_FIELD_NAME not in plugin_config
+            ):
+                raise RuntimeError(
+                    f"Invalid runtime env plugin config {plugin_config}, "
+                    "it should be a object which contains the "
+                    f"{RAY_RUNTIME_ENV_CLASS_FIELD_NAME} field."
+                )
+            plugin_class = import_attr(plugin_config[RAY_RUNTIME_ENV_CLASS_FIELD_NAME])
+            self.validate_plugin_class(plugin_class)
+
+            # The priority should be an integer between 0 and 100.
+            # The default priority is 10. A smaller number indicates a
+            # higher priority and the plugin will be set up first.
+            if RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME in plugin_config:
+                priority = plugin_config[RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME]
+            else:
+                priority = plugin_class.priority
+            self.validate_priority(priority)
+
+            class_instance = plugin_class()
+            self.plugins[plugin_class.name] = PluginSetupContext(
+                plugin_class.name,
+                class_instance,
+                priority,
+                self.create_uri_cache_for_plugin(class_instance),
+            )
+
+    def add_plugin(self, plugin: RuntimeEnvPlugin) -> None:
+        """Add a plugin to the manager and create a URI cache for it.
+
+        Args:
+            plugin: The class instance of the plugin.
+        """
+        plugin_class = type(plugin)
+        self.validate_plugin_class(plugin_class)
+        self.validate_priority(plugin_class.priority)
+        self.plugins[plugin_class.name] = PluginSetupContext(
+            plugin_class.name,
+            plugin,
+            plugin_class.priority,
+            self.create_uri_cache_for_plugin(plugin),
+        )
+
+    def create_uri_cache_for_plugin(self, plugin: RuntimeEnvPlugin) -> URICache:
+        """Create a URI cache for a plugin.
+
+        Args:
+            plugin: The plugin instance whose URIs the cache will manage.
+
+        Returns:
+            The created URI cache for the plugin.
+        """
+        # Set the max size for the cache.  Defaults to 10 GB.
+        cache_size_env_var = f"RAY_RUNTIME_ENV_{plugin.name}_CACHE_SIZE_GB".upper()
+        cache_size_bytes = int(
+            (1024**3) * float(os.environ.get(cache_size_env_var, 10))
+        )
+        return URICache(plugin.delete_uri, cache_size_bytes)
+
+    def sorted_plugin_setup_contexts(self) -> List[PluginSetupContext]:
+        """Get the sorted plugin setup contexts, sorted by increasing priority.
+
+        Returns:
+            The sorted plugin setup contexts.
+        """
+        return sorted(self.plugins.values(), key=lambda x: x.priority)
+
+
+async def create_for_plugin_if_needed(
+    runtime_env: "RuntimeEnv",  # noqa: F821
+    plugin: RuntimeEnvPlugin,
+    uri_cache: URICache,
+    context: RuntimeEnvContext,
+    logger: logging.Logger = default_logger,
+):
+    """Set up the environment using the plugin if not already set up and cached."""
+    if plugin.name not in runtime_env or runtime_env[plugin.name] is None:
+        return
+
+    plugin.validate(runtime_env)
+
+    uris = plugin.get_uris(runtime_env)
+
+    if not uris:
+        logger.debug(
+            f"No URIs for runtime env plugin {plugin.name}; "
+            "create always without checking the cache."
+        )
+        await plugin.create(None, runtime_env, context, logger=logger)
+
+    for uri in uris:
+        if uri not in uri_cache:
+            logger.debug(f"Cache miss for URI {uri}.")
+            size_bytes = await plugin.create(uri, runtime_env, context, logger=logger)
+            uri_cache.add(uri, size_bytes, logger=logger)
+        else:
+            logger.info(
+                f"Runtime env {plugin.name} {uri} is already installed "
+                "and will be reused. Search "
+                "all runtime_env_setup-*.log to find the corresponding setup log."
+            )
+            uri_cache.mark_used(uri, logger=logger)
+
+    plugin.modify_context(uris, runtime_env, context, logger)

python/ray/_private/runtime_env/plugin_schema_manager.py

@@ -0,0 +1,97 @@
+import json
+import logging
+import os
+from typing import List
+
+import jsonschema
+
+from ray._private.runtime_env.constants import (
+    RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX,
+    RAY_RUNTIME_ENV_PLUGIN_SCHEMAS_ENV_VAR,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class RuntimeEnvPluginSchemaManager:
+    """This manager is used to load plugin json schemas."""
+
+    default_schema_path = os.path.join(
+        os.path.dirname(__file__), "../../runtime_env/schemas"
+    )
+    schemas = {}
+    loaded = False
+
+    @classmethod
+    def _load_schemas(cls, schema_paths: List[str]):
+        for schema_path in schema_paths:
+            try:
+                with open(schema_path) as f:
+                    schema = json.load(f)
+            except json.decoder.JSONDecodeError:
+                logger.error("Invalid runtime env schema %s, skip it.", schema_path)
+                continue
+            except OSError:
+                logger.error("Cannot open runtime env schema %s, skip it.", schema_path)
+                continue
+            if "title" not in schema:
+                logger.error(
+                    "No valid title in runtime env schema %s, skip it.", schema_path
+                )
+                continue
+            if schema["title"] in cls.schemas:
+                logger.error(
+                    "The 'title' of runtime env schema %s conflicts with %s, skip it.",
+                    schema_path,
+                    cls.schemas[schema["title"]],
+                )
+                continue
+            cls.schemas[schema["title"]] = schema
+
+    @classmethod
+    def _load_default_schemas(cls):
+        schema_json_files = list()
+        for root, _, files in os.walk(cls.default_schema_path):
+            for f in files:
+                if f.endswith(RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX):
+                    schema_json_files.append(os.path.join(root, f))
+            logger.debug(
+                f"Loading the default runtime env schemas: {schema_json_files}."
+            )
+            cls._load_schemas(schema_json_files)
+
+    @classmethod
+    def _load_schemas_from_env_var(cls):
+        # The format of env var:
+        # "/path/to/env_1_schema.json,/path/to/env_2_schema.json,/path/to/schemas_dir/"
+        schema_paths = os.environ.get(RAY_RUNTIME_ENV_PLUGIN_SCHEMAS_ENV_VAR)
+        if schema_paths:
+            schema_json_files = list()
+            for path in schema_paths.split(","):
+                if path.endswith(RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX):
+                    schema_json_files.append(path)
+                elif os.path.isdir(path):
+                    for root, _, files in os.walk(path):
+                        for f in files:
+                            if f.endswith(RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX):
+                                schema_json_files.append(os.path.join(root, f))
+            logger.info(
+                f"Loading the runtime env schemas from env var: {schema_json_files}."
+            )
+            cls._load_schemas(schema_json_files)
+
+    @classmethod
+    def validate(cls, name, instance):
+        if not cls.loaded:
+            # Load the schemas lazily.
+            cls._load_default_schemas()
+            cls._load_schemas_from_env_var()
+            cls.loaded = True
+        # if no schema matches, skip the validation.
+        if name in cls.schemas:
+            jsonschema.validate(instance=instance, schema=cls.schemas[name])
+
+    @classmethod
+    def clear(cls):
+        cls.schemas.clear()
+        cls.loaded = False

python/ray/_private/runtime_env/protocol.py

@@ -0,0 +1,315 @@
+import enum
+import os
+from urllib.parse import urlparse
+
+RAY_RUNTIME_ENV_HTTP_USER_AGENT_ENV_VAR = "RAY_RUNTIME_ENV_HTTP_USER_AGENT"
+RAY_RUNTIME_ENV_BEARER_TOKEN_ENV_VAR = "RAY_RUNTIME_ENV_BEARER_TOKEN"
+_DEFAULT_HTTP_USER_AGENT = "ray-runtime-env-curl/1.0"
+
+
+class ProtocolsProvider:
+    _MISSING_DEPENDENCIES_WARNING = (
+        "Note that these must be preinstalled "
+        "on all nodes in the Ray cluster; it is not "
+        "sufficient to install them in the runtime_env."
+    )
+
+    @classmethod
+    def get_protocols(cls):
+        return {
+            # For packages dynamically uploaded and managed by the GCS.
+            "gcs",
+            # For conda environments installed locally on each node.
+            "conda",
+            # For pip environments installed locally on each node.
+            "pip",
+            # For uv environments install locally on each node.
+            "uv",
+            # Remote http path, assumes everything packed in one zip file.
+            "http",
+            # Remote https path, assumes everything packed in one zip file.
+            "https",
+            # Remote s3 path, assumes everything packed in one zip file.
+            "s3",
+            # Remote google storage path, assumes everything packed in one zip file.
+            "gs",
+            # Remote azure blob storage path, assumes everything packed in one zip file.
+            "azure",
+            # Remote Azure Blob File System Secure path, assumes everything packed in one zip file.
+            "abfss",
+            # File storage path, assumes everything packed in one zip file.
+            "file",
+        }
+
+    @classmethod
+    def get_remote_protocols(cls):
+        return {"http", "https", "s3", "gs", "azure", "abfss", "file"}
+
+    @classmethod
+    def _handle_s3_protocol(cls):
+        """Set up S3 protocol handling.
+
+        Returns:
+            tuple: (open_file function, transport_params)
+
+        Raises:
+            ImportError: If required dependencies are not installed.
+        """
+        try:
+            import boto3
+            from smart_open import open as open_file
+        except ImportError:
+            raise ImportError(
+                "You must `pip install smart_open[s3]` "
+                "to fetch URIs in s3 bucket. " + cls._MISSING_DEPENDENCIES_WARNING
+            )
+
+        # Create S3 client, falling back to unsigned for public buckets
+        session = boto3.Session()
+        # session.get_credentials() will return None if no credentials can be found.
+        if session.get_credentials():
+            # If credentials are found, use a standard signed client.
+            s3_client = session.client("s3")
+        else:
+            # No credentials found, fall back to an unsigned client for public buckets.
+            from botocore import UNSIGNED
+            from botocore.config import Config
+
+            s3_client = boto3.client("s3", config=Config(signature_version=UNSIGNED))
+
+        transport_params = {"client": s3_client}
+        return open_file, transport_params
+
+    @classmethod
+    def _handle_gs_protocol(cls):
+        """Set up Google Cloud Storage protocol handling.
+
+        Returns:
+            tuple: (open_file function, transport_params)
+
+        Raises:
+            ImportError: If required dependencies are not installed.
+        """
+        try:
+            from google.cloud import storage  # noqa: F401
+            from smart_open import open as open_file
+        except ImportError:
+            raise ImportError(
+                "You must `pip install smart_open[gcs]` "
+                "to fetch URIs in Google Cloud Storage bucket."
+                + cls._MISSING_DEPENDENCIES_WARNING
+            )
+
+        return open_file, None
+
+    @classmethod
+    def _handle_azure_protocol(cls):
+        """Set up Azure blob storage protocol handling.
+
+        Returns:
+            tuple: (open_file function, transport_params)
+
+        Raises:
+            ImportError: If required dependencies are not installed.
+            ValueError: If required environment variables are not set.
+        """
+        try:
+            from azure.identity import DefaultAzureCredential
+            from azure.storage.blob import BlobServiceClient  # noqa: F401
+            from smart_open import open as open_file
+        except ImportError:
+            raise ImportError(
+                "You must `pip install azure-storage-blob azure-identity smart_open[azure]` "
+                "to fetch URIs in Azure Blob Storage. "
+                + cls._MISSING_DEPENDENCIES_WARNING
+            )
+
+        # Define authentication variable
+        azure_storage_account_name = os.getenv("AZURE_STORAGE_ACCOUNT")
+
+        if not azure_storage_account_name:
+            raise ValueError(
+                "Azure Blob Storage authentication requires "
+                "AZURE_STORAGE_ACCOUNT environment variable to be set."
+            )
+
+        account_url = f"https://{azure_storage_account_name}.blob.core.windows.net/"
+        transport_params = {
+            "client": BlobServiceClient(
+                account_url=account_url, credential=DefaultAzureCredential()
+            )
+        }
+
+        return open_file, transport_params
+
+    @classmethod
+    def _handle_abfss_protocol(cls):
+        """Set up Azure Blob File System Secure (ABFSS) protocol handling.
+
+        Returns:
+            tuple: (open_file function, transport_params)
+
+        Raises:
+            ImportError: If required dependencies are not installed.
+            ValueError: If the ABFSS URI format is invalid.
+        """
+        try:
+            import adlfs
+            from azure.identity import DefaultAzureCredential
+        except ImportError:
+            raise ImportError(
+                "You must `pip install adlfs azure-identity` "
+                "to fetch URIs in Azure Blob File System Secure. "
+                + cls._MISSING_DEPENDENCIES_WARNING
+            )
+
+        def open_file(uri, mode, *, transport_params=None):
+            # Parse and validate the ABFSS URI
+            parsed = urlparse(uri)
+
+            # Validate ABFSS URI format: abfss://container@account.dfs.core.windows.net/path
+            if not parsed.netloc or "@" not in parsed.netloc:
+                raise ValueError(
+                    f"Invalid ABFSS URI format - missing container@account: {uri}"
+                )
+
+            container_part, hostname_part = parsed.netloc.split("@", 1)
+
+            # Validate container name (must be non-empty)
+            if not container_part:
+                raise ValueError(
+                    f"Invalid ABFSS URI format - empty container name: {uri}"
+                )
+
+            # Validate hostname format
+            if not hostname_part or not hostname_part.endswith(".dfs.core.windows.net"):
+                raise ValueError(
+                    f"Invalid ABFSS URI format - invalid hostname (must end with .dfs.core.windows.net): {uri}"
+                )
+
+            # Extract and validate account name
+            azure_storage_account_name = hostname_part.split(".")[0]
+            if not azure_storage_account_name:
+                raise ValueError(
+                    f"Invalid ABFSS URI format - empty account name: {uri}"
+                )
+
+            # Handle ABFSS URI with adlfs
+            filesystem = adlfs.AzureBlobFileSystem(
+                account_name=azure_storage_account_name,
+                credential=DefaultAzureCredential(),
+            )
+            return filesystem.open(uri, mode)
+
+        return open_file, None
+
+    @classmethod
+    def _http_headers(cls) -> dict:
+        headers = {
+            "User-Agent": os.environ.get(
+                RAY_RUNTIME_ENV_HTTP_USER_AGENT_ENV_VAR, _DEFAULT_HTTP_USER_AGENT
+            ),
+            "Accept": "*/*",
+        }
+
+        bearer_token = os.environ.get(RAY_RUNTIME_ENV_BEARER_TOKEN_ENV_VAR)
+        if bearer_token:
+            headers["Authorization"] = f"Bearer {bearer_token}"
+
+        return headers
+
+    @classmethod
+    def _handle_http_protocol(cls):
+        """Set up HTTP/HTTPS protocol handling with curl-like headers."""
+
+        try:
+            from smart_open import open as smart_open_open
+        except ImportError:
+            raise ImportError(
+                "You must `pip install smart_open` to fetch HTTP/HTTPS URIs. "
+                + cls._MISSING_DEPENDENCIES_WARNING
+            )
+
+        def open_file(uri, mode, *, transport_params=None):
+            params = {
+                "headers": cls._http_headers(),
+                "timeout": 60,
+            }
+            if transport_params:
+                params.update(transport_params)
+            return smart_open_open(uri, mode, transport_params=params)
+
+        return open_file, None
+
+    @classmethod
+    def download_remote_uri(cls, protocol: str, source_uri: str, dest_file: str):
+        """Download file from remote URI to destination file.
+
+        Args:
+            protocol: The protocol to use for downloading (e.g., 's3', 'https').
+            source_uri: The source URI to download from.
+            dest_file: The destination file path to save to.
+
+        Raises:
+            ImportError: If required dependencies for the protocol are not installed.
+        """
+        assert protocol in cls.get_remote_protocols()
+
+        tp = None
+        open_file = None
+
+        if protocol == "file":
+            source_uri = source_uri[len("file://") :]
+
+            def open_file(uri, mode, *, transport_params=None):
+                return open(uri, mode)
+
+        elif protocol in ("http", "https"):
+            open_file, tp = cls._handle_http_protocol()
+        elif protocol == "s3":
+            open_file, tp = cls._handle_s3_protocol()
+        elif protocol == "gs":
+            open_file, tp = cls._handle_gs_protocol()
+        elif protocol == "azure":
+            open_file, tp = cls._handle_azure_protocol()
+        elif protocol == "abfss":
+            open_file, tp = cls._handle_abfss_protocol()
+        else:
+            try:
+                from smart_open import open as open_file
+            except ImportError:
+                raise ImportError(
+                    "You must `pip install smart_open` "
+                    f"to fetch {protocol.upper()} URIs. "
+                    + cls._MISSING_DEPENDENCIES_WARNING
+                )
+
+        with open_file(source_uri, "rb", transport_params=tp) as fin:
+            with open(dest_file, "wb") as fout:
+                fout.write(fin.read())
+
+
+Protocol = enum.Enum(
+    "Protocol",
+    {protocol.upper(): protocol for protocol in ProtocolsProvider.get_protocols()},
+)
+
+
+@classmethod
+def _remote_protocols(cls):
+    # Returns a list of protocols that support remote storage
+    # These protocols should only be used with paths that end in
+    # ".zip", ".whl", ".tar.gz", or ".tgz"
+    return [
+        cls[protocol.upper()] for protocol in ProtocolsProvider.get_remote_protocols()
+    ]
+
+
+Protocol.remote_protocols = _remote_protocols
+
+
+def _download_remote_uri(self, source_uri, dest_file):
+    return ProtocolsProvider.download_remote_uri(self.value, source_uri, dest_file)
+
+
+Protocol.download_remote_uri = _download_remote_uri

python/ray/_private/runtime_env/py_executable.py

@@ -0,0 +1,45 @@
+import logging
+from typing import List, Optional
+
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+
+default_logger = logging.getLogger(__name__)
+
+
+class PyExecutablePlugin(RuntimeEnvPlugin):
+    """This plugin allows running Ray workers with a custom Python executable.
+
+    You can use it with
+    `ray.init(runtime_env={"py_executable": "<command> <args>"})`. If you specify
+    a `working_dir` in the runtime environment, the executable will have access
+    to the working directory, for example, to a requirements.txt for a package manager,
+    a script for a debugger, or the executable could be a shell script in the
+    working directory. You can also use this plugin to run worker processes
+    in a custom profiler or use a custom Python interpreter or `python` with
+    custom arguments.
+    """
+
+    name = "py_executable"
+
+    def __init__(self):
+        pass
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger = default_logger,
+    ) -> int:
+        return 0
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        logger.info("Running py_executable plugin")
+        context.py_executable = runtime_env.py_executable()

python/ray/_private/runtime_env/py_modules.py

@@ -0,0 +1,242 @@
+import logging
+import os
+from pathlib import Path
+from types import ModuleType
+from typing import Any, Dict, List, Optional
+
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.packaging import (
+    Protocol,
+    delete_package,
+    download_and_unpack_package,
+    get_local_dir_from_uri,
+    get_uri_for_directory,
+    get_uri_for_file,
+    get_uri_for_package,
+    install_wheel_package,
+    is_whl_uri,
+    package_exists,
+    parse_uri,
+    upload_package_if_needed,
+    upload_package_to_gcs,
+)
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.runtime_env.working_dir import set_pythonpath_in_context
+from ray._private.utils import get_directory_size_bytes
+from ray._raylet import GcsClient
+from ray.exceptions import RuntimeEnvSetupError
+
+default_logger = logging.getLogger(__name__)
+
+
+def _check_is_uri(s: str) -> bool:
+    try:
+        protocol, path = parse_uri(s)
+    except ValueError:
+        protocol, path = None, None
+
+    supported_extensions = (".zip", ".whl", ".tar.gz", ".tgz")
+    if protocol in Protocol.remote_protocols() and not any(
+        path.endswith(ext) for ext in supported_extensions
+    ):
+        raise ValueError(
+            "Only .zip, .whl, .tar.gz, and .tgz files supported for remote URIs."
+        )
+
+    return protocol is not None
+
+
+def upload_py_modules_if_needed(
+    runtime_env: Dict[str, Any],
+    include_gitignore: bool,
+    scratch_dir: Optional[str] = None,
+    logger: Optional[logging.Logger] = default_logger,
+    upload_fn=None,
+) -> Dict[str, Any]:
+    """Uploads the entries in py_modules and replaces them with a list of URIs.
+
+    For each entry that is already a URI, this is a no-op.
+    """
+    py_modules = runtime_env.get("py_modules")
+    if py_modules is None:
+        return runtime_env
+
+    if not isinstance(py_modules, list):
+        raise TypeError(
+            "py_modules must be a List of local paths, imported modules, or "
+            f"URIs, got {type(py_modules)}."
+        )
+
+    py_modules_uris = []
+    for module in py_modules:
+        if isinstance(module, str):
+            # module_path is a local path or a URI.
+            module_path = module
+        elif isinstance(module, Path):
+            module_path = str(module)
+        elif isinstance(module, ModuleType):
+            if not hasattr(module, "__path__"):
+                # This is a single-file module.
+                module_path = module.__file__
+            else:
+                # NOTE(edoakes): Python allows some installed Python packages to
+                # be split into multiple directories. We could probably handle
+                # this, but it seems tricky & uncommon. If it's a problem for
+                # users, we can add this support on demand.
+                if len(module.__path__) > 1:
+                    raise ValueError(
+                        "py_modules only supports modules whose __path__"
+                        " has length 1 or those who are single-file."
+                    )
+                [module_path] = module.__path__
+        else:
+            raise TypeError(
+                "py_modules must be a list of file paths, URIs, "
+                f"or imported modules, got {type(module)}."
+            )
+
+        if _check_is_uri(module_path):
+            module_uri = module_path
+        else:
+            # module_path is a local path.
+            if Path(module_path).is_dir() or Path(module_path).suffix == ".py":
+                is_dir = Path(module_path).is_dir()
+                excludes = runtime_env.get("excludes", None)
+                if is_dir:
+                    module_uri = get_uri_for_directory(
+                        module_path,
+                        include_gitignore=include_gitignore,
+                        excludes=excludes,
+                    )
+                else:
+                    module_uri = get_uri_for_file(module_path)
+                if upload_fn is None:
+                    if scratch_dir is None:
+                        scratch_dir = os.getcwd()
+                    try:
+                        upload_package_if_needed(
+                            module_uri,
+                            scratch_dir,
+                            module_path,
+                            include_gitignore=include_gitignore,
+                            include_parent_dir=is_dir,
+                            excludes=excludes,
+                            logger=logger,
+                        )
+                    except Exception as e:
+                        from ray.util.spark.utils import is_in_databricks_runtime
+
+                        if is_in_databricks_runtime():
+                            raise RuntimeEnvSetupError(
+                                f"Failed to upload module {module_path} to the Ray "
+                                f"cluster, please ensure there are only files under "
+                                f"the module path, notebooks under the path are "
+                                f"not allowed, original exception: {e}"
+                            ) from e
+                        raise RuntimeEnvSetupError(
+                            f"Failed to upload module {module_path} to the Ray "
+                            f"cluster: {e}"
+                        ) from e
+                else:
+                    upload_fn(module_path, excludes=excludes)
+            elif Path(module_path).suffix == ".whl":
+                module_uri = get_uri_for_package(Path(module_path))
+                if upload_fn is None:
+                    if not package_exists(module_uri):
+                        try:
+                            upload_package_to_gcs(
+                                module_uri, Path(module_path).read_bytes()
+                            )
+                        except Exception as e:
+                            raise RuntimeEnvSetupError(
+                                f"Failed to upload {module_path} to the Ray "
+                                f"cluster: {e}"
+                            ) from e
+                else:
+                    upload_fn(module_path, excludes=None, is_file=True)
+            else:
+                raise ValueError(
+                    "py_modules entry must be a .py file, "
+                    "a directory, or a .whl file; "
+                    f"got {module_path}"
+                )
+
+        py_modules_uris.append(module_uri)
+
+    # TODO(architkulkarni): Expose a single URI for py_modules.  This plugin
+    # should internally handle the "sub-URIs", the individual modules.
+
+    runtime_env["py_modules"] = py_modules_uris
+    return runtime_env
+
+
+class PyModulesPlugin(RuntimeEnvPlugin):
+
+    name = "py_modules"
+
+    def __init__(self, resources_dir: str, gcs_client: GcsClient):
+        self._resources_dir = os.path.join(resources_dir, "py_modules_files")
+        self._gcs_client = gcs_client
+        try_to_create_directory(self._resources_dir)
+
+    def _get_local_dir_from_uri(self, uri: str):
+        return get_local_dir_from_uri(uri, self._resources_dir)
+
+    def delete_uri(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ) -> int:
+        """Delete URI and return the number of bytes deleted."""
+        logger.info("Got request to delete pymodule URI %s", uri)
+        local_dir = get_local_dir_from_uri(uri, self._resources_dir)
+        local_dir_size = get_directory_size_bytes(local_dir)
+
+        deleted = delete_package(uri, self._resources_dir)
+        if not deleted:
+            logger.warning(f"Tried to delete nonexistent URI: {uri}.")
+            return 0
+
+        return local_dir_size
+
+    def get_uris(self, runtime_env) -> List[str]:
+        return runtime_env.py_modules()
+
+    async def create(
+        self,
+        uri: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ) -> int:
+
+        module_dir = await download_and_unpack_package(
+            uri, self._resources_dir, self._gcs_client, logger=logger
+        )
+
+        if is_whl_uri(uri):
+            wheel_uri = module_dir
+            module_dir = self._get_local_dir_from_uri(uri)
+            await install_wheel_package(
+                wheel_uri=wheel_uri, target_dir=module_dir, logger=logger
+            )
+
+        return get_directory_size_bytes(module_dir)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env_dict: Dict,
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        module_dirs = []
+        for uri in uris:
+            module_dir = self._get_local_dir_from_uri(uri)
+            if not module_dir.exists():
+                raise ValueError(
+                    f"Local directory {module_dir} for URI {uri} does "
+                    "not exist on the cluster. Something may have gone wrong while "
+                    "downloading, unpacking or installing the py_modules files."
+                )
+            module_dirs.append(str(module_dir))
+        set_pythonpath_in_context(os.pathsep.join(module_dirs), context)

python/ray/_private/runtime_env/rocprof_sys.py

@@ -0,0 +1,173 @@
+import asyncio
+import copy
+import logging
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray.exceptions import RuntimeEnvSetupError
+
+default_logger = logging.getLogger(__name__)
+
+# rocprof-sys config used when runtime_env={"_rocprof_sys": "default"}
+# Refer to the following link for more information on rocprof-sys options
+# https://rocm.docs.amd.com/projects/rocprofiler-systems/en/docs-6.4.0/how-to/understanding-rocprof-sys-output.html
+ROCPROFSYS_DEFAULT_CONFIG = {
+    "env": {
+        "ROCPROFSYS_TIME_OUTPUT": "false",
+        "ROCPROFSYS_OUTPUT_PREFIX": "worker_process_%p",
+    },
+    "args": {
+        "F": "true",
+    },
+}
+
+
+def parse_rocprof_sys_config(
+    rocprof_sys_config: Dict[str, str]
+) -> Tuple[List[str], List[str]]:
+    """
+    Function to convert dictionary of rocprof-sys options into
+    rocprof-sys-python command line
+
+    The function returns:
+    - List[str]: rocprof-sys-python cmd line split into list of str
+    """
+    rocprof_sys_cmd = ["rocprof-sys-python"]
+    rocprof_sys_env = {}
+    if "args" in rocprof_sys_config:
+        # Parse rocprof-sys arg options
+        for option, option_val in rocprof_sys_config["args"].items():
+            # option standard based on
+            # https://www.gnu.org/software/libc/manual/html_node/Argument-Syntax.html
+            if len(option) > 1:
+                rocprof_sys_cmd.append(f"--{option}={option_val}")
+            else:
+                rocprof_sys_cmd += [f"-{option}", option_val]
+    if "env" in rocprof_sys_config:
+        rocprof_sys_env = rocprof_sys_config["env"]
+    rocprof_sys_cmd.append("--")
+    return rocprof_sys_cmd, rocprof_sys_env
+
+
+class RocProfSysPlugin(RuntimeEnvPlugin):
+    name = "_rocprof_sys"
+
+    def __init__(self, resources_dir: str):
+        self.rocprof_sys_cmd = []
+        self.rocprof_sys_env = {}
+
+        # replace this with better way to get logs dir
+        session_dir, runtime_dir = os.path.split(resources_dir)
+        self._rocprof_sys_dir = Path(session_dir) / "logs" / "rocprof_sys"
+        try_to_create_directory(self._rocprof_sys_dir)
+
+    async def _check_rocprof_sys_script(
+        self, rocprof_sys_config: Dict[str, str]
+    ) -> Tuple[bool, str]:
+        """
+        Function to validate if rocprof_sys_config is a valid rocprof_sys profile options
+        Args:
+            rocprof_sys_config: dictionary mapping rocprof_sys option to it's value
+        Returns:
+            a tuple consists of a boolean indicating if the rocprof_sys_config
+            is valid option and an error message if the rocprof_sys_config is invalid
+        """
+
+        # use empty as rocprof_sys report test filename
+        test_folder = str(Path(self._rocprof_sys_dir) / "test")
+        rocprof_sys_cmd, rocprof_sys_env = parse_rocprof_sys_config(rocprof_sys_config)
+        rocprof_sys_env_copy = copy.deepcopy(rocprof_sys_env)
+        rocprof_sys_env_copy["ROCPROFSYS_OUTPUT_PATH"] = test_folder
+        rocprof_sys_env_copy.update(os.environ)
+        try_to_create_directory(test_folder)
+
+        # Create a test python file to run rocprof_sys
+        with open(f"{test_folder}/test.py", "w") as f:
+            f.write("import time\n")
+        try:
+            rocprof_sys_cmd = rocprof_sys_cmd + [f"{test_folder}/test.py"]
+            process = await asyncio.create_subprocess_exec(
+                *rocprof_sys_cmd,
+                env=rocprof_sys_env_copy,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            stdout, stderr = await process.communicate()
+            error_msg = stderr.strip() if stderr.strip() != "" else stdout.strip()
+
+            # cleanup temp file
+            clean_up_cmd = ["rm", "-r", test_folder]
+            cleanup_process = await asyncio.create_subprocess_exec(
+                *clean_up_cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            _, _ = await cleanup_process.communicate()
+            if process.returncode == 0:
+                return True, None
+            else:
+                return False, error_msg
+        except FileNotFoundError:
+            return False, ("rocprof_sys is not installed")
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger = default_logger,
+    ) -> int:
+        rocprof_sys_config = runtime_env.rocprof_sys()
+        if not rocprof_sys_config:
+            return 0
+
+        if rocprof_sys_config and sys.platform != "linux":
+            raise RuntimeEnvSetupError("rocprof-sys CLI is only available in Linux.\n")
+
+        if isinstance(rocprof_sys_config, str):
+            if rocprof_sys_config == "default":
+                rocprof_sys_config = ROCPROFSYS_DEFAULT_CONFIG
+            else:
+                raise RuntimeEnvSetupError(
+                    f"Unsupported rocprof_sys config: {rocprof_sys_config}. "
+                    "The supported config is 'default' or "
+                    "Dictionary of rocprof_sys options"
+                )
+
+        is_valid_rocprof_sys_config, error_msg = await self._check_rocprof_sys_script(
+            rocprof_sys_config
+        )
+        if not is_valid_rocprof_sys_config:
+            logger.warning(error_msg)
+            raise RuntimeEnvSetupError(
+                "rocprof-sys profile failed to run with the following "
+                f"error message:\n {error_msg}"
+            )
+        # add set output path to logs dir
+        if "env" not in rocprof_sys_config:
+            rocprof_sys_config["env"] = {}
+        rocprof_sys_config["env"]["ROCPROFSYS_OUTPUT_PATH"] = str(
+            Path(self._rocprof_sys_dir)
+        )
+
+        self.rocprof_sys_cmd, self.rocprof_sys_env = parse_rocprof_sys_config(
+            rocprof_sys_config
+        )
+        return 0
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        logger.info("Running rocprof-sys profiler")
+        context.py_executable = " ".join(self.rocprof_sys_cmd)
+        context.env_vars.update(self.rocprof_sys_env)

python/ray/_private/runtime_env/setup_hook.py

@@ -0,0 +1,261 @@
+import base64
+import logging
+import os
+import traceback
+from typing import Any, Callable, Dict, Optional, Union
+
+import ray
+import ray._private.ray_constants as ray_constants
+import ray.cloudpickle as pickle
+from ray._common.utils import load_class
+from ray._private.function_manager import build_setup_hook_export_entry
+from ray.runtime_env import RuntimeEnv
+
+logger = logging.getLogger(__name__)
+
+RUNTIME_ENV_FUNC_IDENTIFIER = "ray_runtime_env_func::"
+
+
+def get_import_export_timeout():
+    return int(
+        os.environ.get(
+            ray_constants.RAY_WORKER_PROCESS_SETUP_HOOK_LOAD_TIMEOUT_ENV_VAR, "60"
+        )
+    )
+
+
+def decode_function_key(key: bytes) -> str:
+    # b64encode only includes A-Z, a-z, 0-9, + and / characters
+    return RUNTIME_ENV_FUNC_IDENTIFIER + base64.b64encode(key).decode()
+
+
+def _encode_function_key(key: str) -> bytes:
+    assert key.startswith(RUNTIME_ENV_FUNC_IDENTIFIER)
+    return base64.b64decode(key[len(RUNTIME_ENV_FUNC_IDENTIFIER) :])
+
+
+def _raise_setup_hook_conflict(existing_hook_value: str, setup_hook_desc: str) -> None:
+    raise RuntimeError(
+        "Conflicting worker_process_setup_hook: the setup hook env "
+        f"var is already set to '{existing_hook_value}', but "
+        f"runtime_env specifies {setup_hook_desc}."
+    )
+
+
+def export_setup_func_callable(
+    runtime_env: Union[Dict[str, Any], RuntimeEnv],
+    setup_func: Callable,
+    worker: "ray.Worker",
+) -> Union[Dict[str, Any], RuntimeEnv]:
+    assert isinstance(setup_func, Callable)
+    try:
+        key = worker.function_actor_manager.export_setup_func(
+            setup_func, timeout=get_import_export_timeout()
+        )
+    except Exception as e:
+        raise ray.exceptions.RuntimeEnvSetupError(
+            "Failed to export the setup function."
+        ) from e
+    env_vars = runtime_env.get("env_vars", {})
+    assert ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR not in env_vars, (
+        f"The env var, {ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR}, "
+        "is not permitted because it is reserved for the internal use."
+    )
+    env_vars[ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR] = decode_function_key(key)
+    runtime_env["env_vars"] = env_vars
+    # Note: This field is no-op. We don't have a plugin for the setup hook
+    # because we can implement it simply using an env var.
+    # This field is just for the observability purpose, so we store
+    # the name of the method.
+    runtime_env["worker_process_setup_hook"] = setup_func.__name__
+    return runtime_env
+
+
+def export_setup_func_module(
+    runtime_env: Union[Dict[str, Any], RuntimeEnv],
+    setup_func_module: str,
+) -> Union[Dict[str, Any], RuntimeEnv]:
+    assert isinstance(setup_func_module, str)
+    env_vars = runtime_env.get("env_vars", {})
+    assert ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR not in env_vars, (
+        f"The env var, {ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR}, "
+        "is not permitted because it is reserved for the internal use."
+    )
+    env_vars[ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR] = setup_func_module
+    runtime_env["env_vars"] = env_vars
+    return runtime_env
+
+
+def _check_setup_hook_consistency(
+    existing_hook_value: str,
+    setup_func: Union[Callable, str],
+    worker: "ray.Worker",
+) -> None:
+    """Validate that an already-set hook env var is consistent with setup_func.
+
+    When the env var is already populated (e.g. inherited from a job supervisor),
+    we compare it against the `worker_process_setup_hook` field in the runtime_env
+    to detect silent mismatches.
+
+    Args:
+        existing_hook_value: The value of the existing hook env var.
+        setup_func: The setup function or module path.
+        worker: The worker instance.
+
+    Raises:
+        RuntimeError: If a conflict between the existing env var and setup_func is detected.
+    """
+    if isinstance(setup_func, Callable):
+        try:
+            _encode_function_key(existing_hook_value)
+        except Exception:
+            _raise_setup_hook_conflict(
+                existing_hook_value, f"callable '{setup_func.__name__}'"
+            )
+        _check_callable_hooks_match(existing_hook_value, setup_func, worker)
+    elif isinstance(setup_func, str):
+        try:
+            _encode_function_key(existing_hook_value)
+            existing_is_callable_ref = True
+        except Exception:
+            existing_is_callable_ref = False
+        if existing_is_callable_ref or existing_hook_value != setup_func:
+            _raise_setup_hook_conflict(existing_hook_value, f"'{setup_func}'")
+
+
+def _check_callable_hooks_match(
+    existing_hook_value: str,
+    setup_func: Callable,
+    worker: "ray.Worker",
+) -> None:
+    """Verify a callable produces the same GCS key as the existing env var."""
+    _, _, expected_key = build_setup_hook_export_entry(
+        setup_func, worker.current_job_id.binary()
+    )
+    expected_env_value = decode_function_key(expected_key)
+
+    if existing_hook_value != expected_env_value:
+        _raise_setup_hook_conflict(
+            existing_hook_value, f"callable '{setup_func.__name__}'"
+        )
+
+
+def upload_worker_process_setup_hook_if_needed(
+    runtime_env: Union[Dict[str, Any], RuntimeEnv],
+    worker: "ray.Worker",
+) -> Union[Dict[str, Any], RuntimeEnv]:
+    """Uploads the worker_process_setup_hook to GCS with a key.
+
+    runtime_env["worker_process_setup_hook"] is converted to a decoded key
+    that can load the worker setup hook function from GCS.
+    i.e., you can use internalKV.Get(runtime_env["worker_process_setup_hook])
+    to access the worker setup hook from GCS.
+
+    Args:
+        runtime_env: The runtime_env. The value will be modified
+            when returned.
+        worker: ray.worker instance.
+
+    Returns:
+        The modified runtime_env with the setup hook processed into an env var.
+    """
+    setup_func = runtime_env.get("worker_process_setup_hook")
+
+    if setup_func is None:
+        return runtime_env
+
+    env_vars = runtime_env.get("env_vars", {})
+    existing_hook = env_vars.get(ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR)
+
+    if existing_hook is not None:
+        # A setup hook is already populated (e.g. inherited from job supervisor).
+        # Validate that it is consistent with the current worker_process_setup_hook.
+        _check_setup_hook_consistency(existing_hook, setup_func, worker)
+        return runtime_env
+
+    if isinstance(setup_func, Callable):
+        return export_setup_func_callable(runtime_env, setup_func, worker)
+    elif isinstance(setup_func, str):
+        return export_setup_func_module(runtime_env, setup_func)
+    else:
+        raise TypeError(
+            "worker_process_setup_hook must be a function, " f"got {type(setup_func)}."
+        )
+
+
+def load_and_execute_setup_hook(
+    worker_process_setup_hook_key: str,
+) -> Optional[str]:
+    """Load the setup hook from a given key and execute.
+
+    Args:
+        worker_process_setup_hook_key: The key to import the setup hook
+            from GCS.
+    Returns:
+        An error message if it fails. None if it succeeds.
+    """
+    assert worker_process_setup_hook_key is not None
+    if not worker_process_setup_hook_key.startswith(RUNTIME_ENV_FUNC_IDENTIFIER):
+        return load_and_execute_setup_hook_module(worker_process_setup_hook_key)
+    else:
+        return load_and_execute_setup_hook_func(worker_process_setup_hook_key)
+
+
+def load_and_execute_setup_hook_module(
+    worker_process_setup_hook_key: str,
+) -> Optional[str]:
+    try:
+        setup_func = load_class(worker_process_setup_hook_key)
+        setup_func()
+        return None
+    except Exception:
+        error_message = (
+            "Failed to execute the setup hook method, "
+            f"{worker_process_setup_hook_key} "
+            "from ``ray.init(runtime_env="
+            f"{{'worker_process_setup_hook': {worker_process_setup_hook_key}}})``. "
+            "Please make sure the given module exists and is available "
+            "from ray workers. For more details, see the error trace below.\n"
+            f"{traceback.format_exc()}"
+        )
+        return error_message
+
+
+def load_and_execute_setup_hook_func(
+    worker_process_setup_hook_key: str,
+) -> Optional[str]:
+    worker = ray._private.worker.global_worker
+    assert worker.connected
+    func_manager = worker.function_actor_manager
+    try:
+        worker_setup_func_info = func_manager.fetch_registered_method(
+            _encode_function_key(worker_process_setup_hook_key),
+            timeout=get_import_export_timeout(),
+        )
+    except Exception:
+        error_message = (
+            "Failed to import setup hook within "
+            f"{get_import_export_timeout()} seconds.\n"
+            f"{traceback.format_exc()}"
+        )
+        return error_message
+
+    try:
+        setup_func = pickle.loads(worker_setup_func_info.function)
+    except Exception:
+        error_message = (
+            "Failed to deserialize the setup hook method.\n" f"{traceback.format_exc()}"
+        )
+        return error_message
+
+    try:
+        setup_func()
+    except Exception:
+        error_message = (
+            f"Failed to execute the setup hook method. Function name:"
+            f"{worker_setup_func_info.function_name}\n"
+            f"{traceback.format_exc()}"
+        )
+        return error_message
+
+    return None

python/ray/_private/runtime_env/uri_cache.py

@@ -0,0 +1,115 @@
+import logging
+from typing import Callable, Optional, Set
+
+default_logger = logging.getLogger(__name__)
+
+DEFAULT_MAX_URI_CACHE_SIZE_BYTES = (1024**3) * 10  # 10 GB
+
+
+class URICache:
+    """Caches URIs up to a specified total size limit.
+
+    URIs are represented by strings. Each URI has an associated size on disk.
+    When a URI is added to the URICache, it is marked as "in use".
+    When a URI is no longer in use, the user of this class should call
+    `mark_unused` to signal that the URI is safe for deletion.
+
+    URIs in the cache can be marked as "in use" by calling `mark_used`.
+
+    Deletion of URIs on disk does not occur until the size limit is exceeded.
+    When this happens, URIs that are not in use are deleted randomly until the
+    size limit is satisfied, or there are no more URIs that are not in use.
+
+    It is possible for the total size on disk to exceed the size limit if all
+    the URIs are in use.
+
+    """
+
+    def __init__(
+        self,
+        delete_fn: Optional[Callable[[str, logging.Logger], int]] = None,
+        max_total_size_bytes: int = DEFAULT_MAX_URI_CACHE_SIZE_BYTES,
+        debug_mode: bool = False,
+    ):
+        # Maps URIs to the size in bytes of their corresponding disk contents.
+        self._used_uris: Set[str] = set()
+        self._unused_uris: Set[str] = set()
+
+        if delete_fn is None:
+            self._delete_fn = lambda uri, logger: 0
+        else:
+            self._delete_fn = delete_fn
+
+        # Total size of both used and unused URIs in the cache.
+        self._total_size_bytes = 0
+        self.max_total_size_bytes = max_total_size_bytes
+
+        # Used in `self._check_valid()` for testing.
+        self._debug_mode = debug_mode
+
+    def mark_unused(self, uri: str, logger: logging.Logger = default_logger):
+        """Mark a URI as unused and okay to be deleted."""
+        if uri not in self._used_uris:
+            logger.info(f"URI {uri} is already unused.")
+        else:
+            self._unused_uris.add(uri)
+            self._used_uris.remove(uri)
+        logger.info(f"Marked URI {uri} unused.")
+        self._evict_if_needed(logger)
+        self._check_valid()
+
+    def mark_used(self, uri: str, logger: logging.Logger = default_logger):
+        """Mark a URI as in use. URIs in use will not be deleted."""
+        if uri in self._used_uris:
+            return
+        elif uri in self._unused_uris:
+            self._used_uris.add(uri)
+            self._unused_uris.remove(uri)
+        else:
+            raise ValueError(
+                f"Got request to mark URI {uri} used, but this "
+                "URI is not present in the cache."
+            )
+        logger.info(f"Marked URI {uri} used.")
+        self._check_valid()
+
+    def add(self, uri: str, size_bytes: int, logger: logging.Logger = default_logger):
+        """Add a URI to the cache and mark it as in use."""
+        if uri in self._unused_uris:
+            self._unused_uris.remove(uri)
+
+        self._used_uris.add(uri)
+        self._total_size_bytes += size_bytes
+
+        self._evict_if_needed(logger)
+        self._check_valid()
+        logger.info(f"Added URI {uri} with size {size_bytes}")
+
+    def get_total_size_bytes(self) -> int:
+        return self._total_size_bytes
+
+    def _evict_if_needed(self, logger: logging.Logger = default_logger):
+        """Evict unused URIs (if they exist) until total size <= max size."""
+        while (
+            self._unused_uris
+            and self.get_total_size_bytes() > self.max_total_size_bytes
+        ):
+            # TODO(architkulkarni): Evict least recently used URI instead
+            arbitrary_unused_uri = next(iter(self._unused_uris))
+            self._unused_uris.remove(arbitrary_unused_uri)
+            num_bytes_deleted = self._delete_fn(arbitrary_unused_uri, logger)
+            self._total_size_bytes -= num_bytes_deleted
+            logger.info(
+                f"Deleted URI {arbitrary_unused_uri} with size " f"{num_bytes_deleted}."
+            )
+
+    def _check_valid(self):
+        """(Debug mode only) Check "used" and "unused" sets are disjoint."""
+        if self._debug_mode:
+            assert self._used_uris & self._unused_uris == set()
+
+    def __contains__(self, uri):
+        return uri in self._used_uris or uri in self._unused_uris
+
+    def __repr__(self):
+        return str(self.__dict__)

python/ray/_private/runtime_env/utils.py

@@ -0,0 +1,117 @@
+import asyncio
+import itertools
+import logging
+import subprocess
+import textwrap
+import types
+from typing import List
+
+
+class SubprocessCalledProcessError(subprocess.CalledProcessError):
+    """The subprocess.CalledProcessError with stripped stdout."""
+
+    LAST_N_LINES = 50
+
+    def __init__(self, *args, cmd_index=None, **kwargs):
+        self.cmd_index = cmd_index
+        super().__init__(*args, **kwargs)
+
+    @staticmethod
+    def _get_last_n_line(str_data: str, last_n_lines: int) -> str:
+        if last_n_lines < 0:
+            return str_data
+        lines = str_data.strip().split("\n")
+        return "\n".join(lines[-last_n_lines:])
+
+    def __str__(self):
+        str_list = (
+            []
+            if self.cmd_index is None
+            else [f"Run cmd[{self.cmd_index}] failed with the following details."]
+        )
+        str_list.append(super().__str__())
+        out = {
+            "stdout": self.stdout,
+            "stderr": self.stderr,
+        }
+        for name, s in out.items():
+            if s:
+                subtitle = f"Last {self.LAST_N_LINES} lines of {name}:"
+                last_n_line_str = self._get_last_n_line(s, self.LAST_N_LINES).strip()
+                str_list.append(
+                    f"{subtitle}\n{textwrap.indent(last_n_line_str, ' ' * 4)}"
+                )
+        return "\n".join(str_list)
+
+
+async def check_output_cmd(
+    cmd: List[str],
+    *,
+    logger: logging.Logger,
+    cmd_index_gen: types.GeneratorType = itertools.count(1),
+    **kwargs,
+) -> str:
+    """Run command with arguments and return its output.
+
+    If the return code was non-zero it raises a CalledProcessError. The
+    CalledProcessError object will have the return code in the returncode
+    attribute and any output in the output attribute.
+
+    Args:
+        cmd: The cmdline should be a sequence of program arguments or else
+            a single string or path-like object. The program to execute is
+            the first item in cmd.
+        logger: The logger instance.
+        cmd_index_gen: The cmd index generator, default is itertools.count(1).
+        **kwargs: All arguments are passed to the create_subprocess_exec.
+
+    Returns:
+        The stdout of cmd.
+
+    Raises:
+        CalledProcessError: If the return code of cmd is not 0.
+    """
+
+    cmd_index = next(cmd_index_gen)
+    logger.info("Run cmd[%s] %s", cmd_index, repr(cmd))
+
+    proc = None
+    try:
+        proc = await asyncio.create_subprocess_exec(
+            *cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+            **kwargs,
+        )
+        # Use communicate instead of polling stdout:
+        #   * Avoid deadlocks due to streams pausing reading or writing and blocking the
+        #     child process. Please refer to:
+        #     https://docs.python.org/3/library/asyncio-subprocess.html#asyncio.asyncio.subprocess.Process.stderr
+        #   * Avoid mixing multiple outputs of concurrent cmds.
+        stdout, _ = await proc.communicate()
+    except asyncio.exceptions.CancelledError as e:
+        # since Python 3.9, when cancelled, the inner process needs to throw as it is
+        # for asyncio to timeout properly https://bugs.python.org/issue40607
+        raise e
+    except BaseException as e:
+        raise RuntimeError(f"Run cmd[{cmd_index}] got exception.") from e
+    else:
+        stdout = stdout.decode("utf-8")
+        if stdout:
+            logger.info("Output of cmd[%s]: %s", cmd_index, stdout)
+        else:
+            logger.info("No output for cmd[%s]", cmd_index)
+        if proc.returncode != 0:
+            raise SubprocessCalledProcessError(
+                proc.returncode, cmd, output=stdout, cmd_index=cmd_index
+            )
+        return stdout
+    finally:
+        if proc is not None:
+            # Kill process.
+            try:
+                proc.kill()
+            except ProcessLookupError:
+                pass
+            # Wait process exit.
+            await proc.wait()

python/ray/_private/runtime_env/uv.py

@@ -0,0 +1,344 @@
+"""Util class to install packages via uv."""
+
+import asyncio
+import hashlib
+import json
+import logging
+import os
+import shutil
+import sys
+from asyncio import create_task, get_running_loop
+from typing import Dict, List, Optional
+
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env import dependency_utils, virtualenv_utils
+from ray._private.runtime_env.packaging import Protocol, parse_uri
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.runtime_env.utils import check_output_cmd
+from ray._private.utils import get_directory_size_bytes
+
+default_logger = logging.getLogger(__name__)
+
+
+def _get_uv_hash(uv_dict: Dict) -> str:
+    """Get a deterministic hash value for `uv` related runtime envs."""
+    serialized_uv_spec = json.dumps(uv_dict, sort_keys=True)
+    hash_val = hashlib.sha1(serialized_uv_spec.encode("utf-8")).hexdigest()
+    return hash_val
+
+
+def get_uri(runtime_env: Dict) -> Optional[str]:
+    """Return `"uv://<hashed_dependencies>"`, or None if no GC required."""
+    uv = runtime_env.get("uv")
+    if uv is not None:
+        if isinstance(uv, dict):
+            uri = "uv://" + _get_uv_hash(uv_dict=uv)
+        elif isinstance(uv, list):
+            uri = "uv://" + _get_uv_hash(uv_dict=dict(packages=uv))
+        else:
+            raise TypeError(
+                "uv field received by RuntimeEnvAgent must be "
+                f"list or dict, not {type(uv).__name__}."
+            )
+    else:
+        uri = None
+    return uri
+
+
+class UvProcessor:
+    def __init__(
+        self,
+        target_dir: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        try:
+            import virtualenv  # noqa: F401 ensure virtualenv exists.
+        except ImportError:
+            raise RuntimeError(
+                f"Please install virtualenv "
+                f"`{sys.executable} -m pip install virtualenv`"
+                f"to enable uv runtime env."
+            )
+
+        logger.debug("Setting up uv for runtime_env: %s", runtime_env)
+        self._target_dir = target_dir
+        # An empty directory is created to execute cmd.
+        self._exec_cwd = os.path.join(self._target_dir, "exec_cwd")
+        self._runtime_env = runtime_env
+        self._logger = logger
+
+        self._uv_config = self._runtime_env.uv_config()
+        self._uv_env = os.environ.copy()
+        self._uv_env.update(self._runtime_env.env_vars())
+
+    async def _install_uv(
+        self, path: str, cwd: str, pip_env: dict, logger: logging.Logger
+    ):
+        """Before package install, make sure the required version `uv` (if specifieds)
+        is installed.
+        """
+        virtualenv_path = virtualenv_utils.get_virtualenv_path(path)
+        python = virtualenv_utils.get_virtualenv_python(path)
+
+        def _get_uv_exec_to_install() -> str:
+            """Get `uv` executable with version to install."""
+            uv_version = self._uv_config.get("uv_version", None)
+            if uv_version:
+                return f"uv{uv_version}"
+            # Use default version.
+            return "uv"
+
+        uv_install_cmd = [
+            python,
+            "-m",
+            "pip",
+            "install",
+            "--disable-pip-version-check",
+            "--no-cache-dir",
+            _get_uv_exec_to_install(),
+        ]
+        logger.info("Installing package uv to %s", virtualenv_path)
+        await check_output_cmd(uv_install_cmd, logger=logger, cwd=cwd, env=pip_env)
+
+    async def _check_uv_existence(
+        self, path: str, cwd: str, env: dict, logger: logging.Logger
+    ) -> bool:
+        """Check and return the existence of `uv` in virtual env."""
+        python = virtualenv_utils.get_virtualenv_python(path)
+
+        check_existence_cmd = [
+            python,
+            "-m",
+            "uv",
+            "--version",
+        ]
+
+        try:
+            # If `uv` doesn't exist, exception will be thrown.
+            await check_output_cmd(check_existence_cmd, logger=logger, cwd=cwd, env=env)
+            return True
+        except Exception:
+            return False
+
+    async def _uv_check(sef, python: str, cwd: str, logger: logging.Logger) -> None:
+        """Check virtual env dependency compatibility.
+        If any incompatibility detected, exception will be thrown.
+
+        param:
+            python: the path for python executable within virtual environment.
+        """
+        cmd = [python, "-m", "uv", "pip", "check"]
+        await check_output_cmd(
+            cmd,
+            logger=logger,
+            cwd=cwd,
+        )
+
+    async def _install_uv_packages(
+        self,
+        path: str,
+        uv_packages: List[str],
+        cwd: str,
+        pip_env: Dict,
+        logger: logging.Logger,
+    ):
+        """Install required python packages via `uv`."""
+        virtualenv_path = virtualenv_utils.get_virtualenv_path(path)
+        python = virtualenv_utils.get_virtualenv_python(path)
+        # TODO(fyrestone): Support -i, --no-deps, --no-cache-dir, ...
+        requirements_file = dependency_utils.get_requirements_file(path, uv_packages)
+
+        # Check existence for `uv` and see if we could skip `uv` installation.
+        uv_exists = await self._check_uv_existence(path, cwd, pip_env, logger)
+
+        # Install uv, which acts as the default package manager.
+        if (not uv_exists) or (self._uv_config.get("uv_version", None) is not None):
+            await self._install_uv(path, cwd, pip_env, logger)
+
+        # Avoid blocking the event loop.
+        loop = get_running_loop()
+        await loop.run_in_executor(
+            None, dependency_utils.gen_requirements_txt, requirements_file, uv_packages
+        )
+
+        # Install all dependencies.
+        #
+        # Difference with pip:
+        # 1. `--disable-pip-version-check` has no effect for uv.
+        uv_install_cmd = [
+            python,
+            "-m",
+            "uv",
+            "pip",
+            "install",
+            "-r",
+            requirements_file,
+        ]
+
+        uv_opt_list = self._uv_config.get("uv_pip_install_options", ["--no-cache"])
+        if uv_opt_list:
+            uv_install_cmd += uv_opt_list
+
+        logger.info("Installing python requirements to %s", virtualenv_path)
+        await check_output_cmd(uv_install_cmd, logger=logger, cwd=cwd, env=pip_env)
+
+        # Check python environment for conflicts.
+        if self._uv_config.get("uv_check", False):
+            await self._uv_check(python, cwd, logger)
+
+    async def _run(self):
+        path = self._target_dir
+        logger = self._logger
+        uv_packages = self._uv_config["packages"]
+        # We create an empty directory for exec cmd so that the cmd will
+        # run more stable. e.g. if cwd has ray, then checking ray will
+        # look up ray in cwd instead of site packages.
+        os.makedirs(self._exec_cwd, exist_ok=True)
+        try:
+            await virtualenv_utils.create_or_get_virtualenv(
+                path, self._exec_cwd, logger
+            )
+            python = virtualenv_utils.get_virtualenv_python(path)
+            async with dependency_utils.check_ray(python, self._exec_cwd, logger):
+                # Install packages with uv.
+                await self._install_uv_packages(
+                    path,
+                    uv_packages,
+                    self._exec_cwd,
+                    self._uv_env,
+                    logger,
+                )
+        except Exception:
+            logger.info("Delete incomplete virtualenv: %s", path)
+            shutil.rmtree(path, ignore_errors=True)
+            logger.exception("Failed to install uv packages.")
+            raise
+
+    def __await__(self):
+        return self._run().__await__()
+
+
+class UvPlugin(RuntimeEnvPlugin):
+    name = "uv"
+
+    def __init__(self, resources_dir: str):
+        self._uv_resource_dir = os.path.join(resources_dir, "uv")
+        self._creating_task = {}
+        # Maps a URI to a lock that is used to prevent multiple concurrent
+        # installs of the same virtualenv, see #24513
+        self._create_locks: Dict[str, asyncio.Lock] = {}
+        # Key: created hashes. Value: size of the uv dir.
+        self._created_hash_bytes: Dict[str, int] = {}
+        try_to_create_directory(self._uv_resource_dir)
+
+    def _get_path_from_hash(self, hash_val: str) -> str:
+        """Generate a path from the hash of a uv spec.
+
+        Example output:
+            /tmp/ray/session_2021-11-03_16-33-59_356303_41018/runtime_resources
+                /uv/ray-9a7972c3a75f55e976e620484f58410c920db091
+        """
+        return os.path.join(self._uv_resource_dir, hash_val)
+
+    def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]:  # noqa: F821
+        """Return the uv URI from the RuntimeEnv if it exists, else return []."""
+        uv_uri = runtime_env.uv_uri()
+        if uv_uri:
+            return [uv_uri]
+        return []
+
+    def delete_uri(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ) -> int:
+        """Delete URI and return the number of bytes deleted."""
+        logger.info("Got request to delete uv URI %s", uri)
+        protocol, hash_val = parse_uri(uri)
+        if protocol != Protocol.UV:
+            raise ValueError(
+                "UvPlugin can only delete URIs with protocol "
+                f"uv. Received protocol {protocol}, URI {uri}"
+            )
+
+        # Cancel running create task.
+        task = self._creating_task.pop(hash_val, None)
+        if task is not None:
+            task.cancel()
+
+        del self._created_hash_bytes[hash_val]
+
+        uv_env_path = self._get_path_from_hash(hash_val)
+        local_dir_size = get_directory_size_bytes(uv_env_path)
+        del self._create_locks[uri]
+        try:
+            shutil.rmtree(uv_env_path)
+        except OSError as e:
+            logger.warning(f"Error when deleting uv env {uv_env_path}: {str(e)}")
+            return 0
+
+        return local_dir_size
+
+    async def create(
+        self,
+        uri: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: "RuntimeEnvContext",  # noqa: F821
+        logger: Optional[logging.Logger] = default_logger,
+    ) -> int:
+        if not runtime_env.has_uv():
+            return 0
+
+        protocol, hash_val = parse_uri(uri)
+        target_dir = self._get_path_from_hash(hash_val)
+
+        async def _create_for_hash():
+            await UvProcessor(
+                target_dir,
+                runtime_env,
+                logger,
+            )
+
+            loop = get_running_loop()
+            return await loop.run_in_executor(
+                None, get_directory_size_bytes, target_dir
+            )
+
+        if uri not in self._create_locks:
+            # async lock to prevent the same virtualenv being concurrently installed
+            self._create_locks[uri] = asyncio.Lock()
+
+        async with self._create_locks[uri]:
+            if hash_val in self._created_hash_bytes:
+                return self._created_hash_bytes[hash_val]
+            self._creating_task[hash_val] = task = create_task(_create_for_hash())
+            task.add_done_callback(lambda _: self._creating_task.pop(hash_val, None))
+            uv_dir_bytes = await task
+            self._created_hash_bytes[hash_val] = uv_dir_bytes
+            return uv_dir_bytes
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: "RuntimeEnvContext",  # noqa: F821
+        logger: logging.Logger = default_logger,
+    ):
+        if not runtime_env.has_uv():
+            return
+        # UvPlugin only uses a single URI.
+        uri = uris[0]
+        # Update py_executable.
+        protocol, hash_val = parse_uri(uri)
+        target_dir = self._get_path_from_hash(hash_val)
+        virtualenv_python = virtualenv_utils.get_virtualenv_python(target_dir)
+
+        if not os.path.exists(virtualenv_python):
+            raise ValueError(
+                f"Local directory {target_dir} for URI {uri} does "
+                "not exist on the cluster. Something may have gone wrong while "
+                "installing the runtime_env `uv` packages."
+            )
+        context.py_executable = virtualenv_python
+        context.command_prefix += virtualenv_utils.get_virtualenv_activate_command(
+            target_dir
+        )

python/ray/_private/runtime_env/uv_runtime_env_hook.py

@@ -0,0 +1,452 @@
+import argparse
+import copy
+import optparse
+import os
+import pathlib
+import platform
+import sys
+import urllib.parse
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import psutil
+
+
+def _is_path(path_or_uri: str) -> bool:
+    """Returns True if uri_or_path is a path and False otherwise.
+
+    IMPORTANT: This is a duplicate of ray._private.path_utils.is_path().
+
+    Why we can't import from path_utils:
+    - This hook runs via `uv run --no-project uv_runtime_env_hook.py` in test scenarios
+    - UV creates a minimal environment without dependencies installed yet
+    - Importing from ray._private.path_utils triggers the full Ray import chain:
+      ray._private.path_utils → ray/__init__.py → ray._private.worker →
+      ray.widgets → ray.widgets.util → packaging.version
+    - The 'packaging' module is not available in the minimal UV environment,
+      causing: ModuleNotFoundError: No module named 'packaging.version'
+
+    This duplicate implementation uses only stdlib (pathlib, urllib.parse)
+    to avoid the dependency issue. If you modify this function, ensure you
+    also update ray._private.path_utils.is_path() to keep them in sync.
+    """
+    if not isinstance(path_or_uri, str):
+        raise TypeError(f"path_or_uri must be a string, got {type(path_or_uri)}.")
+
+    parsed_path = pathlib.Path(path_or_uri)
+    parsed_uri = urllib.parse.urlparse(path_or_uri)
+
+    if isinstance(parsed_path, pathlib.PurePosixPath):
+        return not parsed_uri.scheme
+    elif isinstance(parsed_path, pathlib.PureWindowsPath):
+        return parsed_uri.scheme == parsed_path.drive.strip(":").lower()
+    else:
+        # this should never happen
+        raise TypeError(f"Unsupported path type: {type(parsed_path).__name__}")
+
+
+def _create_uv_run_parser():
+    """Create and return the argument parser for 'uv run' command."""
+
+    parser = optparse.OptionParser(prog="uv run", add_help_option=False)
+
+    # Disable interspersed args to stop parsing when we hit the first
+    # argument that is not recognized by the parser.
+    parser.disable_interspersed_args()
+
+    # Main options group
+    main_group = optparse.OptionGroup(parser, "Main options")
+    main_group.add_option("--extra", action="append", dest="extras")
+    main_group.add_option("--all-extras", action="store_true")
+    main_group.add_option("--no-extra", action="append", dest="no_extras")
+    main_group.add_option("--no-dev", action="store_true")
+    main_group.add_option("--group", action="append", dest="groups")
+    main_group.add_option("--no-group", action="append", dest="no_groups")
+    main_group.add_option("--no-default-groups", action="store_true")
+    main_group.add_option("--only-group", action="append", dest="only_groups")
+    main_group.add_option("--all-groups", action="store_true")
+    main_group.add_option("-m", "--module")
+    main_group.add_option("--only-dev", action="store_true")
+    main_group.add_option("--no-editable", action="store_true")
+    main_group.add_option("--exact", action="store_true")
+    main_group.add_option("--env-file", action="append", dest="env_files")
+    main_group.add_option("--no-env-file", action="store_true")
+    parser.add_option_group(main_group)
+
+    # With options
+    with_group = optparse.OptionGroup(parser, "With options")
+    with_group.add_option("--with", action="append", dest="with_packages")
+    with_group.add_option("--with-editable", action="append", dest="with_editable")
+    with_group.add_option(
+        "--with-requirements", action="append", dest="with_requirements"
+    )
+    parser.add_option_group(with_group)
+
+    # Environment options
+    env_group = optparse.OptionGroup(parser, "Environment options")
+    env_group.add_option("--isolated", action="store_true")
+    env_group.add_option("--active", action="store_true")
+    env_group.add_option("--no-sync", action="store_true")
+    env_group.add_option("--locked", action="store_true")
+    env_group.add_option("--frozen", action="store_true")
+    parser.add_option_group(env_group)
+
+    # Script options
+    script_group = optparse.OptionGroup(parser, "Script options")
+    script_group.add_option("-s", "--script", action="store_true")
+    script_group.add_option("--gui-script", action="store_true")
+    parser.add_option_group(script_group)
+
+    # Workspace options
+    workspace_group = optparse.OptionGroup(parser, "Workspace options")
+    workspace_group.add_option("--all-packages", action="store_true")
+    workspace_group.add_option("--package")
+    workspace_group.add_option("--no-project", action="store_true")
+    parser.add_option_group(workspace_group)
+
+    # Index options
+    index_group = optparse.OptionGroup(parser, "Index options")
+    index_group.add_option("--index", action="append", dest="indexes")
+    index_group.add_option("--default-index")
+    index_group.add_option("-i", "--index-url")
+    index_group.add_option(
+        "--extra-index-url", action="append", dest="extra_index_urls"
+    )
+    index_group.add_option("-f", "--find-links", action="append", dest="find_links")
+    index_group.add_option("--no-index", action="store_true")
+    index_group.add_option(
+        "--index-strategy",
+        type="choice",
+        choices=["first-index", "unsafe-first-match", "unsafe-best-match"],
+    )
+    index_group.add_option(
+        "--keyring-provider", type="choice", choices=["disabled", "subprocess"]
+    )
+    parser.add_option_group(index_group)
+
+    # Resolver options
+    resolver_group = optparse.OptionGroup(parser, "Resolver options")
+    resolver_group.add_option("-U", "--upgrade", action="store_true")
+    resolver_group.add_option(
+        "-P", "--upgrade-package", action="append", dest="upgrade_packages"
+    )
+    resolver_group.add_option(
+        "--resolution", type="choice", choices=["highest", "lowest", "lowest-direct"]
+    )
+    resolver_group.add_option(
+        "--prerelease",
+        type="choice",
+        choices=[
+            "disallow",
+            "allow",
+            "if-necessary",
+            "explicit",
+            "if-necessary-or-explicit",
+        ],
+    )
+    resolver_group.add_option(
+        "--fork-strategy", type="choice", choices=["fewest", "requires-python"]
+    )
+    resolver_group.add_option("--exclude-newer")
+    resolver_group.add_option("--no-sources", action="store_true")
+    parser.add_option_group(resolver_group)
+
+    # Installer options
+    installer_group = optparse.OptionGroup(parser, "Installer options")
+    installer_group.add_option("--reinstall", action="store_true")
+    installer_group.add_option(
+        "--reinstall-package", action="append", dest="reinstall_packages"
+    )
+    installer_group.add_option(
+        "--link-mode", type="choice", choices=["clone", "copy", "hardlink", "symlink"]
+    )
+    installer_group.add_option("--compile-bytecode", action="store_true")
+    parser.add_option_group(installer_group)
+
+    # Build options
+    build_group = optparse.OptionGroup(parser, "Build options")
+    build_group.add_option(
+        "-C", "--config-setting", action="append", dest="config_settings"
+    )
+    build_group.add_option("--no-build-isolation", action="store_true")
+    build_group.add_option(
+        "--no-build-isolation-package",
+        action="append",
+        dest="no_build_isolation_packages",
+    )
+    build_group.add_option("--no-build", action="store_true")
+    build_group.add_option(
+        "--no-build-package", action="append", dest="no_build_packages"
+    )
+    build_group.add_option("--no-binary", action="store_true")
+    build_group.add_option(
+        "--no-binary-package", action="append", dest="no_binary_packages"
+    )
+    parser.add_option_group(build_group)
+
+    # Cache options
+    cache_group = optparse.OptionGroup(parser, "Cache options")
+    cache_group.add_option("-n", "--no-cache", action="store_true")
+    cache_group.add_option("--cache-dir")
+    cache_group.add_option("--refresh", action="store_true")
+    cache_group.add_option(
+        "--refresh-package", action="append", dest="refresh_packages"
+    )
+    parser.add_option_group(cache_group)
+
+    # Python options
+    python_group = optparse.OptionGroup(parser, "Python options")
+    python_group.add_option("-p", "--python")
+    python_group.add_option("--managed-python", action="store_true")
+    python_group.add_option("--no-managed-python", action="store_true")
+    python_group.add_option("--no-python-downloads", action="store_true")
+    # note: the following is a legacy option and will be removed at some point
+    # https://github.com/astral-sh/uv/pull/12246
+    python_group.add_option(
+        "--python-preference",
+        type="choice",
+        choices=["only-managed", "managed", "system", "only-system"],
+    )
+    parser.add_option_group(python_group)
+
+    # Global options
+    global_group = optparse.OptionGroup(parser, "Global options")
+    global_group.add_option("-q", "--quiet", action="count", default=0)
+    global_group.add_option("-v", "--verbose", action="count", default=0)
+    global_group.add_option(
+        "--color", type="choice", choices=["auto", "always", "never"]
+    )
+    global_group.add_option("--native-tls", action="store_true")
+    global_group.add_option("--offline", action="store_true")
+    global_group.add_option(
+        "--allow-insecure-host", action="append", dest="insecure_hosts"
+    )
+    global_group.add_option("--no-progress", action="store_true")
+    global_group.add_option("--directory")
+    global_group.add_option("--project")
+    global_group.add_option("--config-file")
+    global_group.add_option("--no-config", action="store_true")
+    parser.add_option_group(global_group)
+
+    return parser
+
+
+def _parse_args(
+    parser: optparse.OptionParser, args: List[str]
+) -> Tuple[optparse.Values, List[str]]:
+    """
+    Parse the command-line options found in 'args'.
+
+    Replacement for parser.parse_args that handles unknown arguments
+    by keeping them in the command list instead of erroring and
+    discarding them.
+    """
+    parser.rargs = args
+    parser.largs = []
+    options = parser.get_default_values()
+    try:
+        parser._process_args(parser.largs, parser.rargs, options)
+    except optparse.BadOptionError as err:
+        # If we hit an argument that is not recognized, we put it
+        # back into the unconsumed arguments
+        parser.rargs = [err.opt_str] + parser.rargs
+    return options, parser.rargs
+
+
+def _check_working_dir_files(
+    uv_run_args: optparse.Values, runtime_env: Dict[str, Any]
+) -> None:
+    """
+    Check that the files required by uv are local to the working_dir. This catches
+    the most common cases of how things are different in Ray, i.e. not the whole file
+    system will be available on the workers, only the working_dir.
+
+    The function won't return anything, it just raises a RuntimeError if there is an error.
+    """
+    working_dir = Path(runtime_env["working_dir"]).resolve()
+
+    # Check if the requirements.txt file is in the working_dir
+    if uv_run_args.with_requirements:
+        for requirements_file in uv_run_args.with_requirements:
+            if not Path(requirements_file).resolve().is_relative_to(working_dir):
+                raise RuntimeError(
+                    f"You specified --with-requirements={uv_run_args.with_requirements} but "
+                    f"the requirements file is not in the working_dir {runtime_env['working_dir']}, "
+                    "so the workers will not have access to the file. Make sure "
+                    "the requirements file is in the working directory. "
+                    "You can do so by specifying --directory in 'uv run', by changing the current "
+                    "working directory before running 'uv run', or by using the 'working_dir' "
+                    "parameter of the runtime_environment."
+                )
+
+    # Check if the pyproject.toml file is in the working_dir
+    pyproject = None
+    if uv_run_args.no_project:
+        pyproject = None
+    elif uv_run_args.project:
+        pyproject = Path(uv_run_args.project)
+    else:
+        # Walk up the directory tree until pyproject.toml is found
+        current_path = Path.cwd().resolve()
+        while current_path != current_path.parent:
+            if (current_path / "pyproject.toml").exists():
+                pyproject = Path(current_path / "pyproject.toml")
+                break
+            current_path = current_path.parent
+
+    if pyproject and not pyproject.resolve().is_relative_to(working_dir):
+        raise RuntimeError(
+            f"Your {pyproject.resolve()} is not in the working_dir {runtime_env['working_dir']}, "
+            "so the workers will not have access to the file. Make sure "
+            "the pyproject.toml file is in the working directory. "
+            "You can do so by specifying --directory in 'uv run', by changing the current "
+            "working directory before running 'uv run', or by using the 'working_dir' "
+            "parameter of the runtime_environment."
+        )
+
+
+def _get_uv_run_cmdline() -> Optional[List[str]]:
+    """
+    Return the command line of the first ancestor process that was run with
+    "uv run" and None if there is no such ancestor.
+
+    uv spawns the python process as a child process, so we first check the
+    parent process command line. We also check our parent's parents since
+    the Ray driver might be run as a subprocess of the 'uv run' process.
+    """
+    parents = psutil.Process().parents()
+    for parent in parents:
+        try:
+            cmdline = parent.cmdline()
+            if (
+                len(cmdline) > 1
+                and os.path.basename(cmdline[0]) == "uv"
+                and cmdline[1] == "run"
+            ):
+                return cmdline
+        except psutil.NoSuchProcess:
+            continue
+        except psutil.AccessDenied:
+            continue
+    return None
+
+
+def hook(runtime_env: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+    """Hook that detects if the driver is run in 'uv run' and sets the runtime environment accordingly."""
+
+    runtime_env = copy.deepcopy(runtime_env) or {}
+
+    cmdline = _get_uv_run_cmdline()
+    if not cmdline:
+        # This means the driver was not run in a 'uv run' environment -- in this case
+        # we leave the runtime environment unchanged
+        return runtime_env
+
+    # First check that the "uv" and "pip" runtime environments are not used.
+    if "uv" in runtime_env or "pip" in runtime_env:
+        raise RuntimeError(
+            "You are using the 'pip' or 'uv' runtime environments together with "
+            "'uv run'. These are not compatible since 'uv run' will run the workers "
+            "in an isolated environment -- please add the 'pip' or 'uv' dependencies to your "
+            "'uv run' environment e.g. by including them in your pyproject.toml."
+        )
+
+    # Extract the arguments uv_run_args of 'uv run' that are not part of the command.
+    args_to_parse = cmdline[2:]  # Remove 'uv run' prefix
+    original_length = len(
+        args_to_parse
+    )  # Save before parsing (parser modifies in-place)
+
+    parser = _create_uv_run_parser()
+    (options, command) = _parse_args(parser, args_to_parse)
+
+    # Calculate how many arguments were consumed by the parser.
+    # Since disable_interspersed_args() is set, parsing stops at the first
+    # unrecognized argument (the command), so all consumed args are uv options.
+    args_consumed = original_length - len(command)
+    uv_run_args = cmdline[: 2 + args_consumed]
+
+    # Remove the "--directory" argument since it has already been taken into
+    # account when setting the current working directory of the current process.
+    # Also remove the "--module" argument, since the default_worker.py is
+    # invoked as a script and not as a module.
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--directory")
+    parser.add_argument("-m", "--module")
+    _, remaining_uv_run_args = parser.parse_known_args(uv_run_args)
+
+    # Pin the worker Python for `uv run` unless the driver already specified one.
+    #
+    # Without this, uv may resolve a different Python than the Ray driver (e.g. 3.12 instead of 3.11),
+    # which causes Ray to fail with a Python version mismatch. (See https://github.com/ray-project/ray/issues/59639)
+    # order of precedence:
+    # 1. options.python (driver specified)
+    # 2. env_vars["UV_PYTHON"]
+    # 3. current os.environ["UV_PYTHON"]
+    # 4. platform.python_version() (since uv run uses the same Python as the driver)
+    if not options.python:
+        env_vars = runtime_env.get("env_vars") or {}
+        uv_python = (
+            env_vars.get("UV_PYTHON")
+            or os.environ.get("UV_PYTHON")
+            or platform.python_version()
+        )
+        remaining_uv_run_args = remaining_uv_run_args + ["--python", uv_python]
+
+    # Append "python" to the end so that when Ray adds "-m default_worker.py",
+    # it becomes "uv run --python X.Y.Z python -m default_worker.py"
+    remaining_uv_run_args = remaining_uv_run_args + ["python"]
+
+    runtime_env["py_executable"] = " ".join(remaining_uv_run_args)
+
+    # If the user specified a working_dir, we always honor it, otherwise
+    # use the same working_dir that uv run would use
+    if "working_dir" not in runtime_env:
+        runtime_env["working_dir"] = os.getcwd()
+
+    # Validate that pyproject.toml and requirements files are within working_dir
+    # This prevents runtime errors on workers when files are not accessible
+    # Only validate for local paths - remote URIs will be downloaded by Ray
+    working_dir = runtime_env["working_dir"]
+    if _is_path(working_dir):
+        _check_working_dir_files(options, runtime_env)
+
+    return runtime_env
+
+
+# This __main__ is used for unit testing if the runtime_env_hook picks up the
+# right settings.
+if __name__ == "__main__":
+    import json
+
+    test_parser = argparse.ArgumentParser()
+    test_parser.add_argument("--extra-args", action="store_true")
+    test_parser.add_argument("runtime_env")
+    args = test_parser.parse_args()
+
+    # If the env variable is set, add one more level of subprocess indirection
+    if os.environ.get("RAY_TEST_UV_ADD_SUBPROCESS_INDIRECTION") == "1":
+        import subprocess
+
+        env = os.environ.copy()
+        env.pop("RAY_TEST_UV_ADD_SUBPROCESS_INDIRECTION")
+        subprocess.check_call([sys.executable] + sys.argv, env=env)
+        sys.exit(0)
+
+    # If the following env variable is set, we use multiprocessing
+    # spawn to start the subprocess, since it uses a different way to
+    # modify the command line than subprocess.check_call
+    if os.environ.get("RAY_TEST_UV_MULTIPROCESSING_SPAWN") == "1":
+        import multiprocessing
+
+        multiprocessing.set_start_method("spawn")
+        pool = multiprocessing.Pool(processes=1)
+        runtime_env = json.loads(args.runtime_env)
+        print(json.dumps(pool.apply(hook, (runtime_env,))))
+        sys.exit(0)
+
+    # We purposefully modify sys.argv here to make sure the hook is robust
+    # against such modification.
+    sys.argv.pop(1)
+    runtime_env = json.loads(args.runtime_env)
+    print(json.dumps(hook(runtime_env)))

python/ray/_private/runtime_env/validation.py

@@ -0,0 +1,466 @@
+import logging
+import sys
+from collections import OrderedDict
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import yaml
+
+from ray._private.path_utils import is_path
+from ray._private.runtime_env.packaging import parse_path
+
+logger = logging.getLogger(__name__)
+
+
+def validate_path(path: str) -> None:
+    """Parse the path to ensure it is well-formed and exists."""
+    parse_path(path)
+
+
+def validate_uri(uri: str):
+    try:
+        from ray._private.runtime_env.packaging import Protocol, parse_uri
+
+        protocol, path = parse_uri(uri)
+    except ValueError:
+        raise ValueError(
+            f"{uri} is not a valid URI. Passing directories or modules to "
+            "be dynamically uploaded is only supported at the job level "
+            "(i.e., passed to `ray.init`)."
+        )
+
+    supported_extensions = (".zip", ".whl", ".tar.gz", ".tgz")
+    if protocol in Protocol.remote_protocols() and not any(
+        path.endswith(ext) for ext in supported_extensions
+    ):
+        raise ValueError(
+            "Only .zip, .whl, .tar.gz, and .tgz files supported for remote URIs."
+        )
+
+
+def _handle_local_deps_requirement_file(requirements_file: str):
+    """Read the given [requirements_file], and return all required dependencies."""
+    requirements_path = Path(requirements_file)
+    if not requirements_path.is_file():
+        raise ValueError(f"{requirements_path} is not a valid file")
+    return requirements_path.read_text().strip().split("\n")
+
+
+def validate_py_modules_uris(py_modules_uris: List[str]) -> List[str]:
+    """Parses and validates a 'py_modules' option.
+
+    Expects py_modules to be a list of URIs.
+    """
+    if not isinstance(py_modules_uris, list):
+        raise TypeError(
+            "`py_modules` must be a list of strings, got " f"{type(py_modules_uris)}."
+        )
+
+    for module in py_modules_uris:
+
+        if not isinstance(module, str):
+            raise TypeError("`py_module` must be a string, got " f"{type(module)}.")
+
+        validate_uri(module)
+
+
+def parse_and_validate_py_modules(py_modules: List[str]) -> List[str]:
+    """Parses and validates a 'py_modules' option.
+
+    Expects py_modules to be a list of local paths or URIs.
+    """
+    if not isinstance(py_modules, list):
+        raise TypeError(
+            "`py_modules` must be a list of strings, got " f"{type(py_modules)}."
+        )
+
+    for module in py_modules:
+
+        if not isinstance(module, str):
+            raise TypeError("`py_module` must be a string, got " f"{type(module)}.")
+
+        if is_path(module):
+            validate_path(module)
+        else:
+            validate_uri(module)
+
+    return py_modules
+
+
+def validate_working_dir_uri(working_dir_uri: str) -> str:
+    """Parses and validates a 'working_dir' option."""
+    if not isinstance(working_dir_uri, str):
+        raise TypeError(
+            "`working_dir` must be a string, got " f"{type(working_dir_uri)}."
+        )
+
+    validate_uri(working_dir_uri)
+
+
+def parse_and_validate_working_dir(working_dir: str) -> str:
+    """Parses and validates a 'working_dir' option.
+
+    This can be a URI or a path.
+    """
+    assert working_dir is not None
+
+    if not isinstance(working_dir, str):
+        raise TypeError("`working_dir` must be a string, got " f"{type(working_dir)}.")
+
+    if is_path(working_dir):
+        validate_path(working_dir)
+    else:
+        validate_uri(working_dir)
+
+    return working_dir
+
+
+def parse_and_validate_conda(conda: Union[str, dict]) -> Union[str, dict]:
+    """Parses and validates a user-provided 'conda' option.
+
+    Conda can be one of three cases:
+        1) A dictionary describing the env. This is passed through directly.
+        2) A string referring to the name of a preinstalled conda env.
+        3) A string pointing to a local conda YAML file. This is detected
+           by looking for a '.yaml' or '.yml' suffix. In this case, the file
+           will be read as YAML and passed through as a dictionary.
+    """
+    assert conda is not None
+
+    if sys.platform == "win32":
+        logger.warning(
+            "runtime environment support is experimental on Windows. "
+            "If you run into issues please file a report at "
+            "https://github.com/ray-project/ray/issues."
+        )
+
+    result = conda
+    if isinstance(conda, str):
+        file_path = Path(conda)
+        if file_path.suffix in (".yaml", ".yml"):
+            if not file_path.is_file():
+                raise ValueError(f"Can't find conda YAML file {file_path}.")
+            try:
+                result = yaml.safe_load(file_path.read_text())
+            except Exception as e:
+                raise ValueError(f"Failed to read conda file {file_path}: {e}.")
+        elif file_path.is_absolute():
+            if not file_path.is_dir():
+                raise ValueError(f"Can't find conda env directory {file_path}.")
+            result = str(file_path)
+    elif isinstance(conda, dict):
+        result = conda
+    else:
+        raise TypeError(
+            "runtime_env['conda'] must be of type str or " f"dict, got {type(conda)}."
+        )
+
+    return result
+
+
+def parse_and_validate_uv(uv: Union[str, List[str], Dict]) -> Optional[Dict]:
+    """Parses and validates a user-provided 'uv' option.
+
+    The value of the input 'uv' field can be one of two cases:
+        1) A List[str] describing the requirements. This is passed through.
+           Example usage: ["tensorflow", "requests"]
+        2) a string containing the path to a local pip “requirements.txt” file.
+        3) A python dictionary that has one field:
+            a) packages (required, List[str]): a list of uv packages, it same as 1).
+            b) uv_check (optional, bool): whether to enable pip check at the end of uv
+               install, default to False.
+            c) uv_version (optional, str): user provides a specific uv to use; if
+               unspecified, default version of uv will be used.
+            d) uv_pip_install_options (optional, List[str]): user-provided options for
+              `uv pip install` command, default to ["--no-cache"].
+
+    The returned parsed value will be a list of packages. If a Ray library
+    (e.g. "ray[serve]") is specified, it will be deleted and replaced by its
+    dependencies (e.g. "uvicorn", "requests").
+    """
+    assert uv is not None
+    if sys.platform == "win32":
+        logger.warning(
+            "runtime environment support is experimental on Windows. "
+            "If you run into issues please file a report at "
+            "https://github.com/ray-project/ray/issues."
+        )
+
+    result: str = ""
+    if isinstance(uv, str):
+        uv_list = _handle_local_deps_requirement_file(uv)
+        result = dict(packages=uv_list, uv_check=False)
+    elif isinstance(uv, list) and all(isinstance(dep, str) for dep in uv):
+        result = dict(packages=uv, uv_check=False)
+    elif isinstance(uv, dict):
+        if set(uv.keys()) - {
+            "packages",
+            "uv_check",
+            "uv_version",
+            "uv_pip_install_options",
+        }:
+            raise ValueError(
+                "runtime_env['uv'] can only have these fields: "
+                "packages, uv_check, uv_version and uv_pip_install_options, but got: "
+                f"{list(uv.keys())}"
+            )
+        if "packages" not in uv:
+            raise ValueError(
+                f"runtime_env['uv'] must include field 'packages', but got {uv}"
+            )
+        if "uv_check" in uv and not isinstance(uv["uv_check"], bool):
+            raise TypeError(
+                "runtime_env['uv']['uv_check'] must be of type bool, "
+                f"got {type(uv['uv_check'])}"
+            )
+        if "uv_version" in uv and not isinstance(uv["uv_version"], str):
+            raise TypeError(
+                "runtime_env['uv']['uv_version'] must be of type str, "
+                f"got {type(uv['uv_version'])}"
+            )
+        if "uv_pip_install_options" in uv:
+            if not isinstance(uv["uv_pip_install_options"], list):
+                raise TypeError(
+                    "runtime_env['uv']['uv_pip_install_options'] must be of type "
+                    f"list[str] got {type(uv['uv_pip_install_options'])}"
+                )
+            # Check each item in installation option.
+            for idx, cur_opt in enumerate(uv["uv_pip_install_options"]):
+                if not isinstance(cur_opt, str):
+                    raise TypeError(
+                        "runtime_env['uv']['uv_pip_install_options'] must be of type "
+                        f"list[str] got {type(cur_opt)} for {idx}-th item."
+                    )
+
+        result = uv.copy()
+        result["uv_check"] = uv.get("uv_check", False)
+        result["uv_pip_install_options"] = uv.get(
+            "uv_pip_install_options", ["--no-cache"]
+        )
+        if not isinstance(uv["packages"], list):
+            raise ValueError(
+                "runtime_env['uv']['packages'] must be of type list, "
+                f"got: {type(uv['packages'])}"
+            )
+    else:
+        raise TypeError(
+            "runtime_env['uv'] must be of type " f"List[str], or dict, got {type(uv)}"
+        )
+
+    # Deduplicate packages for package lists.
+    result["packages"] = list(OrderedDict.fromkeys(result["packages"]))
+
+    if len(result["packages"]) == 0:
+        result = None
+    logger.debug(f"Rewrote runtime_env `uv` field from {uv} to {result}.")
+    return result
+
+
+def parse_and_validate_pip(pip: Union[str, List[str], Dict]) -> Optional[Dict]:
+    """Parses and validates a user-provided 'pip' option.
+
+    The value of the input 'pip' field can be one of two cases:
+        1) A List[str] describing the requirements. This is passed through.
+        2) A string pointing to a local requirements file. In this case, the
+           file contents will be read split into a list.
+        3) A python dictionary that has three fields:
+            a) packages (required, List[str]): a list of pip packages, it same as 1).
+            b) pip_check (optional, bool): whether to enable pip check at the end of pip
+               install, default to False.
+            c) pip_version (optional, str): the version of pip, ray will spell
+               the package name 'pip' in front of the `pip_version` to form the final
+               requirement string, the syntax of a requirement specifier is defined in
+               full in PEP 508.
+            d) pip_install_options (optional, List[str]): user-provided options for
+              `pip install` command, defaults to ["--disable-pip-version-check", "--no-cache-dir"].
+
+    The returned parsed value will be a list of pip packages. If a Ray library
+    (e.g. "ray[serve]") is specified, it will be deleted and replaced by its
+    dependencies (e.g. "uvicorn", "requests").
+    """
+    assert pip is not None
+    result = None
+
+    if sys.platform == "win32":
+        logger.warning(
+            "runtime environment support is experimental on Windows. "
+            "If you run into issues please file a report at "
+            "https://github.com/ray-project/ray/issues."
+        )
+    if isinstance(pip, str):
+        # We have been given a path to a requirements.txt file.
+        pip_list = _handle_local_deps_requirement_file(pip)
+        result = dict(
+            packages=pip_list,
+            pip_check=False,
+        )
+    elif isinstance(pip, list) and all(isinstance(dep, str) for dep in pip):
+        result = dict(packages=pip, pip_check=False)
+    elif isinstance(pip, dict):
+        if set(pip.keys()) - {
+            "packages",
+            "pip_check",
+            "pip_install_options",
+            "pip_version",
+        }:
+            raise ValueError(
+                "runtime_env['pip'] can only have these fields: "
+                "packages, pip_check, pip_install_options and pip_version, but got: "
+                f"{list(pip.keys())}"
+            )
+
+        if "pip_check" in pip and not isinstance(pip["pip_check"], bool):
+            raise TypeError(
+                "runtime_env['pip']['pip_check'] must be of type bool, "
+                f"got {type(pip['pip_check'])}"
+            )
+        if "pip_version" in pip:
+            if not isinstance(pip["pip_version"], str):
+                raise TypeError(
+                    "runtime_env['pip']['pip_version'] must be of type str, "
+                    f"got {type(pip['pip_version'])}"
+                )
+        if "pip_install_options" in pip:
+            if not isinstance(pip["pip_install_options"], list):
+                raise TypeError(
+                    "runtime_env['pip']['pip_install_options'] must be of type "
+                    f"list[str] got {type(pip['pip_install_options'])}"
+                )
+            # Check each item in installation option.
+            for idx, cur_opt in enumerate(pip["pip_install_options"]):
+                if not isinstance(cur_opt, str):
+                    raise TypeError(
+                        "runtime_env['pip']['pip_install_options'] must be of type "
+                        f"list[str] got {type(cur_opt)} for {idx}-th item."
+                    )
+
+        result = pip.copy()
+        # Contrary to pip_check, we do not insert the default value of pip_install_options.
+        # This is to maintain backwards compatibility with ray==2.0.1
+        result["pip_check"] = pip.get("pip_check", False)
+
+        if "packages" not in pip:
+            raise ValueError(
+                f"runtime_env['pip'] must include field 'packages', but got {pip}"
+            )
+        elif isinstance(pip["packages"], str):
+            result["packages"] = _handle_local_deps_requirement_file(pip["packages"])
+        elif not isinstance(pip["packages"], list):
+            raise ValueError(
+                "runtime_env['pip']['packages'] must be of type str of list, "
+                f"got: {type(pip['packages'])}"
+            )
+    else:
+        raise TypeError(
+            "runtime_env['pip'] must be of type str or " f"List[str], got {type(pip)}"
+        )
+
+    # Eliminate duplicates to prevent `pip install` from erroring. Use
+    # OrderedDict to preserve the order of the list.  This makes the output
+    # deterministic and easier to debug, because pip install can have
+    # different behavior depending on the order of the input.
+    result["packages"] = list(OrderedDict.fromkeys(result["packages"]))
+
+    if len(result["packages"]) == 0:
+        result = None
+
+    logger.debug(f"Rewrote runtime_env `pip` field from {pip} to {result}.")
+
+    return result
+
+
+def parse_and_validate_container(container: List[str]) -> List[str]:
+    """Parses and validates a user-provided 'container' option.
+
+    This is passed through without validation (for now).
+    """
+    assert container is not None
+    return container
+
+
+def parse_and_validate_excludes(excludes: List[str]) -> List[str]:
+    """Parses and validates a user-provided 'excludes' option.
+
+    This is validated to verify that it is of type List[str].
+
+    If an empty list is passed, we return `None` for consistency.
+    """
+    assert excludes is not None
+
+    if isinstance(excludes, list) and len(excludes) == 0:
+        return None
+
+    if isinstance(excludes, list) and all(isinstance(path, str) for path in excludes):
+        return excludes
+    else:
+        raise TypeError(
+            "runtime_env['excludes'] must be of type "
+            f"List[str], got {type(excludes)}"
+        )
+
+
+def parse_and_validate_env_vars(env_vars: Dict[str, str]) -> Optional[Dict[str, str]]:
+    """Parses and validates a user-provided 'env_vars' option.
+
+    This is validated to verify that all keys and vals are strings.
+
+    If an empty dictionary is passed, we return `None` for consistency.
+
+    Args:
+        env_vars: A dictionary of environment variables to set in the
+            runtime environment.
+
+    Returns:
+        The validated env_vars dictionary, or None if it was empty.
+
+    Raises:
+        TypeError: If the env_vars is not a dictionary of strings. The error message
+            will include the type of the invalid value.
+    """
+    assert env_vars is not None
+    if len(env_vars) == 0:
+        return None
+
+    if not isinstance(env_vars, dict):
+        raise TypeError(
+            "runtime_env['env_vars'] must be of type "
+            f"Dict[str, str], got {type(env_vars)}"
+        )
+
+    for key, val in env_vars.items():
+        if not isinstance(key, str):
+            raise TypeError(
+                "runtime_env['env_vars'] must be of type "
+                f"Dict[str, str], but the key {key} is of type {type(key)}"
+            )
+        if not isinstance(val, str):
+            raise TypeError(
+                "runtime_env['env_vars'] must be of type "
+                f"Dict[str, str], but the value {val} is of type {type(val)}"
+            )
+
+    return env_vars
+
+
+# Dictionary mapping runtime_env options with the function to parse and
+# validate them.
+OPTION_TO_VALIDATION_FN = {
+    "py_modules": parse_and_validate_py_modules,
+    "working_dir": parse_and_validate_working_dir,
+    "excludes": parse_and_validate_excludes,
+    "conda": parse_and_validate_conda,
+    "pip": parse_and_validate_pip,
+    "uv": parse_and_validate_uv,
+    "env_vars": parse_and_validate_env_vars,
+    "container": parse_and_validate_container,
+}
+
+# RuntimeEnv can be created with local paths
+# for these options. However, after the packages
+# for these options have been uploaded to GCS,
+# they must be URIs. These functions provide the ability
+# to validate that these options only contain well-formed URIs.
+OPTION_TO_NO_PATH_VALIDATION_FN = {
+    "working_dir": validate_working_dir_uri,
+    "py_modules": validate_py_modules_uris,
+}

python/ray/_private/runtime_env/virtualenv_utils.py

@@ -0,0 +1,110 @@
+"""Utils to detect runtime environment."""
+
+import logging
+import os
+import sys
+from typing import List
+
+from ray._private.runtime_env.utils import check_output_cmd
+
+_WIN32 = os.name == "nt"
+
+
+def is_in_virtualenv() -> bool:
+    # virtualenv <= 16.7.9 sets the real_prefix,
+    # virtualenv > 16.7.9 & venv set the base_prefix.
+    # So, we check both of them here.
+    # https://github.com/pypa/virtualenv/issues/1622#issuecomment-586186094
+    return hasattr(sys, "real_prefix") or (
+        hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix
+    )
+
+
+def get_virtualenv_path(target_dir: str) -> str:
+    """Get virtual environment path."""
+    return os.path.join(target_dir, "virtualenv")
+
+
+def get_virtualenv_python(target_dir: str) -> str:
+    virtualenv_path = get_virtualenv_path(target_dir)
+    if _WIN32:
+        return os.path.join(virtualenv_path, "Scripts", "python.exe")
+    else:
+        return os.path.join(virtualenv_path, "bin", "python")
+
+
+def get_virtualenv_activate_command(target_dir: str) -> List[str]:
+    """Get the command to activate virtual environment."""
+    virtualenv_path = get_virtualenv_path(target_dir)
+    if _WIN32:
+        cmd = [os.path.join(virtualenv_path, "Scripts", "activate.bat")]
+    else:
+        cmd = ["source", os.path.join(virtualenv_path, "bin/activate")]
+    return cmd + ["1>&2", "&&"]
+
+
+async def create_or_get_virtualenv(path: str, cwd: str, logger: logging.Logger):
+    """Create or get a virtualenv from path."""
+    python = sys.executable
+    virtualenv_path = os.path.join(path, "virtualenv")
+    virtualenv_app_data_path = os.path.join(path, "virtualenv_app_data")
+
+    if _WIN32:
+        current_python_dir = sys.prefix
+        env = os.environ.copy()
+    else:
+        current_python_dir = os.path.abspath(
+            os.path.join(os.path.dirname(python), "..")
+        )
+        env = {}
+
+    if is_in_virtualenv():
+        # virtualenv-clone homepage:
+        # https://github.com/edwardgeorge/virtualenv-clone
+        # virtualenv-clone Usage:
+        # virtualenv-clone /path/to/existing/venv /path/to/cloned/ven
+        # or
+        # python -m clonevirtualenv /path/to/existing/venv /path/to/cloned/ven
+        clonevirtualenv = os.path.join(os.path.dirname(__file__), "_clonevirtualenv.py")
+        create_venv_cmd = [
+            python,
+            clonevirtualenv,
+            current_python_dir,
+            virtualenv_path,
+        ]
+        logger.info("Cloning virtualenv %s to %s", current_python_dir, virtualenv_path)
+    else:
+        # virtualenv options:
+        # https://virtualenv.pypa.io/en/latest/cli_interface.html
+        #
+        # --app-data
+        # --reset-app-data
+        #   Set an empty separated app data folder for current virtualenv.
+        #
+        # --no-periodic-update
+        #   Disable the periodic (once every 14 days) update of the embedded
+        #   wheels.
+        #
+        # --system-site-packages
+        #   Inherit site packages.
+        #
+        # --no-download
+        #   Never download the latest pip/setuptools/wheel from PyPI.
+        create_venv_cmd = [
+            python,
+            "-m",
+            "virtualenv",
+            "--app-data",
+            virtualenv_app_data_path,
+            "--reset-app-data",
+            "--no-periodic-update",
+            "--system-site-packages",
+            "--no-download",
+            virtualenv_path,
+        ]
+        logger.info(
+            "Creating virtualenv at %s, current python dir %s",
+            virtualenv_path,
+            virtualenv_path,
+        )
+    await check_output_cmd(create_venv_cmd, logger=logger, cwd=cwd, env=env)

python/ray/_private/runtime_env/working_dir.py

@@ -0,0 +1,278 @@
+import logging
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+import ray._private.ray_constants as ray_constants
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.packaging import (
+    Protocol,
+    delete_package,
+    download_and_unpack_package,
+    get_local_dir_from_uri,
+    get_uri_for_directory,
+    get_uri_for_package,
+    parse_uri,
+    upload_package_if_needed,
+    upload_package_to_gcs,
+)
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.utils import get_directory_size_bytes
+from ray._raylet import GcsClient
+from ray.exceptions import RuntimeEnvSetupError
+from ray.util.debug import log_once
+
+default_logger = logging.getLogger(__name__)
+
+_WIN32 = os.name == "nt"
+_LOG_ONCE_DEFAULT_EXCLUDE_PREFIX = "runtime_env_default_exclude:"
+
+
+def upload_working_dir_if_needed(
+    runtime_env: Dict[str, Any],
+    include_gitignore: bool,
+    scratch_dir: Optional[str] = None,
+    logger: Optional[logging.Logger] = default_logger,
+    upload_fn: Optional[Callable[[str, Optional[List[str]], bool], None]] = None,
+) -> Dict[str, Any]:
+    """Uploads the working_dir and replaces it with a URI.
+
+    If the working_dir is already a URI, this is a no-op.
+
+    Excludes are combined from:
+      - .gitignore and .rayignore files in the working_dir
+      - runtime_env["excludes"] field
+      - RAY_RUNTIME_ENV_DEFAULT_EXCLUDES constant, overridable via
+        RAY_OVERRIDE_RUNTIME_ENV_DEFAULT_EXCLUDES environment variable
+    """
+    working_dir = runtime_env.get("working_dir")
+    if working_dir is None:
+        return runtime_env
+
+    if not isinstance(working_dir, str) and not isinstance(working_dir, Path):
+        raise TypeError(
+            "working_dir must be a string or Path (either a local path "
+            f"or remote URI), got {type(working_dir)}."
+        )
+
+    if isinstance(working_dir, Path):
+        working_dir = str(working_dir)
+
+    # working_dir is already a URI -- just pass it through.
+    try:
+        protocol, path = parse_uri(working_dir)
+    except ValueError:
+        protocol, path = None, None
+
+    if protocol is not None:
+        supported_extensions = (".zip", ".tar.gz", ".tgz")
+        if protocol in Protocol.remote_protocols() and not any(
+            path.endswith(ext) for ext in supported_extensions
+        ):
+            raise ValueError(
+                "Only .zip, .tar.gz, and .tgz files supported for remote URIs."
+            )
+        return runtime_env
+
+    default_excludes = ray_constants.get_runtime_env_default_excludes()
+    user_excludes = runtime_env.get("excludes") or []
+    excludes = default_excludes + list(user_excludes)
+    # TODO(ricardo): 2026-01-07 Remove these warnings in a few releases. Added in
+    # case users rely on these directories being uploaded with their working_dir
+    # since this change would be difficult to debug.
+    logger = logger or default_logger
+    working_dir_path = Path(working_dir)
+    for d in default_excludes:
+        if (working_dir_path / d).exists() and log_once(
+            f"{_LOG_ONCE_DEFAULT_EXCLUDE_PREFIX}{d}"
+        ):
+            logger.warning(
+                "Directory %r is now ignored by default when packaging the working "
+                "directory. To disable this behavior, set "
+                "the `RAY_OVERRIDE_RUNTIME_ENV_DEFAULT_EXCLUDES=''` environment "
+                "variable.",
+                d,
+            )
+    try:
+        working_dir_uri = get_uri_for_directory(
+            working_dir,
+            include_gitignore=include_gitignore,
+            excludes=excludes,
+        )
+    except ValueError:  # working_dir is not a directory
+        package_path = Path(working_dir)
+        supported_local = (
+            package_path.suffix == ".zip"
+            or package_path.suffix == ".tgz"
+            or package_path.name.endswith(".tar.gz")
+        )
+        if not package_path.exists() or not supported_local:
+            raise ValueError(
+                f"directory {package_path} must be an existing "
+                "directory or a supported archive (.zip, .tar.gz, .tgz)"
+            )
+
+        pkg_uri = get_uri_for_package(package_path)
+        if upload_fn is not None:
+            upload_fn(working_dir, excludes=excludes, is_file=True)
+        else:
+            try:
+                upload_package_to_gcs(pkg_uri, package_path.read_bytes())
+            except Exception as e:
+                raise RuntimeEnvSetupError(
+                    f"Failed to upload package {package_path} to the Ray cluster: {e}"
+                ) from e
+        runtime_env["working_dir"] = pkg_uri
+        return runtime_env
+    if upload_fn is None:
+        if scratch_dir is None:
+            scratch_dir = os.getcwd()
+        try:
+            upload_package_if_needed(
+                working_dir_uri,
+                scratch_dir,
+                working_dir,
+                include_parent_dir=False,
+                excludes=excludes,
+                include_gitignore=include_gitignore,
+                logger=logger,
+            )
+        except Exception as e:
+            raise RuntimeEnvSetupError(
+                f"Failed to upload working_dir {working_dir} to the Ray cluster: {e}"
+            ) from e
+    else:
+        upload_fn(working_dir, excludes=excludes)
+
+    runtime_env["working_dir"] = working_dir_uri
+    return runtime_env
+
+
+def set_pythonpath_in_context(python_path: str, context: RuntimeEnvContext):
+    """Insert the path as the first entry in PYTHONPATH in the runtime env.
+
+    This is compatible with users providing their own PYTHONPATH in env_vars,
+    and is also compatible with the existing PYTHONPATH in the cluster.
+
+    The import priority is as follows:
+    this python_path arg > env_vars PYTHONPATH > existing cluster env PYTHONPATH.
+    """
+    if "PYTHONPATH" in context.env_vars:
+        python_path += os.pathsep + context.env_vars["PYTHONPATH"]
+    if "PYTHONPATH" in os.environ:
+        python_path += os.pathsep + os.environ["PYTHONPATH"]
+    context.env_vars["PYTHONPATH"] = python_path
+
+
+class WorkingDirPlugin(RuntimeEnvPlugin):
+    name = "working_dir"
+
+    # Note working_dir is not following the priority order of other plugins. Instead
+    # it's specially treated to happen before all other plugins.
+    priority = 5
+
+    def __init__(self, resources_dir: str, gcs_client: GcsClient):
+        self._resources_dir = os.path.join(resources_dir, "working_dir_files")
+        self._gcs_client = gcs_client
+        try_to_create_directory(self._resources_dir)
+
+    def delete_uri(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ) -> int:
+        """Delete URI and return the number of bytes deleted."""
+        logger.info("Got request to delete working dir URI %s", uri)
+        local_dir = get_local_dir_from_uri(uri, self._resources_dir)
+        local_dir_size = get_directory_size_bytes(local_dir)
+
+        deleted = delete_package(uri, self._resources_dir)
+        if not deleted:
+            logger.warning(f"Tried to delete nonexistent URI: {uri}.")
+            return 0
+
+        return local_dir_size
+
+    def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]:  # noqa: F821
+        working_dir_uri = runtime_env.working_dir()
+        if working_dir_uri != "":
+            return [working_dir_uri]
+        return []
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: dict,
+        context: RuntimeEnvContext,
+        logger: logging.Logger = default_logger,
+    ) -> int:
+        local_dir = await download_and_unpack_package(
+            uri,
+            self._resources_dir,
+            self._gcs_client,
+            logger=logger,
+            overwrite=True,
+        )
+        return get_directory_size_bytes(local_dir)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env_dict: Dict,
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        if not uris:
+            return
+
+        # WorkingDirPlugin uses a single URI.
+        uri = uris[0]
+        local_dir = get_local_dir_from_uri(uri, self._resources_dir)
+        if not local_dir.exists():
+            raise ValueError(
+                f"Local directory {local_dir} for URI {uri} does "
+                "not exist on the cluster. Something may have gone wrong while "
+                "downloading or unpacking the working_dir."
+            )
+
+        if not _WIN32:
+            context.command_prefix += ["cd", str(local_dir), "&&"]
+        else:
+            # Include '/d' incase temp folder is on different drive than Ray install.
+            context.command_prefix += ["cd", "/d", f"{local_dir}", "&&"]
+        set_pythonpath_in_context(python_path=str(local_dir), context=context)
+
+    @contextmanager
+    def with_working_dir_env(self, uri):
+        """
+        If uri is not None, add the local working directory to the environment variable
+        as "RAY_RUNTIME_ENV_CREATE_WORKING_DIR". This is useful for other plugins to
+        create their environment with reference to the working directory. For example
+        `pip -r ${RAY_RUNTIME_ENV_CREATE_WORKING_DIR}/requirements.txt`
+
+        The environment variable is removed after the context manager exits.
+        """
+        if uri is None:
+            yield
+        else:
+            local_dir = get_local_dir_from_uri(uri, self._resources_dir)
+            if not local_dir.exists():
+                raise ValueError(
+                    f"Local directory {local_dir} for URI {uri} does "
+                    "not exist on the cluster. Something may have gone wrong while "
+                    "downloading or unpacking the working_dir."
+                )
+            key = ray_constants.RAY_RUNTIME_ENV_CREATE_WORKING_DIR_ENV_VAR
+            prev = os.environ.get(key)
+            # Windows backslash paths are weird. When it's passed to the env var, and
+            # when Pip expands it, the backslashes are interpreted as escape characters
+            # and messes up the whole path. So we convert it to forward slashes.
+            # This works at least for all Python applications, including pip.
+            os.environ[key] = local_dir.as_posix()
+            try:
+                yield
+            finally:
+                if prev is None:
+                    del os.environ[key]
+                else:
+                    os.environ[key] = prev