[Core][RuntimeEnv] Add archives plugin

doc/source/ray-core/api/runtime-env.rst

@@ -7,3 +7,4 @@ Runtime Env API
 
     ray.runtime_env.RuntimeEnvConfig
     ray.runtime_env.RuntimeEnv
+    ray.runtime_env.get_archive_paths

doc/source/ray-core/handling-dependencies.rst

@@ -632,6 +632,39 @@ The ``runtime_env`` is a Python dictionary or a Python class :class:`ray.runtime
   ``ray list runtime-envs`` and the Python SDK still receive the plaintext values. See
   :ref:`Runtime environment redaction <runtime-env-redaction>` to change this behavior.
 
+- ``archives`` (str | Dict[str, str]): Specifies one or more remote archives to download and unpack on every node that runs a worker with this runtime environment. Supported formats are ``.zip``, ``.tar.gz``, and ``.tgz``. Unlike ``working_dir``, this field doesn't change the worker's current directory. Unlike ``py_modules``, it doesn't add the archive contents to ``PYTHONPATH``.
+
+  - Example: ``{"archives": "https://example.com/resources.zip"}``
+
+  - Example: ``{"archives": {"model": "s3://bucket/model.tar.gz", "config": "https://example.com/config.zip"}}``
+
+  Use :func:`ray.runtime_env.get_archive_paths` in the worker to get the local unpacked paths. A string input returns one string path. A dictionary input returns a dictionary with the same keys and local paths as values.
+
+  .. code-block:: python
+
+    import os
+
+    import ray
+    from ray.runtime_env import get_archive_paths
+
+    @ray.remote
+    def read_resource():
+        archive_path = get_archive_paths()
+        with open(os.path.join(archive_path, "resource.json")) as resource_file:
+            return resource_file.read()
+
+    result = ray.get(
+        read_resource.options(
+            runtime_env={
+                "archives": "https://example.com/resources.zip",
+            }
+        ).remote()
+    )
+
+  The returned paths are local to the current node and remain valid only while the runtime environment is in use. Treat the contents as read-only because workers using the same URI share a cache directory. ``archives`` isn't compatible with ``container`` or ``image_uri``.
+
+  Don't embed passwords or tokens in archive URIs because runtime environment fields and setup logs may expose the URI. Use the authentication mechanisms described in :ref:`runtime-env-auth` instead.
+
 - ``nsight`` (Union[str, Dict[str, str]]): specifies the config for the Nsight System Profiler. The value is either (1) "default", which refers to the `default config <https://github.com/ray-project/ray/blob/master/python/ray/_private/runtime_env/nsight.py#L20>`_, or (2) a dict of Nsight System Profiler options and their values.
   See :ref:`here <profiling-nsight-profiler>` for more details on setup and usage.
 

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

@@ -16,6 +16,7 @@
     DEFAULT_RUNTIME_ENV_TIMEOUT_SECONDS,
 )
 from ray._private.ray_logging import setup_component_logger
+from ray._private.runtime_env.archives import ArchivesPlugin
 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
@@ -228,9 +229,10 @@ def __init__(
         # 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)
+        # URI reference table. A URI may be cached independently by multiple plugins.
+        # The key includes the plugin type so each cache receives its own lifecycle
+        # notifications.
+        self._uri_reference: Dict[Tuple[str, UriType], int] = defaultdict(int)
         self._uris_parser = uris_parser
         self._unused_uris_callback = unused_uris_callback
         self._unused_runtime_env_callback = unused_runtime_env_callback
@@ -243,20 +245,23 @@ def __init__(
 
     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
+        for uri, uri_type in uris:
+            self._uri_reference[(uri, uri_type)] += 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:
+            uri_key = (uri, uri_type)
+            if self._uri_reference.get(uri_key, 0) > 0:
+                self._uri_reference[uri_key] -= 1
+                if self._uri_reference[uri_key] == 0:
                     unused_uris.append((uri, uri_type))
-                    del self._uri_reference[uri]
+                    del self._uri_reference[uri_key]
             else:
-                default_logger.warning(f"URI {uri} does not exist.")
+                default_logger.warning(
+                    f"URI {uri} for plugin {uri_type} does not exist."
+                )
         if unused_uris:
             default_logger.info(f"Unused uris {unused_uris}.")
             self._unused_uris_callback(unused_uris)
@@ -383,6 +388,7 @@ def __init__(
         self._working_dir_plugin = WorkingDirPlugin(
             self._runtime_env_dir, self._gcs_client
         )
+        self._archives_plugin = ArchivesPlugin(self._runtime_env_dir)
         self._container_plugin = ContainerPlugin(temp_dir)
         # TODO(jonathan-anyscale): change the plugin to ProfilerPlugin
         # and unify with nsight and other profilers.
@@ -395,6 +401,7 @@ def __init__(
         # self._xxx_plugin, we should just iterate through self._plugins.
         self._base_plugins: List[RuntimeEnvPlugin] = [
             self._working_dir_plugin,
+            self._archives_plugin,
             self._uv_plugin,
             self._pip_plugin,
             self._conda_plugin,
@@ -476,7 +483,7 @@ async def _setup_runtime_env(
             with self._setup_logger_factory.setup_logger(
                 request.job_id.decode(), log_files
             ) as per_job_logger:
-                context = RuntimeEnvContext(env_vars=runtime_env.env_vars())
+                context = RuntimeEnvContext(env_vars=dict(runtime_env.env_vars()))
 
                 # Warn about unrecognized fields in the runtime env.
                 for name, _ in runtime_env.plugins():

python/ray/_private/runtime_env/archives.py

@@ -0,0 +1,112 @@
+import json
+import logging
+import os
+from typing import Dict, List, Optional, Union
+
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env.constants import (
+    RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR,
+)
+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,
+)
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.runtime_env.validation import parse_and_validate_archives
+from ray._private.utils import get_directory_size_bytes
+
+default_logger = logging.getLogger(__name__)
+
+
+class ArchivesPlugin(RuntimeEnvPlugin):
+    name = "archives"
+
+    def __init__(self, resources_dir: str):
+        self._resources_dir = os.path.join(resources_dir, "archives_files")
+        try_to_create_directory(self._resources_dir)
+
+    @staticmethod
+    def validate(runtime_env_dict: dict) -> None:
+        parse_and_validate_archives(runtime_env_dict[ArchivesPlugin.name])
+        if RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR in runtime_env_dict.get(
+            "env_vars", {}
+        ):
+            raise ValueError(
+                f"{RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR!r} is managed by the "
+                "archives runtime environment and cannot be set in env_vars."
+            )
+
+    def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]:  # noqa: F821
+        archives = runtime_env.archives()
+        if isinstance(archives, str):
+            return [archives] if archives else []
+        if isinstance(archives, dict):
+            uris = []
+            seen_uris = set()
+            for uri in archives.values():
+                if isinstance(uri, str) and uri and uri not in seen_uris:
+                    uris.append(uri)
+                    seen_uris.add(uri)
+            return uris
+        return []
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger = default_logger,
+    ) -> int:
+        local_dir = await download_and_unpack_package(
+            uri,
+            self._resources_dir,
+            gcs_client=None,
+            logger=logger,
+            overwrite=True,
+        )
+        return get_directory_size_bytes(local_dir)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger = default_logger,
+    ) -> None:
+        if not uris:
+            return
+
+        archives = runtime_env.archives()
+        local_paths: Union[str, Dict[str, str]]
+        if isinstance(archives, str):
+            local_paths = self._get_local_dir(archives)
+        else:
+            local_paths = {
+                name: self._get_local_dir(uri) for name, uri in archives.items()
+            }
+
+        logger.info("Adding archives paths to the worker context.")
+        context.env_vars[RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR] = json.dumps(
+            local_paths, sort_keys=True
+        )
+
+    def delete_uri(self, uri: str, logger: logging.Logger = default_logger) -> int:
+        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("Tried to delete nonexistent archives URI: %s", uri)
+            return 0
+        return local_dir_size
+
+    def _get_local_dir(self, uri: str) -> str:
+        local_dir = get_local_dir_from_uri(uri, self._resources_dir)
+        if not local_dir.exists():
+            raise ValueError(
+                f"Local directory {local_dir} for archives URI {uri} does not "
+                "exist on the cluster. Something may have gone wrong while "
+                "downloading or unpacking the archive."
+            )
+        return str(local_dir)

python/ray/_private/runtime_env/constants.py

@@ -26,3 +26,6 @@
 
 # The file suffix of runtime env plugin schemas.
 RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX = ".json"
+
+# The env var used to expose unpacked archive paths to workers.
+RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR = "RAY_RUNTIME_ENV_ARCHIVES_PATHS"

python/ray/_private/runtime_env/plugin.py

@@ -250,16 +250,19 @@ async def create_for_plugin_if_needed(
         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)
+        async with uri_cache.get_uri_lock(uri):
+            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/uri_cache.py

@@ -1,4 +1,6 @@
+import asyncio
 import logging
+import weakref
 from typing import Callable, Optional, Set
 
 default_logger = logging.getLogger(__name__)
@@ -34,6 +36,9 @@ def __init__(
         # Maps URIs to the size in bytes of their corresponding disk contents.
         self._used_uris: Set[str] = set()
         self._unused_uris: Set[str] = set()
+        # Avoid serializing setup for unrelated URIs while ensuring that concurrent
+        # runtime environments only create and account for a shared URI once.
+        self._uri_locks = weakref.WeakValueDictionary()
 
         if delete_fn is None:
             self._delete_fn = lambda uri, logger: 0
@@ -58,6 +63,13 @@ def mark_unused(self, uri: str, logger: logging.Logger = default_logger):
         self._evict_if_needed(logger)
         self._check_valid()
 
+    def get_uri_lock(self, uri: str) -> asyncio.Lock:
+        lock = self._uri_locks.get(uri)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._uri_locks[uri] = lock
+        return lock
+
     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:

python/ray/_private/runtime_env/validation.py

@@ -116,6 +116,63 @@ def parse_and_validate_working_dir(working_dir: str) -> str:
     return working_dir
 
 
+def parse_and_validate_archives(
+    archives: Union[str, Dict[str, str]],
+) -> Union[str, Dict[str, str]]:
+    """Validate remote archives made available to Ray workers."""
+    assert archives is not None
+
+    if isinstance(archives, str):
+        archive_uris = [archives]
+    elif isinstance(archives, dict):
+        if not archives:
+            raise ValueError("runtime_env['archives'] must not be empty.")
+        for name, uri in archives.items():
+            if not isinstance(name, str) or not name:
+                raise TypeError(
+                    "runtime_env['archives'] keys must be non-empty strings, "
+                    f"got {name!r} of type {type(name)}."
+                )
+            if not isinstance(uri, str):
+                raise TypeError(
+                    "runtime_env['archives'] values must be strings, "
+                    f"got {uri!r} of type {type(uri)} for key {name!r}."
+                )
+        archive_uris = list(archives.values())
+    else:
+        raise TypeError(
+            "runtime_env['archives'] must be of type str or Dict[str, str], "
+            f"got {type(archives)}."
+        )
+
+    from ray._common.runtime_env_uri import parse_uri
+    from ray._private.runtime_env.protocol import Protocol
+
+    supported_extensions = (".zip", ".tar.gz", ".tgz")
+    for uri in archive_uris:
+        if not uri:
+            raise ValueError("runtime_env['archives'] URIs must not be empty.")
+        try:
+            protocol, path = parse_uri(uri)
+        except ValueError as exc:
+            raise ValueError(
+                f"runtime_env['archives'] entries must be remote URIs, got {uri!r}."
+            ) from exc
+
+        if protocol not in Protocol.remote_protocols():
+            raise ValueError(
+                "runtime_env['archives'] only supports remote URI protocols, "
+                f"got {protocol.value!r} for {uri!r}."
+            )
+        if not any(path.endswith(ext) for ext in supported_extensions):
+            raise ValueError(
+                "runtime_env['archives'] only supports .zip, .tar.gz, and .tgz "
+                f"files, got {uri!r}."
+            )
+
+    return archives
+
+
 def parse_and_validate_conda(conda: Union[str, dict]) -> Union[str, dict]:
     """Parses and validates a user-provided 'conda' option.
 
@@ -446,6 +503,7 @@ def parse_and_validate_env_vars(env_vars: Dict[str, str]) -> Optional[Dict[str,
 # Dictionary mapping runtime_env options with the function to parse and
 # validate them.
 OPTION_TO_VALIDATION_FN = {
+    "archives": parse_and_validate_archives,
     "py_modules": parse_and_validate_py_modules,
     "working_dir": parse_and_validate_working_dir,
     "excludes": parse_and_validate_excludes,

python/ray/runtime_env/__init__.py

@@ -1,6 +1,8 @@
+from ray.runtime_env.archives import get_archive_paths
 from ray.runtime_env.runtime_env import RuntimeEnv, RuntimeEnvConfig  # noqa: E402,F401
 
 __all__ = [
+    "get_archive_paths",
     "RuntimeEnvConfig",
     "RuntimeEnv",
 ]