wehub-resource-sync
423 lines
15 KiB
Python
423 lines
15 KiB
Python
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
|
|
)
|