chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,183 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user