wehub-resource-sync
129 lines
4.9 KiB
Python
129 lines
4.9 KiB
Python
import os
|
|
import sys
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
import sklearn.linear_model as logreg_module
|
|
from sklearn import datasets
|
|
|
|
import mlflow.sklearn
|
|
import mlflow.utils.model_utils as mlflow_model_utils
|
|
from mlflow.environment_variables import MLFLOW_RECORD_ENV_VARS_IN_MODEL_LOGGING
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models import Model
|
|
from mlflow.utils.file_utils import TempDir
|
|
from mlflow.utils.model_utils import env_var_tracker
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def sklearn_knn_model():
|
|
iris = datasets.load_iris()
|
|
X = iris.data[:, :2] # we only take the first two features.
|
|
y = iris.target
|
|
logreg_model = logreg_module.LogisticRegression()
|
|
logreg_model.fit(X, y)
|
|
return logreg_model
|
|
|
|
|
|
@pytest.fixture
|
|
def model_path(tmp_path):
|
|
return os.path.join(tmp_path, "model")
|
|
|
|
|
|
def test_get_flavor_configuration_throws_exception_when_requested_flavor_is_missing(
|
|
model_path, sklearn_knn_model
|
|
):
|
|
mlflow.sklearn.save_model(sk_model=sklearn_knn_model, path=model_path)
|
|
|
|
# The saved model contains the "sklearn" flavor, so this call should succeed
|
|
sklearn_flavor_config = mlflow_model_utils._get_flavor_configuration(
|
|
model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME
|
|
)
|
|
assert sklearn_flavor_config is not None
|
|
|
|
|
|
def test_get_flavor_configuration_with_present_flavor_returns_expected_configuration(
|
|
sklearn_knn_model, model_path
|
|
):
|
|
mlflow.sklearn.save_model(sk_model=sklearn_knn_model, path=model_path)
|
|
|
|
sklearn_flavor_config = mlflow_model_utils._get_flavor_configuration(
|
|
model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME
|
|
)
|
|
model_config = Model.load(os.path.join(model_path, "MLmodel"))
|
|
assert sklearn_flavor_config == model_config.flavors[mlflow.sklearn.FLAVOR_NAME]
|
|
|
|
|
|
def test_add_code_to_system_path(sklearn_knn_model, model_path):
|
|
mlflow.sklearn.save_model(
|
|
sk_model=sklearn_knn_model,
|
|
path=model_path,
|
|
code_paths=[
|
|
"tests/utils/test_resources/dummy_module.py",
|
|
"tests/utils/test_resources/dummy_package",
|
|
],
|
|
)
|
|
|
|
sklearn_flavor_config = mlflow_model_utils._get_flavor_configuration(
|
|
model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME
|
|
)
|
|
with TempDir(chdr=True):
|
|
# Load the model from a new directory that is not a parent of the source code path to
|
|
# verify that source code paths and their subdirectories are resolved correctly
|
|
with pytest.raises(ModuleNotFoundError, match="No module named 'dummy_module'"):
|
|
import dummy_module
|
|
|
|
mlflow_model_utils._add_code_from_conf_to_system_path(model_path, sklearn_flavor_config)
|
|
import dummy_module # noqa: F401
|
|
|
|
# If this raises an exception it's because dummy_package.test imported
|
|
# dummy_package.operator and not the built-in operator module. This only
|
|
# happens if MLflow is misconfiguring the system path.
|
|
from dummy_package import base # noqa: F401
|
|
|
|
# Ensure that the custom tests/utils/test_resources/dummy_package/pandas.py is not
|
|
# overwriting the 3rd party `pandas` package
|
|
assert "dummy_package" in sys.modules
|
|
assert "pandas" in sys.modules
|
|
assert "site-packages" in sys.modules["pandas"].__file__
|
|
|
|
|
|
def test_add_code_to_system_path_not_copyable_file(sklearn_knn_model, model_path):
|
|
with mock.patch("builtins.open", side_effect=OSError("[Errno 95] Operation not supported")):
|
|
with pytest.raises(MlflowException, match=r"Failed to copy the specified code path"):
|
|
mlflow.sklearn.save_model(
|
|
sk_model=sklearn_knn_model,
|
|
path=model_path,
|
|
code_paths=["tests/utils/test_resources/dummy_module.py"],
|
|
)
|
|
|
|
|
|
def test_env_var_tracker(monkeypatch):
|
|
monkeypatch.setenv("DATABRICKS_HOST", "host")
|
|
assert "DATABRICKS_HOST" in os.environ
|
|
assert "TEST_API_KEY" not in os.environ
|
|
|
|
with env_var_tracker() as tracked_env_names:
|
|
assert os.environ["DATABRICKS_HOST"] == "host"
|
|
monkeypatch.setenv("TEST_API_KEY", "key")
|
|
# accessed env var is tracked
|
|
assert os.environ.get("TEST_API_KEY") == "key"
|
|
# test non-existing env vars fetched by `get` are not tracked
|
|
os.environ.get("INVALID_API_KEY", "abc")
|
|
# test non-existing env vars are not tracked
|
|
try:
|
|
os.environ["ANOTHER_API_KEY"]
|
|
except KeyError:
|
|
pass
|
|
assert all(x in tracked_env_names for x in ["DATABRICKS_HOST", "TEST_API_KEY"])
|
|
assert all(x not in tracked_env_names for x in ["INVALID_API_KEY", "ANOTHER_API_KEY"])
|
|
|
|
assert isinstance(os.environ, os._Environ)
|
|
assert all(x in os.environ for x in ["DATABRICKS_HOST", "TEST_API_KEY"])
|
|
assert all(x not in os.environ for x in ["INVALID_API_KEY", "ANOTHER_API_KEY"])
|
|
|
|
monkeypatch.setenv(MLFLOW_RECORD_ENV_VARS_IN_MODEL_LOGGING.name, "false")
|
|
with env_var_tracker() as env:
|
|
os.environ.get("API_KEY")
|
|
assert env == set()
|