def test_signature_and_examples_are_saved_correctly(ols_model):
model = ols_model.model
X = ols_model.inference_dataframe
signature_ = infer_signature(X)
example_ = X[0:3, :]
for signature in (None, signature_):
for example in (None, example_):
with TempDir() as tmp:
path = tmp.path("model")
mlflow.statsmodels.save_model(
model, path=path, signature=signature, input_example=example
)
mlflow_model = Model.load(path)
assert signature == mlflow_model.signature
if example is None:
assert mlflow_model.saved_input_example_info is None
else:
assert np.array_equal(_read_example(mlflow_model, path), example)
def test_model_export_with_schema_and_examples(spacy_model_with_data):
spacy_model = spacy_model_with_data.model
signature_ = infer_signature(spacy_model_with_data.inference_data)
example_ = spacy_model_with_data.inference_data.head(3)
for signature in (None, signature_):
for example in (None, example_):
with TempDir() as tmp:
path = tmp.path("model")
mlflow.spacy.save_model(spacy_model,
path=path,
signature=signature,
input_example=example)
mlflow_model = Model.load(path)
assert signature == mlflow_model.signature
if example is None:
assert mlflow_model.saved_input_example_info is None
else:
assert all((_read_example(mlflow_model,
path) == example).all())
def test_signature_and_examples_are_saved_correctly(fastai_model):
model = fastai_model.model
signature_ = infer_signature(fastai_model.inference_dataframe)
example_ = fastai_model.inference_dataframe.head(3)
for signature in (None, signature_):
for example in (None, example_):
with TempDir() as tmp:
path = tmp.path("model")
mlflow.fastai.save_model(model,
path=path,
signature=signature,
input_example=example)
mlflow_model = Model.load(path)
assert signature == mlflow_model.signature
if example is None:
assert mlflow_model.saved_input_example_info is None
else:
assert all((_read_example(mlflow_model,
path) == example).all())
def test_diviner_signature_and_examples_saved_correctly(
grouped_prophet, diviner_data, model_path, use_signature, use_example
):
prediction = grouped_prophet.forecast(horizon=20, frequency="D")
signature = infer_signature(diviner_data.df, prediction) if use_signature else None
example = diviner_data.df[0:5].copy(deep=False) if use_example else None
mlflow.diviner.save_model(
grouped_prophet, 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)
# NB: datetime values are implicitly cast, so this needs to be reverted.
r_example["ds"] = pd.to_datetime(r_example["ds"], format=DS_FORMAT)
np.testing.assert_array_equal(r_example, example)
def test_autolog_logs_signature_and_input_example(data_type):
mlflow.sklearn.autolog(log_input_examples=True, log_model_signatures=True)
X, y = get_iris()
X = data_type(X)
y = data_type(y)
model = sklearn.linear_model.LinearRegression()
with mlflow.start_run() as run:
model.fit(X, y)
model_path = os.path.join(run.info.artifact_uri, MODEL_DIR)
model_conf = get_model_conf(run.info.artifact_uri)
input_example = _read_example(model_conf, model_path)
pyfunc_model = mlflow.pyfunc.load_model(model_path)
assert model_conf.signature == infer_signature(X, model.predict(X[:5]))
np.testing.assert_array_equal(pyfunc_model.predict(input_example),
model.predict(X[:5]))
def test_signature_and_examples_are_saved_correctly(sequential_model, data):
model = sequential_model
signature_ = infer_signature(*data)
example_ = data[0].head(3)
for signature in (None, signature_):
for example in (None, example_):
with TempDir() as tmp:
path = tmp.path("model")
mlflow.pytorch.save_model(model,
path=path,
signature=signature,
input_example=example)
mlflow_model = Model.load(path)
assert signature == mlflow_model.signature
if example is None:
assert mlflow_model.saved_input_example_info is None
else:
assert all((_read_example(mlflow_model,
path) == example).all())
def test_signature_and_examples_are_saved_correctly(gluon_model, model_data):
model = gluon_model
signature_ = infer_signature(model_data[0].asnumpy())
example_ = model_data[0].asnumpy()[:3, ]
for signature in (None, signature_):
for example in (None, example_):
with TempDir() as tmp:
path = tmp.path("model")
mlflow.gluon.save_model(model,
path=path,
signature=signature,
input_example=example)
mlflow_model = Model.load(path)
assert signature == mlflow_model.signature
if example is None:
assert mlflow_model.saved_input_example_info is None
else:
assert np.array_equal(_read_example(mlflow_model, path),
example)
def test_model_export_with_signature_and_examples(iris_df, spark_model_iris):
_, _, iris_spark_df = iris_df
signature_ = infer_signature(iris_spark_df)
example_ = iris_spark_df.toPandas().head(3)
for signature in (None, signature_):
for example in (None, example_):
with TempDir() as tmp:
path = tmp.path("model")
sparkm.save_model(spark_model_iris.model,
path=path,
signature=signature,
input_example=example)
mlflow_model = Model.load(path)
assert signature == mlflow_model.signature
if example is None:
assert mlflow_model.saved_input_example_info is None
else:
assert all((_read_example(mlflow_model,
path) == example).all())
def test_signature_and_examples_are_saved_correctly(sklearn_knn_model):
data = sklearn_knn_model.inference_data
model = sklearn_knn_model.model
signature_ = infer_signature(data)
example_ = data[:3, ]
for signature in (None, signature_):
for example in (None, example_):
with TempDir() as tmp:
path = tmp.path("model")
mlflow.sklearn.save_model(model,
path=path,
signature=signature,
input_example=example)
mlflow_model = Model.load(path)
assert signature == mlflow_model.signature
if example is None:
assert mlflow_model.saved_input_example_info is None
else:
assert all((_read_example(mlflow_model,
path) == example).all())
def test_signature_and_examples_are_saved_correctly(onnx_model, data, onnx_custom_env):
model = onnx_model
signature_ = infer_signature(*data)
example_ = data[0].head(3)
for signature in (None, signature_):
for example in (None, example_):
with TempDir() as tmp:
path = tmp.path("model")
mlflow.onnx.save_model(
model,
path=path,
conda_env=onnx_custom_env,
signature=signature,
input_example=example,
)
mlflow_model = Model.load(path)
assert signature == mlflow_model.signature
if example is None:
assert mlflow_model.saved_input_example_info is None
else:
assert all((_read_example(mlflow_model, path) == example).all())
def test_log_model_with_signature_and_examples(iris_df, spark_model_iris):
_, _, iris_spark_df = iris_df
signature_ = infer_signature(iris_spark_df)
example_ = iris_spark_df.toPandas().head(3)
artifact_path = "model"
for signature in (None, signature_):
for example in (None, example_):
with mlflow.start_run():
sparkm.log_model(spark_model_iris.model,
artifact_path=artifact_path,
signature=signature,
input_example=example)
artifact_uri = mlflow.get_artifact_uri()
model_path = os.path.join(artifact_uri, artifact_path)
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:
assert all((_read_example(mlflow_model,
model_path) == example).all())
def test_signature_and_examples_are_saved_correctly(iris_data, main_scoped_model_class):
def test_predict(sk_model, model_input):
return sk_model.predict(model_input) * 2
data = iris_data
signature_ = infer_signature(*data)
example_ = data[0][:3, ]
for signature in (None, signature_):
for example in (None, example_):
with TempDir() as tmp:
path = tmp.path("model")
mlflow.pyfunc.save_model(path=path,
artifacts={},
python_model=main_scoped_model_class(test_predict),
signature=signature,
input_example=example)
mlflow_model = Model.load(path)
assert signature == mlflow_model.signature
if example is None:
assert mlflow_model.saved_input_example_info is None
else:
assert all((_read_example(mlflow_model, path) == example).all())
def test_pmdarima_signature_and_examples_saved_correctly(
auto_arima_model, test_data, model_path, use_signature, use_example):
# NB: Signature inference will only work on the first element of the tuple return
prediction = auto_arima_model.predict(n_periods=20,
return_conf_int=True,
alpha=0.05)
signature = infer_signature(test_data,
prediction[0]) if use_signature else None
example = test_data[0:5].copy(deep=False) if use_example else None
mlflow.pmdarima.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)
def test_signature_and_examples_saved_correctly(prophet_model):
data = prophet_model.data
model = prophet_model.model
horizon_df = future_horizon_df(model, FORECAST_HORIZON)
signature_ = infer_signature(data, model.predict(horizon_df))
example_ = data[0:5].copy(deep=False)
example_["y"] = pd.to_numeric(example_["y"]) # cast to appropriate precision
for signature in (None, signature_):
for example in (None, example_):
with TempDir() as tmp:
path = tmp.path("model")
mlflow.prophet.save_model(
model, path=path, signature=signature, input_example=example
)
mlflow_model = Model.load(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, path).copy(deep=False)
r_example["ds"] = pd.to_datetime(r_example["ds"], format=DS_FORMAT)
np.testing.assert_array_equal(r_example, example)
def test_schema_and_examples_are_save_correctly(saved_tf_iris_model):
train_x, train_y = iris_data_utils.load_data()[0]
X = pd.DataFrame(train_x)
y = pd.Series(train_y)
for signature in (None, infer_signature(X, y)):
for example in (None, X.head(3)):
with TempDir() as tmp:
path = tmp.path("model")
mlflow.tensorflow.save_model(
tf_saved_model_dir=saved_tf_iris_model.path,
tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags,
tf_signature_def_key=saved_tf_iris_model.signature_def_key,
path=path,
signature=signature,
input_example=example,
)
mlflow_model = Model.load(path)
assert signature == mlflow_model.signature
if example is None:
assert mlflow_model.saved_input_example_info is None
else:
assert all((_read_example(mlflow_model, path) == example).all())
def test_autolog_logs_signature_and_input_example(data_type):
mlflow.sklearn.autolog(log_input_examples=True, log_model_signatures=True)
X, y = get_iris()
X = data_type(X)
y = data_type(y)
model = sklearn.linear_model.LinearRegression()
with mlflow.start_run() as run:
model.fit(X, y)
model_path = os.path.join(run.info.artifact_uri, MODEL_DIR)
model_conf = get_model_conf(run.info.artifact_uri)
input_example = _read_example(model_conf, model_path)
pyfunc_model = mlflow.pyfunc.load_model(model_path)
assert model_conf.signature == infer_signature(X, model.predict(X[:5]))
# On GitHub Actions, `pyfunc_model.predict` and `model.predict` sometimes return
# slightly different results:
#
# >>> pyfunc_model.predict(input_example)
# [[0.171504346208176 ]
# [0.34346150441640155] <- diff
# [0.06895096846585114] <- diff
# [0.05925789882165455]
# [0.03424907823290102]]
#
# >>> model.predict(X[:5])
# [[0.171504346208176 ]
# [0.3434615044164018 ] <- diff
# [0.06895096846585136] <- diff
# [0.05925789882165455]
# [0.03424907823290102]]
#
# As a workaround, use `assert_array_almost_equal` instead of `assert_array_equal`
np.testing.assert_array_almost_equal(pyfunc_model.predict(input_example),
model.predict(X[:5]))
def test_lgb_autolog_gets_input_example(bst_params):
# we need to check the example input against the initial input given to train function.
# we can't use the train_set fixture for this as it defines free_raw_data=False but this
# feature should work even if it is True
iris = datasets.load_iris()
X = pd.DataFrame(iris.data[:, :2], columns=iris.feature_names[:2])
y = iris.target
dataset = lgb.Dataset(X, y, free_raw_data=True)
mlflow.lightgbm.autolog(log_input_examples=True)
lgb.train(bst_params, dataset)
run = get_latest_run()
model_path = os.path.join(run.info.artifact_uri, "model")
model_conf = Model.load(os.path.join(model_path, "MLmodel"))
input_example = _read_example(model_conf, model_path)
assert input_example.equals(X[:5])
pyfunc_model = mlflow.pyfunc.load_model(os.path.join(run.info.artifact_uri, "model"))
# make sure reloading the input_example and predicting on it does not error
pyfunc_model.predict(input_example)
def test_xgb_autolog_gets_input_example(bst_params):
mlflow.xgboost.autolog(log_input_example=True)
# we cannot use dtrain fixture, as the dataset must be constructed
# after the call to autolog() in order to get the input example
iris = datasets.load_iris()
X = pd.DataFrame(iris.data[:, :2], columns=iris.feature_names[:2])
y = iris.target
dataset = xgb.DMatrix(X, y)
xgb.train(bst_params, dataset)
run = get_latest_run()
model_path = os.path.join(run.info.artifact_uri, "model")
model_conf = Model.load(os.path.join(model_path, "MLmodel"))
input_example = _read_example(model_conf, model_path)
assert input_example.equals(X[:5])
pyfunc_model = mlflow.pyfunc.load_model(os.path.join(run.info.artifact_uri, "model"))
# make sure reloading the input_example and predicting on it does not error
pyfunc_model.predict(input_example)
def test_parameter_search_estimators_produce_expected_outputs(
cv_class, search_space, backend):
mlflow.sklearn.autolog(log_input_examples=True, log_model_signatures=True)
svc = sklearn.svm.SVC()
cv_model = cv_class(svc, search_space, n_jobs=5, return_train_score=True)
X, y = get_iris()
def train_cv_model():
if backend is None:
cv_model.fit(X, y)
else:
with sklearn.utils.parallel_backend(backend=backend):
cv_model.fit(X, y)
with mlflow.start_run() as run:
train_cv_model()
run_id = run.info.run_id
params, metrics, tags, artifacts = get_run_data(run_id)
expected_cv_params = truncate_dict(
stringify_dict_values(cv_model.get_params(deep=False)))
expected_cv_params.update({
"best_{}".format(param_name): str(param_value)
for param_name, param_value in cv_model.best_params_.items()
})
assert params == expected_cv_params
assert {
TRAINING_SCORE: cv_model.score(X, y),
"best_cv_score": cv_model.best_score_,
}.items() <= metrics.items()
assert tags == get_expected_class_tags(cv_model)
assert MODEL_DIR in artifacts
assert "best_estimator" in artifacts
assert "cv_results.csv" in artifacts
best_estimator = mlflow.sklearn.load_model(
"runs:/{}/best_estimator".format(run_id))
assert isinstance(best_estimator, sklearn.svm.SVC)
cv_model = mlflow.sklearn.load_model("runs:/{}/{}".format(
run_id, MODEL_DIR))
assert isinstance(cv_model, cv_class)
# Ensure that a signature and input example are produced for the best estimator
best_estimator_conf = get_model_conf(run.info.artifact_uri,
"best_estimator")
assert best_estimator_conf.signature == infer_signature(
X, best_estimator.predict(X[:5]))
best_estimator_path = os.path.join(run.info.artifact_uri, "best_estimator")
input_example = _read_example(best_estimator_conf, best_estimator_path)
best_estimator.predict(
input_example) # Ensure that input example evaluation succeeds
client = mlflow.tracking.MlflowClient()
child_runs = client.search_runs(
run.info.experiment_id,
"tags.`mlflow.parentRunId` = '{}'".format(run_id))
cv_results = pd.DataFrame.from_dict(cv_model.cv_results_)
# We expect to have created a child run for each point in the parameter search space
assert len(child_runs) == len(cv_results)
# Verify that each set of parameter search results has a corresponding MLflow run
# with the expected data
for _, result in cv_results.iterrows():
result_params = result.get("params", {})
params_search_clause = " and ".join([
"params.`{}` = '{}'".format(key, value)
for key, value in result_params.items()
])
search_filter = "tags.`mlflow.parentRunId` = '{}' and {}".format(
run_id, params_search_clause)
child_runs = client.search_runs(run.info.experiment_id, search_filter)
assert len(child_runs) == 1
child_run = child_runs[0]
assert child_run.info.status == RunStatus.to_string(RunStatus.FINISHED)
_, child_metrics, child_tags, _ = get_run_data(child_run.info.run_id)
assert child_tags == get_expected_class_tags(svc)
assert "mean_test_score" in child_metrics.keys()
assert "std_test_score" in child_metrics.keys()
# Ensure that we do not capture separate metrics for each cross validation split, which
# would produce very noisy metrics results
assert len([
metric for metric in child_metrics.keys()
if metric.startswith("split")
]) == 0
