import os import re import shutil import sys import uuid from pathlib import Path import pytest import mlflow from mlflow import cli from mlflow.utils import process from mlflow.utils.virtualenv import _get_mlflow_virtualenv_root from tests.helper_functions import clear_hub_cache, flaky, start_mock_openai_server from tests.integration.utils import invoke_cli_runner EXAMPLES_DIR = "examples" def find_python_env_yaml(directory: Path) -> Path: return next(filter(lambda p: p.name == "python_env.yaml", Path(directory).iterdir())) def replace_mlflow_with_dev_version(yml_path: Path) -> None: old_src = yml_path.read_text() mlflow_dir = Path(mlflow.__path__[0]).parent new_src = re.sub(r"- mlflow.*\n", f"- {mlflow_dir}\n", old_src) yml_path.write_text(new_src) @pytest.fixture(autouse=True) def clean_up_mlflow_virtual_environments(): yield venv_root = Path(_get_mlflow_virtualenv_root()) if not venv_root.exists(): return for path in venv_root.iterdir(): if path.is_dir(): shutil.rmtree(path) @pytest.fixture(scope="module", autouse=True) def mock_openai(): # Some examples includes OpenAI API calls, so we start a mock server. with start_mock_openai_server() as base_url: with pytest.MonkeyPatch.context() as mp: mp.setenv("OPENAI_API_BASE", base_url) mp.setenv("OPENAI_API_KEY", "test") yield @pytest.mark.notrackingurimock @flaky() @pytest.mark.parametrize( ("directory", "params"), [ ("h2o", []), # TODO: Fix the hyperparam example and re-enable it # ("hyperparam", ["-e", "train", "-P", "epochs=1"]), # ("hyperparam", ["-e", "random", "-P", "epochs=1"]), # ("hyperparam", ["-e", "hyperopt", "-P", "epochs=1"]), ( "lightgbm/lightgbm_native", ["-P", "learning_rate=0.1", "-P", "colsample_bytree=0.8", "-P", "subsample=0.9"], ), ("lightgbm/lightgbm_sklearn", []), ("statsmodels", ["-P", "inverse_method=qr"]), ("pytorch", ["-P", "epochs=2"]), ("sklearn_logistic_regression", []), ("sklearn_elasticnet_wine", ["-P", "alpha=0.5"]), ("sklearn_elasticnet_diabetes/linux", []), ("spacy", []), ( "xgboost/xgboost_native", ["-P", "learning_rate=0.3", "-P", "colsample_bytree=0.8", "-P", "subsample=0.9"], ), ("xgboost/xgboost_sklearn", []), ("pytorch/MNIST", ["-P", "max_epochs=1"]), ("pytorch/HPOExample", ["-P", "n_trials=2", "-P", "max_epochs=1"]), ("pytorch/CaptumExample", ["-P", "max_epochs=50"]), ("supply_chain_security", []), ("tensorflow", []), ("sktime", []), ], ) def test_mlflow_run_example(directory, params, tmp_path): # Use tmp_path+uuid as tmp directory to avoid the same # directory being reused when re-trying the test since # tmp_path is named as the test name random_tmp_path = tmp_path / str(uuid.uuid4()) example_dir = Path(EXAMPLES_DIR, directory) tmp_example_dir = random_tmp_path.joinpath(example_dir) shutil.copytree(example_dir, tmp_example_dir) mlflow.set_tracking_uri(f"sqlite:///{random_tmp_path / 'mlruns.db'}") python_env_path = find_python_env_yaml(tmp_example_dir) replace_mlflow_with_dev_version(python_env_path) cli_run_list = [tmp_example_dir] + params invoke_cli_runner(cli.run, list(map(str, cli_run_list))) @pytest.mark.notrackingurimock @pytest.mark.parametrize( ("directory", "command"), [ ("docker", ["docker", "build", "-t", "mlflow-docker-example", "-f", "Dockerfile", "."]), ("keras", [sys.executable, "train.py"]), ( "lightgbm/lightgbm_native", [ sys.executable, "train.py", "--learning-rate", "0.2", "--colsample-bytree", "0.8", "--subsample", "0.9", ], ), ("lightgbm/lightgbm_sklearn", [sys.executable, "train.py"]), ("statsmodels", [sys.executable, "train.py", "--inverse-method", "qr"]), ("quickstart", [sys.executable, "mlflow_tracking.py"]), ("remote_store", [sys.executable, "remote_server.py"]), ( "xgboost/xgboost_native", [ sys.executable, "train.py", "--learning-rate", "0.2", "--colsample-bytree", "0.8", "--subsample", "0.9", ], ), ("xgboost/xgboost_sklearn", [sys.executable, "train.py"]), ("catboost", [sys.executable, "train.py"]), ("prophet", [sys.executable, "train.py"]), ("sklearn_autolog", [sys.executable, "linear_regression.py"]), ("sklearn_autolog", [sys.executable, "pipeline.py"]), ("sklearn_autolog", [sys.executable, "grid_search_cv.py"]), ("pyspark_ml_autologging", [sys.executable, "logistic_regression.py"]), ("pyspark_ml_autologging", [sys.executable, "one_vs_rest.py"]), ("pyspark_ml_autologging", [sys.executable, "pipeline.py"]), ("shap", [sys.executable, "regression.py"]), ("shap", [sys.executable, "binary_classification.py"]), ("shap", [sys.executable, "multiclass_classification.py"]), ("shap", [sys.executable, "explainer_logging.py"]), ("ray_serve", [sys.executable, "train_model.py"]), ("pip_requirements", [sys.executable, "pip_requirements.py"]), ("pmdarima", [sys.executable, "train.py"]), ("evaluation", [sys.executable, "evaluate_on_binary_classifier.py"]), ("evaluation", [sys.executable, "evaluate_on_multiclass_classifier.py"]), ("evaluation", [sys.executable, "evaluate_on_regressor.py"]), ("evaluation", [sys.executable, "evaluate_with_custom_metrics.py"]), ("evaluation", [sys.executable, "evaluate_with_custom_metrics_comprehensive.py"]), ("evaluation", [sys.executable, "evaluate_with_model_validation.py"]), ("spark_udf", [sys.executable, "spark_udf_datetime.py"]), ("pyfunc", [sys.executable, "train.py"]), ("tensorflow", [sys.executable, "train.py"]), ("transformers", [sys.executable, "conversational.py"]), ("transformers", [sys.executable, "load_components.py"]), ("transformers", [sys.executable, "simple.py"]), ("transformers", [sys.executable, "sentence_transformer.py"]), ("transformers", [sys.executable, "whisper.py"]), ("sentence_transformers", [sys.executable, "simple.py"]), ("tracing", [sys.executable, "fluent.py"]), ("tracing", [sys.executable, "client.py"]), ("llama_index", [sys.executable, "simple_index.py"]), ("llama_index", [sys.executable, "autolog.py"]), ], ) def test_command_example(directory, command): cwd_dir = Path(EXAMPLES_DIR, directory) assert os.environ.get("MLFLOW_HOME") is not None if directory == "transformers": # NB: Clearing the huggingface_hub cache is to lower the disk storage pressure for CI clear_hub_cache() process._exec_cmd(command, cwd=cwd_dir, env=os.environ)