wehub-resource-sync
454 lines
16 KiB
Python
454 lines
16 KiB
Python
import json
|
|
import math
|
|
from io import StringIO
|
|
from unittest import mock
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
import sklearn.linear_model as logreg_module
|
|
from scipy.sparse import csc_matrix, csr_matrix
|
|
from sklearn import datasets
|
|
from sklearn.base import BaseEstimator, ClassifierMixin
|
|
|
|
import mlflow
|
|
from mlflow.models import Model
|
|
from mlflow.models.signature import ModelSignature, infer_signature
|
|
from mlflow.models.utils import (
|
|
_Example,
|
|
_read_sparse_matrix_from_json,
|
|
parse_inputs_data,
|
|
)
|
|
from mlflow.types import DataType
|
|
from mlflow.types.schema import ColSpec, Schema, TensorSpec
|
|
from mlflow.types.utils import TensorsNotSupportedException
|
|
from mlflow.utils.file_utils import TempDir
|
|
from mlflow.utils.proto_json_utils import dataframe_from_raw_json
|
|
|
|
|
|
@pytest.fixture
|
|
def pandas_df_with_all_types():
|
|
df = pd.DataFrame({
|
|
"boolean": [True, False, True],
|
|
"integer": np.array([1, 2, 3], np.int32),
|
|
"long": np.array([1, 2, 3], np.int64),
|
|
"float": np.array([math.pi, 2 * math.pi, 3 * math.pi], np.float32),
|
|
"double": [math.pi, 2 * math.pi, 3 * math.pi],
|
|
"binary": [bytes([1, 2, 3]), bytes([4, 5, 6]), bytes([7, 8, 9])],
|
|
"string": ["a", "b", "c"],
|
|
"boolean_ext": [True, False, True],
|
|
"integer_ext": [1, 2, 3],
|
|
"string_ext": ["a", "b", "c"],
|
|
"array": np.array(["a", "b", "c"]),
|
|
})
|
|
df["boolean_ext"] = df["boolean_ext"].astype("boolean")
|
|
df["integer_ext"] = df["integer_ext"].astype("Int64")
|
|
df["string_ext"] = df["string_ext"].astype("string")
|
|
return df
|
|
|
|
|
|
@pytest.fixture
|
|
def df_without_columns():
|
|
return pd.DataFrame({0: [1, 2, 3], 1: [4, 5, 6], 2: [7, 8, 9]})
|
|
|
|
|
|
@pytest.fixture
|
|
def df_with_nan():
|
|
return pd.DataFrame({
|
|
"boolean": [True, False, True],
|
|
"integer": np.array([1, 2, 3], np.int32),
|
|
"long": np.array([1, 2, 3], np.int64),
|
|
"float": np.array([np.nan, 2 * math.pi, 3 * math.pi], np.float32),
|
|
"double": [math.pi, np.nan, 3 * math.pi],
|
|
"binary": [bytes([1, 2, 3]), bytes([4, 5, 6]), bytes([7, 8, 9])],
|
|
"string": ["a", "b", "c"],
|
|
})
|
|
|
|
|
|
@pytest.fixture
|
|
def dict_of_ndarrays():
|
|
return {
|
|
"1D": np.arange(0, 12, 0.5),
|
|
"2D": np.arange(0, 12, 0.5).reshape(3, 8),
|
|
"3D": np.arange(0, 12, 0.5).reshape(2, 3, 4),
|
|
"4D": np.arange(0, 12, 0.5).reshape(3, 2, 2, 2),
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def dict_of_ndarrays_with_nans():
|
|
return {
|
|
"1D": np.array([0.5, np.nan, 2.0]),
|
|
"2D": np.array([[0.1, 0.2], [np.nan, 0.5]]),
|
|
"3D": np.array([[[0.1, np.nan], [0.3, 0.4]], [[np.nan, 0.6], [0.7, np.nan]]]),
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def dict_of_sparse_matrix():
|
|
return {
|
|
"sparse_matrix_csc": csc_matrix(np.arange(0, 12, 0.5).reshape(3, 8)),
|
|
"sparse_matrix_csr": csr_matrix(np.arange(0, 12, 0.5).reshape(3, 8)),
|
|
}
|
|
|
|
|
|
def test_input_examples(pandas_df_with_all_types, dict_of_ndarrays):
|
|
sig = infer_signature(pandas_df_with_all_types)
|
|
# test setting example with data frame with all supported data types
|
|
with TempDir() as tmp:
|
|
example = _Example(pandas_df_with_all_types)
|
|
example.save(tmp.path())
|
|
filename = example.info["artifact_path"]
|
|
with open(tmp.path(filename)) as f:
|
|
data = json.load(f)
|
|
assert set(data.keys()) == {"columns", "data"}
|
|
parsed_df = dataframe_from_raw_json(tmp.path(filename), schema=sig.inputs)
|
|
pd.testing.assert_frame_equal(pandas_df_with_all_types, parsed_df, check_dtype=False)
|
|
# the frame read without schema should match except for the binary values
|
|
pd.testing.assert_frame_equal(
|
|
parsed_df.drop(columns=["binary"]),
|
|
dataframe_from_raw_json(tmp.path(filename)).drop(columns=["binary"]),
|
|
check_dtype=False,
|
|
)
|
|
|
|
# NB: Drop columns that cannot be encoded by proto_json_utils.pyNumpyEncoder
|
|
new_df = pandas_df_with_all_types.drop(columns=["boolean_ext", "integer_ext", "string_ext"])
|
|
|
|
# pass the input as dictionary instead
|
|
with TempDir() as tmp:
|
|
d = {name: new_df[name].values for name in new_df.columns}
|
|
example = _Example(d)
|
|
example.save(tmp.path())
|
|
filename = example.info["artifact_path"]
|
|
parsed_dict = parse_inputs_data(tmp.path(filename))
|
|
assert d.keys() == parsed_dict.keys()
|
|
# Asserting binary will fail since it is converted to base64 encoded strings.
|
|
# The check above suffices that the binary input is stored.
|
|
del d["binary"]
|
|
for key in d:
|
|
np.testing.assert_array_equal(d[key], parsed_dict[key])
|
|
|
|
# input passed as numpy array
|
|
new_df = pandas_df_with_all_types.drop(columns=["binary"])
|
|
for col in new_df:
|
|
input_example = new_df[col].to_numpy()
|
|
with TempDir() as tmp:
|
|
example = _Example(input_example)
|
|
example.save(tmp.path())
|
|
filename = example.info["artifact_path"]
|
|
parsed_ary = parse_inputs_data(tmp.path(filename))
|
|
np.testing.assert_array_equal(parsed_ary, input_example)
|
|
|
|
# pass multidimensional array
|
|
for col in dict_of_ndarrays:
|
|
input_example = dict_of_ndarrays[col]
|
|
with TempDir() as tmp:
|
|
example = _Example(input_example)
|
|
example.save(tmp.path())
|
|
filename = example.info["artifact_path"]
|
|
parsed_ary = parse_inputs_data(tmp.path(filename))
|
|
np.testing.assert_array_equal(parsed_ary, input_example)
|
|
|
|
# pass multidimensional array as a list
|
|
example = np.array([[1, 2, 3]])
|
|
with pytest.raises(
|
|
TensorsNotSupportedException,
|
|
match=r"Numpy arrays in list are not supported as input examples.",
|
|
):
|
|
_Example([example, example])
|
|
|
|
# pass dict with scalars
|
|
with TempDir() as tmp:
|
|
example = {"a": 1, "b": "abc"}
|
|
x = _Example(example)
|
|
x.save(tmp.path())
|
|
filename = x.info["artifact_path"]
|
|
with open(tmp.path(filename)) as f:
|
|
parsed_data = json.load(f)
|
|
assert example == parsed_data
|
|
|
|
|
|
def test_pandas_orients_for_input_examples(
|
|
pandas_df_with_all_types, df_without_columns, dict_of_ndarrays
|
|
):
|
|
# test setting example with data frame with all supported data types
|
|
with TempDir() as tmp:
|
|
example = _Example(pandas_df_with_all_types)
|
|
example.save(tmp.path())
|
|
filename = example.info["artifact_path"]
|
|
assert example.info["type"] == "dataframe"
|
|
assert example.info["pandas_orient"] == "split"
|
|
with open(tmp.path(filename)) as f:
|
|
data = json.load(f)
|
|
dataframe = pd.read_json(
|
|
StringIO(json.dumps(data)), orient=example.info["pandas_orient"], precise_float=True
|
|
)
|
|
pd.testing.assert_frame_equal(
|
|
pandas_df_with_all_types.drop(columns=["binary"]),
|
|
dataframe.drop(columns=["binary"]),
|
|
check_dtype=False,
|
|
)
|
|
|
|
with TempDir() as tmp:
|
|
example = _Example(df_without_columns)
|
|
example.save(tmp.path())
|
|
filename = example.info["artifact_path"]
|
|
assert example.info["type"] == "dataframe"
|
|
assert example.info["pandas_orient"] == "values"
|
|
with open(tmp.path(filename)) as f:
|
|
data = json.load(f)
|
|
assert set(data.keys()) == {"data"}
|
|
# NOTE: when no column names are provided (i.e. values orient),
|
|
# saving an example adds a "data" key rather than directly storing the plain data
|
|
data = data["data"]
|
|
dataframe = pd.read_json(
|
|
StringIO(json.dumps(data)), orient=example.info["pandas_orient"]
|
|
)
|
|
pd.testing.assert_frame_equal(dataframe, df_without_columns, check_dtype=False)
|
|
|
|
# pass dict with scalars
|
|
with TempDir() as tmp:
|
|
example = {"a": 1, "b": "abc"}
|
|
x = _Example(example)
|
|
x.save(tmp.path())
|
|
filename = x.info["artifact_path"]
|
|
assert x.info["type"] == "json_object"
|
|
with open(tmp.path(filename)) as f:
|
|
parsed_json = json.load(f)
|
|
assert parsed_json == example
|
|
|
|
|
|
def test_sparse_matrix_input_examples(dict_of_sparse_matrix):
|
|
for example_type, input_example in dict_of_sparse_matrix.items():
|
|
with TempDir() as tmp:
|
|
example = _Example(input_example)
|
|
example.save(tmp.path())
|
|
filename = example.info["artifact_path"]
|
|
assert example.info["type"] == example_type
|
|
parsed_matrix = _read_sparse_matrix_from_json(tmp.path(filename), example_type)
|
|
np.testing.assert_array_equal(parsed_matrix.toarray(), input_example.toarray())
|
|
|
|
|
|
def test_input_examples_with_nan(df_with_nan, dict_of_ndarrays_with_nans):
|
|
# test setting example with data frame with NaN values in it
|
|
sig = infer_signature(df_with_nan)
|
|
with TempDir() as tmp:
|
|
example = _Example(df_with_nan)
|
|
example.save(tmp.path())
|
|
filename = example.info["artifact_path"]
|
|
assert example.info["type"] == "dataframe"
|
|
assert example.info["pandas_orient"] == "split"
|
|
with open(tmp.path(filename)) as f:
|
|
data = json.load(f)
|
|
assert set(data.keys()) == {"columns", "data"}
|
|
pd.read_json(StringIO(json.dumps(data)), orient=example.info["pandas_orient"])
|
|
|
|
parsed_df = dataframe_from_raw_json(tmp.path(filename), schema=sig.inputs)
|
|
|
|
# by definition of NaN, NaN == NaN is False but NaN != NaN is True
|
|
pd.testing.assert_frame_equal(df_with_nan, parsed_df, check_dtype=False)
|
|
# the frame read without schema should match except for the binary values
|
|
no_schema_df = dataframe_from_raw_json(tmp.path(filename))
|
|
a = parsed_df.drop(columns=["binary"])
|
|
b = no_schema_df.drop(columns=["binary"])
|
|
pd.testing.assert_frame_equal(a, b, check_dtype=False)
|
|
|
|
# pass multidimensional array
|
|
for col in dict_of_ndarrays_with_nans:
|
|
input_example = dict_of_ndarrays_with_nans[col]
|
|
sig = infer_signature(input_example)
|
|
with TempDir() as tmp:
|
|
example = _Example(input_example)
|
|
example.save(tmp.path())
|
|
filename = example.info["artifact_path"]
|
|
assert example.info["type"] == "ndarray"
|
|
parsed_ary = parse_inputs_data(tmp.path(filename), schema=sig.inputs)
|
|
assert np.array_equal(parsed_ary, input_example, equal_nan=True)
|
|
|
|
# without a schema/dtype specified, the resulting tensor will keep the None type
|
|
no_schema_df = parse_inputs_data(tmp.path(filename))
|
|
np.testing.assert_array_equal(
|
|
no_schema_df, np.where(np.isnan(input_example), None, input_example)
|
|
)
|
|
|
|
|
|
class DummySklearnModel(BaseEstimator, ClassifierMixin):
|
|
def __init__(self, output_shape=(1,)):
|
|
self.output_shape = output_shape
|
|
|
|
def fit(self, X, y=None):
|
|
return self
|
|
|
|
def predict(self, X):
|
|
n_samples = X.shape[0]
|
|
full_output_shape = (n_samples,) + self.output_shape
|
|
return np.zeros(full_output_shape, dtype=np.dtype("int64"))
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("input_is_tabular", "output_shape", "expected_signature"),
|
|
[
|
|
# When the input example is column-based, output 1D numpy arrays are interpreted `ColSpec`s
|
|
(
|
|
True,
|
|
(),
|
|
ModelSignature(
|
|
inputs=Schema([ColSpec(name="feature", type=DataType.string)]),
|
|
outputs=Schema([ColSpec(type=DataType.long)]),
|
|
),
|
|
),
|
|
# But if the output numpy array has higher dimensions, fallback to interpreting the model
|
|
# output as `TensorSpec`s.
|
|
(
|
|
True,
|
|
(2,),
|
|
ModelSignature(
|
|
inputs=Schema([ColSpec(name="feature", type=DataType.string)]),
|
|
outputs=Schema([TensorSpec(np.dtype("int64"), (-1, 2))]),
|
|
),
|
|
),
|
|
# If the input example is tensor-based, interpret output numpy arrays as `TensorSpec`s
|
|
(
|
|
False,
|
|
(),
|
|
ModelSignature(
|
|
inputs=Schema([TensorSpec(np.dtype("int64"), (-1, 1))]),
|
|
outputs=Schema([TensorSpec(np.dtype("int64"), (-1,))]),
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_infer_signature_with_input_example(input_is_tabular, output_shape, expected_signature):
|
|
model = DummySklearnModel(output_shape=output_shape)
|
|
artifact_path = "model"
|
|
example = pd.DataFrame({"feature": ["value"]}) if input_is_tabular else np.array([[1]])
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.sklearn.log_model(
|
|
model,
|
|
name=artifact_path,
|
|
input_example=example,
|
|
serialization_format="cloudpickle",
|
|
)
|
|
|
|
mlflow_model = Model.load(model_info.model_uri)
|
|
assert mlflow_model.signature == expected_signature
|
|
|
|
|
|
def test_infer_signature_from_example_can_be_disabled():
|
|
artifact_path = "model"
|
|
with mlflow.start_run():
|
|
model_info = mlflow.sklearn.log_model(
|
|
DummySklearnModel(output_shape=()),
|
|
name=artifact_path,
|
|
input_example=np.array([[1]]),
|
|
signature=False,
|
|
serialization_format="cloudpickle",
|
|
)
|
|
|
|
mlflow_model = Model.load(model_info.model_uri)
|
|
assert mlflow_model.signature is None
|
|
|
|
|
|
def test_infer_signature_raises_if_predict_on_input_example_fails(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_TESTING", "false")
|
|
|
|
class ErrorModel(BaseEstimator, ClassifierMixin):
|
|
def fit(self, X, y=None):
|
|
return self
|
|
|
|
def predict(self, X):
|
|
raise Exception("oh no!")
|
|
|
|
with mock.patch("mlflow.models.model._logger.warning") as mock_warning:
|
|
with mlflow.start_run():
|
|
mlflow.sklearn.log_model(
|
|
ErrorModel(),
|
|
name="model",
|
|
input_example=np.array([[1]]),
|
|
serialization_format="cloudpickle",
|
|
)
|
|
assert any(
|
|
"Failed to validate serving input example" in call[0][0]
|
|
for call in mock_warning.call_args_list
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def iris_model():
|
|
X, y = datasets.load_iris(return_X_y=True, as_frame=True)
|
|
return logreg_module.LogisticRegression().fit(X, y)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_example",
|
|
[
|
|
{
|
|
"sepal length (cm)": 5.1,
|
|
"sepal width (cm)": 3.5,
|
|
"petal length (cm)": 1.4,
|
|
"petal width (cm)": 0.2,
|
|
},
|
|
pd.DataFrame([[5.1, 3.5, 1.4, 0.2]]),
|
|
pd.DataFrame(
|
|
{
|
|
"sepal length (cm)": 5.1,
|
|
"sepal width (cm)": 3.5,
|
|
"petal length (cm)": 1.4,
|
|
"petal width (cm)": 0.2,
|
|
},
|
|
index=[0],
|
|
),
|
|
],
|
|
)
|
|
def test_infer_signature_on_multi_column_input_examples(input_example, iris_model):
|
|
artifact_path = "model"
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.sklearn.log_model(
|
|
iris_model, name=artifact_path, input_example=input_example
|
|
)
|
|
|
|
mlflow_model = Model.load(model_info.model_uri)
|
|
input_columns = mlflow_model.signature.inputs.inputs
|
|
assert len(input_columns) == 4
|
|
assert all(col.type == DataType.double for col in input_columns)
|
|
assert mlflow_model.signature.outputs == Schema([ColSpec(type=DataType.long)])
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_example",
|
|
["some string", bytes([1, 2, 3])],
|
|
)
|
|
def test_infer_signature_on_scalar_input_examples(input_example):
|
|
class IdentitySklearnModel(BaseEstimator, ClassifierMixin):
|
|
def fit(self, X, y=None):
|
|
return self
|
|
|
|
def predict(self, X):
|
|
if isinstance(X, pd.DataFrame):
|
|
return X
|
|
raise Exception("Unsupported input type")
|
|
|
|
artifact_path = "model"
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.sklearn.log_model(
|
|
IdentitySklearnModel(),
|
|
name=artifact_path,
|
|
input_example=input_example,
|
|
serialization_format="cloudpickle",
|
|
)
|
|
|
|
mlflow_model = Model.load(model_info.model_uri)
|
|
signature = mlflow_model.signature
|
|
assert isinstance(signature, ModelSignature)
|
|
assert signature.inputs.inputs[0].name is None
|
|
t = DataType.string if isinstance(input_example, str) else DataType.binary
|
|
assert signature == ModelSignature(
|
|
inputs=Schema([ColSpec(type=t)]),
|
|
outputs=Schema([ColSpec(name=0, type=t)]),
|
|
)
|
|
# test that a single string still passes pyfunc schema enforcement
|
|
mlflow.pyfunc.load_model(model_info.model_uri).predict(input_example)
|