wehub-resource-sync
388 lines
15 KiB
Python
388 lines
15 KiB
Python
import os
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
import boto3
|
|
import flavor
|
|
import moto
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
from botocore.config import Config
|
|
from sktime.datasets import load_airline, load_longley
|
|
from sktime.datatypes import convert
|
|
from sktime.forecasting.arima import AutoARIMA
|
|
from sktime.forecasting.model_selection import temporal_train_test_split
|
|
from sktime.forecasting.naive import NaiveForecaster
|
|
|
|
import mlflow
|
|
from mlflow import pyfunc
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models import Model, infer_signature
|
|
from mlflow.models.utils import _read_example
|
|
from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
|
|
from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
|
|
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
|
|
from mlflow.utils.environment import _mlflow_conda_env
|
|
|
|
FH = [1, 2, 3]
|
|
COVERAGE = [0.1, 0.5, 0.9]
|
|
ALPHA = [0.1, 0.5, 0.9]
|
|
COV = False
|
|
|
|
|
|
@pytest.fixture
|
|
def model_path(tmp_path):
|
|
"""Create a temporary path to save/log model."""
|
|
return tmp_path.joinpath("model")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def mock_s3_bucket():
|
|
"""Create a mock S3 bucket using moto.
|
|
|
|
Returns
|
|
-------
|
|
string with name of mock S3 bucket
|
|
"""
|
|
with moto.mock_s3():
|
|
bucket_name = "mock-bucket"
|
|
my_config = Config(region_name="us-east-1")
|
|
s3_client = boto3.client("s3", config=my_config)
|
|
s3_client.create_bucket(Bucket=bucket_name)
|
|
yield bucket_name
|
|
|
|
|
|
@pytest.fixture
|
|
def sktime_custom_env(tmp_path):
|
|
"""Create a conda environment and returns path to conda environment yml file."""
|
|
conda_env = tmp_path.joinpath("conda_env.yml")
|
|
_mlflow_conda_env(conda_env, additional_pip_deps=["sktime"])
|
|
return conda_env
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def data_airline():
|
|
"""Create sample data for univariate model without exogenous regressor."""
|
|
return load_airline()
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def data_longley():
|
|
"""Create sample data for univariate model with exogenous regressor."""
|
|
y, X = load_longley()
|
|
y_train, y_test, X_train, X_test = temporal_train_test_split(y, X)
|
|
return y_train, y_test, X_train, X_test
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def auto_arima_model(data_airline):
|
|
"""Create instance of fitted auto arima model."""
|
|
return AutoARIMA(sp=12, d=0, max_p=2, max_q=2, suppress_warnings=True).fit(data_airline)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def naive_forecaster_model_with_regressor(data_longley):
|
|
"""Create instance of fitted naive forecaster model."""
|
|
y_train, _, X_train, _ = data_longley
|
|
model = NaiveForecaster()
|
|
return model.fit(y_train, X_train)
|
|
|
|
|
|
@pytest.mark.parametrize("serialization_format", ["pickle", "cloudpickle"])
|
|
def test_auto_arima_model_save_and_load(auto_arima_model, model_path, serialization_format):
|
|
flavor.save_model(
|
|
sktime_model=auto_arima_model,
|
|
path=model_path,
|
|
serialization_format=serialization_format,
|
|
)
|
|
loaded_model = flavor.load_model(
|
|
model_uri=model_path,
|
|
)
|
|
|
|
np.testing.assert_array_equal(auto_arima_model.predict(fh=FH), loaded_model.predict(fh=FH))
|
|
|
|
|
|
@pytest.mark.parametrize("serialization_format", ["pickle", "cloudpickle"])
|
|
def test_auto_arima_model_pyfunc_output(auto_arima_model, model_path, serialization_format):
|
|
flavor.save_model(
|
|
sktime_model=auto_arima_model,
|
|
path=model_path,
|
|
serialization_format=serialization_format,
|
|
)
|
|
loaded_pyfunc = flavor.pyfunc.load_model(model_uri=model_path)
|
|
|
|
model_predict = auto_arima_model.predict(fh=FH)
|
|
predict_conf = pd.DataFrame([{"fh": FH, "predict_method": "predict"}])
|
|
pyfunc_predict = loaded_pyfunc.predict(predict_conf)
|
|
np.testing.assert_array_equal(model_predict, pyfunc_predict)
|
|
|
|
model_predict_interval = auto_arima_model.predict_interval(fh=FH, coverage=COVERAGE)
|
|
predict_interval_conf = pd.DataFrame([
|
|
{
|
|
"fh": FH,
|
|
"predict_method": "predict_interval",
|
|
"coverage": COVERAGE,
|
|
}
|
|
])
|
|
pyfunc_predict_interval = loaded_pyfunc.predict(predict_interval_conf)
|
|
np.testing.assert_array_equal(model_predict_interval.values, pyfunc_predict_interval.values)
|
|
|
|
model_predict_quantiles = auto_arima_model.predict_quantiles(fh=FH, alpha=ALPHA)
|
|
predict_quantiles_conf = pd.DataFrame([
|
|
{
|
|
"fh": FH,
|
|
"predict_method": "predict_quantiles",
|
|
"alpha": ALPHA,
|
|
}
|
|
])
|
|
pyfunc_predict_quantiles = loaded_pyfunc.predict(predict_quantiles_conf)
|
|
np.testing.assert_array_equal(model_predict_quantiles.values, pyfunc_predict_quantiles.values)
|
|
|
|
model_predict_var = auto_arima_model.predict_var(fh=FH, cov=COV)
|
|
predict_var_conf = pd.DataFrame([{"fh": FH, "predict_method": "predict_var", "cov": COV}])
|
|
pyfunc_predict_var = loaded_pyfunc.predict(predict_var_conf)
|
|
np.testing.assert_array_equal(model_predict_var.values, pyfunc_predict_var.values)
|
|
|
|
|
|
def test_naive_forecaster_model_with_regressor_pyfunc_output(
|
|
naive_forecaster_model_with_regressor, model_path, data_longley
|
|
):
|
|
_, _, _, X_test = data_longley
|
|
|
|
flavor.save_model(sktime_model=naive_forecaster_model_with_regressor, path=model_path)
|
|
loaded_pyfunc = flavor.pyfunc.load_model(model_uri=model_path)
|
|
|
|
X_test_array = convert(X_test, "pd.DataFrame", "np.ndarray")
|
|
|
|
model_predict = naive_forecaster_model_with_regressor.predict(fh=FH, X=X_test)
|
|
predict_conf = pd.DataFrame([{"fh": FH, "predict_method": "predict", "X": X_test_array}])
|
|
pyfunc_predict = loaded_pyfunc.predict(predict_conf)
|
|
np.testing.assert_array_equal(model_predict, pyfunc_predict)
|
|
|
|
model_predict_interval = naive_forecaster_model_with_regressor.predict_interval(
|
|
fh=FH, coverage=COVERAGE, X=X_test
|
|
)
|
|
predict_interval_conf = pd.DataFrame([
|
|
{
|
|
"fh": FH,
|
|
"predict_method": "predict_interval",
|
|
"coverage": COVERAGE,
|
|
"X": X_test_array,
|
|
}
|
|
])
|
|
pyfunc_predict_interval = loaded_pyfunc.predict(predict_interval_conf)
|
|
np.testing.assert_array_equal(model_predict_interval.values, pyfunc_predict_interval.values)
|
|
|
|
model_predict_quantiles = naive_forecaster_model_with_regressor.predict_quantiles(
|
|
fh=FH, alpha=ALPHA, X=X_test
|
|
)
|
|
predict_quantiles_conf = pd.DataFrame([
|
|
{
|
|
"fh": FH,
|
|
"predict_method": "predict_quantiles",
|
|
"alpha": ALPHA,
|
|
"X": X_test_array,
|
|
}
|
|
])
|
|
pyfunc_predict_quantiles = loaded_pyfunc.predict(predict_quantiles_conf)
|
|
np.testing.assert_array_equal(model_predict_quantiles.values, pyfunc_predict_quantiles.values)
|
|
|
|
model_predict_var = naive_forecaster_model_with_regressor.predict_var(fh=FH, cov=COV, X=X_test)
|
|
predict_var_conf = pd.DataFrame([
|
|
{
|
|
"fh": FH,
|
|
"predict_method": "predict_var",
|
|
"cov": COV,
|
|
"X": X_test_array,
|
|
}
|
|
])
|
|
pyfunc_predict_var = loaded_pyfunc.predict(predict_var_conf)
|
|
np.testing.assert_array_equal(model_predict_var.values, pyfunc_predict_var.values)
|
|
|
|
|
|
@pytest.mark.parametrize("use_signature", [True, False])
|
|
@pytest.mark.parametrize("use_example", [True, False])
|
|
def test_signature_and_examples_saved_correctly(
|
|
auto_arima_model, data_airline, model_path, use_signature, use_example
|
|
):
|
|
# Note: Signature inference fails on native model predict_interval/predict_quantiles
|
|
prediction = auto_arima_model.predict(fh=FH)
|
|
signature = infer_signature(data_airline, prediction) if use_signature else None
|
|
example = pd.DataFrame(data_airline[0:5].copy(deep=False)) if use_example else None
|
|
flavor.save_model(auto_arima_model, path=model_path, signature=signature, input_example=example)
|
|
mlflow_model = Model.load(model_path)
|
|
assert signature == mlflow_model.signature
|
|
if example is None:
|
|
assert mlflow_model.saved_input_example_info is None
|
|
else:
|
|
r_example = _read_example(mlflow_model, model_path).copy(deep=False)
|
|
np.testing.assert_array_equal(r_example, example)
|
|
|
|
|
|
@pytest.mark.parametrize("use_signature", [True, False])
|
|
def test_predict_var_signature_saved_correctly(
|
|
auto_arima_model, data_airline, model_path, use_signature
|
|
):
|
|
prediction = auto_arima_model.predict_var(fh=FH)
|
|
signature = infer_signature(data_airline, prediction) if use_signature else None
|
|
flavor.save_model(auto_arima_model, path=model_path, signature=signature)
|
|
mlflow_model = Model.load(model_path)
|
|
assert signature == mlflow_model.signature
|
|
|
|
|
|
@pytest.mark.parametrize("use_signature", [True, False])
|
|
@pytest.mark.parametrize("use_example", [True, False])
|
|
def test_signature_and_example_for_pyfunc_predict_interval(
|
|
auto_arima_model, model_path, data_airline, use_signature, use_example
|
|
):
|
|
model_path_primary = model_path.joinpath("primary")
|
|
model_path_secondary = model_path.joinpath("secondary")
|
|
flavor.save_model(sktime_model=auto_arima_model, path=model_path_primary)
|
|
loaded_pyfunc = flavor.pyfunc.load_model(model_uri=model_path_primary)
|
|
predict_conf = pd.DataFrame([
|
|
{
|
|
"fh": FH,
|
|
"predict_method": "predict_interval",
|
|
"coverage": COVERAGE,
|
|
}
|
|
])
|
|
forecast = loaded_pyfunc.predict(predict_conf)
|
|
signature = infer_signature(data_airline, forecast) if use_signature else None
|
|
example = pd.DataFrame(data_airline[0:5].copy(deep=False)) if use_example else None
|
|
flavor.save_model(
|
|
auto_arima_model,
|
|
path=model_path_secondary,
|
|
signature=signature,
|
|
input_example=example,
|
|
)
|
|
mlflow_model = Model.load(model_path_secondary)
|
|
assert signature == mlflow_model.signature
|
|
if example is None:
|
|
assert mlflow_model.saved_input_example_info is None
|
|
else:
|
|
r_example = _read_example(mlflow_model, model_path_secondary).copy(deep=False)
|
|
np.testing.assert_array_equal(r_example, example)
|
|
|
|
|
|
@pytest.mark.parametrize("use_signature", [True, False])
|
|
def test_signature_for_pyfunc_predict_quantiles(
|
|
auto_arima_model, model_path, data_airline, use_signature
|
|
):
|
|
model_path_primary = model_path.joinpath("primary")
|
|
model_path_secondary = model_path.joinpath("secondary")
|
|
flavor.save_model(sktime_model=auto_arima_model, path=model_path_primary)
|
|
loaded_pyfunc = flavor.pyfunc.load_model(model_uri=model_path_primary)
|
|
predict_conf = pd.DataFrame([
|
|
{
|
|
"fh": FH,
|
|
"predict_method": "predict_quantiles",
|
|
"alpha": ALPHA,
|
|
}
|
|
])
|
|
forecast = loaded_pyfunc.predict(predict_conf)
|
|
signature = infer_signature(data_airline, forecast) if use_signature else None
|
|
flavor.save_model(auto_arima_model, path=model_path_secondary, signature=signature)
|
|
mlflow_model = Model.load(model_path_secondary)
|
|
assert signature == mlflow_model.signature
|
|
|
|
|
|
def test_load_from_remote_uri_succeeds(auto_arima_model, model_path, mock_s3_bucket):
|
|
flavor.save_model(sktime_model=auto_arima_model, path=model_path)
|
|
|
|
artifact_root = f"s3://{mock_s3_bucket}"
|
|
artifact_path = "model"
|
|
artifact_repo = S3ArtifactRepository(artifact_root)
|
|
artifact_repo.log_artifacts(model_path, artifact_path=artifact_path)
|
|
|
|
model_uri = os.path.join(artifact_root, artifact_path)
|
|
reloaded_sktime_model = flavor.load_model(model_uri=model_uri)
|
|
|
|
np.testing.assert_array_equal(
|
|
auto_arima_model.predict(fh=FH),
|
|
reloaded_sktime_model.predict(fh=FH),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("should_start_run", [True, False])
|
|
@pytest.mark.parametrize("serialization_format", ["pickle", "cloudpickle"])
|
|
def test_log_model(auto_arima_model, tmp_path, should_start_run, serialization_format):
|
|
try:
|
|
if should_start_run:
|
|
mlflow.start_run()
|
|
artifact_path = "sktime"
|
|
conda_env = tmp_path.joinpath("conda_env.yaml")
|
|
_mlflow_conda_env(conda_env, additional_pip_deps=["sktime"])
|
|
model_info = flavor.log_model(
|
|
sktime_model=auto_arima_model,
|
|
artifact_path=artifact_path,
|
|
conda_env=str(conda_env),
|
|
serialization_format=serialization_format,
|
|
)
|
|
model_uri = f"runs:/{mlflow.active_run().info.run_id}/{artifact_path}"
|
|
assert model_info.model_uri == model_uri
|
|
reloaded_model = flavor.load_model(
|
|
model_uri=model_uri,
|
|
)
|
|
np.testing.assert_array_equal(auto_arima_model.predict(), reloaded_model.predict())
|
|
model_path = Path(_download_artifact_from_uri(artifact_uri=model_uri))
|
|
model_config = Model.load(str(model_path.joinpath("MLmodel")))
|
|
assert pyfunc.FLAVOR_NAME in model_config.flavors
|
|
finally:
|
|
mlflow.end_run()
|
|
|
|
|
|
def test_log_model_calls_register_model(auto_arima_model, tmp_path):
|
|
artifact_path = "sktime"
|
|
register_model_patch = mock.patch("mlflow.register_model")
|
|
with mlflow.start_run(), register_model_patch:
|
|
conda_env = tmp_path.joinpath("conda_env.yaml")
|
|
_mlflow_conda_env(conda_env, additional_pip_deps=["sktime"])
|
|
flavor.log_model(
|
|
sktime_model=auto_arima_model,
|
|
artifact_path=artifact_path,
|
|
conda_env=str(conda_env),
|
|
registered_model_name="SktimeModel",
|
|
)
|
|
model_uri = f"runs:/{mlflow.active_run().info.run_id}/{artifact_path}"
|
|
mlflow.register_model.assert_called_once_with(
|
|
model_uri,
|
|
"SktimeModel",
|
|
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
|
|
)
|
|
|
|
|
|
def test_log_model_no_registered_model_name(auto_arima_model, tmp_path):
|
|
artifact_path = "sktime"
|
|
register_model_patch = mock.patch("mlflow.register_model")
|
|
with mlflow.start_run(), register_model_patch:
|
|
conda_env = tmp_path.joinpath("conda_env.yaml")
|
|
_mlflow_conda_env(conda_env, additional_pip_deps=["sktime"])
|
|
flavor.log_model(
|
|
sktime_model=auto_arima_model,
|
|
artifact_path=artifact_path,
|
|
conda_env=str(conda_env),
|
|
)
|
|
mlflow.register_model.assert_not_called()
|
|
|
|
|
|
def test_sktime_pyfunc_raises_invalid_df_input(auto_arima_model, model_path):
|
|
flavor.save_model(sktime_model=auto_arima_model, path=model_path)
|
|
loaded_pyfunc = flavor.pyfunc.load_model(model_uri=model_path)
|
|
|
|
with pytest.raises(MlflowException, match="The provided prediction pd.DataFrame "):
|
|
loaded_pyfunc.predict(pd.DataFrame([{"predict_method": "predict"}, {"fh": FH}]))
|
|
|
|
with pytest.raises(MlflowException, match="The provided prediction configuration "):
|
|
loaded_pyfunc.predict(pd.DataFrame([{"invalid": True}]))
|
|
|
|
with pytest.raises(MlflowException, match="Invalid `predict_method` value"):
|
|
loaded_pyfunc.predict(pd.DataFrame([{"predict_method": "predict_proba"}]))
|
|
|
|
|
|
def test_sktime_save_model_raises_invalid_serialization_format(auto_arima_model, model_path):
|
|
with pytest.raises(MlflowException, match="Unrecognized serialization format: "):
|
|
flavor.save_model(
|
|
sktime_model=auto_arima_model, path=model_path, serialization_format="json"
|
|
)
|