def test_sparkml_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies(
spark_model_iris):
artifact_path = "model"
with kiwi.start_run():
sparkm.log_model(
spark_model=spark_model_iris.model, artifact_path=artifact_path, conda_env=None)
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=artifact_path)
model_path = _download_artifact_from_uri(artifact_uri=model_uri)
pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
with open(conda_env_path, "r") as f:
conda_env = yaml.safe_load(f)
assert conda_env == sparkm.get_default_conda_env()
def test_model_save_persists_specified_conda_env_in_mlflow_model_directory(
spacy_model_with_data, model_path, spacy_custom_env):
kiwi.spacy.save_model(spacy_model=spacy_model_with_data.model,
path=model_path,
conda_env=spacy_custom_env)
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])
assert os.path.exists(saved_conda_env_path)
assert saved_conda_env_path != spacy_custom_env
with open(spacy_custom_env, "r") as f:
spacy_custom_env_text = f.read()
with open(saved_conda_env_path, "r") as f:
saved_conda_env_text = f.read()
assert saved_conda_env_text == spacy_custom_env_text
def test_model_save_persists_specified_conda_env_in_mlflow_model_directory(
sklearn_knn_model, model_path, sklearn_custom_env):
kiwi.sklearn.save_model(sk_model=sklearn_knn_model.model,
path=model_path,
conda_env=sklearn_custom_env)
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])
assert os.path.exists(saved_conda_env_path)
assert saved_conda_env_path != sklearn_custom_env
with open(sklearn_custom_env, "r") as f:
sklearn_custom_env_parsed = yaml.safe_load(f)
with open(saved_conda_env_path, "r") as f:
saved_conda_env_parsed = yaml.safe_load(f)
assert saved_conda_env_parsed == sklearn_custom_env_parsed
def test_save_model_accepts_conda_env_as_dict(saved_tf_iris_model, model_path):
conda_env = dict(kiwi.tensorflow.get_default_conda_env())
conda_env["dependencies"].append("pytest")
kiwi.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=model_path,
conda_env=conda_env)
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])
assert os.path.exists(saved_conda_env_path)
with open(saved_conda_env_path, "r") as f:
saved_conda_env_parsed = yaml.safe_load(f)
assert saved_conda_env_parsed == conda_env
def test_save_model_persists_specified_conda_env_in_mlflow_model_directory(
saved_tf_iris_model, model_path, tf_custom_env):
kiwi.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=model_path,
conda_env=tf_custom_env)
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])
assert os.path.exists(saved_conda_env_path)
assert saved_conda_env_path != tf_custom_env
with open(tf_custom_env, "r") as f:
tf_custom_env_text = f.read()
with open(saved_conda_env_path, "r") as f:
saved_conda_env_text = f.read()
assert saved_conda_env_text == tf_custom_env_text
def test_model_log_persists_specified_conda_env_in_mlflow_model_directory(
sequential_model, pytorch_custom_env):
artifact_path = "model"
with kiwi.start_run():
kiwi.pytorch.log_model(pytorch_model=sequential_model,
artifact_path=artifact_path,
conda_env=pytorch_custom_env)
model_path = _download_artifact_from_uri("runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id, artifact_path=artifact_path))
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])
assert os.path.exists(saved_conda_env_path)
assert saved_conda_env_path != pytorch_custom_env
with open(pytorch_custom_env, "r") as f:
pytorch_custom_env_text = f.read()
with open(saved_conda_env_path, "r") as f:
saved_conda_env_text = f.read()
assert saved_conda_env_text == pytorch_custom_env_text
def _get_and_parse_flavor_configuration(model_path):
"""
:param path: Local filesystem path to the MLflow Model with the ``tensorflow`` flavor.
:return: A triple containing the following elements:
- ``tf_saved_model_dir``: The local filesystem path to the underlying TensorFlow
SavedModel directory.
- ``tf_meta_graph_tags``: A list of tags identifying the TensorFlow model's metagraph
within the serialized ``SavedModel`` object.
- ``tf_signature_def_key``: A string identifying the input/output signature associated
with the model. This is a key within the serialized
``SavedModel``'s signature definition mapping.
"""
flavor_conf = _get_flavor_configuration(model_path=model_path,
flavor_name=FLAVOR_NAME)
tf_saved_model_dir = os.path.join(model_path,
flavor_conf['saved_model_dir'])
tf_meta_graph_tags = flavor_conf['meta_graph_tags']
tf_signature_def_key = flavor_conf['signature_def_key']
return tf_saved_model_dir, tf_meta_graph_tags, tf_signature_def_key
def test_model_log_persists_specified_conda_env_in_mlflow_model_directory(onnx_model,
onnx_custom_env):
import kiwi.onnx
artifact_path = "model"
with kiwi.start_run():
kiwi.onnx.log_model(
onnx_model=onnx_model, artifact_path=artifact_path, conda_env=onnx_custom_env)
model_path = _download_artifact_from_uri("runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id, artifact_path=artifact_path))
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])
assert os.path.exists(saved_conda_env_path)
assert saved_conda_env_path != onnx_custom_env
with open(onnx_custom_env, "r") as f:
onnx_custom_env_parsed = yaml.safe_load(f)
with open(saved_conda_env_path, "r") as f:
saved_conda_env_parsed = yaml.safe_load(f)
assert saved_conda_env_parsed == onnx_custom_env_parsed
def test_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies(
sklearn_knn_model):
artifact_path = "model"
knn_model = sklearn_knn_model.model
with kiwi.start_run():
kiwi.sklearn.log_model(
sk_model=knn_model,
artifact_path=artifact_path,
conda_env=None,
serialization_format=kiwi.sklearn.SERIALIZATION_FORMAT_PICKLE)
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id, artifact_path=artifact_path)
model_path = _download_artifact_from_uri(artifact_uri=model_uri)
pyfunc_conf = _get_flavor_configuration(model_path=model_path,
flavor_name=pyfunc.FLAVOR_NAME)
conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
with open(conda_env_path, "r") as f:
conda_env = yaml.safe_load(f)
assert conda_env == kiwi.sklearn.get_default_conda_env()
def test_save_model_without_specified_conda_env_uses_default_env_with_expected_dependencies(
sklearn_logreg_model, main_scoped_model_class, tmpdir):
sklearn_model_path = os.path.join(str(tmpdir), "sklearn_model")
kiwi.sklearn.save_model(sk_model=sklearn_logreg_model,
path=sklearn_model_path)
pyfunc_model_path = os.path.join(str(tmpdir), "pyfunc_model")
kiwi.pyfunc.save_model(
path=pyfunc_model_path,
artifacts={"sk_model": sklearn_model_path},
python_model=main_scoped_model_class(predict_fn=None),
conda_env=_conda_env())
pyfunc_conf = _get_flavor_configuration(
model_path=pyfunc_model_path, flavor_name=kiwi.pyfunc.FLAVOR_NAME)
conda_env_path = os.path.join(pyfunc_model_path,
pyfunc_conf[kiwi.pyfunc.ENV])
with open(conda_env_path, "r") as f:
conda_env = yaml.safe_load(f)
assert conda_env == _conda_env()
def load_model(model_uri):
"""
Load an XGBoost model from a local file or a run.
:param model_uri: The location, in URI format, of the MLflow model. For example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/tracking.html#
artifact-locations>`_.
:return: An XGBoost model (an instance of `xgboost.Booster`_)
"""
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri)
flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME)
xgb_model_file_path = os.path.join(local_model_path, flavor_conf.get("data", "model.xgb"))
return _load_model(path=xgb_model_file_path)
def test_log_model_without_specified_conda_env_uses_default_env_with_expected_dependencies(
saved_tf_iris_model):
artifact_path = "model"
with kiwi.start_run():
kiwi.tensorflow.log_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,
artifact_path=artifact_path,
conda_env=None)
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id, artifact_path=artifact_path)
model_path = _download_artifact_from_uri(artifact_uri=model_uri)
pyfunc_conf = _get_flavor_configuration(model_path=model_path,
flavor_name=pyfunc.FLAVOR_NAME)
conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
with open(conda_env_path, "r") as f:
conda_env = yaml.safe_load(f)
assert conda_env == kiwi.tensorflow.get_default_conda_env()
def load_model(model_uri):
"""
Load a scikit-learn model from a local file or a run.
:param model_uri: The location, in URI format, of the MLflow model, for example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
- ``models:/<model_name>/<model_version>``
- ``models:/<model_name>/<stage>``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
artifact-locations>`_.
:return: A scikit-learn model.
.. code-block:: python
:caption: Example
import mlflow.sklearn
sk_model = mlflow.sklearn.load_model("runs:/96771d893a5e46159d9f3b49bf9013e2/sk_models")
# use Pandas DataFrame to make predictions
pandas_df = ...
predictions = sk_model.predict(pandas_df)
"""
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri)
flavor_conf = _get_flavor_configuration(model_path=local_model_path,
flavor_name=FLAVOR_NAME)
sklearn_model_artifacts_path = os.path.join(local_model_path,
flavor_conf['pickled_model'])
serialization_format = flavor_conf.get('serialization_format',
SERIALIZATION_FORMAT_PICKLE)
return _load_model_from_local_file(
path=sklearn_model_artifacts_path,
serialization_format=serialization_format)
def test_load_model_raises_exception_when_pickle_module_cannot_be_imported(
main_scoped_subclassed_model, model_path):
kiwi.pytorch.save_model(
path=model_path,
pytorch_model=main_scoped_subclassed_model,
conda_env=None)
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 kiwi.pytorch._PICKLE_MODULE_INFO_FILE_NAME in os.listdir(model_data_path)
with open(
os.path.join(model_data_path, kiwi.pytorch._PICKLE_MODULE_INFO_FILE_NAME), "w") as f:
f.write(bad_pickle_module_name)
with pytest.raises(MlflowException) as exc_info:
kiwi.pytorch.load_model(model_uri=model_path)
assert "Failed to import the pickle module" in str(exc_info)
assert bad_pickle_module_name in str(exc_info)
def load_model(model_uri, **kwargs):
"""
Load a Keras model from a local file or a run.
Extra arguments are passed through to keras.load_model.
:param model_uri: The location, in URI format, of the MLflow model. For example:
- ``/Users/me/path/to/local/model``
- ``relative/path/to/local/model``
- ``s3://my_bucket/path/to/model``
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
- ``models:/<model_name>/<model_version>``
- ``models:/<model_name>/<stage>``
For more information about supported URI schemes, see
`Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
artifact-locations>`_.
:return: A Keras model instance.
.. code-block:: python
:caption: Example
# Load persisted model as a Keras model or as a PyFunc, call predict() on a pandas DataFrame
keras_model = mlflow.keras.load_model("runs:/96771d893a5e46159d9f3b49bf9013e2" + "/models")
predictions = keras_model.predict(x_test)
"""
local_model_path = _download_artifact_from_uri(artifact_uri=model_uri)
flavor_conf = _get_flavor_configuration(model_path=local_model_path,
flavor_name=FLAVOR_NAME)
keras_module = importlib.import_module(
flavor_conf.get("keras_module", "keras"))
keras_model_artifacts_path = os.path.join(
local_model_path, flavor_conf.get("data", _MODEL_SAVE_PATH))
return _load_model(model_path=keras_model_artifacts_path,
keras_module=keras_module,
**kwargs)
def test_sparkml_model_log_persists_specified_conda_env_in_mlflow_model_directory(
spark_model_iris, model_path, spark_custom_env):
artifact_path = "model"
with kiwi.start_run():
sparkm.log_model(
spark_model=spark_model_iris.model,
artifact_path=artifact_path,
conda_env=spark_custom_env)
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=artifact_path)
model_path = _download_artifact_from_uri(artifact_uri=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])
assert os.path.exists(saved_conda_env_path)
assert saved_conda_env_path != spark_custom_env
with open(spark_custom_env, "r") as f:
spark_custom_env_parsed = yaml.safe_load(f)
with open(saved_conda_env_path, "r") as f:
saved_conda_env_parsed = yaml.safe_load(f)
assert saved_conda_env_parsed == spark_custom_env_parsed
def _load_pyfunc(model_path):
pyfunc_config = _get_flavor_configuration(
model_path=model_path, flavor_name=kiwi.pyfunc.FLAVOR_NAME)
python_model_cloudpickle_version = pyfunc_config.get(
CONFIG_KEY_CLOUDPICKLE_VERSION, None)
if python_model_cloudpickle_version is None:
kiwi.pyfunc._logger.warning(
"The version of CloudPickle used to save the model could not be found in the MLmodel"
" configuration")
elif python_model_cloudpickle_version != cloudpickle.__version__:
# CloudPickle does not have a well-defined cross-version compatibility policy. Micro version
# releases have been known to cause incompatibilities. Therefore, we match on the full
# library version
kiwi.pyfunc._logger.warning(
"The version of CloudPickle that was used to save the model, `CloudPickle %s`, differs"
" from the version of CloudPickle that is currently running, `CloudPickle %s`, and may"
" be incompatible", python_model_cloudpickle_version,
cloudpickle.__version__)
python_model_subpath = pyfunc_config.get(CONFIG_KEY_PYTHON_MODEL, None)
if python_model_subpath is None:
raise MlflowException(
"Python model path was not specified in the model configuration")
with open(os.path.join(model_path, python_model_subpath), "rb") as f:
python_model = cloudpickle.load(f)
artifacts = {}
for saved_artifact_name, saved_artifact_info in\
pyfunc_config.get(CONFIG_KEY_ARTIFACTS, {}).items():
artifacts[saved_artifact_name] = os.path.join(
model_path, saved_artifact_info[CONFIG_KEY_ARTIFACT_RELATIVE_PATH])
context = PythonModelContext(artifacts=artifacts)
python_model.load_context(context=context)
return _PythonModelPyfuncWrapper(python_model=python_model,
context=context)
def test_log_model_persists_specified_conda_env_in_mlflow_model_directory(
sklearn_knn_model, main_scoped_model_class, pyfunc_custom_env):
sklearn_artifact_path = "sk_model"
with kiwi.start_run():
kiwi.sklearn.log_model(sk_model=sklearn_knn_model,
artifact_path=sklearn_artifact_path)
sklearn_run_id = kiwi.active_run().info.run_id
pyfunc_artifact_path = "pyfunc_model"
with kiwi.start_run():
kiwi.pyfunc.log_model(
artifact_path=pyfunc_artifact_path,
artifacts={
"sk_model":
utils_get_artifact_uri(artifact_path=sklearn_artifact_path,
run_id=sklearn_run_id)
},
python_model=main_scoped_model_class(predict_fn=None),
conda_env=pyfunc_custom_env)
pyfunc_model_path = _download_artifact_from_uri(
"runs:/{run_id}/{artifact_path}".format(
run_id=kiwi.active_run().info.run_id,
artifact_path=pyfunc_artifact_path))
pyfunc_conf = _get_flavor_configuration(
model_path=pyfunc_model_path, flavor_name=kiwi.pyfunc.FLAVOR_NAME)
saved_conda_env_path = os.path.join(pyfunc_model_path,
pyfunc_conf[kiwi.pyfunc.ENV])
assert os.path.exists(saved_conda_env_path)
assert saved_conda_env_path != pyfunc_custom_env
with open(pyfunc_custom_env, "r") as f:
pyfunc_custom_env_parsed = yaml.safe_load(f)
with open(saved_conda_env_path, "r") as f:
saved_conda_env_parsed = yaml.safe_load(f)
assert saved_conda_env_parsed == pyfunc_custom_env_parsed
def test_get_flavor_configuration_throws_exception_when_model_configuration_does_not_exist(
model_path):
with pytest.raises(MlflowException) as exc:
mlflow_model_utils._get_flavor_configuration(
model_path=model_path, flavor_name=kiwi.mleap.FLAVOR_NAME)
assert exc.error_code == RESOURCE_DOES_NOT_EXIST
