wehub-resource-sync
1470 lines
52 KiB
Python
1470 lines
52 KiB
Python
import importlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import pickle
|
|
import re
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
import torch
|
|
import yaml
|
|
from packaging.version import Version
|
|
from sklearn import datasets
|
|
from torch import nn
|
|
from torch.utils.data import DataLoader
|
|
|
|
import mlflow.pyfunc.scoring_server as pyfunc_scoring_server
|
|
import mlflow.pytorch
|
|
from mlflow import pyfunc
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models import Model, ModelSignature
|
|
from mlflow.models.utils import _read_example, load_serving_example
|
|
from mlflow.pytorch import pickle_module as mlflow_pytorch_pickle_module
|
|
from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
|
|
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
|
|
from mlflow.types.schema import DataType, Schema, TensorSpec
|
|
from mlflow.utils.environment import _mlflow_conda_env
|
|
from mlflow.utils.file_utils import TempDir
|
|
from mlflow.utils.model_utils import _get_flavor_configuration
|
|
|
|
from tests.helper_functions import (
|
|
_assert_pip_requirements,
|
|
_compare_conda_env_requirements,
|
|
_compare_logged_code_paths,
|
|
_is_available_on_pypi,
|
|
_is_importable,
|
|
_mlflow_major_version_string,
|
|
assert_array_almost_equal,
|
|
assert_register_model_called_with_local_model_path,
|
|
)
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
# This test suite is included as a code dependency when testing PyTorch model scoring in new
|
|
# processes and docker containers. In these environments, the `tests` module is not available.
|
|
# Therefore, we attempt to import from `tests` and gracefully emit a warning if it's unavailable.
|
|
try:
|
|
from tests.helper_functions import pyfunc_serve_and_score_model
|
|
except ImportError:
|
|
_logger.warning(
|
|
"Failed to import test helper functions. Tests depending on these functions may fail!"
|
|
)
|
|
|
|
EXTRA_PYFUNC_SERVING_TEST_ARGS = (
|
|
[] if _is_available_on_pypi("torch") else ["--env-manager", "local"]
|
|
)
|
|
|
|
# in pytorch >= 2.6.0, the `weights_only` kwarg default has been changed from
|
|
# `False` to `True`. this can cause pickle deserialization errors when loading
|
|
# models, unless the model classes have been explicitly marked as safe using
|
|
# `torch.serialization.add_safe_globals()`
|
|
ENABLE_LEGACY_DESERIALIZATION = Version(torch.__version__) >= Version("2.6.0")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def data():
|
|
iris = datasets.load_iris()
|
|
data = pd.DataFrame(
|
|
data=np.c_[iris["data"], iris["target"]], columns=iris["feature_names"] + ["target"]
|
|
)
|
|
y = data["target"]
|
|
x = data.drop("target", axis=1)
|
|
return x, y
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def iris_tensor_spec():
|
|
return ModelSignature(
|
|
inputs=Schema([TensorSpec(np.dtype("float32"), (-1, 4))]),
|
|
outputs=Schema([TensorSpec(np.dtype("float32"), (-1, 1))]),
|
|
)
|
|
|
|
|
|
def get_dataset(data):
|
|
x, y = data
|
|
return [(xi.astype(np.float32), yi.astype(np.float32)) for xi, yi in zip(x.values, y.values)]
|
|
|
|
|
|
def train_model(model, data):
|
|
dataset = get_dataset(data)
|
|
criterion = nn.MSELoss()
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
|
|
batch_size = 16
|
|
num_workers = 4
|
|
dataloader = DataLoader(
|
|
dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True, drop_last=False
|
|
)
|
|
|
|
model.train()
|
|
for _ in range(5):
|
|
for batch in dataloader:
|
|
optimizer.zero_grad()
|
|
batch_size = batch[0].shape[0]
|
|
y_pred = model(batch[0]).squeeze(dim=1)
|
|
loss = criterion(y_pred, batch[1])
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
|
|
def get_sequential_model():
|
|
return nn.Sequential(nn.Linear(4, 3), nn.ReLU(), nn.Linear(3, 1))
|
|
|
|
|
|
@pytest.fixture
|
|
def sequential_model(data, scripted_model):
|
|
model = get_sequential_model()
|
|
if scripted_model:
|
|
model = torch.jit.script(model)
|
|
|
|
train_model(model=model, data=data)
|
|
return model
|
|
|
|
|
|
def get_subclassed_model_definition():
|
|
"""
|
|
Defines a PyTorch model class that inherits from ``torch.nn.Module``. This method can be invoked
|
|
within a pytest fixture to define the model class in the ``__main__`` scope. Alternatively, it
|
|
can be invoked within a module to define the class in the module's scope.
|
|
"""
|
|
|
|
class SubclassedModel(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear = torch.nn.Linear(4, 1)
|
|
|
|
def forward(self, x):
|
|
return self.linear(x)
|
|
|
|
return SubclassedModel
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def main_scoped_subclassed_model(data):
|
|
"""
|
|
A custom PyTorch model inheriting from ``torch.nn.Module`` whose class is defined in the
|
|
"__main__" scope.
|
|
"""
|
|
model_class = get_subclassed_model_definition()
|
|
model = model_class()
|
|
train_model(model=model, data=data)
|
|
return model
|
|
|
|
|
|
class ModuleScopedSubclassedModel(get_subclassed_model_definition()):
|
|
"""
|
|
A custom PyTorch model class defined in the test module scope. This is a subclass of
|
|
``torch.nn.Module``.
|
|
"""
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def module_scoped_subclassed_model(data):
|
|
"""
|
|
A custom PyTorch model inheriting from ``torch.nn.Module`` whose class is defined in the test
|
|
module scope.
|
|
"""
|
|
model = ModuleScopedSubclassedModel()
|
|
train_model(model=model, data=data)
|
|
return model
|
|
|
|
|
|
@pytest.fixture
|
|
def model_path(tmp_path):
|
|
return os.path.join(tmp_path, "model")
|
|
|
|
|
|
@pytest.fixture
|
|
def pytorch_custom_env(tmp_path):
|
|
conda_env = os.path.join(tmp_path, "conda_env.yml")
|
|
_mlflow_conda_env(conda_env, additional_pip_deps=["pytorch", "torchvision", "pytest"])
|
|
return conda_env
|
|
|
|
|
|
def _predict(model, data):
|
|
from torch.fx import GraphModule
|
|
|
|
dataset = get_dataset(data)
|
|
batch_size = 16
|
|
num_workers = 4
|
|
dataloader = DataLoader(
|
|
dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False, drop_last=False
|
|
)
|
|
predictions = np.zeros((len(dataloader.sampler),))
|
|
|
|
if not isinstance(model, GraphModule):
|
|
model.eval()
|
|
with torch.no_grad():
|
|
for i, batch in enumerate(dataloader):
|
|
y_preds = model(batch[0]).squeeze(dim=1).numpy()
|
|
predictions[i * batch_size : (i + 1) * batch_size] = y_preds
|
|
return predictions
|
|
|
|
|
|
@pytest.fixture
|
|
def sequential_predicted(sequential_model, data):
|
|
return _predict(sequential_model, data)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_signature_and_examples_are_saved_correctly(sequential_model, data, iris_tensor_spec):
|
|
model = sequential_model
|
|
example_ = data[0].head(3).values.astype(np.float32)
|
|
for signature in (None, iris_tensor_spec):
|
|
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,
|
|
serialization_format="pickle",
|
|
)
|
|
mlflow_model = Model.load(path)
|
|
if signature is None and example is None:
|
|
assert mlflow_model.signature is None
|
|
else:
|
|
assert mlflow_model.signature == iris_tensor_spec
|
|
if example is None:
|
|
assert mlflow_model.saved_input_example_info is None
|
|
else:
|
|
np.testing.assert_allclose(_read_example(mlflow_model, path), example)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_log_model(sequential_model, data, sequential_predicted):
|
|
try:
|
|
artifact_path = "pytorch"
|
|
model_info = mlflow.pytorch.log_model(
|
|
sequential_model, name=artifact_path, serialization_format="pickle"
|
|
)
|
|
|
|
sequential_model_loaded = mlflow.pytorch.load_model(model_uri=model_info.model_uri)
|
|
test_predictions = _predict(sequential_model_loaded, data)
|
|
np.testing.assert_array_equal(test_predictions, sequential_predicted)
|
|
finally:
|
|
mlflow.end_run()
|
|
|
|
|
|
def test_log_model_calls_register_model(module_scoped_subclassed_model):
|
|
custom_pickle_module = pickle
|
|
artifact_path = "model"
|
|
register_model_patch = mock.patch("mlflow.tracking._model_registry.fluent._register_model")
|
|
with mlflow.start_run(), register_model_patch:
|
|
model_info = mlflow.pytorch.log_model(
|
|
module_scoped_subclassed_model,
|
|
name=artifact_path,
|
|
pickle_module=custom_pickle_module,
|
|
registered_model_name="AdsModel1",
|
|
serialization_format="pickle",
|
|
)
|
|
assert_register_model_called_with_local_model_path(
|
|
register_model_mock=mlflow.tracking._model_registry.fluent._register_model,
|
|
model_uri=model_info.model_uri,
|
|
registered_model_name="AdsModel1",
|
|
)
|
|
|
|
|
|
def test_log_model_no_registered_model_name(module_scoped_subclassed_model):
|
|
custom_pickle_module = pickle
|
|
artifact_path = "model"
|
|
register_model_patch = mock.patch("mlflow.tracking._model_registry.fluent._register_model")
|
|
with mlflow.start_run(), register_model_patch:
|
|
mlflow.pytorch.log_model(
|
|
module_scoped_subclassed_model,
|
|
name=artifact_path,
|
|
pickle_module=custom_pickle_module,
|
|
serialization_format="pickle",
|
|
)
|
|
mlflow.tracking._model_registry.fluent._register_model.assert_not_called()
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_raise_exception(sequential_model):
|
|
with TempDir(chdr=True, remove_on_exit=True) as tmp:
|
|
path = tmp.path("model")
|
|
with pytest.raises(MlflowException, match="No such artifact"):
|
|
mlflow.pytorch.load_model(path)
|
|
|
|
with pytest.raises(TypeError, match="Argument 'pytorch_model' should be a torch.nn.Module"):
|
|
mlflow.pytorch.save_model([1, 2, 3], path)
|
|
|
|
mlflow.pytorch.save_model(sequential_model, path, serialization_format="pickle")
|
|
with pytest.raises(MlflowException, match=f"Path '{os.path.abspath(path)}' already exists"):
|
|
mlflow.pytorch.save_model(sequential_model, path, serialization_format="pickle")
|
|
|
|
import sklearn.neighbors as knn
|
|
|
|
from mlflow import sklearn
|
|
|
|
path = tmp.path("knn.pkl")
|
|
knn = knn.KNeighborsClassifier()
|
|
with open(path, "wb") as f:
|
|
pickle.dump(knn, f)
|
|
path = tmp.path("knn")
|
|
sklearn.save_model(knn, path=path)
|
|
with pytest.raises(MlflowException, match='Model does not have the "pytorch" flavor'):
|
|
mlflow.pytorch.load_model(path)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_save_and_load_model(sequential_model, model_path, data, sequential_predicted):
|
|
mlflow.pytorch.save_model(sequential_model, model_path, serialization_format="pickle")
|
|
|
|
# Loading pytorch model
|
|
sequential_model_loaded = mlflow.pytorch.load_model(model_path)
|
|
np.testing.assert_array_equal(_predict(sequential_model_loaded, data), sequential_predicted)
|
|
|
|
# Loading pyfunc model
|
|
pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
|
|
np.testing.assert_array_almost_equal(
|
|
pyfunc_loaded.predict(data[0]).values[:, 0], sequential_predicted, decimal=4
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_pyfunc_model_works_with_np_input_type(
|
|
sequential_model, model_path, data, sequential_predicted
|
|
):
|
|
mlflow.pytorch.save_model(sequential_model, model_path, serialization_format="pickle")
|
|
|
|
# Loading pyfunc model
|
|
pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
|
|
|
|
# predict works with dataframes
|
|
df_result = pyfunc_loaded.predict(data[0])
|
|
assert type(df_result) == pd.DataFrame
|
|
np.testing.assert_array_almost_equal(df_result.values[:, 0], sequential_predicted, decimal=4)
|
|
|
|
# predict works with numpy ndarray
|
|
np_result = pyfunc_loaded.predict(data[0].values.astype(np.float32))
|
|
assert type(np_result) == np.ndarray
|
|
np.testing.assert_array_almost_equal(np_result[:, 0], sequential_predicted, decimal=4)
|
|
|
|
# predict does not work with lists
|
|
with pytest.raises(
|
|
TypeError, match="The PyTorch flavor does not support List or Dict input types"
|
|
):
|
|
pyfunc_loaded.predict([1, 2, 3, 4])
|
|
|
|
# predict does not work with scalars
|
|
with pytest.raises(TypeError, match="Input data should be pandas.DataFrame or numpy.ndarray"):
|
|
pyfunc_loaded.predict(4)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_load_model_from_remote_uri_succeeds(
|
|
sequential_model, model_path, mock_s3_bucket, data, sequential_predicted
|
|
):
|
|
mlflow.pytorch.save_model(sequential_model, model_path, serialization_format="pickle")
|
|
|
|
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 = artifact_root + "/" + artifact_path
|
|
sequential_model_loaded = mlflow.pytorch.load_model(model_uri=model_uri)
|
|
np.testing.assert_array_equal(_predict(sequential_model_loaded, data), sequential_predicted)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_model_save_persists_specified_conda_env_in_mlflow_model_directory(
|
|
sequential_model, model_path, pytorch_custom_env
|
|
):
|
|
mlflow.pytorch.save_model(
|
|
pytorch_model=sequential_model,
|
|
path=model_path,
|
|
conda_env=pytorch_custom_env,
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
|
|
saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]["conda"])
|
|
assert os.path.exists(saved_conda_env_path)
|
|
assert saved_conda_env_path != pytorch_custom_env
|
|
|
|
with open(pytorch_custom_env) as f:
|
|
pytorch_custom_env_text = f.read()
|
|
with open(saved_conda_env_path) as f:
|
|
saved_conda_env_text = f.read()
|
|
assert saved_conda_env_text == pytorch_custom_env_text
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_model_save_persists_requirements_in_mlflow_model_directory(
|
|
sequential_model, model_path, pytorch_custom_env
|
|
):
|
|
mlflow.pytorch.save_model(
|
|
pytorch_model=sequential_model,
|
|
path=model_path,
|
|
conda_env=pytorch_custom_env,
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
saved_pip_req_path = os.path.join(model_path, "requirements.txt")
|
|
_compare_conda_env_requirements(pytorch_custom_env, saved_pip_req_path)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [False])
|
|
def test_save_model_with_pip_requirements(sequential_model, tmp_path):
|
|
expected_mlflow_version = _mlflow_major_version_string()
|
|
# Path to a requirements file
|
|
tmpdir1 = tmp_path.joinpath("1")
|
|
req_file = tmp_path.joinpath("requirements.txt")
|
|
req_file.write_text("a")
|
|
mlflow.pytorch.save_model(
|
|
sequential_model, tmpdir1, pip_requirements=str(req_file), serialization_format="pickle"
|
|
)
|
|
_assert_pip_requirements(tmpdir1, [expected_mlflow_version, "a"], strict=True)
|
|
|
|
# List of requirements
|
|
tmpdir2 = tmp_path.joinpath("2")
|
|
mlflow.pytorch.save_model(
|
|
sequential_model,
|
|
tmpdir2,
|
|
pip_requirements=[f"-r {req_file}", "b"],
|
|
serialization_format="pickle",
|
|
)
|
|
_assert_pip_requirements(tmpdir2, [expected_mlflow_version, "a", "b"], strict=True)
|
|
|
|
# Constraints file
|
|
tmpdir3 = tmp_path.joinpath("3")
|
|
mlflow.pytorch.save_model(
|
|
sequential_model,
|
|
tmpdir3,
|
|
pip_requirements=[f"-c {req_file}", "b"],
|
|
serialization_format="pickle",
|
|
)
|
|
_assert_pip_requirements(
|
|
tmpdir3, [expected_mlflow_version, "b", "-c constraints.txt"], ["a"], strict=True
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [False])
|
|
def test_save_model_with_extra_pip_requirements(sequential_model, tmp_path):
|
|
expected_mlflow_version = _mlflow_major_version_string()
|
|
default_reqs = mlflow.pytorch.get_default_pip_requirements()
|
|
|
|
# Path to a requirements file
|
|
tmpdir1 = tmp_path.joinpath("1")
|
|
req_file = tmp_path.joinpath("requirements.txt")
|
|
req_file.write_text("a")
|
|
mlflow.pytorch.save_model(
|
|
sequential_model,
|
|
tmpdir1,
|
|
extra_pip_requirements=str(req_file),
|
|
serialization_format="pickle",
|
|
)
|
|
_assert_pip_requirements(tmpdir1, [expected_mlflow_version, *default_reqs, "a"])
|
|
|
|
# List of requirements
|
|
tmpdir2 = tmp_path.joinpath("2")
|
|
mlflow.pytorch.save_model(
|
|
sequential_model,
|
|
tmpdir2,
|
|
extra_pip_requirements=[f"-r {req_file}", "b"],
|
|
serialization_format="pickle",
|
|
)
|
|
_assert_pip_requirements(tmpdir2, [expected_mlflow_version, *default_reqs, "a", "b"])
|
|
|
|
# Constraints file
|
|
tmpdir3 = tmp_path.joinpath("3")
|
|
mlflow.pytorch.save_model(
|
|
sequential_model,
|
|
tmpdir3,
|
|
extra_pip_requirements=[f"-c {req_file}", "b"],
|
|
serialization_format="pickle",
|
|
)
|
|
_assert_pip_requirements(
|
|
tmpdir3, [expected_mlflow_version, *default_reqs, "b", "-c constraints.txt"], ["a"]
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_model_save_accepts_conda_env_as_dict(sequential_model, model_path):
|
|
conda_env = dict(mlflow.pytorch.get_default_conda_env())
|
|
conda_env["dependencies"].append("pytest")
|
|
mlflow.pytorch.save_model(
|
|
pytorch_model=sequential_model,
|
|
path=model_path,
|
|
conda_env=conda_env,
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
|
|
saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]["conda"])
|
|
assert os.path.exists(saved_conda_env_path)
|
|
|
|
with open(saved_conda_env_path) as f:
|
|
saved_conda_env_parsed = yaml.safe_load(f)
|
|
assert saved_conda_env_parsed == conda_env
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_model_log_persists_specified_conda_env_in_mlflow_model_directory(
|
|
sequential_model, pytorch_custom_env
|
|
):
|
|
artifact_path = "model"
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
sequential_model,
|
|
name=artifact_path,
|
|
conda_env=pytorch_custom_env,
|
|
serialization_format="pickle",
|
|
)
|
|
model_path = _download_artifact_from_uri(model_info.model_uri)
|
|
|
|
pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
|
|
saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]["conda"])
|
|
assert os.path.exists(saved_conda_env_path)
|
|
assert saved_conda_env_path != pytorch_custom_env
|
|
|
|
with open(pytorch_custom_env) as f:
|
|
pytorch_custom_env_text = f.read()
|
|
with open(saved_conda_env_path) as f:
|
|
saved_conda_env_text = f.read()
|
|
assert saved_conda_env_text == pytorch_custom_env_text
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_model_log_persists_requirements_in_mlflow_model_directory(
|
|
sequential_model, pytorch_custom_env
|
|
):
|
|
artifact_path = "model"
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
sequential_model,
|
|
name=artifact_path,
|
|
conda_env=pytorch_custom_env,
|
|
serialization_format="pickle",
|
|
)
|
|
model_path = _download_artifact_from_uri(model_info.model_uri)
|
|
|
|
saved_pip_req_path = os.path.join(model_path, "requirements.txt")
|
|
_compare_conda_env_requirements(pytorch_custom_env, saved_pip_req_path)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
|
|
sequential_model, model_path
|
|
):
|
|
mlflow.pytorch.save_model(
|
|
pytorch_model=sequential_model, path=model_path, serialization_format="pickle"
|
|
)
|
|
_assert_pip_requirements(model_path, mlflow.pytorch.get_default_pip_requirements())
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies(
|
|
sequential_model,
|
|
):
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
sequential_model, name="model", serialization_format="pickle"
|
|
)
|
|
|
|
_assert_pip_requirements(model_info.model_uri, mlflow.pytorch.get_default_pip_requirements())
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_load_model_with_differing_pytorch_version_logs_warning(sequential_model, model_path):
|
|
mlflow.pytorch.save_model(
|
|
pytorch_model=sequential_model, path=model_path, serialization_format="pickle"
|
|
)
|
|
saver_pytorch_version = "1.0"
|
|
model_config_path = os.path.join(model_path, "MLmodel")
|
|
model_config = Model.load(model_config_path)
|
|
model_config.flavors[mlflow.pytorch.FLAVOR_NAME]["pytorch_version"] = saver_pytorch_version
|
|
model_config.save(model_config_path)
|
|
|
|
log_messages = []
|
|
|
|
def custom_warn(message_text, *args, **kwargs):
|
|
log_messages.append(message_text % args % kwargs)
|
|
|
|
loader_pytorch_version = "0.8.2"
|
|
with (
|
|
mock.patch("mlflow.pytorch._logger.warning") as warn_mock,
|
|
mock.patch("torch.__version__", loader_pytorch_version),
|
|
):
|
|
warn_mock.side_effect = custom_warn
|
|
mlflow.pytorch.load_model(model_uri=model_path)
|
|
|
|
assert any(
|
|
"does not match installed PyTorch version" in log_message
|
|
and saver_pytorch_version in log_message
|
|
and loader_pytorch_version in log_message
|
|
for log_message in log_messages
|
|
)
|
|
|
|
|
|
def test_pyfunc_model_serving_with_module_scoped_subclassed_model_and_default_conda_env(
|
|
module_scoped_subclassed_model, data
|
|
):
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
module_scoped_subclassed_model,
|
|
name="pytorch_model",
|
|
code_paths=[__file__],
|
|
input_example=data[0],
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
inference_payload = load_serving_example(model_info.model_uri)
|
|
scoring_response = pyfunc_serve_and_score_model(
|
|
model_uri=model_info.model_uri,
|
|
data=inference_payload,
|
|
content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
|
|
extra_args=["--env-manager", "local"],
|
|
)
|
|
assert scoring_response.status_code == 200
|
|
|
|
deployed_model_preds = pd.DataFrame(json.loads(scoring_response.content)["predictions"])
|
|
np.testing.assert_array_almost_equal(
|
|
deployed_model_preds.values[:, 0],
|
|
_predict(model=module_scoped_subclassed_model, data=data),
|
|
decimal=4,
|
|
)
|
|
|
|
|
|
def test_save_model_with_wrong_codepaths_fails_correctly(
|
|
module_scoped_subclassed_model, model_path, data
|
|
):
|
|
with pytest.raises(TypeError, match="Argument code_paths should be a list, not <class 'str'>"):
|
|
mlflow.pytorch.save_model(
|
|
path=model_path, pytorch_model=module_scoped_subclassed_model, code_paths="some string"
|
|
)
|
|
|
|
|
|
def test_pyfunc_model_serving_with_main_scoped_subclassed_model_and_custom_pickle_module(
|
|
main_scoped_subclassed_model, data
|
|
):
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
main_scoped_subclassed_model,
|
|
name="pytorch_model",
|
|
pickle_module=mlflow_pytorch_pickle_module,
|
|
input_example=data[0],
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
inference_payload = load_serving_example(model_info.model_uri)
|
|
scoring_response = pyfunc_serve_and_score_model(
|
|
model_uri=model_info.model_uri,
|
|
data=inference_payload,
|
|
content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
|
|
extra_args=["--env-manager", "local"],
|
|
)
|
|
assert scoring_response.status_code == 200
|
|
|
|
deployed_model_preds = pd.DataFrame(json.loads(scoring_response.content)["predictions"])
|
|
np.testing.assert_array_almost_equal(
|
|
deployed_model_preds.values[:, 0],
|
|
_predict(model=main_scoped_subclassed_model, data=data),
|
|
decimal=4,
|
|
)
|
|
|
|
|
|
def test_load_model_succeeds_with_dependencies_specified_via_code_paths(
|
|
module_scoped_subclassed_model, model_path, data
|
|
):
|
|
# Save a PyTorch model whose class is defined in the current test suite. Because the
|
|
# `tests` module is not available when the model is deployed for local scoring, we include
|
|
# the test suite file as a code dependency
|
|
mlflow.pytorch.save_model(
|
|
path=model_path,
|
|
pytorch_model=module_scoped_subclassed_model,
|
|
code_paths=[__file__],
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
# Define a custom pyfunc model that loads a PyTorch model artifact using
|
|
# `mlflow.pytorch.load_model`
|
|
class TorchValidatorModel(pyfunc.PythonModel):
|
|
def load_context(self, context):
|
|
self.pytorch_model = mlflow.pytorch.load_model(context.artifacts["pytorch_model"])
|
|
|
|
def predict(self, context, model_input, params=None):
|
|
with torch.no_grad():
|
|
input_tensor = torch.from_numpy(model_input.values.astype(np.float32))
|
|
output_tensor = self.pytorch_model(input_tensor)
|
|
return pd.DataFrame(output_tensor.numpy())
|
|
|
|
pyfunc_artifact_path = "pyfunc_model"
|
|
with mlflow.start_run():
|
|
model_info = pyfunc.log_model(
|
|
pyfunc_artifact_path,
|
|
python_model=TorchValidatorModel(),
|
|
artifacts={"pytorch_model": model_path},
|
|
input_example=data[0],
|
|
# save file into code_paths, otherwise after first model loading (happens when
|
|
# validating input_example) then we can not load the model again
|
|
code_paths=[__file__],
|
|
)
|
|
|
|
# Deploy the custom pyfunc model and ensure that it is able to successfully load its
|
|
# constituent PyTorch model via `mlflow.pytorch.load_model`
|
|
inference_payload = load_serving_example(model_info.model_uri)
|
|
scoring_response = pyfunc_serve_and_score_model(
|
|
model_uri=model_info.model_uri,
|
|
data=inference_payload,
|
|
content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON,
|
|
extra_args=["--env-manager", "local"],
|
|
)
|
|
assert scoring_response.status_code == 200
|
|
|
|
deployed_model_preds = pd.DataFrame(json.loads(scoring_response.content)["predictions"])
|
|
np.testing.assert_array_almost_equal(
|
|
deployed_model_preds.values[:, 0],
|
|
_predict(model=module_scoped_subclassed_model, data=data),
|
|
decimal=4,
|
|
)
|
|
|
|
|
|
def test_load_pyfunc_loads_torch_model_using_pickle_module_specified_at_save_time(
|
|
module_scoped_subclassed_model, model_path
|
|
):
|
|
custom_pickle_module = pickle
|
|
|
|
mlflow.pytorch.save_model(
|
|
path=model_path,
|
|
pytorch_model=module_scoped_subclassed_model,
|
|
pickle_module=custom_pickle_module,
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
import_module_fn = importlib.import_module
|
|
imported_modules = []
|
|
|
|
def track_module_imports(module_name):
|
|
imported_modules.append(module_name)
|
|
return import_module_fn(module_name)
|
|
|
|
with (
|
|
mock.patch("importlib.import_module") as import_mock,
|
|
mock.patch("torch.load") as torch_load_mock,
|
|
):
|
|
import_mock.side_effect = track_module_imports
|
|
pyfunc.load_model(model_path)
|
|
|
|
expected_kwargs = {"pickle_module": custom_pickle_module}
|
|
if ENABLE_LEGACY_DESERIALIZATION:
|
|
expected_kwargs["weights_only"] = False
|
|
|
|
torch_load_mock.assert_called_with(mock.ANY, **expected_kwargs)
|
|
assert custom_pickle_module.__name__ in imported_modules
|
|
|
|
|
|
def test_load_model_loads_torch_model_using_pickle_module_specified_at_save_time(
|
|
module_scoped_subclassed_model,
|
|
):
|
|
custom_pickle_module = pickle
|
|
|
|
artifact_path = "pytorch_model"
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
module_scoped_subclassed_model,
|
|
name=artifact_path,
|
|
pickle_module=custom_pickle_module,
|
|
serialization_format="pickle",
|
|
)
|
|
model_uri = model_info.model_uri
|
|
|
|
import_module_fn = importlib.import_module
|
|
imported_modules = []
|
|
|
|
def track_module_imports(module_name):
|
|
imported_modules.append(module_name)
|
|
return import_module_fn(module_name)
|
|
|
|
with (
|
|
mock.patch("importlib.import_module") as import_mock,
|
|
mock.patch("torch.load") as torch_load_mock,
|
|
):
|
|
import_mock.side_effect = track_module_imports
|
|
pyfunc.load_model(model_uri=model_uri)
|
|
|
|
expected_kwargs = {"pickle_module": custom_pickle_module}
|
|
if ENABLE_LEGACY_DESERIALIZATION:
|
|
expected_kwargs["weights_only"] = False
|
|
|
|
torch_load_mock.assert_called_with(mock.ANY, **expected_kwargs)
|
|
assert custom_pickle_module.__name__ in imported_modules
|
|
|
|
|
|
def test_load_pyfunc_succeeds_when_data_is_model_file_instead_of_directory(
|
|
module_scoped_subclassed_model, model_path, data
|
|
):
|
|
"""
|
|
This test verifies that PyTorch models saved in older versions of MLflow are loaded successfully
|
|
by ``mlflow.pytorch.load_model``. The ``data`` path associated with these older models is
|
|
serialized PyTorch model file, as opposed to the current format: a directory containing a
|
|
serialized model file and pickle module information.
|
|
"""
|
|
mlflow.pytorch.save_model(
|
|
path=model_path, pytorch_model=module_scoped_subclassed_model, serialization_format="pickle"
|
|
)
|
|
|
|
model_conf_path = os.path.join(model_path, "MLmodel")
|
|
model_conf = Model.load(model_conf_path)
|
|
pyfunc_conf = model_conf.flavors.get(pyfunc.FLAVOR_NAME)
|
|
assert pyfunc_conf is not None
|
|
model_data_path = os.path.join(model_path, pyfunc_conf[pyfunc.DATA])
|
|
assert os.path.exists(model_data_path)
|
|
assert mlflow.pytorch._SERIALIZED_TORCH_MODEL_FILE_NAME in os.listdir(model_data_path)
|
|
pyfunc_conf[pyfunc.DATA] = os.path.join(
|
|
model_data_path, mlflow.pytorch._SERIALIZED_TORCH_MODEL_FILE_NAME
|
|
)
|
|
model_conf.save(model_conf_path)
|
|
|
|
loaded_pyfunc = pyfunc.load_model(model_path)
|
|
|
|
np.testing.assert_array_almost_equal(
|
|
loaded_pyfunc.predict(data[0]),
|
|
pd.DataFrame(_predict(model=module_scoped_subclassed_model, data=data)),
|
|
decimal=4,
|
|
)
|
|
|
|
|
|
def test_load_model_succeeds_when_data_is_model_file_instead_of_directory(
|
|
module_scoped_subclassed_model, model_path, data
|
|
):
|
|
"""
|
|
This test verifies that PyTorch models saved in older versions of MLflow are loaded successfully
|
|
by ``mlflow.pytorch.load_model``. The ``data`` path associated with these older models is
|
|
serialized PyTorch model file, as opposed to the current format: a directory containing a
|
|
serialized model file and pickle module information.
|
|
"""
|
|
artifact_path = "pytorch_model"
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
module_scoped_subclassed_model, name=artifact_path, serialization_format="pickle"
|
|
)
|
|
model_path = _download_artifact_from_uri(model_info.model_uri)
|
|
|
|
model_conf_path = os.path.join(model_path, "MLmodel")
|
|
model_conf = Model.load(model_conf_path)
|
|
pyfunc_conf = model_conf.flavors.get(pyfunc.FLAVOR_NAME)
|
|
assert pyfunc_conf is not None
|
|
model_data_path = os.path.join(model_path, pyfunc_conf[pyfunc.DATA])
|
|
assert os.path.exists(model_data_path)
|
|
assert mlflow.pytorch._SERIALIZED_TORCH_MODEL_FILE_NAME in os.listdir(model_data_path)
|
|
pyfunc_conf[pyfunc.DATA] = os.path.join(
|
|
model_data_path, mlflow.pytorch._SERIALIZED_TORCH_MODEL_FILE_NAME
|
|
)
|
|
model_conf.save(model_conf_path)
|
|
|
|
loaded_pyfunc = pyfunc.load_model(model_path)
|
|
|
|
np.testing.assert_array_almost_equal(
|
|
loaded_pyfunc.predict(data[0]),
|
|
pd.DataFrame(_predict(model=module_scoped_subclassed_model, data=data)),
|
|
decimal=4,
|
|
)
|
|
|
|
|
|
def test_load_model_allows_user_to_override_pickle_module_via_keyword_argument(
|
|
module_scoped_subclassed_model, model_path
|
|
):
|
|
mlflow.pytorch.save_model(
|
|
path=model_path,
|
|
pytorch_model=module_scoped_subclassed_model,
|
|
pickle_module=pickle,
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
with (
|
|
mock.patch("torch.load") as torch_load_mock,
|
|
mock.patch("mlflow.pytorch._logger.warning") as warn_mock,
|
|
):
|
|
mlflow.pytorch.load_model(model_uri=model_path, pickle_module=mlflow_pytorch_pickle_module)
|
|
torch_load_mock.assert_called_with(mock.ANY, pickle_module=mlflow_pytorch_pickle_module)
|
|
warn_mock.assert_any_call(mock.ANY, mlflow_pytorch_pickle_module.__name__, pickle.__name__)
|
|
|
|
|
|
def test_load_model_raises_exception_when_pickle_module_cannot_be_imported(
|
|
main_scoped_subclassed_model, model_path
|
|
):
|
|
mlflow.pytorch.save_model(
|
|
path=model_path, pytorch_model=main_scoped_subclassed_model, serialization_format="pickle"
|
|
)
|
|
|
|
bad_pickle_module_name = "not.a.real.module"
|
|
|
|
pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
|
|
model_data_path = os.path.join(model_path, pyfunc_conf[pyfunc.DATA])
|
|
assert os.path.exists(model_data_path)
|
|
assert mlflow.pytorch._PICKLE_MODULE_INFO_FILE_NAME in os.listdir(model_data_path)
|
|
with open(
|
|
os.path.join(model_data_path, mlflow.pytorch._PICKLE_MODULE_INFO_FILE_NAME), "w"
|
|
) as f:
|
|
f.write(bad_pickle_module_name)
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Failed to import the pickle module.+" + re.escape(bad_pickle_module_name),
|
|
):
|
|
mlflow.pytorch.load_model(model_uri=model_path)
|
|
|
|
|
|
def test_pyfunc_serve_and_score(data):
|
|
model = torch.nn.Linear(4, 1)
|
|
train_model(model=model, data=data)
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
model, name="model", input_example=data[0], serialization_format="pickle"
|
|
)
|
|
|
|
inference_payload = load_serving_example(model_info.model_uri)
|
|
resp = pyfunc_serve_and_score_model(
|
|
model_info.model_uri,
|
|
inference_payload,
|
|
pyfunc_scoring_server.CONTENT_TYPE_JSON,
|
|
extra_args=EXTRA_PYFUNC_SERVING_TEST_ARGS,
|
|
)
|
|
from mlflow.deployments import PredictionsResponse
|
|
|
|
scores = PredictionsResponse.from_json(resp.content).get_predictions()
|
|
np.testing.assert_array_almost_equal(scores.values[:, 0], _predict(model=model, data=data))
|
|
|
|
|
|
@pytest.mark.skipif(not _is_importable("transformers"), reason="This test requires transformers")
|
|
def test_pyfunc_serve_and_score_transformers():
|
|
from transformers import BertConfig, BertModel
|
|
|
|
from mlflow.deployments import PredictionsResponse
|
|
|
|
class MyBertModel(BertModel):
|
|
def forward(self, *args, **kwargs):
|
|
return super().forward(*args, **kwargs).last_hidden_state
|
|
|
|
model = MyBertModel(
|
|
BertConfig(
|
|
vocab_size=16,
|
|
hidden_size=2,
|
|
num_hidden_layers=2,
|
|
num_attention_heads=2,
|
|
intermediate_size=2,
|
|
)
|
|
)
|
|
model.eval()
|
|
|
|
input_ids = model.dummy_inputs["input_ids"]
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
model,
|
|
name="model",
|
|
input_example=np.array(input_ids.tolist()),
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
inference_payload = load_serving_example(model_info.model_uri)
|
|
resp = pyfunc_serve_and_score_model(
|
|
model_info.model_uri,
|
|
inference_payload,
|
|
pyfunc_scoring_server.CONTENT_TYPE_JSON,
|
|
extra_args=EXTRA_PYFUNC_SERVING_TEST_ARGS,
|
|
)
|
|
|
|
scores = PredictionsResponse.from_json(resp.content.decode("utf-8")).get_predictions(
|
|
predictions_format="ndarray"
|
|
)
|
|
assert_array_almost_equal(scores, model(input_ids).detach().numpy(), rtol=1e-6)
|
|
|
|
|
|
@pytest.fixture
|
|
def create_requirements_file(tmp_path):
|
|
requirement_file_name = "requirements.txt"
|
|
fp = tmp_path.joinpath(requirement_file_name)
|
|
test_string = "mlflow"
|
|
fp.write_text(test_string)
|
|
return str(fp), test_string
|
|
|
|
|
|
@pytest.fixture
|
|
def create_extra_files(tmp_path):
|
|
fp1 = tmp_path.joinpath("extra1.txt")
|
|
fp2 = tmp_path.joinpath("extra2.txt")
|
|
fp1.write_text("1")
|
|
fp2.write_text("2")
|
|
return [str(fp1), str(fp2)], ["1", "2"]
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_extra_files_log_model(create_extra_files, sequential_model):
|
|
extra_files, contents_expected = create_extra_files
|
|
with mlflow.start_run():
|
|
mlflow.pytorch.log_model(
|
|
sequential_model, name="models", extra_files=extra_files, serialization_format="pickle"
|
|
)
|
|
|
|
model_uri = "runs:/{run_id}/{model_path}".format(
|
|
run_id=mlflow.active_run().info.run_id, model_path="models"
|
|
)
|
|
with TempDir(remove_on_exit=True) as tmp:
|
|
model_path = _download_artifact_from_uri(model_uri, tmp.path())
|
|
model_config_path = os.path.join(model_path, "MLmodel")
|
|
model_config = Model.load(model_config_path)
|
|
flavor_config = model_config.flavors["pytorch"]
|
|
|
|
assert "extra_files" in flavor_config
|
|
loaded_extra_files = flavor_config["extra_files"]
|
|
|
|
for loaded_extra_file, content_expected in zip(loaded_extra_files, contents_expected):
|
|
assert "path" in loaded_extra_file
|
|
extra_file_path = os.path.join(model_path, loaded_extra_file["path"])
|
|
with open(extra_file_path) as fp:
|
|
assert fp.read() == content_expected
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_extra_files_save_model(create_extra_files, sequential_model):
|
|
extra_files, contents_expected = create_extra_files
|
|
with TempDir(remove_on_exit=True) as tmp:
|
|
model_path = os.path.join(tmp.path(), "models")
|
|
mlflow.pytorch.save_model(
|
|
pytorch_model=sequential_model,
|
|
path=model_path,
|
|
extra_files=extra_files,
|
|
serialization_format="pickle",
|
|
)
|
|
model_config_path = os.path.join(model_path, "MLmodel")
|
|
model_config = Model.load(model_config_path)
|
|
flavor_config = model_config.flavors["pytorch"]
|
|
|
|
assert "extra_files" in flavor_config
|
|
loaded_extra_files = flavor_config["extra_files"]
|
|
|
|
for loaded_extra_file, content_expected in zip(loaded_extra_files, contents_expected):
|
|
assert "path" in loaded_extra_file
|
|
extra_file_path = os.path.join(model_path, loaded_extra_file["path"])
|
|
with open(extra_file_path) as fp:
|
|
assert fp.read() == content_expected
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_log_model_invalid_extra_file_path(sequential_model):
|
|
with (
|
|
mlflow.start_run(),
|
|
pytest.raises(MlflowException, match="No such artifact: 'non_existing_file.txt'"),
|
|
):
|
|
mlflow.pytorch.log_model(
|
|
sequential_model,
|
|
name="models",
|
|
extra_files=["non_existing_file.txt"],
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_log_model_invalid_extra_file_type(sequential_model):
|
|
with (
|
|
mlflow.start_run(),
|
|
pytest.raises(TypeError, match="Extra files argument should be a list"),
|
|
):
|
|
mlflow.pytorch.log_model(
|
|
sequential_model,
|
|
name="models",
|
|
extra_files="non_existing_file.txt",
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
|
|
def state_dict_equal(state_dict1, state_dict2):
|
|
for key1 in state_dict1:
|
|
if key1 not in state_dict2:
|
|
return False
|
|
|
|
value1 = state_dict1[key1]
|
|
value2 = state_dict2[key1]
|
|
|
|
if type(value1) != type(value2):
|
|
return False
|
|
elif isinstance(value1, dict):
|
|
if not state_dict_equal(value1, value2):
|
|
return False
|
|
elif isinstance(value1, torch.Tensor):
|
|
if not torch.equal(value1, value2):
|
|
return False
|
|
elif value1 != value2:
|
|
return False
|
|
else:
|
|
continue
|
|
|
|
return True
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_save_state_dict(sequential_model, model_path, data):
|
|
state_dict = sequential_model.state_dict()
|
|
mlflow.pytorch.save_state_dict(state_dict, model_path)
|
|
|
|
loaded_state_dict = mlflow.pytorch.load_state_dict(model_path)
|
|
assert state_dict_equal(loaded_state_dict, state_dict)
|
|
model = get_sequential_model()
|
|
model.load_state_dict(loaded_state_dict)
|
|
np.testing.assert_array_almost_equal(
|
|
_predict(model, data),
|
|
_predict(sequential_model, data),
|
|
decimal=4,
|
|
)
|
|
|
|
|
|
def test_save_state_dict_can_save_nested_state_dict(model_path):
|
|
"""
|
|
This test ensures that `save_state_dict` supports a use case described in the page below
|
|
where a user bundles multiple objects (e.g., model, optimizer, learning-rate scheduler)
|
|
into a single nested state_dict and loads it back later for inference or re-training:
|
|
https://pytorch.org/tutorials/recipes/recipes/saving_and_loading_a_general_checkpoint.html
|
|
"""
|
|
model = get_sequential_model()
|
|
optim = torch.optim.Adam(model.parameters())
|
|
state_dict = {"model": model.state_dict(), "optim": optim.state_dict()}
|
|
mlflow.pytorch.save_state_dict(state_dict, model_path)
|
|
|
|
loaded_state_dict = mlflow.pytorch.load_state_dict(model_path)
|
|
assert state_dict_equal(loaded_state_dict, state_dict)
|
|
model.load_state_dict(loaded_state_dict["model"])
|
|
optim.load_state_dict(loaded_state_dict["optim"])
|
|
|
|
|
|
def test_load_state_dict_disallows_pickle_deserialization(model_path, monkeypatch):
|
|
model = get_sequential_model()
|
|
mlflow.pytorch.save_state_dict(model.state_dict(), model_path)
|
|
|
|
monkeypatch.setenv("MLFLOW_ALLOW_PICKLE_DESERIALIZATION", "false")
|
|
with pytest.raises(MlflowException, match="MLFLOW_ALLOW_PICKLE_DESERIALIZATION"):
|
|
mlflow.pytorch.load_state_dict(model_path)
|
|
|
|
|
|
@pytest.mark.parametrize("not_state_dict", [0, "", get_sequential_model()])
|
|
def test_save_state_dict_throws_for_invalid_object_type(not_state_dict, model_path):
|
|
with pytest.raises(TypeError, match="Invalid object type for `state_dict`"):
|
|
mlflow.pytorch.save_state_dict(not_state_dict, model_path)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_log_state_dict(sequential_model, data):
|
|
artifact_path = "model"
|
|
state_dict = sequential_model.state_dict()
|
|
with mlflow.start_run():
|
|
mlflow.pytorch.log_state_dict(state_dict, artifact_path)
|
|
state_dict_uri = mlflow.get_artifact_uri(artifact_path)
|
|
|
|
loaded_state_dict = mlflow.pytorch.load_state_dict(state_dict_uri)
|
|
assert state_dict_equal(loaded_state_dict, state_dict)
|
|
model = get_sequential_model()
|
|
model.load_state_dict(loaded_state_dict)
|
|
np.testing.assert_array_almost_equal(
|
|
_predict(model, data),
|
|
_predict(sequential_model, data),
|
|
decimal=4,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_log_model_with_code_paths(sequential_model):
|
|
artifact_path = "model"
|
|
with (
|
|
mlflow.start_run(),
|
|
mock.patch("mlflow.pytorch._add_code_from_conf_to_system_path") as add_mock,
|
|
):
|
|
model_info = mlflow.pytorch.log_model(
|
|
sequential_model,
|
|
name=artifact_path,
|
|
code_paths=[__file__],
|
|
serialization_format="pickle",
|
|
)
|
|
_compare_logged_code_paths(__file__, model_info.model_uri, mlflow.pytorch.FLAVOR_NAME)
|
|
mlflow.pytorch.load_model(model_info.model_uri)
|
|
add_mock.assert_called()
|
|
|
|
|
|
def test_virtualenv_subfield_points_to_correct_path(model_path):
|
|
model = get_sequential_model()
|
|
mlflow.pytorch.save_model(model, path=model_path, serialization_format="pickle")
|
|
pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
|
|
python_env_path = Path(model_path, pyfunc_conf[pyfunc.ENV]["virtualenv"])
|
|
assert python_env_path.exists()
|
|
assert python_env_path.is_file()
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_model_save_load_with_metadata(sequential_model, model_path):
|
|
mlflow.pytorch.save_model(
|
|
sequential_model,
|
|
path=model_path,
|
|
metadata={"metadata_key": "metadata_value"},
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
reloaded_model = mlflow.pyfunc.load_model(model_uri=model_path)
|
|
assert reloaded_model.metadata.metadata["metadata_key"] == "metadata_value"
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_model_log_with_metadata(sequential_model):
|
|
artifact_path = "model"
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
sequential_model,
|
|
name=artifact_path,
|
|
metadata={"metadata_key": "metadata_value"},
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
reloaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
|
|
assert reloaded_model.metadata.metadata["metadata_key"] == "metadata_value"
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [True, False])
|
|
def test_model_log_with_signature_inference(sequential_model, data):
|
|
artifact_path = "model"
|
|
example_ = data[0].head(3).values.astype(np.float32)
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
sequential_model,
|
|
name=artifact_path,
|
|
input_example=example_,
|
|
serialization_format="pickle",
|
|
)
|
|
|
|
assert model_info.signature == ModelSignature(
|
|
inputs=Schema([TensorSpec(np.dtype("float32"), (-1, 4))]),
|
|
outputs=Schema([TensorSpec(np.dtype("float32"), (-1, 1))]),
|
|
)
|
|
inference_payload = load_serving_example(model_info.model_uri)
|
|
response = pyfunc_serve_and_score_model(
|
|
model_info.model_uri,
|
|
inference_payload,
|
|
pyfunc_scoring_server.CONTENT_TYPE_JSON,
|
|
extra_args=["--env-manager", "local"],
|
|
)
|
|
assert response.status_code == 200
|
|
deployed_model_preds = pd.DataFrame(json.loads(response.content)["predictions"])
|
|
np.testing.assert_array_almost_equal(
|
|
deployed_model_preds.values[:, 0],
|
|
_predict(model=sequential_model, data=(data[0].head(3), data[1].head(3))),
|
|
decimal=4,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("scripted_model", [False])
|
|
def test_load_model_to_device(sequential_model):
|
|
with mock.patch("mlflow.pytorch._load_model") as load_model_mock:
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
sequential_model, name="pytorch", serialization_format="pickle"
|
|
)
|
|
model_config = {"device": "cuda"}
|
|
if ENABLE_LEGACY_DESERIALIZATION:
|
|
model_config["weights_only"] = False
|
|
|
|
mlflow.pyfunc.load_model(model_uri=model_info.model_uri, model_config=model_config)
|
|
|
|
load_model_mock.assert_called_with(mock.ANY, **model_config)
|
|
mlflow.pytorch.load_model(model_uri=model_info.model_uri, **model_config)
|
|
load_model_mock.assert_called_with(path=mock.ANY, **model_config)
|
|
|
|
|
|
def test_passing_params_to_model(data):
|
|
class CustomModel(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear = torch.nn.Linear(4, 1)
|
|
|
|
def forward(self, x, y):
|
|
if not torch.is_tensor(x):
|
|
x = torch.from_numpy(x)
|
|
y = torch.tensor(y)
|
|
combined = x * y
|
|
return self.linear(combined)
|
|
|
|
model = CustomModel()
|
|
x = np.random.randn(8, 4).astype(np.float32)
|
|
|
|
signature = mlflow.models.infer_signature(x, None, {"y": 1})
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pytorch.log_model(
|
|
model, name="model", signature=signature, serialization_format="pickle"
|
|
)
|
|
|
|
pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
with torch.no_grad():
|
|
np.testing.assert_array_almost_equal(pyfunc_model.predict(x), model(x, 1), decimal=4)
|
|
np.testing.assert_array_almost_equal(
|
|
pyfunc_model.predict(x, {"y": 2}), model(x, 2), decimal=4
|
|
)
|
|
|
|
|
|
def test_log_model_with_datetime_input():
|
|
df = pd.DataFrame({
|
|
"datetime": pd.date_range("2022-01-01", periods=5, freq="D"),
|
|
"x": np.random.uniform(20, 30, 5),
|
|
"y": np.random.uniform(2, 4, 5),
|
|
"z": np.random.uniform(0, 10, 5),
|
|
})
|
|
model = get_sequential_model()
|
|
model_info = mlflow.pytorch.log_model(
|
|
model, name="pytorch", input_example=df, serialization_format="pickle"
|
|
)
|
|
assert model_info.signature.inputs.inputs[0].type == DataType.datetime
|
|
pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
with torch.no_grad():
|
|
input_tensor = torch.from_numpy(df.to_numpy(dtype=np.float32))
|
|
expected_result = model(input_tensor)
|
|
with torch.no_grad():
|
|
np.testing.assert_array_almost_equal(pyfunc_model.predict(df), expected_result, decimal=4)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
Version(torch.__version__) < Version("2.4"), reason="This test requires torch>=2.4"
|
|
)
|
|
@pytest.mark.parametrize("scripted_model", [False])
|
|
def test_save_and_load_exported_model(sequential_model, model_path, data, sequential_predicted):
|
|
input_example = data[0].to_numpy(dtype=np.float32)
|
|
|
|
mlflow.pytorch.save_model(
|
|
sequential_model,
|
|
model_path,
|
|
serialization_format="pt2",
|
|
input_example=input_example,
|
|
)
|
|
|
|
# Loading pytorch model
|
|
sequential_model_loaded = mlflow.pytorch.load_model(model_path)
|
|
np.testing.assert_array_equal(_predict(sequential_model_loaded, data), sequential_predicted)
|
|
|
|
# Loading pyfunc model
|
|
pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
|
|
np.testing.assert_array_almost_equal(
|
|
pyfunc_loaded.predict(input_example)[:, 0], sequential_predicted, decimal=4
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
Version(torch.__version__) < Version("2.4"), reason="This test requires torch>=2.4"
|
|
)
|
|
def test_exported_model_infer_dynamic_dim(tmp_path):
|
|
class MyModule(torch.nn.Module):
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return torch.sin(x)
|
|
|
|
origin_model = MyModule()
|
|
|
|
input_example = torch.randn(3, 4, 5).numpy()
|
|
|
|
save_path1 = tmp_path / "model1"
|
|
|
|
# test exporting model with auto inferred signature,
|
|
# which sets the first dim (batch dim) of input data as dynamic dim.
|
|
mlflow.pytorch.save_model(
|
|
origin_model,
|
|
save_path1,
|
|
serialization_format="pt2",
|
|
input_example=input_example,
|
|
)
|
|
|
|
# Test the exported model works with test data that changes the first dim (batch dim) size.
|
|
loaded_model1 = mlflow.pytorch.load_model(save_path1)
|
|
|
|
test_data1 = torch.randn(6, 4, 5)
|
|
np.testing.assert_array_almost_equal(
|
|
loaded_model1(test_data1),
|
|
origin_model(test_data1),
|
|
decimal=4,
|
|
)
|
|
|
|
save_path2 = tmp_path / "model2"
|
|
# test exporting model with provided signature,
|
|
# which sets the second dim of input data as dynamic dim.
|
|
mlflow.pytorch.save_model(
|
|
origin_model,
|
|
save_path2,
|
|
serialization_format="pt2",
|
|
input_example=input_example,
|
|
signature=ModelSignature(
|
|
inputs=Schema([TensorSpec(np.dtype("float32"), (3, -1, 5))]),
|
|
),
|
|
)
|
|
|
|
# Test the exported model works with test data that changes the second dim (batch dim) size.
|
|
loaded_model2 = mlflow.pytorch.load_model(save_path2)
|
|
|
|
test_data2 = torch.randn(3, 2, 5)
|
|
np.testing.assert_array_almost_equal(
|
|
loaded_model2(test_data2),
|
|
origin_model(test_data2),
|
|
decimal=4,
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
Version(torch.__version__) < Version("2.4"), reason="This test requires torch>=2.4"
|
|
)
|
|
@pytest.mark.parametrize("scripted_model", [False])
|
|
def test_load_exported_model_check_device_mismatch(sequential_model, model_path):
|
|
mlflow.pytorch.save_model(
|
|
sequential_model,
|
|
model_path,
|
|
serialization_format="pt2",
|
|
input_example=torch.randn(3, 4).numpy(),
|
|
)
|
|
|
|
# test loading model to CPU works
|
|
mlflow.pytorch.load_model(model_path, device="cpu")
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="it can't be loaded on 'cuda' device.",
|
|
):
|
|
mlflow.pytorch.load_model(model_path, device="cuda")
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
Version(torch.__version__) < Version("2.4"), reason="This test requires torch>=2.4"
|
|
)
|
|
def test_save_and_load_exported_model_with_multi_inputs(model_path):
|
|
|
|
class CustomModel(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.linear = torch.nn.Linear(4, 1)
|
|
|
|
def forward(self, x, y):
|
|
with torch.no_grad():
|
|
return self.linear(x + y)
|
|
|
|
model = CustomModel()
|
|
input_example = (torch.randn(10, 4), torch.randn(10, 4))
|
|
|
|
mlflow.pytorch.save_model(
|
|
model,
|
|
model_path,
|
|
serialization_format="pt2",
|
|
input_example=input_example,
|
|
signature=ModelSignature(
|
|
inputs=Schema([
|
|
TensorSpec(np.dtype("float32"), (-1, 4), "v1"),
|
|
TensorSpec(np.dtype("float32"), (-1, 4), "v2"),
|
|
]),
|
|
),
|
|
)
|
|
|
|
model_loaded = mlflow.pytorch.load_model(model_path)
|
|
|
|
np.testing.assert_array_almost_equal(
|
|
model(*input_example),
|
|
model_loaded(*input_example),
|
|
decimal=4,
|
|
)
|