diff --git a/doc/source/ray-core/api/runtime-env.rst b/doc/source/ray-core/api/runtime-env.rst index 35db69c7f908..82a352a8a3c0 100644 --- a/doc/source/ray-core/api/runtime-env.rst +++ b/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 diff --git a/doc/source/ray-core/handling-dependencies.rst b/doc/source/ray-core/handling-dependencies.rst index dcc5b3de3039..5eaaaf9dd1cd 100644 --- a/doc/source/ray-core/handling-dependencies.rst +++ b/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 ` 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 `_, or (2) a dict of Nsight System Profiler options and their values. See :ref:`here ` for more details on setup and usage. diff --git a/python/ray/_private/runtime_env/agent/runtime_env_agent.py b/python/ray/_private/runtime_env/agent/runtime_env_agent.py index d20f4b32eb73..9bce43b1c38f 100644 --- a/python/ray/_private/runtime_env/agent/runtime_env_agent.py +++ b/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(): diff --git a/python/ray/_private/runtime_env/archives.py b/python/ray/_private/runtime_env/archives.py new file mode 100644 index 000000000000..43bf644bb069 --- /dev/null +++ b/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) diff --git a/python/ray/_private/runtime_env/constants.py b/python/ray/_private/runtime_env/constants.py index 3c6096b5993e..a7252097c73a 100644 --- a/python/ray/_private/runtime_env/constants.py +++ b/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" diff --git a/python/ray/_private/runtime_env/plugin.py b/python/ray/_private/runtime_env/plugin.py index 29ce2d3a90f6..f39355fb0f2c 100644 --- a/python/ray/_private/runtime_env/plugin.py +++ b/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) diff --git a/python/ray/_private/runtime_env/uri_cache.py b/python/ray/_private/runtime_env/uri_cache.py index 430bcf6d7e4e..b31771a1f50f 100644 --- a/python/ray/_private/runtime_env/uri_cache.py +++ b/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: diff --git a/python/ray/_private/runtime_env/validation.py b/python/ray/_private/runtime_env/validation.py index f7bbb41f3995..7e7db57a7f01 100644 --- a/python/ray/_private/runtime_env/validation.py +++ b/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, diff --git a/python/ray/runtime_env/__init__.py b/python/ray/runtime_env/__init__.py index f3cd30f708d0..288256f192c2 100644 --- a/python/ray/runtime_env/__init__.py +++ b/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", ] diff --git a/python/ray/runtime_env/archives.py b/python/ray/runtime_env/archives.py new file mode 100644 index 000000000000..b35abbf1cbc0 --- /dev/null +++ b/python/ray/runtime_env/archives.py @@ -0,0 +1,55 @@ +import json +import os +from typing import Dict, Union + +from ray._private.runtime_env.constants import ( + RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR, +) +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +def get_archive_paths() -> Union[str, Dict[str, str]]: + """Return the local paths of archives in the current runtime environment. + + The return value has the same shape as the ``archives`` runtime environment + field: a string input produces a string path, while a dictionary input + produces a dictionary containing the same keys and local paths as values. + + The returned paths are local to the current node and should be treated as + read-only shared cache directories. + + Returns: + The local archive path, or a dictionary mapping configured names to + local archive paths. + + Raises: + RuntimeError: If the current process has no ``archives`` runtime + environment or its internal path metadata is invalid. + """ + serialized_paths = os.environ.get(RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR) + if serialized_paths is None: + raise RuntimeError( + "No archives are available in this process. Configure the 'archives' " + "field in the runtime environment before calling get_archive_paths()." + ) + + try: + paths = json.loads(serialized_paths) + except json.JSONDecodeError as exc: + raise RuntimeError("The archives path metadata is not valid JSON.") from exc + + if isinstance(paths, str): + if paths: + return paths + elif ( + isinstance(paths, dict) + and paths + and all( + isinstance(name, str) and name and isinstance(path, str) and path + for name, path in paths.items() + ) + ): + return paths + + raise RuntimeError("The archives path metadata has an invalid value.") diff --git a/python/ray/runtime_env/runtime_env.py b/python/ray/runtime_env/runtime_env.py index 7b015833355c..733a4d67a8d9 100644 --- a/python/ray/runtime_env/runtime_env.py +++ b/python/ray/runtime_env/runtime_env.py @@ -8,6 +8,9 @@ import ray from ray._private.ray_constants import DEFAULT_RUNTIME_ENV_TIMEOUT_SECONDS from ray._private.runtime_env.conda import get_uri as get_conda_uri +from ray._private.runtime_env.constants import ( + RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR, +) from ray._private.runtime_env.default_impl import get_image_uri_plugin_cls from ray._private.runtime_env.pip import get_uri as get_pip_uri from ray._private.runtime_env.plugin_schema_manager import RuntimeEnvPluginSchemaManager @@ -243,6 +246,10 @@ class MyClass: working_dir: Local path or remote URI (either in the GCS or external storage) of an archive that Ray unpacks in the directory of each task/actor. Supported formats for remote URIs: ``.zip``, ``.tar.gz``, and ``.tgz``. + archives: A remote archive URI, or a dictionary mapping names to remote + archive URIs. Ray downloads and unpacks each archive on every node that + runs a task or actor with this runtime environment. Supported formats are + ``.zip``, ``.tar.gz``, and ``.tgz``. pip: Either a list of pip packages, a string containing the path to a pip requirements.txt file, or a Python dictionary that has three fields: 1) ``packages`` (required, List[str]): a @@ -299,6 +306,7 @@ class MyClass: "py_executable", "java_jars", "working_dir", + "archives", "conda", "pip", "uv", @@ -327,6 +335,7 @@ def __init__( py_modules: Optional[List[str]] = None, py_executable: Optional[str] = None, working_dir: Optional[str] = None, + archives: Optional[Union[str, Dict[str, str]]] = None, pip: Optional[List[str]] = None, conda: Optional[Union[Dict[str, str], str]] = None, container: Optional[Dict[str, str]] = None, @@ -349,6 +358,8 @@ def __init__( runtime_env["py_executable"] = py_executable if working_dir is not None: runtime_env["working_dir"] = working_dir + if archives is not None: + runtime_env["archives"] = archives if pip is not None: runtime_env["pip"] = pip if uv is not None: @@ -421,6 +432,14 @@ def __init__( del self[option] self[option] = option_val + if self.get("archives") is not None and ( + RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR in self.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." + ) + if "_ray_commit" not in self: if self.get("pip") or self.get("conda"): self["_ray_commit"] = ray.__commit__ @@ -518,6 +537,9 @@ def plugin_uris(self) -> List[str]: def working_dir(self) -> str: return self.get("working_dir", "") + def archives(self) -> Optional[Union[str, Dict[str, str]]]: + return self.get("archives") + def py_modules(self) -> List[str]: if "py_modules" in self: return list(self["py_modules"]) diff --git a/python/ray/runtime_env/schemas/archives_schema.json b/python/ray/runtime_env/schemas/archives_schema.json new file mode 100644 index 000000000000..f5f8af7028aa --- /dev/null +++ b/python/ray/runtime_env/schemas/archives_schema.json @@ -0,0 +1,23 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "$id": "http://github.com/ray-project/ray/runtime_env/archives_schema.json", + "title": "archives", + "description": "Remote archives made available to Ray workers.", + "oneOf": [ + { + "type": "string", + "minLength": 1 + }, + { + "type": "object", + "minProperties": 1, + "propertyNames": { + "minLength": 1 + }, + "additionalProperties": { + "type": "string", + "minLength": 1 + } + } + ] +} diff --git a/python/ray/tests/BUILD.bazel b/python/ray/tests/BUILD.bazel index 831e54f49d29..4eda22cffca9 100644 --- a/python/ray/tests/BUILD.bazel +++ b/python/ray/tests/BUILD.bazel @@ -1061,6 +1061,7 @@ py_test_module_list( data = ["pip_install_test-0.5-py3-none-any.whl"], files = [ "test_runtime_env.py", + "test_runtime_env_archives.py", "test_runtime_env_working_dir_2.py", "test_runtime_env_working_dir_3.py", ], diff --git a/python/ray/tests/test_runtime_env_agent.py b/python/ray/tests/test_runtime_env_agent.py index 55592ae0d5ba..c8a80f9ea9c8 100644 --- a/python/ray/tests/test_runtime_env_agent.py +++ b/python/ray/tests/test_runtime_env_agent.py @@ -100,6 +100,30 @@ def unused_runtime_env_processor(unused_runtime_env: str) -> None: assert not expected_unused_runtime_env +def test_reference_table_distinguishes_uri_types(): + uri = "s3://shared/archive.zip" + unused_uris = [] + + def uris_parser(runtime_env): + return [ + (runtime_env.working_dir(), "working_dir"), + (runtime_env.archives(), "archives"), + ] + + reference_table = ReferenceTable( + uris_parser, + unused_uris.extend, + lambda _: None, + ) + runtime_env = RuntimeEnv(working_dir=uri, archives=uri) + serialized_runtime_env = runtime_env.serialize() + + reference_table.increase_reference(runtime_env, serialized_runtime_env, "raylet") + reference_table.decrease_reference(runtime_env, serialized_runtime_env, "raylet") + + assert unused_uris == [(uri, "working_dir"), (uri, "archives")] + + def search_agent(processes): for p in processes: try: diff --git a/python/ray/tests/test_runtime_env_archives.py b/python/ray/tests/test_runtime_env_archives.py new file mode 100644 index 000000000000..fa4a20eb88a8 --- /dev/null +++ b/python/ray/tests/test_runtime_env_archives.py @@ -0,0 +1,255 @@ +import asyncio +import json +import os +import sys +import zipfile +from pathlib import Path +from typing import Dict + +import pytest + +import ray +from ray._private.runtime_env.archives import ArchivesPlugin +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 get_local_dir_from_uri +from ray._private.runtime_env.plugin import ( + RuntimeEnvPlugin, + create_for_plugin_if_needed, +) +from ray._private.runtime_env.uri_cache import URICache +from ray.runtime_env import RuntimeEnv, get_archive_paths + + +class ArchivesRetryPlugin(RuntimeEnvPlugin): + name = "archives_retry_plugin" + priority = 20 + + def __init__(self): + self._attempts = 0 + + def modify_context(self, uris, runtime_env, context, logger): + self._attempts += 1 + if self._attempts == 1: + raise ValueError("Retry archives runtime environment setup.") + + +def _create_zip(path: Path, files: Dict[str, str]) -> str: + with zipfile.ZipFile(path, "w") as archive: + for name, content in files.items(): + archive.writestr(name, content) + return path.as_uri() + + +def test_get_archive_paths(monkeypatch): + monkeypatch.delenv(RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR, raising=False) + with pytest.raises(RuntimeError, match="No archives are available"): + get_archive_paths() + + monkeypatch.setenv( + RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR, json.dumps("/tmp/archive") + ) + assert get_archive_paths() == "/tmp/archive" + + expected = {"model": "/tmp/model", "config": "/tmp/config"} + monkeypatch.setenv(RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR, json.dumps(expected)) + assert get_archive_paths() == expected + + monkeypatch.setenv(RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR, "not-json") + with pytest.raises(RuntimeError, match="not valid JSON"): + get_archive_paths() + + +@pytest.mark.asyncio +async def test_archives_plugin_lifecycle(tmp_path): + archive_uri = _create_zip( + tmp_path / "resources.zip", + {"resource.txt": "resource-data"}, + ) + plugin = ArchivesPlugin(str(tmp_path / "runtime_resources")) + runtime_env = RuntimeEnv( + archives={"primary": archive_uri, "duplicate": archive_uri} + ) + context = RuntimeEnvContext() + + assert plugin.get_uris(runtime_env) == [archive_uri] + size_bytes = await plugin.create(archive_uri, runtime_env, context) + assert size_bytes > 0 + + plugin.modify_context(plugin.get_uris(runtime_env), runtime_env, context) + local_paths = json.loads(context.env_vars[RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR]) + assert local_paths["primary"] == local_paths["duplicate"] + local_dir = Path(local_paths["primary"]) + assert (local_dir / "resource.txt").read_text() == "resource-data" + + uri_cache = URICache(plugin.delete_uri, max_total_size_bytes=0) + uri_cache.add(archive_uri, size_bytes) + uri_cache.mark_unused(archive_uri) + assert uri_cache.get_total_size_bytes() == 0 + assert not local_dir.exists() + assert plugin.delete_uri(archive_uri) == 0 + + +@pytest.mark.asyncio +async def test_archives_plugin_overwrites_incomplete_directory(tmp_path): + archive_uri = _create_zip( + tmp_path / "resources.zip", + {"resource.txt": "resource-data"}, + ) + plugin = ArchivesPlugin(str(tmp_path / "runtime_resources")) + local_dir = get_local_dir_from_uri(archive_uri, plugin._resources_dir) + local_dir.mkdir() + (local_dir / "incomplete.txt").write_text("incomplete") + + await plugin.create( + archive_uri, + RuntimeEnv(archives=archive_uri), + RuntimeEnvContext(), + ) + + assert (local_dir / "resource.txt").read_text() == "resource-data" + assert not (local_dir / "incomplete.txt").exists() + + +@pytest.mark.asyncio +async def test_archives_concurrent_setup_is_single_flight(tmp_path, monkeypatch): + archive_uri = _create_zip( + tmp_path / "resources.zip", + {"resource.txt": "resource-data"}, + ) + plugin = ArchivesPlugin(str(tmp_path / "runtime_resources")) + uri_cache = URICache(plugin.delete_uri) + contexts = [RuntimeEnvContext(), RuntimeEnvContext()] + created_sizes = [] + original_create = plugin.create + + async def counted_create(*args, **kwargs): + size_bytes = await original_create(*args, **kwargs) + created_sizes.append(size_bytes) + return size_bytes + + monkeypatch.setattr(plugin, "create", counted_create) + await asyncio.gather( + create_for_plugin_if_needed( + RuntimeEnv(archives=archive_uri), plugin, uri_cache, contexts[0] + ), + create_for_plugin_if_needed( + RuntimeEnv(archives={"resource": archive_uri}), + plugin, + uri_cache, + contexts[1], + ), + ) + + assert len(created_sizes) == 1 + assert uri_cache.get_total_size_bytes() == created_sizes[0] + assert all( + RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR in context.env_vars + for context in contexts + ) + + +def test_archives_runtime_env(tmp_path): + resources_uri = _create_zip( + tmp_path / "resources.zip", + {"resource.txt": "resource-data"}, + ) + config_uri = _create_zip( + tmp_path / "config.zip", + {"config.json": '{"enabled": true}'}, + ) + + try: + ray.init(num_cpus=1, include_dashboard=False) + + @ray.remote + def read_single_archive(): + archive_path = Path(get_archive_paths()) + return (archive_path / "resource.txt").read_text() + + @ray.remote + class NamedArchiveReader: + def read(self): + archive_paths = get_archive_paths() + return { + "resource": Path( + archive_paths["resource"], "resource.txt" + ).read_text(), + "duplicate_path": archive_paths["resource"] + == archive_paths["duplicate"], + "config": Path(archive_paths["config"], "config.json").read_text(), + } + + assert ( + ray.get( + read_single_archive.options( + runtime_env={"archives": resources_uri} + ).remote() + ) + == "resource-data" + ) + + reader = NamedArchiveReader.options( + runtime_env={ + "archives": { + "resource": resources_uri, + "duplicate": resources_uri, + "config": config_uri, + } + } + ).remote() + assert ray.get(reader.read.remote()) == { + "resource": "resource-data", + "duplicate_path": True, + "config": '{"enabled": true}', + } + finally: + ray.shutdown() + + +@pytest.mark.parametrize("set_runtime_env_retry_times", ["2"], indirect=True) +@pytest.mark.parametrize( + "set_runtime_env_plugins", + [ + json.dumps( + [{"class": ("ray.tests.test_runtime_env_archives.ArchivesRetryPlugin")}] + ) + ], + indirect=True, +) +def test_archives_runtime_env_retry( + tmp_path, set_runtime_env_retry_times, set_runtime_env_plugins +): + resources_uri = _create_zip( + tmp_path / "resources.zip", + {"resource.txt": "resource-data"}, + ) + + try: + ray.init(num_cpus=1, include_dashboard=False) + + @ray.remote + def read_archive(): + archive_path = Path(get_archive_paths()) + return ( + (archive_path / "resource.txt").read_text(), + os.environ["USER_ENV_VAR"], + ) + + assert ray.get( + read_archive.options( + runtime_env={ + "archives": resources_uri, + "env_vars": {"USER_ENV_VAR": "user-value"}, + ArchivesRetryPlugin.name: {}, + } + ).remote() + ) == ("resource-data", "user-value") + finally: + ray.shutdown() + + +if __name__ == "__main__": + sys.exit(pytest.main(["-sv", __file__])) diff --git a/python/ray/tests/unit/test_runtime_env_validation.py b/python/ray/tests/unit/test_runtime_env_validation.py index 36941f163ced..6c97ff0b6757 100644 --- a/python/ray/tests/unit/test_runtime_env_validation.py +++ b/python/ray/tests/unit/test_runtime_env_validation.py @@ -9,9 +9,13 @@ from ray import job_config from ray._private.runtime_env import validation +from ray._private.runtime_env.constants import ( + RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR, +) from ray._private.runtime_env.pip import _get_pip_hash from ray._private.runtime_env.plugin_schema_manager import RuntimeEnvPluginSchemaManager from ray._private.runtime_env.validation import ( + parse_and_validate_archives, parse_and_validate_conda, parse_and_validate_excludes, parse_and_validate_py_modules, @@ -111,6 +115,53 @@ def test_validate_path_valid_input(self, test_directory): assert working_dir == valid_working_dir_path +class TestValidateArchives: + @pytest.mark.parametrize( + "archives", + [ + "https://example.com/resources.zip", + "s3://bucket/resources.tar.gz", + {"config": "https://example.com/config.tgz"}, + ], + ) + def test_valid(self, archives): + assert parse_and_validate_archives(archives) == archives + assert RuntimeEnv(archives=archives).archives() == archives + + @pytest.mark.parametrize( + "archives", + [ + "", + {}, + {"": "https://example.com/resources.zip"}, + {"data": ""}, + {"data": 1}, + ["https://example.com/resources.zip"], + "/tmp/resources.zip", + "gcs://resources.zip", + "https://example.com/resources.whl", + ], + ) + def test_invalid(self, archives): + with pytest.raises((TypeError, ValueError, jsonschema.ValidationError)): + RuntimeEnv(archives=archives) + + def test_serialize(self): + archives = { + "model": "s3://bucket/model.tar.gz", + "config": "https://example.com/config.zip", + } + runtime_env = RuntimeEnv(archives=archives) + assert RuntimeEnv.deserialize(runtime_env.serialize()).archives() == archives + + def test_reserved_env_var(self): + with pytest.raises(ValueError, match="is managed by the archives"): + RuntimeEnv( + archives="https://example.com/resources.zip", + env_vars={RAY_RUNTIME_ENV_ARCHIVES_PATHS_ENV_VAR: "user-value"}, + ) + + class TestValidatePyModules: def test_validate_not_a_list(self): with pytest.raises(TypeError, match="must be a list of strings"): @@ -347,6 +398,15 @@ def test_validate_working_dir(self, set_runtime_env_plugin_schemas): with pytest.raises(jsonschema.exceptions.ValidationError, match="working_dir"): runtime_env["working_dir"] = ["https://abc/file.zip"] + def test_validate_archives(self, set_runtime_env_plugin_schemas): + runtime_env = RuntimeEnv() + runtime_env.set("archives", "https://abc/file.zip") + runtime_env.set("archives", {"data": "https://abc/file.tar.gz"}) + with pytest.raises(jsonschema.exceptions.ValidationError): + runtime_env.set("archives", {}) + with pytest.raises(jsonschema.exceptions.ValidationError): + runtime_env["archives"] = {"data": 1} + def test_validate_test_env_1(self, set_runtime_env_plugin_schemas): runtime_env = RuntimeEnv() runtime_env.set("test_env_1", {"array": ["123"], "bool": True}) diff --git a/python/setup.py b/python/setup.py index bb3b1d5bfa33..b1e209f30afc 100644 --- a/python/setup.py +++ b/python/setup.py @@ -858,6 +858,7 @@ def get_tag(self): "ray": [ "includes/*.pxd", "*.pxd", + "runtime_env/schemas/*.json", "serve/_private/ingress_request_router.lua.tmpl", ], },