chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)
|
||||
Reference in New Issue
Block a user