def test_log_model_without_specified_conda_env_uses_default_env_with_expected_dependencies(
saved_tf_iris_model):
artifact_path = "model"
with mlflow.start_run():
mlflow.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=mlflow.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 == mlflow.tensorflow.get_default_conda_env()
def test_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies(
onnx_model):
import mlflow.onnx
artifact_path = "model"
with mlflow.start_run():
mlflow.onnx.log_model(onnx_model=onnx_model,
artifact_path=artifact_path,
conda_env=None)
model_path = _download_artifact_from_uri(
"runs:/{run_id}/{artifact_path}".format(
run_id=mlflow.active_run().info.run_id,
artifact_path=artifact_path))
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 == mlflow.onnx.get_default_conda_env()
def load_model(path, run_id=None):
"""
Load an H2O model from a local file (if ``run_id`` is ``None``) or a run.
This function expects there is an H2O instance initialised with ``h2o.init``.
:param path: Local filesystem path or run-relative artifact path to the model saved
by :py:func:`mlflow.h2o.save_model`.
:param run_id: Run ID. If provided, combined with ``path`` to identify the model.
"""
if run_id is not None:
path = mlflow.tracking.artifact_utils._get_model_log_dir(
model_name=path, run_id=run_id)
path = os.path.abspath(path)
flavor_conf = _get_flavor_configuration(model_path=path,
flavor_name=FLAVOR_NAME)
# Flavor configurations for models saved in MLflow version <= 0.8.0 may not contain a
# `data` key; in this case, we assume the model artifact path to be `model.h2o`
h2o_model_file_path = os.path.join(path,
flavor_conf.get("data", "model.h2o"))
return _load_model(path=h2o_model_file_path)
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")
mlflow.sklearn.save_model(sk_model=sklearn_logreg_model,
path=sklearn_model_path)
pyfunc_model_path = os.path.join(str(tmpdir), "pyfunc_model")
mlflow.pyfunc.save_model(
dst_path=pyfunc_model_path,
artifacts={"sk_model": sklearn_model_path},
python_model=main_scoped_model_class(predict_fn=None))
pyfunc_conf = _get_flavor_configuration(
model_path=pyfunc_model_path, flavor_name=mlflow.pyfunc.FLAVOR_NAME)
conda_env_path = os.path.join(pyfunc_model_path,
pyfunc_conf[mlflow.pyfunc.ENV])
with open(conda_env_path, "r") as f:
conda_env = yaml.safe_load(f)
assert conda_env == mlflow.pyfunc.model.DEFAULT_CONDA_ENV
def test_model_log_persists_specified_conda_env_in_mlflow_model_directory(
sklearn_knn_model, sklearn_custom_env):
artifact_path = "model"
with mlflow.start_run():
mlflow.sklearn.log_model(sk_model=sklearn_knn_model.model,
artifact_path=artifact_path,
conda_env=sklearn_custom_env)
run_id = mlflow.active_run().info.run_uuid
model_path = _get_model_log_dir(artifact_path, run_id)
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_log_model_without_specified_conda_env_uses_default_env_with_expected_dependencies(
saved_tf_iris_model, model_path):
artifact_path = "model"
with mlflow.start_run():
mlflow.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)
run_id = mlflow.active_run().info.run_id
model_path = _get_model_log_dir(artifact_path, run_id)
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 == mlflow.tensorflow.DEFAULT_CONDA_ENV
def test_model_log_persists_specified_conda_env_in_mlflow_model_directory(
model, pytorch_custom_env):
artifact_path = "model"
with mlflow.start_run():
mlflow.pytorch.log_model(pytorch_model=model,
artifact_path=artifact_path,
conda_env=pytorch_custom_env)
run_id = mlflow.active_run().info.run_uuid
model_path = tracking.utils._get_model_log_dir(artifact_path, run_id)
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 test_save_model_persists_specified_conda_env_in_mlflow_model_directory(
saved_tf_iris_model, model_path, tf_custom_env):
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=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(onnx_model,
onnx_custom_env):
import mlflow.onnx
artifact_path = "model"
with mlflow.start_run():
mlflow.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=mlflow.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_persists_specified_conda_env_in_mlflow_model_directory(
h2o_iris_model, h2o_custom_env):
artifact_path = "model"
with mlflow.start_run():
mlflow.h2o.log_model(h2o_model=h2o_iris_model.model,
artifact_path=artifact_path,
conda_env=h2o_custom_env)
model_path = _download_artifact_from_uri("runs:/{run_id}/{artifact_path}".format(
run_id=mlflow.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 != h2o_custom_env
with open(h2o_custom_env, "r") as f:
h2o_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 == h2o_custom_env_text
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 mlflow.start_run():
mlflow.sklearn.log_model(
sk_model=knn_model,
artifact_path=artifact_path,
conda_env=None,
serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE)
run_id = mlflow.active_run().info.run_uuid
model_path = _get_model_log_dir(artifact_path, run_id)
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 == mlflow.sklearn.DEFAULT_CONDA_ENV
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 load_model(model_uri):
"""
Load a fastai 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: A fastai model (an instance of `fastai.Learner`_).
.. code-block:: python
:caption: Example
import mlflow.fastai
# Define the Learner model
model = ...
# log the fastai Leaner model
with mlflow.start_run() as run:
model.fit(epochs, learning_rate)
mlflow.fastai.log_model(model, "model")
# Load the model for scoring
model_uri = "runs:/{}/model".format(run.info.run_id)
loaded_model = mlflow.fastai.load_model(model_uri)
results = loaded_model.predict(predict_data)
"""
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)
model_file_path = os.path.join(local_model_path,
flavor_conf.get("data", "model.fastai"))
return _load_model(path=model_file_path)
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))
# For backwards compatibility, we assume h5 when the save_format is absent
save_format = flavor_conf.get("save_format", "h5")
return _load_model(model_path=keras_model_artifacts_path,
keras_module=keras_module,
save_format=save_format,
**kwargs)
def test_sklearn_log_explainer_self_serialization():
"""
Tests mlflow.shap log_explainer with SHAP internal serialization of the underlying model
"""
with mlflow.start_run() as run:
run_id = run.info.run_id
X, y = shap.datasets.boston()
model = sklearn.ensemble.RandomForestRegressor(n_estimators=100)
model.fit(X, y)
explainer_original = shap.Explainer(model.predict,
X,
algorithm="permutation")
shap_values_original = explainer_original(X[:5])
mlflow.shap.log_explainer(explainer_original,
"test_explainer",
serialize_model_using_mlflow=False)
explainer_uri = "runs:/" + run_id + "/test_explainer"
explainer_loaded = mlflow.shap.load_explainer("runs:/" + run_id +
"/test_explainer")
shap_values_new = explainer_loaded(X[:5])
explainer_path = _download_artifact_from_uri(
artifact_uri=explainer_uri)
flavor_conf = _get_flavor_configuration(
model_path=explainer_path, flavor_name=mlflow.shap.FLAVOR_NAME)
underlying_model_flavor = flavor_conf["underlying_model_flavor"]
assert underlying_model_flavor is None
np.testing.assert_array_equal(shap_values_original.base_values,
shap_values_new.base_values)
np.testing.assert_allclose(shap_values_original.values,
shap_values_new.values,
rtol=100,
atol=100)
def _load_model(path):
"""
Load Model Implementation.
:param path: Local filesystem path to
the MLflow Model with the ``xgboost`` flavor (MLflow < 1.22.0) or
the top-level MLflow Model directory (MLflow >= 1.22.0).
"""
model_dir = os.path.dirname(path) if os.path.isfile(path) else path
flavor_conf = _get_flavor_configuration(model_path=model_dir, flavor_name=FLAVOR_NAME)
# XGBoost models saved in MLflow >=1.22.0 have `model_class`
# in the XGBoost flavor configuration to specify its XGBoost model class.
# When loading models, we first get the XGBoost model from
# its flavor configuration and then create an instance based on its class.
model_class = flavor_conf.get("model_class", "xgboost.core.Booster")
xgb_model_path = os.path.join(model_dir, flavor_conf.get("data"))
model = _get_class_from_string(model_class)()
model.load_model(xgb_model_path)
return model
def load_pyfunc(path, run_id=None, suppress_warnings=False):
"""
Load a model stored in Python function format.
:param path: Path to the model.
:param run_id: MLflow run ID.
:param suppress_warnings: If True, non-fatal warning messages associated with the model
loading process will be suppressed. If False, these warning messages
will be emitted.
"""
if run_id is not None:
path = tracking.utils._get_model_log_dir(path, run_id)
conf = _get_flavor_configuration(model_path=path, flavor_name=FLAVOR_NAME)
model_py_version = conf.get(PY_VERSION)
if not suppress_warnings:
_warn_potentially_incompatible_py_version_if_necessary(model_py_version=model_py_version)
if CODE in conf and conf[CODE]:
code_path = os.path.join(path, conf[CODE])
mlflow.pyfunc.utils._add_code_to_system_path(code_path=code_path)
data_path = os.path.join(path, conf[DATA]) if (DATA in conf) else path
return importlib.import_module(conf[MAIN])._load_pyfunc(data_path)
def load_model(model_uri):
"""
Load a LightGBM 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: A LightGBM model (an instance of `lightgbm.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)
lgb_model_file_path = os.path.join(local_model_path, flavor_conf.get("data", "model.lgb"))
return _load_model(path=lgb_model_file_path)
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)
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) as exc_info:
mlflow.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_explainer(model_uri):
"""
Load a SHAP explainer 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 SHAP explainer.
"""
explainer_path = _download_artifact_from_uri(artifact_uri=model_uri)
flavor_conf = _get_flavor_configuration(model_path=explainer_path,
flavor_name=FLAVOR_NAME)
_add_code_from_conf_to_system_path(explainer_path, flavor_conf)
explainer_artifacts_path = os.path.join(
explainer_path, flavor_conf["serialized_explainer"])
underlying_model_flavor = flavor_conf["underlying_model_flavor"]
model = None
if underlying_model_flavor != _UNKNOWN_MODEL_FLAVOR:
underlying_model_path = os.path.join(explainer_path,
_UNDERLYING_MODEL_SUBPATH)
if underlying_model_flavor == mlflow.sklearn.FLAVOR_NAME:
model = mlflow.sklearn._load_pyfunc(underlying_model_path).predict
elif underlying_model_flavor == mlflow.pytorch.FLAVOR_NAME:
model = mlflow.pytorch._load_model(
os.path.join(underlying_model_path, "data"))
return _load_explainer(explainer_file=explainer_artifacts_path,
model=model)
def _load_pyfunc(model_path):
pyfunc_config = _get_flavor_configuration(
model_path=model_path, flavor_name=mlflow.pyfunc.FLAVOR_NAME
)
python_model_cloudpickle_version = pyfunc_config.get(CONFIG_KEY_CLOUDPICKLE_VERSION, None)
if python_model_cloudpickle_version is None:
mlflow.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
mlflow.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 load_model(model_uri):
"""
Load an ONNX 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 the
`Artifacts Documentation <https://www.mlflow.org/docs/latest/
tracking.html#artifact-stores>`_.
:return: An ONNX model instance.
"""
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)
onnx_model_artifacts_path = os.path.join(local_model_path, flavor_conf["data"])
return _load_model(model_file=onnx_model_artifacts_path)
def test_log_model_persists_specified_conda_env_in_mlflow_model_directory(
saved_tf_iris_model, tf_custom_env):
artifact_path = "model"
with mlflow.start_run():
mlflow.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=tf_custom_env)
run_id = mlflow.active_run().info.run_uuid
model_path = _get_model_log_dir(artifact_path, run_id)
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_without_specified_conda_env_uses_default_env_with_expected_dependencies(
sklearn_knn_model):
artifact_path = "model"
knn_model = sklearn_knn_model.model
with mlflow.start_run():
mlflow.sklearn.log_model(
sk_model=knn_model,
artifact_path=artifact_path,
conda_env=None,
serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE)
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=mlflow.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 == mlflow.sklearn.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 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.Booster model object
<https://xgboost.readthedocs.io/en/latest/python/python_api.html#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_persists_specified_conda_env_in_mlflow_model_directory(
sklearn_knn_model, main_scoped_model_class, pyfunc_custom_env
):
sklearn_artifact_path = "sk_model"
with mlflow.start_run():
mlflow.sklearn.log_model(sk_model=sklearn_knn_model, artifact_path=sklearn_artifact_path)
sklearn_run_id = mlflow.active_run().info.run_id
pyfunc_artifact_path = "pyfunc_model"
with mlflow.start_run():
mlflow.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=mlflow.active_run().info.run_id, artifact_path=pyfunc_artifact_path
)
)
pyfunc_conf = _get_flavor_configuration(
model_path=pyfunc_model_path, flavor_name=mlflow.pyfunc.FLAVOR_NAME
)
saved_conda_env_path = os.path.join(pyfunc_model_path, pyfunc_conf[mlflow.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 load_model(path, run_id=None):
"""
Load a Keras model from a local file (if ``run_id`` is None) or a run.
:param path: Local filesystem path or run-relative artifact path to the model saved
by :py:func:`mlflow.keras.log_model`.
:param run_id: Run ID. If provided, combined with ``path`` to identify the model.
>>> # Load persisted model as a Keras model or as a PyFunc, call predict() on a Pandas DataFrame
>>> keras_model = mlflow.keras.load_model("models", run_id="96771d893a5e46159d9f3b49bf9013e2")
>>> predictions = keras_model.predict(x_test)
"""
if run_id is not None:
path = mlflow.tracking.artifact_utils._get_model_log_dir(
model_name=path, run_id=run_id)
path = os.path.abspath(path)
flavor_conf = _get_flavor_configuration(model_path=path,
flavor_name=FLAVOR_NAME)
# Flavor configurations for models saved in MLflow version <= 0.8.0 may not contain a
# `data` key; in this case, we assume the model artifact path to be `model.h5`
keras_model_artifacts_path = os.path.join(
path, flavor_conf.get("data", "model.h5"))
return _load_model(model_file=keras_model_artifacts_path)
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 mlflow.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=mlflow.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_model(path, run_id=None):
"""
Load a scikit-learn model from a local file (if ``run_id`` is None) or a run.
:param path: Local filesystem path or run-relative artifact path to the model saved
by :py:func:`mlflow.sklearn.save_model`.
:param run_id: Run ID. If provided, combined with ``path`` to identify the model.
>>> import mlflow.sklearn
>>> sk_model = mlflow.sklearn.load_model("sk_models", run_id="96771d893a5e46159d9f3b49bf9013e2")
>>> #use Pandas DataFrame to make predictions
>>> pandas_df = ...
>>> predictions = sk_model.predict(pandas_df)
"""
if run_id is not None:
path = mlflow.tracking.utils._get_model_log_dir(model_name=path,
run_id=run_id)
path = os.path.abspath(path)
flavor_conf = _get_flavor_configuration(model_path=path,
flavor_name=FLAVOR_NAME)
sklearn_model_artifacts_path = os.path.join(path,
flavor_conf['pickled_model'])
return _load_model_from_local_file(path=sklearn_model_artifacts_path)
