import functools import json import logging import numbers import os import random import signal import socket import subprocess import sys import tempfile import time import uuid from contextlib import contextmanager from functools import wraps from pathlib import Path from typing import Iterator from unittest import mock import pytest import requests import mlflow from mlflow.entities.logged_model import LoggedModel from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS from mlflow.tracking.artifact_utils import _download_artifact_from_uri from mlflow.utils.os import is_windows AWS_METADATA_IP = "169.254.169.254" # Used to fetch AWS Instance and User metadata. def kill_process_tree(pid: int) -> None: """ Gracefully terminate or kill a process tree (children first, then parent). """ import psutil try: parent = psutil.Process(pid) except psutil.NoSuchProcess: return # Kill children first to prevent the parent from spawning new processes children = parent.children(recursive=True) for child in children: try: child.terminate() except psutil.NoSuchProcess: pass # Wait for children to terminate, then force-kill any that remain _, still_alive = psutil.wait_procs(children, timeout=5) for p in still_alive: try: p.kill() except psutil.NoSuchProcess: pass # Finally, kill the parent try: parent.terminate() parent.wait(timeout=5) except psutil.NoSuchProcess: pass except psutil.TimeoutExpired: parent.kill() LOCALHOST = "127.0.0.1" PROTOBUF_REQUIREMENT = "protobuf<4.0.0" _logger = logging.getLogger(__name__) def get_safe_port(): """Returns an ephemeral port that is very likely to be free to bind to.""" sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.bind((LOCALHOST, 0)) port = sock.getsockname()[1] sock.close() return port def random_int(lo=1, hi=1e10): return random.randint(int(lo), int(hi)) def random_str(size=12): msg = ( "UUID4 generated strings have a high potential for collision at small sizes. " "10 is set as the lower bounds for random string generation to prevent non-deterministic " "test failures." ) assert size >= 10, msg return uuid.uuid4().hex[:size] def random_file(ext): return f"temp_test_{random_int()}.{ext}" def expect_status_code(http_response, expected_code): assert http_response.status_code == expected_code, ( f"Unexpected status code. {http_response.status_code} != {expected_code}, " f"body: {http_response.text}" ) def score_model_in_sagemaker_docker_container( model_uri, data, content_type, flavor="python_function", activity_polling_timeout_seconds=500, ): """ Args: model_uri: URI to the model to be served. data: The data to send to the docker container for testing. This is either a Pandas dataframe or string of the format specified by `content_type`. content_type: The type of the data to send to the docker container for testing. This is one of `mlflow.pyfunc.scoring_server.CONTENT_TYPES`. flavor: Model flavor to be deployed. activity_polling_timeout_seconds: The amount of time, in seconds, to wait before declaring the scoring process to have failed. """ env = dict(os.environ) env.update(LC_ALL="en_US.UTF-8", LANG="en_US.UTF-8") port = get_safe_port() scoring_cmd = ( f"mlflow deployments run-local -t sagemaker --name test -m {model_uri}" f" -C image=mlflow-pyfunc -C port={port} --flavor {flavor}" ) proc = _start_scoring_proc( cmd=scoring_cmd.split(" "), env=env, ) with RestEndpoint( proc, port, activity_polling_timeout_seconds, validate_version=False ) as endpoint: return endpoint.invoke(data, content_type) def pyfunc_generate_dockerfile(output_directory, model_uri=None, extra_args=None, env=None): """ Builds a dockerfile for the specified model. Args: output_directory: Output directory to generate Dockerfile and model artifacts model_uri: URI of model, e.g. runs:/some-run-id/run-relative/path/to/model extra_args: List of extra args to pass to `mlflow models build-docker` command env: Environment variables to use. """ cmd = [ "mlflow", "models", "generate-dockerfile", *(["-m", model_uri] if model_uri else []), "-d", output_directory, ] if mlflow_home := os.environ.get("MLFLOW_HOME"): cmd += ["--mlflow-home", mlflow_home] if extra_args: cmd += extra_args subprocess.run(cmd, check=True, env=env) def pyfunc_build_image(model_uri=None, extra_args=None, env=None): """ Builds a docker image containing the specified model, returning the name of the image. Args: model_uri: URI of model, e.g. runs:/some-run-id/run-relative/path/to/model extra_args: List of extra args to pass to `mlflow models build-docker` command env: Environment variables to pass to the subprocess building the image. """ name = uuid.uuid4().hex cmd = [ sys.executable, "-m", "mlflow", "models", "build-docker", *(["-m", model_uri] if model_uri else []), "-n", name, ] if mlflow_home := os.environ.get("MLFLOW_HOME"): cmd += ["--mlflow-home", mlflow_home] if extra_args: cmd += extra_args # Docker image build occasionally fails on GitHub Actions while running `apt-get` due to # transient network issues. Retry the build a few times as a workaround. for _ in range(3): p = subprocess.Popen(cmd, env=env) if p.wait() == 0: return name time.sleep(5) raise RuntimeError(f"Failed to build docker image to serve model from {model_uri}") def pyfunc_serve_from_docker_image(image_name, host_port, extra_args=None): """ Serves a model from a docker container, exposing it as an endpoint at the specified port on the host machine. Returns a handle (Popen object) to the server process. """ env = dict(os.environ) env.update(LC_ALL="en_US.UTF-8", LANG="en_US.UTF-8") scoring_cmd = ["docker", "run", "-p", f"{host_port}:8080", image_name] if extra_args is not None: scoring_cmd += extra_args return _start_scoring_proc(cmd=scoring_cmd, env=env) def pyfunc_serve_from_docker_image_with_env_override( image_name, host_port, extra_args=None, extra_docker_run_options=None ): """ Serves a model from a docker container, exposing it as an endpoint at the specified port on the host machine. Returns a handle (Popen object) to the server process. """ env = dict(os.environ) env.update(LC_ALL="en_US.UTF-8", LANG="en_US.UTF-8") scoring_cmd = [ "docker", "run", "-p", f"{host_port}:8080", *(extra_docker_run_options or []), image_name, ] if extra_args is not None: scoring_cmd += extra_args return _start_scoring_proc(cmd=scoring_cmd, env=env) def pyfunc_serve_and_score_model( model_uri, data, content_type, activity_polling_timeout_seconds=500, extra_args=None, stdout=sys.stdout, ): with pyfunc_scoring_endpoint( model_uri, extra_args=extra_args, activity_polling_timeout_seconds=activity_polling_timeout_seconds, stdout=stdout, ) as endpoint: return endpoint.invoke(data, content_type) @contextmanager def pyfunc_scoring_endpoint( model_uri, activity_polling_timeout_seconds=500, extra_args=None, stdout=sys.stdout ): """ Args: model_uri: URI to the model to be served. activity_polling_timeout_seconds: The amount of time, in seconds, to wait before declaring the scoring process to have failed. extra_args: A list of extra arguments to pass to the pyfunc scoring server command. For example, passing ``extra_args=["--env-manager", "local"]`` will pass the ``--env-manager local`` flag to the scoring server to ensure that conda environment activation is skipped. stdout: The output stream to which standard output is redirected. Defaults to `sys.stdout`. """ env = dict(os.environ) env.update(LC_ALL="en_US.UTF-8", LANG="en_US.UTF-8") env.update(MLFLOW_TRACKING_URI=mlflow.get_tracking_uri()) env.update(MLFLOW_HOME=_get_mlflow_home()) port = get_safe_port() scoring_cmd = [ sys.executable, "-m", "mlflow", "models", "serve", "-m", model_uri, "-p", str(port), "--install-mlflow", ] + (extra_args or []) with _start_scoring_proc(cmd=scoring_cmd, env=env, stdout=stdout, stderr=stdout) as proc: try: with RestEndpoint(proc, port, activity_polling_timeout_seconds) as endpoint: yield endpoint finally: proc.terminate() def _get_mlflow_home(): """ Returns: The path to the MLflow installation root directory. """ mlflow_module_path = os.path.dirname(os.path.abspath(mlflow.__file__)) # The MLflow root directory is one level about the mlflow module location return os.path.join(mlflow_module_path, os.pardir) def _start_scoring_proc(cmd, env, stdout=sys.stdout, stderr=sys.stderr): if not is_windows(): return subprocess.Popen( cmd, stdout=stdout, stderr=stderr, text=True, env=env, # Assign the scoring process to a process group. All child processes of the # scoring process will be assigned to this group as well. This allows child # processes of the scoring process to be terminated successfully preexec_fn=os.setsid, ) else: return subprocess.Popen( cmd, stdout=stdout, stderr=stderr, text=True, env=env, # On Windows, `os.setsid` and `preexec_fn` are unavailable creationflags=subprocess.CREATE_NEW_PROCESS_GROUP, ) class RestEndpoint: def __init__(self, proc, port, activity_polling_timeout_seconds=60 * 8, validate_version=True): self._proc = proc self._port = port self._activity_polling_timeout_seconds = activity_polling_timeout_seconds self._validate_version = validate_version def __enter__(self): ping_status = None for i in range(self._activity_polling_timeout_seconds): assert self._proc.poll() is None, "scoring process died" time.sleep(1) # noinspection PyBroadException try: ping_status = requests.get(url=f"http://localhost:{self._port}/ping") _logger.info(f"connection attempt {i} server is up! ping status {ping_status}") if ping_status.status_code == 200: break except Exception: _logger.info(f"connection attempt {i} failed, server is not up yet") if ping_status is None or ping_status.status_code != 200: raise Exception("ping failed, server is not happy") _logger.info(f"server up, ping status {ping_status}") if self._validate_version: resp_status = requests.get(url=f"http://localhost:{self._port}/version") version = resp_status.text _logger.info(f"mlflow server version {version}") if version != mlflow.__version__: raise Exception("version path is not returning correct mlflow version") return self def __exit__(self, tp, val, traceback): if self._proc.poll() is None: # Terminate the process group containing the scoring process. # This will terminate all child processes of the scoring process if not is_windows(): pgrp = os.getpgid(self._proc.pid) os.killpg(pgrp, signal.SIGTERM) else: # https://stackoverflow.com/questions/47016723/windows-equivalent-for-spawning-and-killing-separate-process-group-in-python-3 self._proc.send_signal(signal.CTRL_BREAK_EVENT) self._proc.kill() def invoke(self, data, content_type): import pandas as pd from mlflow.pyfunc import scoring_server as pyfunc_scoring_server if isinstance(data, pd.DataFrame): if content_type == pyfunc_scoring_server.CONTENT_TYPE_CSV: data = data.to_csv(index=False) else: assert content_type == pyfunc_scoring_server.CONTENT_TYPE_JSON data = json.dumps({"dataframe_split": data.to_dict(orient="split")}) elif type(data) not in {str, dict}: data = json.dumps({"instances": data}) return requests.post( url=f"http://localhost:{self._port}/invocations", data=data, headers={"Content-Type": content_type}, ) @pytest.fixture(autouse=True) def set_boto_credentials(monkeypatch): monkeypatch.setenv("AWS_ACCESS_KEY_ID", "NotARealAccessKey") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "NotARealSecretAccessKey") monkeypatch.setenv("AWS_SESSION_TOKEN", "NotARealSessionToken") def create_mock_response(status_code, text): """ Create a mock response object with the status_code and text Args: status_code: HTTP status code. text: Message from the response. Returns: Mock HTTP Response. """ response = mock.MagicMock() response.status_code = status_code response.text = text return response def _read_lines(path): with open(path) as f: return f.read().splitlines() def _compare_logged_code_paths(code_path: str, model_uri: str, flavor_name: str) -> None: from mlflow.utils.model_utils import FLAVOR_CONFIG_CODE, _get_flavor_configuration model_path = _download_artifact_from_uri(model_uri) pyfunc_conf = _get_flavor_configuration( model_path=model_path, flavor_name=mlflow.pyfunc.FLAVOR_NAME ) flavor_conf = _get_flavor_configuration(model_path, flavor_name=flavor_name) assert pyfunc_conf[mlflow.pyfunc.CODE] == flavor_conf[FLAVOR_CONFIG_CODE] saved_code_path = os.path.join(model_path, pyfunc_conf[mlflow.pyfunc.CODE]) assert os.path.exists(saved_code_path) with open(os.path.join(saved_code_path, os.path.basename(code_path))) as f1: with open(code_path) as f2: assert f1.read() == f2.read() def _compare_conda_env_requirements(env_path, req_path): from mlflow.utils.environment import _get_pip_deps from mlflow.utils.yaml_utils import read_yaml assert os.path.exists(req_path) env_root, env_path = os.path.split(env_path) custom_env_parsed = read_yaml(env_root, env_path) requirements = _read_lines(req_path) assert _get_pip_deps(custom_env_parsed) == requirements def _get_deps_from_requirement_file(model_uri): """ Returns a list of pip dependencies for the model at `model_uri` and truncate the version number. """ from mlflow.utils.environment import _REQUIREMENTS_FILE_NAME local_path = _download_artifact_from_uri(model_uri) pip_packages = _read_lines(os.path.join(local_path, _REQUIREMENTS_FILE_NAME)) return [req.split("==")[0] if "==" in req else req for req in pip_packages] def assert_register_model_called_with_local_model_path( register_model_mock, model_uri, registered_model_name ): register_model_call_args = register_model_mock.call_args assert register_model_call_args.args == (model_uri, registered_model_name) assert ( register_model_call_args.kwargs["await_registration_for"] == DEFAULT_AWAIT_MAX_SLEEP_SECONDS ) local_model_path = register_model_call_args.kwargs["local_model_path"] assert local_model_path.startswith(tempfile.gettempdir()) def _assert_pip_requirements(model_uri, requirements, constraints=None, strict=False): """ Loads the pip requirements (and optionally constraints) from `model_uri` and compares them to `requirements` (and `constraints`). If `strict` is True, evaluate `set(requirements) == set(loaded_requirements)`. Otherwise, evaluate `set(requirements) <= set(loaded_requirements)`. """ from mlflow.utils.environment import ( _CONDA_ENV_FILE_NAME, _CONSTRAINTS_FILE_NAME, _REQUIREMENTS_FILE_NAME, _get_pip_deps, ) from mlflow.utils.yaml_utils import read_yaml local_path = _download_artifact_from_uri(model_uri) txt_reqs = _read_lines(os.path.join(local_path, _REQUIREMENTS_FILE_NAME)) conda_reqs = _get_pip_deps(read_yaml(local_path, _CONDA_ENV_FILE_NAME)) compare_func = set.__eq__ if strict else set.__le__ requirements = set(requirements) assert compare_func(requirements, set(txt_reqs)) assert compare_func(requirements, set(conda_reqs)) if constraints is not None: assert f"-c {_CONSTRAINTS_FILE_NAME}" in txt_reqs assert f"-c {_CONSTRAINTS_FILE_NAME}" in conda_reqs cons = _read_lines(os.path.join(local_path, _CONSTRAINTS_FILE_NAME)) assert compare_func(set(constraints), set(cons)) def _is_available_on_pypi(package, version=None, module=None): """ Returns True if the specified package version is available on PyPI. Args: package: The name of the package. 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`. """ from mlflow.utils.requirements_utils import _get_installed_version url = f"https://pypi.python.org/pypi/{package}/json" for sec in range(3): try: time.sleep(sec) resp = requests.get(url) except requests.exceptions.ConnectionError: continue if resp.status_code == 404: return False if resp.status_code == 200: break else: raise Exception(f"Failed to connect to {url}") version = version or _get_installed_version(module or package) dist_files = resp.json()["releases"].get(version) return ( dist_files is not None # specified version exists and (len(dist_files) > 0) # at least one distribution file exists and not dist_files[0].get("yanked", False) # specified version is not yanked ) def _is_importable(module_name): try: __import__(module_name) return True except ImportError: return False def allow_infer_pip_requirements_fallback_if(condition): def decorator(f): return pytest.mark.allow_infer_pip_requirements_fallback(f) if condition else f return decorator def mock_method_chain(mock_obj, methods, return_value=None, side_effect=None): """ Mock a chain of methods. Examples -------- >>> from unittest import mock >>> m = mock.MagicMock() >>> mock_method_chain(m, ["a", "b"], return_value=0) >>> m.a().b() 0 >>> mock_method_chain(m, ["c.d", "e"], return_value=1) >>> m.c.d().e() 1 >>> mock_method_chain(m, ["f"], side_effect=Exception("side_effect")) >>> m.f() Traceback (most recent call last): ... Exception: side_effect """ length = len(methods) for idx, method in enumerate(methods): mock_obj = functools.reduce(getattr, method.split("."), mock_obj) if idx != length - 1: mock_obj = mock_obj.return_value else: mock_obj.return_value = return_value mock_obj.side_effect = side_effect class StartsWithMatcher: def __init__(self, prefix): self.prefix = prefix def __eq__(self, other): return isinstance(other, str) and other.startswith(self.prefix) class AnyStringWith(str): def __eq__(self, other): return self in other def assert_array_almost_equal(actual_array, desired_array, rtol=1e-6): import numpy as np elem0 = actual_array[0] if isinstance(elem0, numbers.Number) or ( isinstance(elem0, (list, np.ndarray)) and isinstance(elem0[0], numbers.Number) ): np.testing.assert_allclose(actual_array, desired_array, rtol=rtol) else: np.testing.assert_array_equal(actual_array, desired_array) def _mlflow_major_version_string(): from mlflow.utils.environment import _generate_mlflow_version_pinning return _generate_mlflow_version_pinning() @contextmanager def mock_http_request_200(): with mock.patch( "mlflow.utils.rest_utils.http_request", return_value=mock.MagicMock(status_code=200, text="{}"), ) as m: yield m def mock_http_200(f): @functools.wraps(f) @mock.patch( "mlflow.utils.rest_utils.http_request", return_value=mock.MagicMock(status_code=200, text="{}"), ) def wrapper(*args, **kwargs): return f(*args, **kwargs) return wrapper @contextmanager def mock_http_request_403_200(): with mock.patch( "mlflow.utils.rest_utils.http_request", side_effect=[ mock.MagicMock(status_code=403, text='{"error_code": "ENDPOINT_NOT_FOUND"}'), mock.MagicMock(status_code=200, text="{}"), ], ) as m: yield m def clear_hub_cache(): """ Frees up disk space for cached huggingface transformers models and components. This function will remove all files within the cache if the total size of objects exceeds 1 GB on disk. It is used only in CI testing to alleviate the disk burden on the runners as they have limited allocated space and will terminate if the available disk space drops too low. """ try: from huggingface_hub import scan_cache_dir full_cache = scan_cache_dir() cache_size_in_gb = full_cache.size_on_disk / 1000**3 if cache_size_in_gb > 1: commits_to_purge = [ rev.commit_hash for repo in full_cache.repos for rev in repo.revisions ] delete_strategy = full_cache.delete_revisions(*commits_to_purge) delete_strategy.execute() except ImportError: # Local import check for mlflow-skinny not including huggingface_hub pass except Exception as e: _logger.warning(f"Failed to clear cache: {e}", exc_info=True) def flaky(max_tries=3): """ Annotation decorator for retrying flaky functions up to max_tries times, and raise the Exception if it fails after max_tries attempts. Args: max_tries: Maximum number of times to retry the function. Returns: Decorated function. """ def flaky_test_func(test_func): @wraps(test_func) def decorated_func(*args, **kwargs): for i in range(max_tries): try: return test_func(*args, **kwargs) except Exception as e: _logger.warning(f"Attempt {i + 1} failed with error: {e}") if i == max_tries - 1: raise time.sleep(3) return decorated_func return flaky_test_func @contextmanager def start_mock_openai_server(): """ Start a fake service that mimics the OpenAI endpoints such as /chat/completions. Yields: The base URL of the mock OpenAI server. """ port = get_safe_port() script_path = Path(__file__).parent / "openai" / "mock_openai.py" with subprocess.Popen([ sys.executable, script_path, "--host", "localhost", "--port", str(port), ]) as proc: try: base_url = f"http://localhost:{port}" for _ in range(10): try: resp = requests.get(f"{base_url}/health") except requests.ConnectionError: time.sleep(2) continue if resp.ok: break else: proc.kill() proc.wait() raise RuntimeError("Failed to start mock OpenAI server") yield base_url finally: proc.kill() def _is_hf_hub_healthy() -> bool: """ Check if the Hugging Face Hub is healthy by attempting to load a small dataset. """ try: import datasets from huggingface_hub import HfApi except ImportError: # Cannot import datasets or huggingface_hub, so we assume the hub is healthy. return True try: for dataset in HfApi().list_datasets(filter="size_categories:n<1K", limit=10): # Gated datasets (e.g., https://huggingface.co/datasets/PatronusAI/TRAIL) require # authentication to access. if not dataset.gated: datasets.load_dataset(dataset.id) return True return True except requests.exceptions.RequestException: return False except Exception as e: _logger.warning(f"Unexpected error while checking Hugging Face Hub health: {e}. ") # For any other exceptions, we assume the hub is healthy. return True def _iter_pr_files() -> Iterator[str]: if "GITHUB_ACTIONS" not in os.environ: return if os.environ.get("GITHUB_EVENT_NAME") != "pull_request": return with open(os.environ["GITHUB_EVENT_PATH"]) as f: pr_data = json.load(f) pull_number = pr_data["pull_request"]["number"] repo = pr_data["repository"]["full_name"] page = 1 per_page = 100 headers = {"Authorization": token} if (token := os.environ.get("GH_TOKEN")) else None while True: resp = requests.get( f"https://api.github.com/repos/{repo}/pulls/{pull_number}/files", params={"per_page": per_page, "page": page}, headers=headers, ) try: resp.raise_for_status() except requests.exceptions.HTTPError as e: _logger.warning( f"Failed to fetch PR files: {e}. Skipping the check for Hugging Face Hub health." ) return files = [f["filename"] for f in resp.json()] yield from files if len(files) < per_page: break page += 1 @functools.lru_cache(maxsize=1) def _should_skip_hf_test() -> bool: if "CI" not in os.environ: # This is not a CI run. Do not skip tests. return False if any(("huggingface" in f or "transformers" in f) for f in _iter_pr_files()): # This PR modifies huggingface-related files. Do not skip tests. return False # Skip tests if the Hugging Face Hub is unhealthy. return not _is_hf_hub_healthy() def skip_if_hf_hub_unhealthy(): return pytest.mark.skipif( _should_skip_hf_test(), reason=( "Skipping test because Hugging Face Hub is unhealthy. " "See https://status.huggingface.co/ for more information." ), ) def get_logged_model_by_name(name: str) -> LoggedModel | None: """ Get a logged model by name. If multiple logged models with the same name exist, get the latest one. Args: name: The name of the logged model. Returns: The logged model. """ logged_models = mlflow.search_logged_models( filter_string=f"name='{name}'", output_format="list", max_results=1 ) return logged_models[0] if len(logged_models) >= 1 else None