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