wehub-resource-sync
358 lines
13 KiB
Python
358 lines
13 KiB
Python
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
|
|
import yaml
|
|
|
|
from mlflow.environment_variables import MLFLOW_CONDA_CREATE_ENV_CMD, MLFLOW_CONDA_HOME
|
|
from mlflow.exceptions import ExecutionException
|
|
from mlflow.utils import process
|
|
from mlflow.utils.environment import Environment
|
|
from mlflow.utils.os import is_windows
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
CONDA_EXE = "CONDA_EXE"
|
|
|
|
|
|
def get_conda_command(conda_env_name):
|
|
# Checking for newer conda versions
|
|
if not is_windows() and (CONDA_EXE in os.environ or MLFLOW_CONDA_HOME.defined):
|
|
conda_path = get_conda_bin_executable("conda")
|
|
activate_conda_env = [f"source {os.path.dirname(conda_path)}/../etc/profile.d/conda.sh"]
|
|
activate_conda_env += [f"conda activate {conda_env_name} 1>&2"]
|
|
else:
|
|
activate_path = get_conda_bin_executable("activate")
|
|
# in case os name is not 'nt', we are not running on windows. It introduces
|
|
# bash command otherwise.
|
|
if not is_windows():
|
|
return [f"source {activate_path} {conda_env_name} 1>&2"]
|
|
else:
|
|
return [f"conda activate {conda_env_name}"]
|
|
return activate_conda_env
|
|
|
|
|
|
def get_conda_bin_executable(executable_name):
|
|
"""
|
|
Return path to the specified executable, assumed to be discoverable within the 'bin'
|
|
subdirectory of a conda installation.
|
|
|
|
The conda home directory (expected to contain a 'bin' subdirectory) is configurable via the
|
|
``mlflow.projects.MLFLOW_CONDA_HOME`` environment variable. If
|
|
``mlflow.projects.MLFLOW_CONDA_HOME`` is unspecified, this method simply returns the passed-in
|
|
executable name.
|
|
"""
|
|
if conda_home := MLFLOW_CONDA_HOME.get():
|
|
return os.path.join(conda_home, f"bin/{executable_name}")
|
|
# Use CONDA_EXE as per https://github.com/conda/conda/issues/7126
|
|
if conda_exe := os.environ.get(CONDA_EXE):
|
|
conda_bin_dir = os.path.dirname(conda_exe)
|
|
return os.path.join(conda_bin_dir, executable_name)
|
|
return executable_name
|
|
|
|
|
|
def _get_conda_env_name(conda_env_path, env_id=None, env_root_dir=None):
|
|
if conda_env_path:
|
|
with open(conda_env_path) as f:
|
|
conda_env_contents = f.read()
|
|
else:
|
|
conda_env_contents = ""
|
|
|
|
if env_id:
|
|
conda_env_contents += env_id
|
|
|
|
env_name = "mlflow-{}".format(
|
|
hashlib.sha1(conda_env_contents.encode("utf-8"), usedforsecurity=False).hexdigest()
|
|
)
|
|
if env_root_dir:
|
|
env_root_dir = os.path.normpath(env_root_dir)
|
|
# Generate env name with format "mlflow-{conda_env_contents_hash}-{env_root_dir_hash}"
|
|
# hashing `conda_env_contents` and `env_root_dir` separately helps debugging
|
|
env_name += "-{}".format(
|
|
hashlib.sha1(env_root_dir.encode("utf-8"), usedforsecurity=False).hexdigest()
|
|
)
|
|
|
|
return env_name
|
|
|
|
|
|
def _get_conda_executable_for_create_env():
|
|
"""
|
|
Returns the executable that should be used to create environments. This is "conda"
|
|
by default, but it can be set to something else by setting the environment variable
|
|
|
|
"""
|
|
|
|
return get_conda_bin_executable(MLFLOW_CONDA_CREATE_ENV_CMD.get())
|
|
|
|
|
|
def _list_conda_environments(extra_env=None):
|
|
"""Return a list of names of conda environments.
|
|
|
|
Args:
|
|
extra_env: Extra environment variables for running "conda env list" command.
|
|
|
|
"""
|
|
prc = process._exec_cmd(
|
|
[get_conda_bin_executable("conda"), "env", "list", "--json"], extra_env=extra_env
|
|
)
|
|
return list(map(os.path.basename, json.loads(prc.stdout).get("envs", [])))
|
|
|
|
|
|
_CONDA_ENVS_DIR = "conda_envs"
|
|
_CONDA_CACHE_PKGS_DIR = "conda_cache_pkgs"
|
|
_PIP_CACHE_DIR = "pip_cache_pkgs"
|
|
|
|
|
|
def _create_conda_env(
|
|
conda_env_path,
|
|
conda_env_create_path,
|
|
project_env_name,
|
|
conda_extra_env_vars,
|
|
capture_output,
|
|
):
|
|
if conda_env_path:
|
|
process._exec_cmd(
|
|
[
|
|
conda_env_create_path,
|
|
"env",
|
|
"create",
|
|
"-n",
|
|
project_env_name,
|
|
"--file",
|
|
conda_env_path,
|
|
],
|
|
extra_env=conda_extra_env_vars,
|
|
capture_output=capture_output,
|
|
)
|
|
else:
|
|
process._exec_cmd(
|
|
[
|
|
conda_env_create_path,
|
|
"create",
|
|
"--channel",
|
|
"conda-forge",
|
|
"--yes",
|
|
"--override-channels",
|
|
"-n",
|
|
project_env_name,
|
|
"python",
|
|
],
|
|
extra_env=conda_extra_env_vars,
|
|
capture_output=capture_output,
|
|
)
|
|
|
|
return Environment(get_conda_command(project_env_name), conda_extra_env_vars)
|
|
|
|
|
|
def _create_conda_env_retry(
|
|
conda_env_path, conda_env_create_path, project_env_name, conda_extra_env_vars, _capture_output
|
|
):
|
|
"""
|
|
`conda env create` command can fail due to network issues such as `ConnectionResetError`
|
|
while collecting package metadata. This function retries the command up to 3 times.
|
|
"""
|
|
num_attempts = 3
|
|
retryable_errors = (
|
|
"ConnectionResetError",
|
|
"ChunkedEncodingError",
|
|
"CONNECTION FAILED",
|
|
"CondaHTTPError",
|
|
)
|
|
for attempt in range(num_attempts):
|
|
try:
|
|
return _create_conda_env(
|
|
conda_env_path,
|
|
conda_env_create_path,
|
|
project_env_name,
|
|
conda_extra_env_vars,
|
|
capture_output=True,
|
|
)
|
|
except process.ShellCommandException as e:
|
|
if attempt < num_attempts - 1 and any(err in str(e) for err in retryable_errors):
|
|
_logger.warning("Conda env creation failed due to network issue. Retrying...")
|
|
continue
|
|
raise
|
|
|
|
|
|
def _get_conda_extra_env_vars(env_root_dir=None):
|
|
"""
|
|
Given the `env_root_dir` (See doc of PyFuncBackend constructor argument `env_root_dir`),
|
|
return a dict of environment variables which are used to config conda to generate envs
|
|
under the expected `env_root_dir`.
|
|
"""
|
|
if env_root_dir is None:
|
|
return None
|
|
|
|
# Create isolated conda package cache dir "conda_pkgs" under the env_root_dir
|
|
# for each python process.
|
|
# Note: shared conda package cache dir causes race condition issues:
|
|
# See https://github.com/conda/conda/issues/8870
|
|
# See https://docs.conda.io/projects/conda/en/latest/user-guide/configuration/use-condarc.html#specify-environment-directories-envs-dirs
|
|
# and https://docs.conda.io/projects/conda/en/latest/user-guide/configuration/use-condarc.html#specify-package-directories-pkgs-dirs
|
|
|
|
conda_envs_path = os.path.join(env_root_dir, _CONDA_ENVS_DIR)
|
|
conda_pkgs_path = os.path.join(env_root_dir, _CONDA_CACHE_PKGS_DIR)
|
|
pip_cache_dir = os.path.join(env_root_dir, _PIP_CACHE_DIR)
|
|
|
|
os.makedirs(conda_envs_path, exist_ok=True)
|
|
os.makedirs(conda_pkgs_path, exist_ok=True)
|
|
os.makedirs(pip_cache_dir, exist_ok=True)
|
|
|
|
return {
|
|
"CONDA_ENVS_PATH": conda_envs_path,
|
|
"CONDA_PKGS_DIRS": conda_pkgs_path,
|
|
"PIP_CACHE_DIR": pip_cache_dir,
|
|
# PIP_NO_INPUT=1 makes pip run in non-interactive mode,
|
|
# otherwise pip might prompt "yes or no" and ask stdin input
|
|
"PIP_NO_INPUT": "1",
|
|
}
|
|
|
|
|
|
def get_or_create_conda_env(
|
|
conda_env_path,
|
|
env_id=None,
|
|
capture_output=False,
|
|
env_root_dir=None,
|
|
pip_requirements_override=None,
|
|
extra_envs=None,
|
|
):
|
|
"""Given a `Project`, creates a conda environment containing the project's dependencies if such
|
|
a conda environment doesn't already exist. Returns the name of the conda environment.
|
|
|
|
Args:
|
|
conda_env_path: Path to a conda yaml file.
|
|
env_id: Optional string that is added to the contents of the yaml file before
|
|
calculating the hash. It can be used to distinguish environments that have the
|
|
same conda dependencies but are supposed to be different based on the context.
|
|
For example, when serving the model we may install additional dependencies to the
|
|
environment after the environment has been activated.
|
|
capture_output: Specify the capture_output argument while executing the
|
|
"conda env create" command.
|
|
env_root_dir: See doc of PyFuncBackend constructor argument `env_root_dir`.
|
|
pip_requirements_override: If specified, install the specified python dependencies to
|
|
the environment (upgrade if already installed).
|
|
extra_envs: If specified, a dictionary of extra environment variables will be passed to the
|
|
model inference environment.
|
|
|
|
Returns:
|
|
The name of the conda environment.
|
|
|
|
"""
|
|
|
|
conda_path = get_conda_bin_executable("conda")
|
|
conda_env_create_path = _get_conda_executable_for_create_env()
|
|
|
|
try:
|
|
# Checks if Conda executable exists
|
|
process._exec_cmd([conda_path, "--help"], throw_on_error=False, extra_env=extra_envs)
|
|
except OSError:
|
|
raise ExecutionException(
|
|
f"Could not find Conda executable at {conda_path}. "
|
|
"Ensure Conda is installed as per the instructions at "
|
|
"https://conda.io/projects/conda/en/latest/"
|
|
"user-guide/install/index.html. "
|
|
"You can also configure MLflow to look for a specific "
|
|
f"Conda executable by setting the {MLFLOW_CONDA_HOME} environment variable "
|
|
"to the path of the Conda executable"
|
|
)
|
|
|
|
try:
|
|
# Checks if executable for environment creation exists
|
|
process._exec_cmd(
|
|
[conda_env_create_path, "--help"], throw_on_error=False, extra_env=extra_envs
|
|
)
|
|
except OSError:
|
|
raise ExecutionException(
|
|
f"You have set the env variable {MLFLOW_CONDA_CREATE_ENV_CMD}, but "
|
|
f"{conda_env_create_path} does not exist or it is not working properly. "
|
|
f"Note that {conda_env_create_path} and the conda executable need to be "
|
|
"in the same conda environment. You can change the search path by"
|
|
f"modifying the env variable {MLFLOW_CONDA_HOME}"
|
|
)
|
|
|
|
conda_extra_env_vars = _get_conda_extra_env_vars(env_root_dir)
|
|
if extra_envs:
|
|
conda_extra_env_vars.update(extra_envs)
|
|
|
|
# Include the env_root_dir hash in the project_env_name,
|
|
# this is for avoid conda env name conflicts between different CONDA_ENVS_PATH.
|
|
project_env_name = _get_conda_env_name(conda_env_path, env_id=env_id, env_root_dir=env_root_dir)
|
|
if env_root_dir is not None:
|
|
project_env_path = os.path.join(env_root_dir, _CONDA_ENVS_DIR, project_env_name)
|
|
else:
|
|
project_env_path = project_env_name
|
|
|
|
if project_env_name in _list_conda_environments(conda_extra_env_vars):
|
|
_logger.info("Conda environment %s already exists.", project_env_path)
|
|
return Environment(get_conda_command(project_env_name), conda_extra_env_vars)
|
|
|
|
_logger.info("=== Creating conda environment %s ===", project_env_path)
|
|
try:
|
|
_create_conda_env_func = (
|
|
# Retry conda env creation in a pytest session to avoid flaky test failures
|
|
_create_conda_env_retry if "PYTEST_CURRENT_TEST" in os.environ else _create_conda_env
|
|
)
|
|
conda_env = _create_conda_env_func(
|
|
conda_env_path,
|
|
conda_env_create_path,
|
|
project_env_name,
|
|
conda_extra_env_vars,
|
|
capture_output,
|
|
)
|
|
|
|
if pip_requirements_override:
|
|
_logger.info(
|
|
"Installing additional dependencies specified"
|
|
f"by pip_requirements_override: {pip_requirements_override}"
|
|
)
|
|
cmd = [
|
|
conda_path,
|
|
"install",
|
|
"-n",
|
|
project_env_name,
|
|
"--yes",
|
|
*pip_requirements_override,
|
|
]
|
|
process._exec_cmd(cmd, extra_env=conda_extra_env_vars, capture_output=capture_output)
|
|
|
|
return conda_env
|
|
|
|
except Exception:
|
|
try:
|
|
if project_env_name in _list_conda_environments(conda_extra_env_vars):
|
|
_logger.warning(
|
|
"Encountered unexpected error while creating conda environment. Removing %s.",
|
|
project_env_path,
|
|
)
|
|
process._exec_cmd(
|
|
[
|
|
conda_path,
|
|
"remove",
|
|
"--yes",
|
|
"--name",
|
|
project_env_name,
|
|
"--all",
|
|
],
|
|
extra_env=conda_extra_env_vars,
|
|
capture_output=False,
|
|
)
|
|
except Exception as e:
|
|
_logger.warning(
|
|
"Removing conda environment %s failed (error: %s)",
|
|
project_env_path,
|
|
repr(e),
|
|
)
|
|
raise
|
|
|
|
|
|
def _get_conda_dependencies(conda_yaml_path):
|
|
"""Extracts conda dependencies from a conda yaml file.
|
|
|
|
Args:
|
|
conda_yaml_path: Conda yaml file path.
|
|
"""
|
|
with open(conda_yaml_path) as f:
|
|
conda_yaml = yaml.safe_load(f)
|
|
return [d for d in conda_yaml.get("dependencies", []) if isinstance(d, str)]
|