""" This module provides a set of utilities for interpreting and creating requirements files (e.g. pip's `requirements.txt`), which is useful for managing ML software environments. """ import importlib.metadata import json import logging import os import re import subprocess import sys import tempfile from itertools import chain, filterfalse from pathlib import Path from threading import Timer from typing import NamedTuple import importlib_metadata from packaging.requirements import Requirement from packaging.version import InvalidVersion, Version import mlflow from mlflow.environment_variables import ( _MLFLOW_IN_CAPTURE_MODULE_PROCESS, MLFLOW_REQUIREMENTS_INFERENCE_RAISE_ERRORS, MLFLOW_REQUIREMENTS_INFERENCE_TIMEOUT, ) from mlflow.exceptions import MlflowException from mlflow.tracking.artifact_utils import _download_artifact_from_uri from mlflow.utils.autologging_utils.versioning import _strip_dev_version_suffix from mlflow.utils.databricks_utils import ( get_databricks_env_vars, is_in_databricks_runtime, ) _logger = logging.getLogger(__name__) def _is_comment(line): return line.startswith("#") def _is_empty(line): return line == "" def _strip_inline_comment(line): return line[: line.find(" #")].rstrip() if " #" in line else line def _is_requirements_file(line): return line.startswith("-r ") or line.startswith("--requirement ") def _is_constraints_file(line): return line.startswith("-c ") or line.startswith("--constraint ") def _join_continued_lines(lines): """ Joins lines ending with '\\'. >>> _join_continued_lines["a\\", "b\\", "c"] >>> "abc" """ continued_lines = [] for line in lines: if line.endswith("\\"): continued_lines.append(line.rstrip("\\")) else: continued_lines.append(line) yield "".join(continued_lines) continued_lines.clear() # The last line ends with '\' if continued_lines: yield "".join(continued_lines) class _Requirement(NamedTuple): # A string representation of the requirement. req_str: str # A boolean indicating whether this requirement is a constraint. is_constraint: bool def _parse_requirements(requirements, is_constraint, base_dir=None): """A simplified version of `pip._internal.req.parse_requirements` which performs the following operations on the given requirements file and yields the parsed requirements. - Remove comments and blank lines - Join continued lines - Resolve requirements file references (e.g. '-r requirements.txt') - Resolve constraints file references (e.g. '-c constraints.txt') Args: requirements: A string path to a requirements file on the local filesystem or an iterable of pip requirement strings. is_constraint: Indicates the parsed requirements file is a constraint file. base_dir: If specified, resolve relative file references (e.g. '-r requirements.txt') against the specified directory. Returns: A list of ``_Requirement`` instances. References: - `pip._internal.req.parse_requirements`: https://github.com/pypa/pip/blob/7a77484a492c8f1e1f5ef24eaf71a43df9ea47eb/src/pip/_internal/req/req_file.py#L118 - Requirements File Format: https://pip.pypa.io/en/stable/cli/pip_install/#requirements-file-format - Constraints Files: https://pip.pypa.io/en/stable/user_guide/#constraints-files """ if base_dir is None: if isinstance(requirements, (str, Path)): base_dir = os.path.dirname(requirements) with open(requirements) as f: requirements = f.read().splitlines() else: base_dir = os.getcwd() lines = map(str.strip, requirements) lines = map(_strip_inline_comment, lines) lines = _join_continued_lines(lines) lines = filterfalse(_is_comment, lines) lines = filterfalse(_is_empty, lines) for line in lines: if _is_requirements_file(line): req_file = line.split(maxsplit=1)[1] # If `req_file` is an absolute path, `os.path.join` returns `req_file`: # https://docs.python.org/3/library/os.path.html#os.path.join abs_path = os.path.join(base_dir, req_file) yield from _parse_requirements(abs_path, is_constraint=False) elif _is_constraints_file(line): req_file = line.split(maxsplit=1)[1] abs_path = os.path.join(base_dir, req_file) yield from _parse_requirements(abs_path, is_constraint=True) else: yield _Requirement(line, is_constraint) def _flatten(iterable): return chain.from_iterable(iterable) # https://www.python.org/dev/peps/pep-0508/#names _PACKAGE_NAME_REGEX = re.compile(r"^(\w+|\w+[\w._-]*\w+)") def _get_package_name(requirement): m = _PACKAGE_NAME_REGEX.match(requirement) return m and m.group(1) _NORMALIZE_REGEX = re.compile(r"[-_.]+") def _normalize_package_name(pkg_name): """ Normalizes a package name using the rule defined in PEP 503: https://www.python.org/dev/peps/pep-0503/#normalized-names """ return _NORMALIZE_REGEX.sub("-", pkg_name).lower() def _iter_requires(name: str): """ Iterates over the requirements of the specified package. Args: name: The name of the package. Yields: The names of the required packages. """ try: reqs = importlib.metadata.requires(name) except importlib.metadata.PackageNotFoundError: return if reqs is None: return for req in reqs: # Skip extra dependencies semi_colon_idx = req.find(";") if (semi_colon_idx != -1) and req[semi_colon_idx:].startswith("; extra =="): continue req = Requirement(req) # Skip the requirement if the environment marker is not satisfied if req.marker and not req.marker.evaluate(): continue yield req.name def _get_requires(pkg_name): norm_pkg_name = _normalize_package_name(pkg_name) for req in _iter_requires(norm_pkg_name): yield _normalize_package_name(req) def _get_requires_recursive(pkg_name, seen_before=None): """ Recursively yields both direct and transitive dependencies of the specified package. """ norm_pkg_name = _normalize_package_name(pkg_name) seen_before = seen_before or {norm_pkg_name} for req in _get_requires(pkg_name): # Prevent infinite recursion due to cyclic dependencies if req in seen_before: continue seen_before.add(req) yield req yield from _get_requires_recursive(req, seen_before) def _prune_packages(packages): """ Prunes packages required by other packages. For example, `["scikit-learn", "numpy"]` is pruned to `["scikit-learn"]`. """ packages = set(packages) requires = set(_flatten(map(_get_requires_recursive, packages))) # LlamaIndex have one root "llama-index" package that bundles many sub-packages such as # llama-index-llms-openai. Many of those sub-packages are optional, but some are defined # as dependencies of the root package. However, the root package does not pin the versions # for those sub-packages, resulting in non-deterministic behavior when loading the model # later. To address this issue, we keep all sub-packages within the requirements. # Ref: https://github.com/run-llama/llama_index/issues/14788#issuecomment-2232107585 requires = {req for req in requires if not req.startswith("llama-index-")} # Do not exclude mlflow's dependencies # Do not exclude databricks-connect since it conflicts with pyspark during execution time, # and we need to determine if pyspark needs to be stripped based on the inferred packages return packages - (requires - set(_get_requires("mlflow")) - {"databricks-connect"}) def _run_command(cmd, timeout_seconds, env=None): """ Runs the specified command. If it exits with non-zero status, `MlflowException` is raised. """ proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env) timer = Timer(timeout_seconds, proc.kill) try: timer.start() stdout, stderr = proc.communicate() stdout = stdout.decode("utf-8") stderr = stderr.decode("utf-8") if proc.returncode != 0: msg = "\n".join([ f"Encountered an unexpected error while running {cmd}", f"exit status: {proc.returncode}", f"stdout: {stdout}", f"stderr: {stderr}", ]) raise MlflowException(msg) finally: if timer.is_alive(): timer.cancel() def _get_installed_version(package: str, module: str | None = None) -> str: """ Obtains the installed package version using `importlib_metadata.version`. If it fails, use `__import__(module or package).__version__`. """ if package == "mlflow": # `importlib.metadata.version` may return an incorrect version of MLflow when it's # installed in editable mode (e.g. `pip install -e .`). return mlflow.__version__ try: version = importlib_metadata.version(package) except importlib_metadata.PackageNotFoundError: # Note `importlib_metadata.version(package)` is not necessarily equal to # `__import__(package).__version__`. See the example for pytorch below. # # Example # ------- # $ pip install torch==1.9.0 # $ python -c "import torch; print(torch.__version__)" # 1.9.0+cu102 # $ python -c "import importlib_metadata; print(importlib_metadata.version('torch'))" # 1.9.0 version = __import__(module or package).__version__ # Strip the suffix from `dev` versions of PySpark, which are not available for installation # from Anaconda or PyPI if package == "pyspark": version = _strip_dev_version_suffix(version) return version def _capture_imported_modules(model_uri, flavor, record_full_module=False, extra_env_vars=None): """Runs `_capture_modules.py` in a subprocess and captures modules imported during the model loading procedure. If flavor is `transformers`, `_capture_transformers_modules.py` is run instead. Args: model_uri: The URI of the model. flavor: The flavor name of the model. record_full_module: Whether to capture top level modules for inferring python package purpose. Default to False. extra_env_vars: A dictionary of extra environment variables to pass to the subprocess. Default to None. Returns: A list of captured modules. """ local_model_path = _download_artifact_from_uri(model_uri) process_timeout = MLFLOW_REQUIREMENTS_INFERENCE_TIMEOUT.get() raise_on_error = MLFLOW_REQUIREMENTS_INFERENCE_RAISE_ERRORS.get() extra_env_vars = extra_env_vars or {} # Run `_capture_modules.py` to capture modules imported during the loading procedure with tempfile.TemporaryDirectory() as tmpdir: output_file = os.path.join(tmpdir, "imported_modules.txt") # Pass the main environment variables to the subprocess for environment variable mapping main_env = os.environ.copy() # Reset the path variable from the main process so that the subprocess retains all # main process configuration that a user has. # See: ``https://github.com/mlflow/mlflow/issues/6905`` for context on minio configuration # resolution in a subprocess based on PATH entries. main_env["PATH"] = "/usr/sbin:/sbin:" + main_env["PATH"] # Clear py4j gateway env vars to prevent the subprocess from connecting to the parent's # py4j gateway. If these are inherited, libraries like databricks-sdk may attempt to use # them, which can corrupt the parent process's py4j connection state and cause errors # like "Error while obtaining a new communication channel" after the subprocess exits. main_env.pop("PYSPARK_GATEWAY_PORT", None) main_env.pop("PYSPARK_GATEWAY_SECRET", None) # Add databricks env, for langchain models loading we might need CLI configurations if is_in_databricks_runtime(): main_env.update(get_databricks_env_vars(mlflow.get_tracking_uri())) record_full_module_args = ["--record-full-module"] if record_full_module else [] if flavor == mlflow.transformers.FLAVOR_NAME: # Lazily import `_capture_transformers_module` here to avoid circular imports. from mlflow.utils import _capture_transformers_modules for module_to_throw in ["tensorflow", "torch"]: # NB: Setting USE_TF or USE_TORCH here as Transformers only checks these env # variable on the first import of the library, which could happen anytime during # the model loading process (or even mlflow import). When these variables are not # set, Transformers import some torch/tensorflow modules even if they are not # used by the model, resulting in false positives in the captured modules. transformer_env = ( {"USE_TF": "TRUE"} if module_to_throw == "torch" else {"USE_TORCH": "TRUE"} ) try: _run_command( [ sys.executable, _capture_transformers_modules.__file__, "--model-path", local_model_path, "--flavor", flavor, "--output-file", output_file, "--sys-path", json.dumps(sys.path), "--module-to-throw", module_to_throw, *record_full_module_args, ], timeout_seconds=process_timeout, env={ **main_env, **transformer_env, _MLFLOW_IN_CAPTURE_MODULE_PROCESS.name: "true", **extra_env_vars, }, ) with open(output_file) as f: return f.read().splitlines() except MlflowException: pass # Lazily import `_capture_module` here to avoid circular imports. from mlflow.utils import _capture_modules error_file = os.path.join(tmpdir, "error.txt") _run_command( [ sys.executable, _capture_modules.__file__, "--model-path", local_model_path, "--flavor", flavor, "--output-file", output_file, "--error-file", error_file, "--sys-path", json.dumps(sys.path), *record_full_module_args, ], timeout_seconds=process_timeout, env={ **main_env, _MLFLOW_IN_CAPTURE_MODULE_PROCESS.name: "true", **extra_env_vars, }, ) if os.path.exists(error_file): with open(error_file) as f: errors = f.read() if errors: if raise_on_error: raise MlflowException( f"Encountered an error while capturing imported modules: {errors}" ) _logger.warning(errors) with open(output_file) as f: return f.read().splitlines() DATABRICKS_MODULES_TO_PACKAGES = { "databricks.automl": ["databricks-automl-runtime"], "databricks.automl_runtime": ["databricks-automl-runtime"], "databricks.model_monitoring": ["databricks-model-monitoring"], } MLFLOW_MODULES_TO_PACKAGES = { "mlflow.gateway": ["mlflow[gateway]"], } _MODULES_TO_PACKAGES = None _PACKAGES_TO_MODULES = None def _init_modules_to_packages_map(): global _MODULES_TO_PACKAGES if _MODULES_TO_PACKAGES is None: # Note `importlib_metadata.packages_distributions` only captures packages installed into # Python's site-packages directory via tools such as pip: # https://importlib-metadata.readthedocs.io/en/latest/using.html#using-importlib-metadata _MODULES_TO_PACKAGES = importlib_metadata.packages_distributions() # Add mapping for MLflow extras _MODULES_TO_PACKAGES.update(MLFLOW_MODULES_TO_PACKAGES) # Multiple packages populate the `databricks` module namespace on Databricks; to avoid # bundling extraneous Databricks packages into model dependencies, we scope each module # to its relevant package _MODULES_TO_PACKAGES.update(DATABRICKS_MODULES_TO_PACKAGES) if "databricks" in _MODULES_TO_PACKAGES: _MODULES_TO_PACKAGES["databricks"] = [ package for package in _MODULES_TO_PACKAGES["databricks"] if package not in _flatten(DATABRICKS_MODULES_TO_PACKAGES.values()) ] # In Databricks, `_MODULES_TO_PACKAGES` doesn't contain pyspark since it's not installed # via pip or conda. To work around this issue, manually add pyspark. if is_in_databricks_runtime(): _MODULES_TO_PACKAGES.update({"pyspark": ["pyspark"]}) def _init_packages_to_modules_map(): _init_modules_to_packages_map() global _PACKAGES_TO_MODULES _PACKAGES_TO_MODULES = {} for module, pkg_list in _MODULES_TO_PACKAGES.items(): for pkg_name in pkg_list: _PACKAGES_TO_MODULES[pkg_name] = module def _infer_requirements(model_uri, flavor, raise_on_error=False, extra_env_vars=None): """Infers the pip requirements of the specified model by creating a subprocess and loading the model in it to determine which packages are imported. Args: model_uri: The URI of the model. flavor: The flavor name of the model. raise_on_error: If True, raise an exception if an unrecognized package is encountered. extra_env_vars: A dictionary of extra environment variables to pass to the subprocess. Default to None. Returns: A list of inferred pip requirements. """ _init_modules_to_packages_map() modules = _capture_imported_modules(model_uri, flavor, extra_env_vars=extra_env_vars) packages = _flatten([_MODULES_TO_PACKAGES.get(module, []) for module in modules]) packages = map(_normalize_package_name, packages) packages = _prune_packages(packages) excluded_packages = [ # Certain packages (e.g. scikit-learn 0.24.2) imports `setuptools` or `pkg_resources` # (a module provided by `setuptools`) to process or interact with package metadata. # It should be safe to exclude `setuptools` because it's rare to encounter a python # environment where `setuptools` is not pre-installed. "setuptools", # Exclude a package that provides the mlflow module (e.g. mlflow, mlflow-skinny). # Certain flavors (e.g. pytorch) import mlflow while loading a model, but mlflow should # not be counted as a model requirement. *_MODULES_TO_PACKAGES.get("mlflow", []), ] packages = packages - set(excluded_packages) # Handle pandas incompatibility issue with numpy 2.x https://github.com/pandas-dev/pandas/issues/55519 # pandas == 2.2.*: compatible with numpy >= 2 # pandas >= 2.1.2: incompatible with numpy >= 2, but it pins numpy < 2 # pandas < 2.1.2: incompatible with numpy >= 2 and doesn't pin numpy, so we need to pin numpy if any( package == "pandas" and Version(_get_pinned_requirement(package).split("==")[1]) < Version("2.1.2") for package in packages ): packages.add("numpy") return sorted(map(_get_pinned_requirement, packages)) def _get_local_version_label(version): """Extracts a local version label from `version`. Args: version: A version string. """ try: return Version(version).local except InvalidVersion: return None def _strip_local_version_label(version): """Strips a local version label in `version`. Local version identifiers: https://www.python.org/dev/peps/pep-0440/#local-version-identifiers Args: version: A version string to strip. """ class IgnoreLocal(Version): @property def local(self): return None try: return str(IgnoreLocal(version)) except InvalidVersion: return version def _get_pinned_requirement(req_str, version=None, module=None): """Returns a string representing a pinned pip requirement to install the specified package and version (e.g. 'mlflow==1.2.3'). Args: req_str: The package requirement string (e.g. "mlflow" or "mlflow[gateway]"). version: The version of the package. If None, defaults to the installed version. module: The name of the top-level module provided by the package . For example, if `package` is 'scikit-learn', `module` should be 'sklearn'. If None, defaults to `package`. extras: A list of extra names for the package. """ req = Requirement(req_str) package = req.name if version is None: version_raw = _get_installed_version(package, module) if local_version_label := _get_local_version_label(version_raw): version = _strip_local_version_label(version_raw) if not (is_in_databricks_runtime() and package in ("torch", "torchvision")): msg = ( f"Found {package} version ({version_raw}) contains a local version label " f"(+{local_version_label}). MLflow logged a pip requirement for this package " f"as '{package}=={version}' without the local version label to make it " "installable from PyPI. To specify pip requirements containing local version " "labels, please use `conda_env` or `pip_requirements`." ) _logger.warning(msg) else: version = version_raw if req.extras: return f"{package}[{','.join(req.extras)}]=={version}" return f"{package}=={version}" class _MismatchedPackageInfo(NamedTuple): package_name: str installed_version: str | None requirement: str def __str__(self): current_status = self.installed_version or "uninstalled" return f"{self.package_name} (current: {current_status}, required: {self.requirement})" def _check_requirement_satisfied(requirement_str: str) -> _MismatchedPackageInfo | None: """ Checks whether the current python environment satisfies the given requirement if it is parsable as a package name and a set of version specifiers, and returns a `_MismatchedPackageInfo` object containing the mismatched package name, installed version, and requirement if the requirement is not satisfied. Otherwise, returns None. """ try: req = Requirement(requirement_str) except Exception: # We reach here if the requirement string is a file path or a URL. # Extracting the package name from the requirement string is not trivial, # so we skip the check. return None if req.marker and not req.marker.evaluate(): return None _init_packages_to_modules_map() pkg_name = req.name try: installed_version = _get_installed_version(pkg_name, _PACKAGES_TO_MODULES.get(pkg_name)) except ModuleNotFoundError: return _MismatchedPackageInfo( package_name=pkg_name, installed_version=None, requirement=requirement_str, ) if pkg_name == "mlflow" and "gateway" in req.extras: try: from mlflow import gateway # noqa: F401 except ModuleNotFoundError: return _MismatchedPackageInfo( package_name="mlflow[gateway]", installed_version=None, requirement=requirement_str, ) if pkg_name == "mlflow" and Version(installed_version).is_devrelease: return None if len(req.specifier) > 0 and not req.specifier.contains(installed_version): return _MismatchedPackageInfo( package_name=pkg_name, installed_version=installed_version, requirement=requirement_str, ) return None def warn_dependency_requirement_mismatches(model_requirements: list[str]): """ Inspects the model's dependencies and prints a warning if the current Python environment doesn't satisfy them. """ # Suppress databricks-feature-lookup warning for feature store cases # Suppress databricks-chains, databricks-rag, and databricks-agents warnings for RAG # Studio cases # NB: When a final name has been decided for GA for the aforementioned # "Databricks RAG Studio" product, remove unrelated names from this listing. _DATABRICKS_FEATURE_LOOKUP = "databricks-feature-lookup" _DATABRICKS_AGENTS = "databricks-agents" # List of packages to ignore packages_to_ignore = [ _DATABRICKS_FEATURE_LOOKUP, _DATABRICKS_AGENTS, ] # Normalize package names and create ignore list ignore_packages = list(map(_normalize_package_name, packages_to_ignore)) try: mismatch_infos = [] for req in model_requirements: mismatch_info = _check_requirement_satisfied(req) if mismatch_info is not None: if _normalize_package_name(mismatch_info.package_name) in ignore_packages: continue mismatch_infos.append(str(mismatch_info)) if len(mismatch_infos) > 0: mismatch_str = " - " + "\n - ".join(mismatch_infos) warning_msg = ( "Detected one or more mismatches between the model's dependencies and the current " f"Python environment:\n{mismatch_str}\n" "To fix the mismatches, call `mlflow.pyfunc.get_model_dependencies(model_uri)` " "to fetch the model's environment and install dependencies using the resulting " "environment file." ) _logger.warning(warning_msg) except Exception as e: _logger.warning( f"Encountered an unexpected error ({e!r}) while detecting model dependency " "mismatches. Set logging level to DEBUG to see the full traceback." ) _logger.debug("", exc_info=True)