From a49978f60e4a284a6b7894b65528ee1cd6fa3b7c Mon Sep 17 00:00:00 2001 From: dentiny Date: Wed, 6 Nov 2024 07:42:18 -0800 Subject: [PATCH] [core] [2/N] Implement uv processor (#48486) Signed-off-by: hjiang --- python/ray/_private/runtime_env/BUILD | 33 ++- .../_private/runtime_env/dependency_utils.py | 113 ++++++++ python/ray/_private/runtime_env/pip.py | 255 ++---------------- python/ray/_private/runtime_env/uv.py | 133 +++++++++ python/ray/_private/runtime_env/validation.py | 5 +- .../_private/runtime_env/virtualenv_utils.py | 109 ++++++++ python/ray/runtime_env/runtime_env.py | 1 + python/ray/serve/tests/BUILD | 1 - python/ray/tests/conftest.py | 4 +- .../tests/test_runtime_env_conda_and_pip.py | 22 +- .../tests/test_runtime_env_conda_and_pip_4.py | 14 +- python/ray/tests/unit/BUILD | 9 + python/ray/tests/unit/test_runtime_env_uv.py | 44 +++ 13 files changed, 486 insertions(+), 257 deletions(-) create mode 100644 python/ray/_private/runtime_env/dependency_utils.py create mode 100644 python/ray/_private/runtime_env/virtualenv_utils.py create mode 100644 python/ray/tests/unit/test_runtime_env_uv.py diff --git a/python/ray/_private/runtime_env/BUILD b/python/ray/_private/runtime_env/BUILD index f62818c89bfa7..4ba4f80d5b9e2 100644 --- a/python/ray/_private/runtime_env/BUILD +++ b/python/ray/_private/runtime_env/BUILD @@ -1,4 +1,4 @@ -load("@rules_python//python:defs.bzl", "py_library") +load("@rules_python//python:defs.bzl", "py_library", "py_test") package(default_visibility = ["//visibility:public"]) @@ -6,3 +6,34 @@ py_library( name = "validation", srcs = ["validation.py"], ) + +py_library( + name = "utils", + srcs = ["utils.py"], +) + +py_library( + name = "virtualenv_utils", + srcs = ["virtualenv_utils.py"], + deps = [ + ":utils", + ], +) + +py_library( + name = "dependency_utils", + srcs = ["dependency_utils.py"], + deps = [ + ":utils", + ], +) + +py_library( + name = "uv", + srcs= ["uv.py"], + deps = [ + ":utils", + ":virtualenv_utils", + ":dependency_utils", + ], +) diff --git a/python/ray/_private/runtime_env/dependency_utils.py b/python/ray/_private/runtime_env/dependency_utils.py new file mode 100644 index 0000000000000..baac4a3d2ee40 --- /dev/null +++ b/python/ray/_private/runtime_env/dependency_utils.py @@ -0,0 +1,113 @@ +"""Util functions to manage dependency requirements.""" + +from typing import List, Tuple, Optional +import os +import tempfile +import logging +from contextlib import asynccontextmanager +from ray._private.runtime_env import virtualenv_utils +from ray._private.runtime_env.utils import check_output_cmd + +INTERNAL_PIP_FILENAME = "ray_runtime_env_internal_pip_requirements.txt" +MAX_INTERNAL_PIP_FILENAME_TRIES = 100 + + +def gen_requirements_txt(requirements_file: str, pip_packages: List[str]): + """Dump [pip_packages] to the given [requirements_file] for later env setup.""" + with open(requirements_file, "w") as file: + for line in pip_packages: + file.write(line + "\n") + + +@asynccontextmanager +async def check_ray(python: str, cwd: str, logger: logging.Logger): + """A context manager to check ray is not overwritten. + + Currently, we only check ray version and path. It works for virtualenv, + - ray is in Python's site-packages. + - ray is overwritten during yield. + - ray is in virtualenv's site-packages. + """ + + async def _get_ray_version_and_path() -> Tuple[str, str]: + with tempfile.TemporaryDirectory( + prefix="check_ray_version_tempfile" + ) as tmp_dir: + ray_version_path = os.path.join(tmp_dir, "ray_version.txt") + check_ray_cmd = [ + python, + "-c", + """ +import ray +with open(r"{ray_version_path}", "wt") as f: + f.write(ray.__version__) + f.write(" ") + f.write(ray.__path__[0]) + """.format( + ray_version_path=ray_version_path + ), + ] + if virtualenv_utils._WIN32: + env = os.environ.copy() + else: + env = {} + output = await check_output_cmd( + check_ray_cmd, logger=logger, cwd=cwd, env=env + ) + logger.info(f"try to write ray version information in: {ray_version_path}") + with open(ray_version_path, "rt") as f: + output = f.read() + # print after import ray may have  endings, so we strip them by *_ + ray_version, ray_path, *_ = [s.strip() for s in output.split()] + return ray_version, ray_path + + version, path = await _get_ray_version_and_path() + yield + actual_version, actual_path = await _get_ray_version_and_path() + if actual_version != version or actual_path != path: + raise RuntimeError( + "Changing the ray version is not allowed: \n" + f" current version: {actual_version}, " + f"current path: {actual_path}\n" + f" expect version: {version}, " + f"expect path: {path}\n" + "Please ensure the dependencies in the runtime_env pip field " + "do not install a different version of Ray." + ) + + +def get_requirements_file(target_dir: str, pip_list: Optional[List[str]]) -> str: + """Returns the path to the requirements file to use for this runtime env. + + If pip_list is not None, we will check if the internal pip filename is in any of + the entries of pip_list. If so, we will append numbers to the end of the + filename until we find one that doesn't conflict. This prevents infinite + recursion if the user specifies the internal pip filename in their pip list. + + Args: + target_dir: The directory to store the requirements file in. + pip_list: A list of pip requirements specified by the user. + + Returns: + The path to the requirements file to use for this runtime env. + """ + + def filename_in_pip_list(filename: str) -> bool: + for pip_entry in pip_list: + if filename in pip_entry: + return True + return False + + filename = INTERNAL_PIP_FILENAME + if pip_list is not None: + i = 1 + while filename_in_pip_list(filename) and i < MAX_INTERNAL_PIP_FILENAME_TRIES: + filename = f"{INTERNAL_PIP_FILENAME}.{i}" + i += 1 + if i == MAX_INTERNAL_PIP_FILENAME_TRIES: + raise RuntimeError( + "Could not find a valid filename for the internal " + "pip requirements file. Please specify a different " + "pip list in your runtime env." + ) + return os.path.join(target_dir, filename) diff --git a/python/ray/_private/runtime_env/pip.py b/python/ray/_private/runtime_env/pip.py index f7422094ffe38..e3559721239aa 100644 --- a/python/ray/_private/runtime_env/pip.py +++ b/python/ray/_private/runtime_env/pip.py @@ -5,12 +5,11 @@ import os import shutil import sys -import tempfile -from typing import Dict, List, Optional, Tuple -from contextlib import asynccontextmanager +from typing import Dict, List, Optional from asyncio import create_task, get_running_loop -from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.runtime_env import virtualenv_utils +from ray._private.runtime_env import dependency_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 @@ -18,11 +17,6 @@ default_logger = logging.getLogger(__name__) -_WIN32 = os.name == "nt" - -INTERNAL_PIP_FILENAME = "ray_runtime_env_internal_pip_requirements.txt" -MAX_INTERNAL_PIP_FILENAME_TRIES = 100 - def _get_pip_hash(pip_dict: Dict) -> str: serialized_pip_spec = json.dumps(pip_dict, sort_keys=True) @@ -48,69 +42,6 @@ def get_uri(runtime_env: Dict) -> Optional[str]: return uri -class _PathHelper: - @staticmethod - def get_virtualenv_path(target_dir: str) -> str: - return os.path.join(target_dir, "virtualenv") - - @classmethod - def get_virtualenv_python(cls, target_dir: str) -> str: - virtualenv_path = cls.get_virtualenv_path(target_dir) - if _WIN32: - return os.path.join(virtualenv_path, "Scripts", "python.exe") - else: - return os.path.join(virtualenv_path, "bin", "python") - - @classmethod - def get_virtualenv_activate_command(cls, target_dir: str) -> List[str]: - virtualenv_path = cls.get_virtualenv_path(target_dir) - if _WIN32: - cmd = [os.path.join(virtualenv_path, "Scripts", "activate.bat")] - - else: - cmd = ["source", os.path.join(virtualenv_path, "bin/activate")] - return cmd + ["1>&2", "&&"] - - @staticmethod - def get_requirements_file(target_dir: str, pip_list: Optional[List[str]]) -> str: - """Returns the path to the requirements file to use for this runtime env. - - If pip_list is not None, we will check if the internal pip filename is in any of - the entries of pip_list. If so, we will append numbers to the end of the - filename until we find one that doesn't conflict. This prevents infinite - recursion if the user specifies the internal pip filename in their pip list. - - Args: - target_dir: The directory to store the requirements file in. - pip_list: A list of pip requirements specified by the user. - - Returns: - The path to the requirements file to use for this runtime env. - """ - - def filename_in_pip_list(filename: str) -> bool: - for pip_entry in pip_list: - if filename in pip_entry: - return True - return False - - filename = INTERNAL_PIP_FILENAME - if pip_list is not None: - i = 1 - while ( - filename_in_pip_list(filename) and i < MAX_INTERNAL_PIP_FILENAME_TRIES - ): - filename = f"{INTERNAL_PIP_FILENAME}.{i}" - i += 1 - if i == MAX_INTERNAL_PIP_FILENAME_TRIES: - raise RuntimeError( - "Could not find a valid filename for the internal " - "pip requirements file. Please specify a different " - "pip list in your runtime env." - ) - return os.path.join(target_dir, filename) - - class PipProcessor: def __init__( self, @@ -135,16 +66,6 @@ def __init__( self._pip_env = os.environ.copy() self._pip_env.update(self._runtime_env.env_vars()) - @staticmethod - def _is_in_virtualenv() -> bool: - # virtualenv <= 16.7.9 sets the real_prefix, - # virtualenv > 16.7.9 & venv set the base_prefix. - # So, we check both of them here. - # https://github.com/pypa/virtualenv/issues/1622#issuecomment-586186094 - return hasattr(sys, "real_prefix") or ( - hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix - ) - @classmethod async def _ensure_pip_version( cls, @@ -158,7 +79,7 @@ async def _ensure_pip_version( if not pip_version: return - python = _PathHelper.get_virtualenv_python(path) + python = virtualenv_utils.get_virtualenv_python(path) # Ensure pip version. pip_reinstall_cmd = [ python, @@ -186,7 +107,7 @@ async def _pip_check( if not pip_check: logger.info("Skip pip check.") return - python = _PathHelper.get_virtualenv_python(path) + python = virtualenv_utils.get_virtualenv_python(path) await check_output_cmd( [python, "-m", "pip", "check", "--disable-pip-version-check"], @@ -197,138 +118,6 @@ async def _pip_check( logger.info("Pip check on %s successfully.", path) - @staticmethod - @asynccontextmanager - async def _check_ray(python: str, cwd: str, logger: logging.Logger): - """A context manager to check ray is not overwritten. - - Currently, we only check ray version and path. It works for virtualenv, - - ray is in Python's site-packages. - - ray is overwritten during yield. - - ray is in virtualenv's site-packages. - """ - - async def _get_ray_version_and_path() -> Tuple[str, str]: - with tempfile.TemporaryDirectory( - prefix="check_ray_version_tempfile" - ) as tmp_dir: - ray_version_path = os.path.join(tmp_dir, "ray_version.txt") - check_ray_cmd = [ - python, - "-c", - """ -import ray -with open(r"{ray_version_path}", "wt") as f: - f.write(ray.__version__) - f.write(" ") - f.write(ray.__path__[0]) - """.format( - ray_version_path=ray_version_path - ), - ] - if _WIN32: - env = os.environ.copy() - else: - env = {} - output = await check_output_cmd( - check_ray_cmd, logger=logger, cwd=cwd, env=env - ) - logger.info( - f"try to write ray version information in: {ray_version_path}" - ) - with open(ray_version_path, "rt") as f: - output = f.read() - # print after import ray may have  endings, so we strip them by *_ - ray_version, ray_path, *_ = [s.strip() for s in output.split()] - return ray_version, ray_path - - version, path = await _get_ray_version_and_path() - yield - actual_version, actual_path = await _get_ray_version_and_path() - if actual_version != version or actual_path != path: - raise RuntimeError( - "Changing the ray version is not allowed: \n" - f" current version: {actual_version}, " - f"current path: {actual_path}\n" - f" expect version: {version}, " - f"expect path: {path}\n" - "Please ensure the dependencies in the runtime_env pip field " - "do not install a different version of Ray." - ) - - @classmethod - async def _create_or_get_virtualenv( - cls, path: str, cwd: str, logger: logging.Logger - ): - """Create or get a virtualenv from path.""" - python = sys.executable - virtualenv_path = os.path.join(path, "virtualenv") - virtualenv_app_data_path = os.path.join(path, "virtualenv_app_data") - - if _WIN32: - current_python_dir = sys.prefix - env = os.environ.copy() - else: - current_python_dir = os.path.abspath( - os.path.join(os.path.dirname(python), "..") - ) - env = {} - - if cls._is_in_virtualenv(): - # virtualenv-clone homepage: - # https://github.com/edwardgeorge/virtualenv-clone - # virtualenv-clone Usage: - # virtualenv-clone /path/to/existing/venv /path/to/cloned/ven - # or - # python -m clonevirtualenv /path/to/existing/venv /path/to/cloned/ven - clonevirtualenv = os.path.join( - os.path.dirname(__file__), "_clonevirtualenv.py" - ) - create_venv_cmd = [ - python, - clonevirtualenv, - current_python_dir, - virtualenv_path, - ] - logger.info( - "Cloning virtualenv %s to %s", current_python_dir, virtualenv_path - ) - else: - # virtualenv options: - # https://virtualenv.pypa.io/en/latest/cli_interface.html - # - # --app-data - # --reset-app-data - # Set an empty seperated app data folder for current virtualenv. - # - # --no-periodic-update - # Disable the periodic (once every 14 days) update of the embedded - # wheels. - # - # --system-site-packages - # Inherit site packages. - # - # --no-download - # Never download the latest pip/setuptools/wheel from PyPI. - create_venv_cmd = [ - python, - "-m", - "virtualenv", - "--app-data", - virtualenv_app_data_path, - "--reset-app-data", - "--no-periodic-update", - "--system-site-packages", - "--no-download", - virtualenv_path, - ] - logger.info( - "Creating virtualenv at %s, current python dir %s", - virtualenv_path, - virtualenv_path, - ) - await check_output_cmd(create_venv_cmd, logger=logger, cwd=cwd, env=env) - @classmethod async def _install_pip_packages( cls, @@ -338,19 +127,21 @@ async def _install_pip_packages( pip_env: Dict, logger: logging.Logger, ): - virtualenv_path = _PathHelper.get_virtualenv_path(path) - python = _PathHelper.get_virtualenv_python(path) + 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 = _PathHelper.get_requirements_file(path, pip_packages) - - def _gen_requirements_txt(): - with open(pip_requirements_file, "w") as file: - for line in pip_packages: - file.write(line + "\n") + 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, _gen_requirements_txt) + await loop.run_in_executor( + None, + dependency_utils.gen_requirements_txt, + pip_requirements_file, + pip_packages, + ) # pip options # @@ -385,9 +176,9 @@ async def _run(self): exec_cwd = os.path.join(path, "exec_cwd") os.makedirs(exec_cwd, exist_ok=True) try: - await self._create_or_get_virtualenv(path, exec_cwd, logger) - python = _PathHelper.get_virtualenv_python(path) - async with self._check_ray(python, exec_cwd, logger): + 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, @@ -485,7 +276,7 @@ async def create( self, uri: str, runtime_env: "RuntimeEnv", # noqa: F821 - context: RuntimeEnvContext, + context: "RuntimeEnvContext", # noqa: F821 logger: Optional[logging.Logger] = default_logger, ) -> int: if not runtime_env.has_pip(): @@ -523,7 +314,7 @@ def modify_context( self, uris: List[str], runtime_env: "RuntimeEnv", # noqa: F821 - context: RuntimeEnvContext, + context: "RuntimeEnvContext", # noqa: F821 logger: logging.Logger = default_logger, ): if not runtime_env.has_pip(): @@ -533,7 +324,7 @@ def modify_context( # Update py_executable. protocol, hash_val = parse_uri(uri) target_dir = self._get_path_from_hash(hash_val) - virtualenv_python = _PathHelper.get_virtualenv_python(target_dir) + virtualenv_python = virtualenv_utils.get_virtualenv_python(target_dir) if not os.path.exists(virtualenv_python): raise ValueError( @@ -542,6 +333,6 @@ def modify_context( "installing the runtime_env `pip` packages." ) context.py_executable = virtualenv_python - context.command_prefix += _PathHelper.get_virtualenv_activate_command( + context.command_prefix += virtualenv_utils.get_virtualenv_activate_command( target_dir ) diff --git a/python/ray/_private/runtime_env/uv.py b/python/ray/_private/runtime_env/uv.py index e69de29bb2d1d..5d0c5db980c91 100644 --- a/python/ray/_private/runtime_env/uv.py +++ b/python/ray/_private/runtime_env/uv.py @@ -0,0 +1,133 @@ +"""Util class to install packages via uv. +""" + +# TODO(hjiang): Implement `UvPlugin`, which is the counterpart for `PipPlugin`. + +from typing import Dict, List, Optional +from asyncio import get_running_loop +import os +from ray._private.runtime_env import virtualenv_utils +from ray._private.runtime_env import dependency_utils +from ray._private.runtime_env.utils import check_output_cmd +import shutil +import logging +import sys + +default_logger = logging.getLogger(__name__) + + +class UvProcessor: + 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 uv runtime env." + ) + + logger.debug("Setting up uv for runtime_env: %s", runtime_env) + self._target_dir = target_dir + self._runtime_env = runtime_env + self._logger = logger + + self._uv_config = self._runtime_env.uv_config() + self._uv_env = os.environ.copy() + self._uv_env.update(self._runtime_env.env_vars()) + + # TODO(hjiang): Check `uv` existence before installation, so we don't blindly + # install. + async def _install_uv( + self, path: str, cwd: str, pip_env: dict, logger: logging.Logger + ): + """Before package install, make sure `uv` is installed.""" + virtualenv_path = virtualenv_utils.get_virtualenv_path(path) + python = virtualenv_utils.get_virtualenv_python(path) + + uv_install_cmd = [ + python, + "-m", + "pip", + "install", + "--disable-pip-version-check", + "--no-cache-dir", + "uv", + ] + logger.info("Installing package uv to %s", virtualenv_path) + await check_output_cmd(uv_install_cmd, logger=logger, cwd=cwd, env=pip_env) + + async def _install_uv_packages( + self, + path: str, + uv_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, ... + requirements_file = dependency_utils.get_requirements_file(path, uv_packages) + + # Install uv, which acts as the default package manager. + await self._install_uv(path, cwd, pip_env, logger) + + # Avoid blocking the event loop. + loop = get_running_loop() + await loop.run_in_executor( + None, dependency_utils.gen_requirements_txt, requirements_file, uv_packages + ) + + # Install all dependencies. + # + # Difference with pip: + # 1. `--disable-pip-version-check` has no effect for uv. + # 2. `--no-cache-dir` for `pip` maps to `--no-cache` for uv. + pip_install_cmd = [ + python, + "-m", + "uv", + "pip", + "install", + "--no-cache", + "-r", + requirements_file, + ] + 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 + uv_packages = self._uv_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): + # Install packages with uv. + await self._install_uv_packages( + path, + uv_packages, + exec_cwd, + self._uv_env, + logger, + ) + except Exception: + logger.info("Delete incomplete virtualenv: %s", path) + shutil.rmtree(path, ignore_errors=True) + logger.exception("Failed to install uv packages.") + raise + + def __await__(self): + return self._run().__await__() diff --git a/python/ray/_private/runtime_env/validation.py b/python/ray/_private/runtime_env/validation.py index 14fca127a4e42..ac478df592032 100644 --- a/python/ray/_private/runtime_env/validation.py +++ b/python/ray/_private/runtime_env/validation.py @@ -120,7 +120,7 @@ def parse_and_validate_uv(uv: Union[str, List[str], Dict]) -> Optional[Dict]: The value of the input 'uv' field can be one of two cases: 1) A List[str] describing the requirements. This is passed through. Example usage: ["tensorflow", "requests"] - 2) A python dictionary that has three fields: + 2) A python dictionary that has one field: a) packages (required, List[str]): a list of uv packages, it same as 1). The returned parsed value will be a list of packages. If a Ray library @@ -337,13 +337,14 @@ 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. +# +# TODO(hjiang): Expose `uv` related validation after implementation finished. OPTION_TO_VALIDATION_FN = { "py_modules": parse_and_validate_py_modules, "working_dir": parse_and_validate_working_dir, "excludes": parse_and_validate_excludes, "conda": parse_and_validate_conda, "pip": parse_and_validate_pip, - "uv": parse_and_validate_uv, "env_vars": parse_and_validate_env_vars, "container": parse_and_validate_container, } diff --git a/python/ray/_private/runtime_env/virtualenv_utils.py b/python/ray/_private/runtime_env/virtualenv_utils.py new file mode 100644 index 0000000000000..09bd9f82e7f96 --- /dev/null +++ b/python/ray/_private/runtime_env/virtualenv_utils.py @@ -0,0 +1,109 @@ +"""Utils to detect runtime environment.""" + +import sys +from ray._private.runtime_env.utils import check_output_cmd +import logging +import os +from typing import List + +_WIN32 = os.name == "nt" + + +def is_in_virtualenv() -> bool: + # virtualenv <= 16.7.9 sets the real_prefix, + # virtualenv > 16.7.9 & venv set the base_prefix. + # So, we check both of them here. + # https://github.com/pypa/virtualenv/issues/1622#issuecomment-586186094 + return hasattr(sys, "real_prefix") or ( + hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix + ) + + +def get_virtualenv_path(target_dir: str) -> str: + """Get virtual environment path.""" + return os.path.join(target_dir, "virtualenv") + + +def get_virtualenv_python(target_dir: str) -> str: + virtualenv_path = get_virtualenv_path(target_dir) + if _WIN32: + return os.path.join(virtualenv_path, "Scripts", "python.exe") + else: + return os.path.join(virtualenv_path, "bin", "python") + + +def get_virtualenv_activate_command(target_dir: str) -> List[str]: + """Get the command to activate virtual environment.""" + virtualenv_path = get_virtualenv_path(target_dir) + if _WIN32: + cmd = [os.path.join(virtualenv_path, "Scripts", "activate.bat")] + else: + cmd = ["source", os.path.join(virtualenv_path, "bin/activate")] + return cmd + ["1>&2", "&&"] + + +async def create_or_get_virtualenv(path: str, cwd: str, logger: logging.Logger): + """Create or get a virtualenv from path.""" + python = sys.executable + virtualenv_path = os.path.join(path, "virtualenv") + virtualenv_app_data_path = os.path.join(path, "virtualenv_app_data") + + if _WIN32: + current_python_dir = sys.prefix + env = os.environ.copy() + else: + current_python_dir = os.path.abspath( + os.path.join(os.path.dirname(python), "..") + ) + env = {} + + if is_in_virtualenv(): + # virtualenv-clone homepage: + # https://github.com/edwardgeorge/virtualenv-clone + # virtualenv-clone Usage: + # virtualenv-clone /path/to/existing/venv /path/to/cloned/ven + # or + # python -m clonevirtualenv /path/to/existing/venv /path/to/cloned/ven + clonevirtualenv = os.path.join(os.path.dirname(__file__), "_clonevirtualenv.py") + create_venv_cmd = [ + python, + clonevirtualenv, + current_python_dir, + virtualenv_path, + ] + logger.info("Cloning virtualenv %s to %s", current_python_dir, virtualenv_path) + else: + # virtualenv options: + # https://virtualenv.pypa.io/en/latest/cli_interface.html + # + # --app-data + # --reset-app-data + # Set an empty seperated app data folder for current virtualenv. + # + # --no-periodic-update + # Disable the periodic (once every 14 days) update of the embedded + # wheels. + # + # --system-site-packages + # Inherit site packages. + # + # --no-download + # Never download the latest pip/setuptools/wheel from PyPI. + create_venv_cmd = [ + python, + "-m", + "virtualenv", + "--app-data", + virtualenv_app_data_path, + "--reset-app-data", + "--no-periodic-update", + "--system-site-packages", + "--no-download", + virtualenv_path, + ] + logger.info( + "Creating virtualenv at %s, current python dir %s", + virtualenv_path, + virtualenv_path, + ) + await check_output_cmd(create_venv_cmd, logger=logger, cwd=cwd, env=env) diff --git a/python/ray/runtime_env/runtime_env.py b/python/ray/runtime_env/runtime_env.py index feb4f8bb0e526..ab91e04ef8cda 100644 --- a/python/ray/runtime_env/runtime_env.py +++ b/python/ray/runtime_env/runtime_env.py @@ -148,6 +148,7 @@ def to_dict(self) -> Dict: ] = RuntimeEnvConfig.parse_and_validate_runtime_env_config +# TODO(hjiang): Expose `uv` related fields after implementation finished. @PublicAPI class RuntimeEnv(dict): """This class is used to define a runtime environment for a job, task, diff --git a/python/ray/serve/tests/BUILD b/python/ray/serve/tests/BUILD index 0af13939ec414..42702b576c08c 100644 --- a/python/ray/serve/tests/BUILD +++ b/python/ray/serve/tests/BUILD @@ -466,4 +466,3 @@ py_test_module_list( "//python/ray/serve:serve_lib", ], ) - diff --git a/python/ray/tests/conftest.py b/python/ray/tests/conftest.py index 4ee34d2b73c97..47557bc36a32f 100644 --- a/python/ray/tests/conftest.py +++ b/python/ray/tests/conftest.py @@ -21,7 +21,7 @@ import ray import ray._private.ray_constants as ray_constants from ray._private.conftest_utils import set_override_dashboard_url # noqa: F401 -from ray._private.runtime_env.pip import PipProcessor +from ray._private.runtime_env import virtualenv_utils from ray._private.runtime_env.plugin_schema_manager import RuntimeEnvPluginSchemaManager from ray._private.test_utils import ( @@ -980,7 +980,7 @@ def cloned_virtualenv(): # aviod import `pytest_virtualenv` in test case `Minimal install` from pytest_virtualenv import VirtualEnv - if PipProcessor._is_in_virtualenv(): + if virtualenv_utils.is_in_virtualenv(): raise RuntimeError("Forbid the use of this fixture in virtualenv") venv = VirtualEnv( diff --git a/python/ray/tests/test_runtime_env_conda_and_pip.py b/python/ray/tests/test_runtime_env_conda_and_pip.py index d241e8c0a90d5..df9ba385c9380 100644 --- a/python/ray/tests/test_runtime_env_conda_and_pip.py +++ b/python/ray/tests/test_runtime_env_conda_and_pip.py @@ -9,11 +9,11 @@ check_local_files_gced, generate_runtime_env_dict, ) +from ray._private.runtime_env import dependency_utils from ray._private.runtime_env.conda import _get_conda_dict_with_ray_inserted -from ray._private.runtime_env.pip import ( +from ray._private.runtime_env.dependency_utils import ( INTERNAL_PIP_FILENAME, MAX_INTERNAL_PIP_FILENAME_TRIES, - _PathHelper, ) from ray.runtime_env import RuntimeEnv @@ -231,27 +231,25 @@ def f(): def test_get_requirements_file(): - """Unit test for _PathHelper.get_requirements_file.""" + """Unit test for dependency_utils.get_requirements_file.""" with tempfile.TemporaryDirectory() as tmpdir: - path_helper = _PathHelper() - # If pip_list is None, we should return the internal pip filename. - assert path_helper.get_requirements_file(tmpdir, pip_list=None) == os.path.join( - tmpdir, INTERNAL_PIP_FILENAME - ) + assert dependency_utils.get_requirements_file( + tmpdir, pip_list=None + ) == os.path.join(tmpdir, INTERNAL_PIP_FILENAME) # If the internal pip filename is not in pip_list, we should return the internal # pip filename. - assert path_helper.get_requirements_file( + assert dependency_utils.get_requirements_file( tmpdir, pip_list=["foo", "bar"] ) == os.path.join(tmpdir, INTERNAL_PIP_FILENAME) # If the internal pip filename is in pip_list, we should append numbers to the # end of the filename until we find one that doesn't conflict. - assert path_helper.get_requirements_file( + assert dependency_utils.get_requirements_file( tmpdir, pip_list=["foo", "bar", f"-r {INTERNAL_PIP_FILENAME}"] ) == os.path.join(tmpdir, f"{INTERNAL_PIP_FILENAME}.1") - assert path_helper.get_requirements_file( + assert dependency_utils.get_requirements_file( tmpdir, pip_list=[ "foo", @@ -263,7 +261,7 @@ def test_get_requirements_file(): # If we can't find a valid filename, we should raise an error. with pytest.raises(RuntimeError) as excinfo: - path_helper.get_requirements_file( + dependency_utils.get_requirements_file( tmpdir, pip_list=[ "foo", diff --git a/python/ray/tests/test_runtime_env_conda_and_pip_4.py b/python/ray/tests/test_runtime_env_conda_and_pip_4.py index f2d559e844a16..3eaaad030c2a3 100644 --- a/python/ray/tests/test_runtime_env_conda_and_pip_4.py +++ b/python/ray/tests/test_runtime_env_conda_and_pip_4.py @@ -2,7 +2,7 @@ import pytest import sys -from ray._private.runtime_env.pip import PipProcessor +from ray._private.runtime_env import virtualenv_utils import ray @@ -14,8 +14,9 @@ def test_in_virtualenv(start_cluster): assert ( - PipProcessor._is_in_virtualenv() is False and "IN_VIRTUALENV" not in os.environ - ) or (PipProcessor._is_in_virtualenv() is True and "IN_VIRTUALENV" in os.environ) + virtualenv_utils.is_in_virtualenv() is False + and "IN_VIRTUALENV" not in os.environ + ) or (virtualenv_utils.is_in_virtualenv() is True and "IN_VIRTUALENV" in os.environ) cluster, address = start_cluster runtime_env = {"pip": ["pip-install-test==0.5"]} @@ -25,7 +26,7 @@ def test_in_virtualenv(start_cluster): def f(): import pip_install_test # noqa: F401 - return PipProcessor._is_in_virtualenv() + return virtualenv_utils.is_in_virtualenv() # Ensure that the runtime env has been installed # and virtualenv is activated. @@ -140,12 +141,11 @@ def f(): ) def test_run_in_virtualenv(cloned_virtualenv): python_exe_path = cloned_virtualenv.python - print(python_exe_path) # make sure cloned_virtualenv.run will run in virtualenv. cloned_virtualenv.run( - f"{python_exe_path} -c 'from ray._private.runtime_env.pip import PipProcessor;" - "assert PipProcessor._is_in_virtualenv()'", + f"{python_exe_path} -c 'from ray._private.runtime_env import virtualenv_utils;" + "assert virtualenv_utils.is_in_virtualenv()'", capture=True, ) diff --git a/python/ray/tests/unit/BUILD b/python/ray/tests/unit/BUILD index ec799f0544f24..572e3c4fe37d0 100644 --- a/python/ray/tests/unit/BUILD +++ b/python/ray/tests/unit/BUILD @@ -6,3 +6,12 @@ py_test( "//python/ray/_private/runtime_env:validation", ], ) + +py_test( + name = "test_runtime_env_uv", + srcs = ["test_runtime_env_uv.py"], + tags = ["team:core"], + deps = [ + "//python/ray/_private/runtime_env:uv", + ], +) diff --git a/python/ray/tests/unit/test_runtime_env_uv.py b/python/ray/tests/unit/test_runtime_env_uv.py new file mode 100644 index 0000000000000..b4e210049003c --- /dev/null +++ b/python/ray/tests/unit/test_runtime_env_uv.py @@ -0,0 +1,44 @@ +from ray._private.runtime_env import uv + +import pytest +import sys +from unittest.mock import patch + + +class TestRuntimeEnv: + def uv_config(self): + return {"packages": ["requests"]} + + def env_vars(self): + return {} + + +@pytest.fixture +def mock_install_uv(): + with patch( + "ray._private.runtime_env.uv.UvProcessor._install_uv" + ) as mock_install_uv: + mock_install_uv.return_value = None + yield mock_install_uv + + +@pytest.fixture +def mock_install_uv_packages(): + with patch( + "ray._private.runtime_env.uv.UvProcessor._install_uv_packages" + ) as mock_install_uv_packages: + mock_install_uv_packages.return_value = None + yield mock_install_uv_packages + + +@pytest.mark.asyncio +async def test_run(mock_install_uv, mock_install_uv_packages): + target_dir = "/tmp" + runtime_env = TestRuntimeEnv() + + uv_processor = uv.UvProcessor(target_dir=target_dir, runtime_env=runtime_env) + await uv_processor._run() + + +if __name__ == "__main__": + sys.exit(pytest.main(["-vv", __file__]))