wehub-resource-sync
496 lines
17 KiB
Python
496 lines
17 KiB
Python
import json
|
|
from unittest import mock
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
import shap
|
|
import sklearn
|
|
from numba import njit
|
|
from packaging.version import Version
|
|
from sklearn.datasets import load_diabetes
|
|
|
|
import mlflow
|
|
import mlflow.pyfunc.scoring_server as pyfunc_scoring_server
|
|
from mlflow import MlflowClient
|
|
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
|
|
from mlflow.utils import PYTHON_VERSION
|
|
from mlflow.utils.model_utils import _get_flavor_configuration
|
|
|
|
from tests.helper_functions import (
|
|
_assert_pip_requirements,
|
|
_compare_logged_code_paths,
|
|
_mlflow_major_version_string,
|
|
pyfunc_serve_and_score_model,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def shap_model():
|
|
X, y = load_diabetes(return_X_y=True, as_frame=True)
|
|
model = sklearn.ensemble.RandomForestRegressor(n_estimators=100)
|
|
model.fit(X, y)
|
|
return shap.Explainer(model.predict, X, algorithm="permutation")
|
|
|
|
|
|
def get_test_dataset():
|
|
X, y = load_diabetes(as_frame=True, return_X_y=True)
|
|
return X, y
|
|
|
|
|
|
def test_sklearn_log_explainer():
|
|
"""
|
|
Tests mlflow.shap log_explainer with mlflow serialization of the underlying model
|
|
"""
|
|
|
|
with mlflow.start_run() as run:
|
|
run_id = run.info.run_id
|
|
|
|
X, y = get_test_dataset()
|
|
|
|
model = sklearn.ensemble.RandomForestRegressor(n_estimators=100)
|
|
model.fit(X, y)
|
|
|
|
explainer_original = shap.Explainer(model.predict, X, algorithm="permutation")
|
|
shap_values_original = explainer_original(X[:5])
|
|
|
|
mlflow.shap.log_explainer(explainer_original, "test_explainer")
|
|
|
|
explainer_uri = "runs:/" + run_id + "/test_explainer"
|
|
|
|
explainer_loaded = mlflow.shap.load_explainer(explainer_uri)
|
|
shap_values_new = explainer_loaded(X[:5])
|
|
|
|
explainer_path = _download_artifact_from_uri(artifact_uri=explainer_uri)
|
|
flavor_conf = _get_flavor_configuration(
|
|
model_path=explainer_path, flavor_name=mlflow.shap.FLAVOR_NAME
|
|
)
|
|
underlying_model_flavor = flavor_conf["underlying_model_flavor"]
|
|
|
|
assert underlying_model_flavor == mlflow.sklearn.FLAVOR_NAME
|
|
np.testing.assert_array_equal(shap_values_original.base_values, shap_values_new.base_values)
|
|
np.testing.assert_allclose(
|
|
shap_values_original.values, shap_values_new.values, rtol=100, atol=100
|
|
)
|
|
|
|
|
|
def test_sklearn_log_explainer_self_serialization():
|
|
"""
|
|
Tests mlflow.shap log_explainer with SHAP internal serialization of the underlying model
|
|
"""
|
|
|
|
with mlflow.start_run() as run:
|
|
run_id = run.info.run_id
|
|
|
|
X, y = get_test_dataset()
|
|
|
|
model = sklearn.ensemble.RandomForestRegressor(n_estimators=100)
|
|
model.fit(X, y)
|
|
|
|
explainer_original = shap.Explainer(model.predict, X, algorithm="permutation")
|
|
shap_values_original = explainer_original(X[:5])
|
|
|
|
mlflow.shap.log_explainer(
|
|
explainer_original, "test_explainer", serialize_model_using_mlflow=False
|
|
)
|
|
|
|
explainer_uri = "runs:/" + run_id + "/test_explainer"
|
|
|
|
explainer_loaded = mlflow.shap.load_explainer("runs:/" + run_id + "/test_explainer")
|
|
shap_values_new = explainer_loaded(X[:5])
|
|
|
|
explainer_path = _download_artifact_from_uri(artifact_uri=explainer_uri)
|
|
flavor_conf = _get_flavor_configuration(
|
|
model_path=explainer_path, flavor_name=mlflow.shap.FLAVOR_NAME
|
|
)
|
|
underlying_model_flavor = flavor_conf["underlying_model_flavor"]
|
|
|
|
assert underlying_model_flavor is None
|
|
np.testing.assert_array_equal(shap_values_original.base_values, shap_values_new.base_values)
|
|
np.testing.assert_allclose(
|
|
shap_values_original.values, shap_values_new.values, rtol=100, atol=100
|
|
)
|
|
|
|
|
|
def test_sklearn_log_explainer_pyfunc():
|
|
"""
|
|
Tests mlflow.shap log_explainer with mlflow
|
|
serialization of the underlying model using pyfunc flavor
|
|
"""
|
|
|
|
with mlflow.start_run() as run:
|
|
run_id = run.info.run_id
|
|
|
|
X, y = get_test_dataset()
|
|
|
|
model = sklearn.ensemble.RandomForestRegressor(n_estimators=100)
|
|
model.fit(X, y)
|
|
|
|
explainer_original = shap.Explainer(model.predict, X, algorithm="permutation")
|
|
shap_values_original = explainer_original(X[:2])
|
|
|
|
mlflow.shap.log_explainer(explainer_original, "test_explainer")
|
|
|
|
explainer_pyfunc = mlflow.pyfunc.load_model("runs:/" + run_id + "/test_explainer")
|
|
shap_values_new = explainer_pyfunc.predict(X[:2])
|
|
|
|
np.testing.assert_allclose(shap_values_original.values, shap_values_new, rtol=100, atol=100)
|
|
|
|
|
|
def test_log_explanation_doesnt_create_autologged_run():
|
|
try:
|
|
mlflow.sklearn.autolog(disable=False, exclusive=False)
|
|
X, y = sklearn.datasets.load_diabetes(return_X_y=True, as_frame=True)
|
|
X = X.iloc[:50, :4]
|
|
y = y.iloc[:50]
|
|
model = sklearn.linear_model.LinearRegression()
|
|
model.fit(X, y)
|
|
|
|
with mlflow.start_run() as run:
|
|
mlflow.shap.log_explanation(model.predict, X)
|
|
|
|
run_data = MlflowClient().get_run(run.info.run_id).data
|
|
metrics = run_data.metrics
|
|
params = run_data.params
|
|
tags = run_data.tags
|
|
assert not metrics
|
|
assert not params
|
|
assert all("mlflow." in key for key in tags)
|
|
assert "mlflow.autologging" not in tags
|
|
finally:
|
|
mlflow.sklearn.autolog(disable=True)
|
|
|
|
|
|
def test_load_pyfunc(tmp_path):
|
|
X, y = get_test_dataset()
|
|
|
|
model = sklearn.ensemble.RandomForestRegressor(n_estimators=100)
|
|
model.fit(X, y)
|
|
|
|
explainer_original = shap.Explainer(model.predict, X, algorithm="permutation")
|
|
shap_values_original = explainer_original(X[:2])
|
|
path = str(tmp_path.joinpath("pyfunc_test"))
|
|
mlflow.shap.save_explainer(explainer_original, path)
|
|
|
|
explainer_pyfunc = mlflow.shap._load_pyfunc(path)
|
|
shap_values_new = explainer_pyfunc.predict(X[:2])
|
|
|
|
np.testing.assert_allclose(shap_values_original.values, shap_values_new, rtol=100, atol=100)
|
|
|
|
|
|
def test_merge_environment():
|
|
expected_mlflow_version = _mlflow_major_version_string()
|
|
test_shap_env = {
|
|
"channels": ["conda-forge"],
|
|
"dependencies": ["python=3.8.5", "pip", {"pip": [expected_mlflow_version, "shap==0.38.0"]}],
|
|
}
|
|
|
|
test_model_env = {
|
|
"channels": ["conda-forge"],
|
|
"dependencies": [
|
|
"python=3.8.5",
|
|
"pip",
|
|
{"pip": [expected_mlflow_version, "scikit-learn==0.24.0", "cloudpickle==1.6.0"]},
|
|
],
|
|
}
|
|
|
|
expected_merged_env = {
|
|
"name": "mlflow-env",
|
|
"channels": ["conda-forge"],
|
|
"dependencies": [
|
|
f"python={PYTHON_VERSION}",
|
|
"pip",
|
|
{
|
|
"pip": [
|
|
expected_mlflow_version,
|
|
"scikit-learn==0.24.0",
|
|
"cloudpickle==1.6.0",
|
|
"shap==0.38.0",
|
|
]
|
|
},
|
|
],
|
|
}
|
|
|
|
actual_merged_env = mlflow.shap._merge_environments(test_shap_env, test_model_env)
|
|
|
|
assert sorted(expected_merged_env["channels"]) == sorted(actual_merged_env["channels"])
|
|
|
|
expected_conda_deps, expected_pip_deps = mlflow.shap._get_conda_and_pip_dependencies(
|
|
expected_merged_env
|
|
)
|
|
actual_conda_deps, actual_pip_deps = mlflow.shap._get_conda_and_pip_dependencies(
|
|
actual_merged_env
|
|
)
|
|
|
|
assert sorted(expected_pip_deps) == actual_pip_deps
|
|
assert sorted(expected_conda_deps) == actual_conda_deps
|
|
|
|
|
|
def test_merge_environment_with_duplicates():
|
|
expected_mlflow_version = _mlflow_major_version_string()
|
|
duplicate_dependency = "numpy==1.19.2"
|
|
|
|
# Introduce the duplicate in both environments
|
|
test_shap_env = {
|
|
"channels": ["conda-forge"],
|
|
"dependencies": [
|
|
"python=3.8.5",
|
|
"pip",
|
|
{"pip": [expected_mlflow_version, "shap==0.38.0", duplicate_dependency]},
|
|
],
|
|
}
|
|
|
|
test_model_env = {
|
|
"channels": ["conda-forge"],
|
|
"dependencies": [
|
|
"python=3.8.5",
|
|
"pip",
|
|
{
|
|
"pip": [
|
|
expected_mlflow_version,
|
|
"scikit-learn==0.24.0",
|
|
"cloudpickle==1.6.0",
|
|
duplicate_dependency,
|
|
]
|
|
},
|
|
],
|
|
}
|
|
|
|
# The expected merged environment should not have duplicates
|
|
expected_merged_env = {
|
|
"name": "mlflow-env",
|
|
"channels": ["conda-forge"],
|
|
"dependencies": [
|
|
f"python={PYTHON_VERSION}",
|
|
"pip",
|
|
{
|
|
"pip": [
|
|
expected_mlflow_version,
|
|
"scikit-learn==0.24.0",
|
|
"cloudpickle==1.6.0",
|
|
"shap==0.38.0",
|
|
duplicate_dependency,
|
|
]
|
|
},
|
|
],
|
|
}
|
|
|
|
actual_merged_env = mlflow.shap._merge_environments(test_shap_env, test_model_env)
|
|
|
|
assert sorted(expected_merged_env["channels"]) == sorted(actual_merged_env["channels"])
|
|
|
|
expected_conda_deps, expected_pip_deps = mlflow.shap._get_conda_and_pip_dependencies(
|
|
expected_merged_env
|
|
)
|
|
actual_conda_deps, actual_pip_deps = mlflow.shap._get_conda_and_pip_dependencies(
|
|
actual_merged_env
|
|
)
|
|
|
|
# Check that there are no duplicates in the actual pip dependencies
|
|
assert sorted(actual_pip_deps) == sorted(set(actual_pip_deps))
|
|
assert sorted(expected_pip_deps) == actual_pip_deps
|
|
assert sorted(expected_conda_deps) == actual_conda_deps
|
|
|
|
|
|
def test_log_model_with_pip_requirements(shap_model, tmp_path):
|
|
expected_mlflow_version = _mlflow_major_version_string()
|
|
sklearn_default_reqs = mlflow.sklearn.get_default_pip_requirements(include_cloudpickle=True)
|
|
# Path to a requirements file
|
|
req_file = tmp_path.joinpath("requirements.txt")
|
|
req_file.write_text("a")
|
|
with mlflow.start_run():
|
|
model_info = mlflow.shap.log_explainer(shap_model, "model", pip_requirements=str(req_file))
|
|
_assert_pip_requirements(
|
|
model_info.model_uri,
|
|
[expected_mlflow_version, "a", *sklearn_default_reqs],
|
|
strict=False,
|
|
)
|
|
|
|
# List of requirements
|
|
with mlflow.start_run():
|
|
model_info = mlflow.shap.log_explainer(
|
|
shap_model, "model", pip_requirements=[f"-r {req_file}", "b"]
|
|
)
|
|
_assert_pip_requirements(
|
|
model_info.model_uri,
|
|
[expected_mlflow_version, "a", "b", *sklearn_default_reqs],
|
|
strict=False,
|
|
)
|
|
|
|
# Constraints file
|
|
with mlflow.start_run():
|
|
model_info = mlflow.shap.log_explainer(
|
|
shap_model, "model", pip_requirements=[f"-c {req_file}", "b"]
|
|
)
|
|
_assert_pip_requirements(
|
|
model_info.model_uri,
|
|
[expected_mlflow_version, "b", "-c constraints.txt", *sklearn_default_reqs],
|
|
["a"],
|
|
strict=False,
|
|
)
|
|
|
|
|
|
def test_log_model_with_extra_pip_requirements(shap_model, tmp_path):
|
|
expected_mlflow_version = _mlflow_major_version_string()
|
|
shap_default_reqs = mlflow.shap.get_default_pip_requirements()
|
|
sklearn_default_reqs = mlflow.sklearn.get_default_pip_requirements(include_cloudpickle=True)
|
|
|
|
# Path to a requirements file
|
|
req_file = tmp_path.joinpath("requirements.txt")
|
|
req_file.write_text("a")
|
|
with mlflow.start_run():
|
|
log_info = mlflow.shap.log_explainer(
|
|
shap_model, "model", extra_pip_requirements=str(req_file)
|
|
)
|
|
_assert_pip_requirements(
|
|
log_info.model_uri,
|
|
[expected_mlflow_version, *shap_default_reqs, "a", *sklearn_default_reqs],
|
|
)
|
|
|
|
# List of requirements
|
|
with mlflow.start_run():
|
|
log_info = mlflow.shap.log_explainer(
|
|
shap_model, "model", extra_pip_requirements=[f"-r {req_file}", "b"]
|
|
)
|
|
_assert_pip_requirements(
|
|
log_info.model_uri,
|
|
[expected_mlflow_version, *shap_default_reqs, "a", "b", *sklearn_default_reqs],
|
|
)
|
|
|
|
# Constraints file
|
|
with mlflow.start_run():
|
|
log_info = mlflow.shap.log_explainer(
|
|
shap_model, "model", extra_pip_requirements=[f"-c {req_file}", "b"]
|
|
)
|
|
_assert_pip_requirements(
|
|
log_info.model_uri,
|
|
[
|
|
expected_mlflow_version,
|
|
*shap_default_reqs,
|
|
"b",
|
|
"-c constraints.txt",
|
|
*sklearn_default_reqs,
|
|
],
|
|
["a"],
|
|
)
|
|
|
|
|
|
def create_identity_function():
|
|
def identity(x):
|
|
return x
|
|
|
|
def _identity_inverse(x):
|
|
return x
|
|
|
|
identity.inverse = _identity_inverse
|
|
|
|
return identity
|
|
|
|
|
|
@pytest.mark.skipif(Version(shap.__version__) < Version("0.42.0"), reason="numba njit compatible")
|
|
def test_pyfunc_serve_and_score_njit():
|
|
# Create a numba-compatible identify link function due to breaking changes in shap
|
|
# version 0.42.0. Python functions can no longer be passed to the numba jit compiler
|
|
# with the changes introduced in that version.
|
|
@njit
|
|
def identity_function(x):
|
|
return x
|
|
|
|
X, y = get_test_dataset()
|
|
|
|
reg = sklearn.ensemble.RandomForestRegressor(n_estimators=10).fit(X, y)
|
|
model = shap.Explainer(
|
|
reg.predict,
|
|
masker=X,
|
|
algorithm="permutation",
|
|
# `link` defaults to `shap.links.identity` which is decorated by `numba.jit` and causes
|
|
# the following error when loading the explainer for serving:
|
|
# ```
|
|
# Exception: The passed link function needs to be callable and have a callable
|
|
# .inverse property!
|
|
# ```
|
|
# As a workaround, use an identify function that's NOT decorated by `numba.jit`.
|
|
link=identity_function,
|
|
)
|
|
artifact_path = "model"
|
|
with mlflow.start_run():
|
|
model_info = mlflow.shap.log_explainer(model, artifact_path)
|
|
|
|
resp = pyfunc_serve_and_score_model(
|
|
model_info.model_uri,
|
|
data=pd.DataFrame(X[:3]),
|
|
content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
|
|
)
|
|
decoded_json = json.loads(resp.content.decode("utf-8"))
|
|
scores = pd.DataFrame(data=decoded_json["predictions"]).values
|
|
np.testing.assert_allclose(scores, model(X[:3]).values, rtol=100, atol=100)
|
|
|
|
|
|
@pytest.mark.skipif(Version(shap.__version__) > Version("0.41.0"), reason="numba jit compatible")
|
|
def test_pyfunc_serve_and_score():
|
|
# Note: this implementation of an identify function is only compatible with versions of
|
|
# shap <= 0.41.0. A breaking change was introduced with how numba is used with shap in version
|
|
# 0.42.0.
|
|
X, y = get_test_dataset()
|
|
|
|
reg = sklearn.ensemble.RandomForestRegressor(n_estimators=10).fit(X, y)
|
|
model = shap.Explainer(
|
|
reg.predict,
|
|
masker=X,
|
|
algorithm="permutation",
|
|
# `link` defaults to `shap.links.identity` which is decorated by `numba.jit` and causes
|
|
# the following error when loading the explainer for serving:
|
|
# ```
|
|
# Exception: The passed link function needs to be callable and have a callable
|
|
# .inverse property!
|
|
# ```
|
|
# As a workaround, use an identify function that's NOT decorated by `numba.jit`.
|
|
link=create_identity_function(),
|
|
)
|
|
artifact_path = "model"
|
|
with mlflow.start_run():
|
|
model_info = mlflow.shap.log_explainer(model, artifact_path)
|
|
|
|
resp = pyfunc_serve_and_score_model(
|
|
model_info.model_uri,
|
|
data=pd.DataFrame(X[:3]),
|
|
content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
|
|
)
|
|
decoded_json = json.loads(resp.content.decode("utf-8"))
|
|
scores = pd.DataFrame(data=decoded_json["predictions"]).values
|
|
np.testing.assert_allclose(scores, model(X[:3]).values, rtol=100, atol=100)
|
|
|
|
|
|
def test_log_model_with_code_paths(shap_model):
|
|
artifact_path = "model"
|
|
with (
|
|
mlflow.start_run(),
|
|
mock.patch("mlflow.shap._add_code_from_conf_to_system_path") as add_mock,
|
|
):
|
|
model_info = mlflow.shap.log_explainer(shap_model, artifact_path, code_paths=[__file__])
|
|
_compare_logged_code_paths(__file__, model_info.model_uri, mlflow.shap.FLAVOR_NAME)
|
|
mlflow.shap.load_explainer(model_info.model_uri)
|
|
add_mock.assert_called()
|
|
|
|
|
|
def test_model_save_load_with_metadata(shap_model, tmp_path):
|
|
model_path = str(tmp_path.joinpath("pyfunc_test"))
|
|
mlflow.shap.save_explainer(
|
|
shap_model, path=model_path, metadata={"metadata_key": "metadata_value"}
|
|
)
|
|
|
|
reloaded_model = mlflow.pyfunc.load_model(model_uri=model_path)
|
|
assert reloaded_model.metadata.metadata["metadata_key"] == "metadata_value"
|
|
|
|
|
|
def test_model_log_with_metadata(shap_model):
|
|
artifact_path = "model"
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.shap.log_explainer(
|
|
shap_model, artifact_path=artifact_path, metadata={"metadata_key": "metadata_value"}
|
|
)
|
|
|
|
reloaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
|
|
assert reloaded_model.metadata.metadata["metadata_key"] == "metadata_value"
|