def save_model(path, loader_module=None, data_path=None, code_path=None, conda_env=None,
mlflow_model=Model(), python_model=None, artifacts=None, **kwargs):
"""
save_model(path, loader_module=None, data_path=None, code_path=None, conda_env=None,\
mlflow_model=Model(), python_model=None, artifacts=None)
Save a Pyfunc model with custom inference logic and optional data dependencies to a path on the
local filesystem.
For information about the workflows that this method supports, please see :ref:`"workflows for
creating custom pyfunc models" <pyfunc-create-custom-workflows>` and
:ref:`"which workflow is right for my use case?" <pyfunc-create-custom-selecting-workflow>`.
Note that the parameters for the second workflow: ``loader_module``, ``data_path`` and the
parameters for the first workflow: ``python_model``, ``artifacts``, cannot be
specified together.
:param path: The path to which to save the Python model.
:param loader_module: The name of the Python module that is used to load the model
from ``data_path``. This module must define a method with the prototype
``_load_pyfunc(data_path)``. If not ``None``, this module and its
dependencies must be included in one of the following locations:
- The MLflow library.
- Package(s) listed in the model's Conda environment, specified by
the ``conda_env`` parameter.
- One or more of the files specified by the ``code_path`` parameter.
:param data_path: Path to a file or directory containing model data.
:param code_path: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path before the model is loaded.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. This decribes the environment this model should
be run in. If ``python_model`` is not ``None``, the Conda environment must
at least specify the dependencies contained in
:func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is added to the
model. The following is an *example* dictionary representation of a Conda
environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'cloudpickle==0.5.8'
]
}
:param mlflow_model: :py:mod:`mlflow.models.Model` configuration to which to add the
**python_function** flavor.
:param python_model: An instance of a subclass of :class:`~PythonModel`. This class is
serialized using the CloudPickle library. Any dependencies of the class
should be included in one of the following locations:
- The MLflow library.
- Package(s) listed in the model's Conda environment, specified by
the ``conda_env`` parameter.
- One or more of the files specified by the ``code_path`` parameter.
Note: If the class is imported from another module, as opposed to being
defined in the ``__main__`` scope, the defining module should also be
included in one of the listed locations.
:param artifacts: A dictionary containing ``<name, artifact_uri>`` entries. Remote artifact URIs
are resolved to absolute filesystem paths, producing a dictionary of
``<name, absolute_path>`` entries. ``python_model`` can reference these
resolved entries as the ``artifacts`` property of the ``context`` parameter
in :func:`PythonModel.load_context() <mlflow.pyfunc.PythonModel.load_context>`
and :func:`PythonModel.predict() <mlflow.pyfunc.PythonModel.predict>`.
For example, consider the following ``artifacts`` dictionary::
{
"my_file": "s3://my-bucket/path/to/my/file"
}
In this case, the ``"my_file"`` artifact is downloaded from S3. The
``python_model`` can then refer to ``"my_file"`` as an absolute filesystem
path via ``context.artifacts["my_file"]``.
If ``None``, no artifacts are added to the model.
"""
mlflow_model = kwargs.pop('model', mlflow_model)
if len(kwargs) > 0:
raise TypeError("save_model() got unexpected keyword arguments: {}".format(kwargs))
first_argument_set = {
"loader_module": loader_module,
"data_path": data_path,
}
second_argument_set = {
"artifacts": artifacts,
"python_model": python_model,
}
first_argument_set_specified = any([item is not None for item in first_argument_set.values()])
second_argument_set_specified = any([item is not None for item in second_argument_set.values()])
if first_argument_set_specified and second_argument_set_specified:
raise MlflowException(
message=(
"The following sets of parameters cannot be specified together: {first_set_keys}"
" and {second_set_keys}. All parameters in one set must be `None`. Instead, found"
" the following values: {first_set_entries} and {second_set_entries}".format(
first_set_keys=first_argument_set.keys(),
second_set_keys=second_argument_set.keys(),
first_set_entries=first_argument_set,
second_set_entries=second_argument_set)),
error_code=INVALID_PARAMETER_VALUE)
elif (loader_module is None) and (python_model is None):
raise MlflowException(
message="Either `loader_module` or `python_model` must be specified!",
error_code=INVALID_PARAMETER_VALUE)
if first_argument_set_specified:
return _save_model_with_loader_module_and_data_path(
path=path, loader_module=loader_module, data_path=data_path,
code_paths=code_path, conda_env=conda_env, mlflow_model=mlflow_model)
elif second_argument_set_specified:
return mlflow.pyfunc.model._save_model_with_class_artifacts_params(
path=path, python_model=python_model, artifacts=artifacts, conda_env=conda_env,
code_paths=code_path, mlflow_model=mlflow_model)
def log_model(
pd_model,
artifact_path,
training=False,
conda_env=None,
registered_model_name=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Log a paddle model as an MLflow artifact for the current run. Produces an MLflow Model
containing the following flavors:
- :py:mod:`mlflow.paddle`
- :py:mod:`mlflow.pyfunc`. NOTE: This flavor is only included for paddle models
that define `predict()`, since `predict()` is required for pyfunc model inference.
:param pd_model: paddle model to be saved.
:param artifact_path: Run-relative artifact path.
:param training: Only valid when saving a model trained using the PaddlePaddle high level API.
If set to True, the saved model supports both re-training and
inference. If set to False, it only supports inference.
:param conda_env: {{ conda_env }}
:param registered_model_name: If given, create a model version under
``registered_model_name``, also creating a registered model if one
with the given name does not exist.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param await_registration_for: Number of seconds to wait for the model version to finish
being created and is in ``READY`` status. By default, the function
waits for five minutes. Specify 0 or None to skip waiting.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
.. code-block:: python
:caption: Example
import mlflow.paddle
def load_data():
...
class Regressor():
...
model = Regressor()
model.train()
training_data, test_data = load_data()
opt = paddle.optimizer.SGD(learning_rate=0.01, parameters=model.parameters())
EPOCH_NUM = 10
BATCH_SIZE = 10
for epoch_id in range(EPOCH_NUM):
...
mlflow.log_param('learning_rate', 0.01)
mlflow.paddle.log_model(model, "model")
sk_path_dir = ...
mlflow.paddle.save_model(model, sk_path_dir)
"""
return Model.log(
artifact_path=artifact_path,
flavor=mlflow.paddle,
pd_model=pd_model,
conda_env=conda_env,
registered_model_name=registered_model_name,
signature=signature,
input_example=input_example,
await_registration_for=await_registration_for,
training=training,
pip_requirements=pip_requirements,
extra_pip_requirements=extra_pip_requirements,
)
def log_model(keras_model,
artifact_path,
conda_env=None,
custom_objects=None,
keras_module=None,
registered_model_name=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
**kwargs):
"""
Log a Keras model as an MLflow artifact for the current run.
:param keras_model: Keras model to be saved.
:param artifact_path: Run-relative artifact path.
:param conda_env: Either a dictionary representation of a Conda environment or
the path to a Conda environment yaml file.
If provided, this describes the environment this model should be
run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`mlflow.keras.get_default_conda_env()` environment is added to
the model. The following is an *example* dictionary representation of a
Conda environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'keras=2.2.4',
'tensorflow=1.8.0'
]
}
:param custom_objects: A Keras ``custom_objects`` dictionary mapping names (strings) to
custom classes or functions associated with the Keras model. MLflow saves
these custom layers using CloudPickle and restores them automatically
when the model is loaded with :py:func:`mlflow.keras.load_model` and
:py:func:`mlflow.pyfunc.load_model`.
:param keras_module: Keras module to be used to save / load the model
(``keras`` or ``tf.keras``). If not provided, MLflow will
attempt to infer the Keras module based on the given model.
:param registered_model_name: (Experimental) If given, create a model version under
``registered_model_name``, also creating a registered model if one
with the given name does not exist.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param await_registration_for: Number of seconds to wait for the model version to finish
being created and is in ``READY`` status. By default, the function
waits for five minutes. Specify 0 or None to skip waiting.
:param kwargs: kwargs to pass to ``keras_model.save`` method.
.. code-block:: python
:caption: Example
from keras import Dense, layers
import mlflow
# Build, compile, and train your model
keras_model = ...
keras_model.compile(optimizer="rmsprop", loss="mse", metrics=["accuracy"])
results = keras_model.fit(
x_train, y_train, epochs=20, batch_size = 128, validation_data=(x_val, y_val))
# Log metrics and log the model
with mlflow.start_run() as run:
mlflow.keras.log_model(keras_model, "models")
"""
Model.log(artifact_path=artifact_path,
flavor=mlflow.keras,
keras_model=keras_model,
conda_env=conda_env,
custom_objects=custom_objects,
keras_module=keras_module,
registered_model_name=registered_model_name,
signature=signature,
input_example=input_example,
await_registration_for=await_registration_for,
**kwargs)
def save_model(
gluon_model,
path,
mlflow_model=None,
conda_env=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Save a Gluon model to a path on the local file system.
:param gluon_model: Gluon model to be saved. Must be already hybridized.
:param path: Local path where the model is to be saved.
:param mlflow_model: MLflow model config this flavor is being added to.
:param conda_env: Either a dictionary representation of a Conda environment or
the path to a Conda environment yaml file.
If provided, this decribes the environment this model should be
run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`mlflow.gluon.get_default_conda_env()` environment is added to
the model. The following is an *example* dictionary representation of a
Conda environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'mxnet=1.5.0'
]
}
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example can be a Pandas DataFrame where the given
example will be serialized to json using the Pandas split-oriented
format, or a numpy array where the example will be serialized to json
by converting it to a list. Bytes are base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
.. code-block:: python
:caption: Example
from mxnet.gluon import Trainer
from mxnet.gluon.contrib import estimator
from mxnet.gluon.loss import SoftmaxCrossEntropyLoss
from mxnet.gluon.nn import HybridSequential
from mxnet.metric import Accuracy
import mlflow
# Build, compile, and train your model
gluon_model_path = ...
net = HybridSequential()
with net.name_scope():
...
net.hybridize()
net.collect_params().initialize()
softmax_loss = SoftmaxCrossEntropyLoss()
trainer = Trainer(net.collect_params())
est = estimator.Estimator(net=net, loss=softmax_loss, metrics=Accuracy(), trainer=trainer)
est.fit(train_data=train_data, epochs=100, val_data=validation_data)
# Save the model as an MLflow Model
mlflow.gluon.save_model(net, gluon_model_path)
"""
_validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
data_subpath = "data"
data_path = os.path.join(path, data_subpath)
os.makedirs(data_path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# The epoch argument of the export method does not play any role in selecting
# a specific epoch's parameters, and is there only for display purposes.
gluon_model.export(os.path.join(data_path, _MODEL_SAVE_PATH))
conda_env, pip_requirements, pip_constraints = (
_process_pip_requirements(
get_default_pip_requirements(), pip_requirements, extra_pip_requirements,
)
if conda_env is None
else _process_conda_env(conda_env)
)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
pyfunc.add_to_model(mlflow_model, loader_module="mlflow.gluon", env=_CONDA_ENV_FILE_NAME)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(spark_model,
path,
mlflow_model=Model(),
conda_env=None,
dfs_tmpdir=None,
sample_input=None):
"""
Save a Spark MLlib Model to a local path.
By default, this function saves models using the Spark MLlib persistence mechanism.
Additionally, if a sample input is specified using the ``sample_input`` parameter, the model
is also serialized in MLeap format and the MLeap flavor is added.
:param spark_model: Spark model to be saved - MLFlow can only save descendants of
pyspark.ml.Model which implement MLReadable and MLWritable.
:param path: Local path where the model is to be saved.
:param mlflow_model: MLflow model config this flavor is being added to.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If `None`, the default
:func:`get_default_conda_env()` environment is added to the model.
The following is an *example* dictionary representation of a Conda
environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'pyspark=2.3.0'
]
}
:param dfs_tmpdir: Temporary directory path on Distributed (Hadoop) File System (DFS) or local
filesystem if running in local mode. The model is be written in this
destination and then copied to the requested local path. This is necessary
as Spark ML models read from and write to DFS if running on a cluster. All
temporary files created on the DFS are removed if this operation
completes successfully. Defaults to ``/tmp/mlflow``.
:param sample_input: A sample input that is used to add the MLeap flavor to the model.
This must be a PySpark DataFrame that the model can evaluate. If
``sample_input`` is ``None``, the MLeap flavor is not added.
>>> from mlflow import spark
>>> from pyspark.ml.pipeline.PipelineModel
>>>
>>> #your pyspark.ml.pipeline.PipelineModel type
>>> model = ...
>>> mlflow.spark.save_model(model, "spark-model")
"""
_validate_model(spark_model)
from pyspark.ml import PipelineModel
if not isinstance(spark_model, PipelineModel):
spark_model = PipelineModel([spark_model])
# Spark ML stores the model on DFS if running on a cluster
# Save it to a DFS temp dir first and copy it to local path
if dfs_tmpdir is None:
dfs_tmpdir = DFS_TMP
tmp_path = _tmp_path(dfs_tmpdir)
spark_model.save(tmp_path)
sparkml_data_path = os.path.abspath(
os.path.join(path, _SPARK_MODEL_PATH_SUB))
_HadoopFileSystem.copy_to_local_file(tmp_path,
sparkml_data_path,
remove_src=True)
_save_model_metadata(dst_dir=path,
spark_model=spark_model,
mlflow_model=mlflow_model,
sample_input=sample_input,
conda_env=conda_env)
def log_model(spark_model,
artifact_path,
conda_env=None,
dfs_tmpdir=None,
sample_input=None,
registered_model_name=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None):
"""
Log a Spark MLlib model as an MLflow artifact for the current run. This uses the
MLlib persistence format and produces an MLflow Model with the Spark flavor.
Note: If no run is active, it will instantiate a run to obtain a run_id.
:param spark_model: Spark model to be saved - MLflow can only save descendants of
pyspark.ml.Model which implement MLReadable and MLWritable.
:param artifact_path: Run relative artifact path.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If `None`, the default
:func:`get_default_conda_env()` environment is added to the model.
The following is an *example* dictionary representation of a Conda
environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'pyspark=2.3.0'
]
}
:param dfs_tmpdir: Temporary directory path on Distributed (Hadoop) File System (DFS) or local
filesystem if running in local mode. The model is written in this
destination and then copied into the model's artifact directory. This is
necessary as Spark ML models read from and write to DFS if running on a
cluster. If this operation completes successfully, all temporary files
created on the DFS are removed. Defaults to ``/tmp/mlflow``.
:param sample_input: A sample input used to add the MLeap flavor to the model.
This must be a PySpark DataFrame that the model can evaluate. If
``sample_input`` is ``None``, the MLeap flavor is not added.
:param registered_model_name: (Experimental) If given, create a model version under
``registered_model_name``, also creating a registered model if one
with the given name does not exist.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
.. code-block:: python
:caption: Example
from pyspark.ml import Pipeline
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.feature import HashingTF, Tokenizer
training = spark.createDataFrame([
(0, "a b c d e spark", 1.0),
(1, "b d", 0.0),
(2, "spark f g h", 1.0),
(3, "hadoop mapreduce", 0.0) ], ["id", "text", "label"])
tokenizer = Tokenizer(inputCol="text", outputCol="words")
hashingTF = HashingTF(inputCol=tokenizer.getOutputCol(), outputCol="features")
lr = LogisticRegression(maxIter=10, regParam=0.001)
pipeline = Pipeline(stages=[tokenizer, hashingTF, lr])
model = pipeline.fit(training)
mlflow.spark.log_model(model, "spark-model")
"""
from py4j.protocol import Py4JJavaError
_validate_model(spark_model)
from pyspark.ml import PipelineModel
if not isinstance(spark_model, PipelineModel):
spark_model = PipelineModel([spark_model])
run_id = mlflow.tracking.fluent._get_or_start_run().info.run_id
run_root_artifact_uri = mlflow.get_artifact_uri()
# If the artifact URI is a local filesystem path, defer to Model.log() to persist the model,
# since Spark may not be able to write directly to the driver's filesystem. For example,
# writing to `file:/uri` will write to the local filesystem from each executor, which will
# be incorrect on multi-node clusters - to avoid such issues we just use the Model.log() path
# here.
if is_local_uri(run_root_artifact_uri):
return Model.log(artifact_path=artifact_path,
flavor=mlflow.spark,
spark_model=spark_model,
conda_env=conda_env,
dfs_tmpdir=dfs_tmpdir,
sample_input=sample_input,
registered_model_name=registered_model_name,
signature=signature,
input_example=input_example)
# If Spark cannot write directly to the artifact repo, defer to Model.log() to persist the
# model
model_dir = os.path.join(run_root_artifact_uri, artifact_path)
try:
spark_model.save(os.path.join(model_dir, _SPARK_MODEL_PATH_SUB))
except Py4JJavaError:
return Model.log(artifact_path=artifact_path,
flavor=mlflow.spark,
spark_model=spark_model,
conda_env=conda_env,
dfs_tmpdir=dfs_tmpdir,
sample_input=sample_input,
registered_model_name=registered_model_name,
signature=signature,
input_example=input_example)
# Otherwise, override the default model log behavior and save model directly to artifact repo
mlflow_model = Model(artifact_path=artifact_path, run_id=run_id)
with TempDir() as tmp:
tmp_model_metadata_dir = tmp.path()
_save_model_metadata(tmp_model_metadata_dir,
spark_model,
mlflow_model,
sample_input,
conda_env,
signature=signature,
input_example=input_example)
mlflow.tracking.fluent.log_artifacts(tmp_model_metadata_dir,
artifact_path)
if registered_model_name is not None:
mlflow.register_model("runs:/%s/%s" % (run_id, artifact_path),
registered_model_name)
def log_model(
onnx_model,
artifact_path,
conda_env=None,
code_paths=None,
registered_model_name=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
pip_requirements=None,
extra_pip_requirements=None,
onnx_execution_providers=None,
):
"""
Log an ONNX model as an MLflow artifact for the current run.
:param onnx_model: ONNX model to be saved.
:param artifact_path: Run-relative artifact path.
:param conda_env: {{ conda_env }}
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param registered_model_name: If given, create a model version under
``registered_model_name``, also creating a registered model if one
with the given name does not exist.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example can be a Pandas DataFrame where the given
example will be serialized to json using the Pandas split-oriented
format, or a numpy array where the example will be serialized to json
by converting it to a list. Bytes are base64-encoded.
:param await_registration_for: Number of seconds to wait for the model version to finish
being created and is in ``READY`` status. By default, the function
waits for five minutes. Specify 0 or None to skip waiting.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
:param onnx_execution_providers: List of strings defining onnxruntime execution providers.
Defaults to example:
['CUDAExecutionProvider', 'CPUExecutionProvider']
This uses GPU preferentially over CPU.
See onnxruntime API for further descriptions:
https://onnxruntime.ai/docs/execution-providers/
:return: A :py:class:`ModelInfo <mlflow.models.model.ModelInfo>` instance that contains the
metadata of the logged model.
"""
return Model.log(
artifact_path=artifact_path,
flavor=mlflow.onnx,
onnx_model=onnx_model,
conda_env=conda_env,
code_paths=code_paths,
registered_model_name=registered_model_name,
signature=signature,
input_example=input_example,
await_registration_for=await_registration_for,
pip_requirements=pip_requirements,
extra_pip_requirements=extra_pip_requirements,
onnx_execution_providers=onnx_execution_providers,
)
def get_model_conf(artifact_uri, model_subpath="model"):
model_conf_path = os.path.join(artifact_uri, model_subpath, "MLmodel")
return Model.load(model_conf_path)
def log_model(
lgb_model,
artifact_path,
conda_env=None,
registered_model_name=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
pip_requirements=None,
extra_pip_requirements=None,
**kwargs,
):
"""
Log a LightGBM model as an MLflow artifact for the current run.
:param lgb_model: LightGBM model (an instance of `lightgbm.Booster`_) or
models that implement the `scikit-learn API`_ to be saved.
:param artifact_path: Run-relative artifact path.
:param conda_env: {{ conda_env }}
:param registered_model_name: If given, create a model version under
``registered_model_name``, also creating a registered model if one
with the given name does not exist.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param await_registration_for: Number of seconds to wait for the model version to finish
being created and is in ``READY`` status. By default, the function
waits for five minutes. Specify 0 or None to skip waiting.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
:param kwargs: kwargs to pass to `lightgbm.Booster.save_model`_ method.
:return: A :py:class:`ModelInfo <mlflow.models.model.ModelInfo>` instance that contains the
metadata of the logged model.
"""
return Model.log(
artifact_path=artifact_path,
flavor=mlflow.lightgbm,
registered_model_name=registered_model_name,
lgb_model=lgb_model,
conda_env=conda_env,
signature=signature,
input_example=input_example,
await_registration_for=await_registration_for,
pip_requirements=pip_requirements,
extra_pip_requirements=extra_pip_requirements,
**kwargs,
)
def log_model(
tf_saved_model_dir,
tf_meta_graph_tags,
tf_signature_def_key,
artifact_path,
conda_env=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
registered_model_name=None,
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
):
"""
Log a *serialized* collection of TensorFlow graphs and variables as an MLflow model
for the current run. This method operates on TensorFlow variables and graphs that have been
serialized in TensorFlow's ``SavedModel`` format. For more information about ``SavedModel``
format, see the TensorFlow documentation:
https://www.tensorflow.org/guide/saved_model#save_and_restore_models.
This method saves a model with both ``python_function`` and ``tensorflow`` flavors.
If loaded back using the ``python_function`` flavor, the model can be used to predict on
pandas DataFrames, producing a pandas DataFrame whose output columns correspond to the
TensorFlow model's outputs. The python_function model will flatten outputs that are length-one,
one-dimensional tensors of a single scalar value (e.g.
``{"predictions": [[1.0], [2.0], [3.0]]}``) into the scalar values (e.g.
``{"predictions": [1, 2, 3]}``), so that the resulting output column is a column of scalars
rather than lists of length one. All other model output types are included as-is in the output
DataFrame.
:param tf_saved_model_dir: Path to the directory containing serialized TensorFlow variables and
graphs in ``SavedModel`` format.
:param tf_meta_graph_tags: A list of tags identifying the model's metagraph within the
serialized ``SavedModel`` object. For more information, see the
``tags`` parameter of the
``tf.saved_model.builder.SavedModelBuilder`` method.
:param tf_signature_def_key: A string identifying the input/output signature associated with the
model. This is a key within the serialized ``SavedModel`` signature
definition mapping. For more information, see the
``signature_def_map`` parameter of the
``tf.saved_model.builder.SavedModelBuilder`` method.
:param artifact_path: The run-relative path to which to log model artifacts.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is added to the model. The
following is an *example* dictionary representation of a Conda environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'tensorflow=1.8.0'
]
}
:param registered_model_name: (Experimental) If given, create a model version under
``registered_model_name``, also creating a registered model if one
with the given name does not exist.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param await_registration_for: Number of seconds to wait for the model version to finish
being created and is in ``READY`` status. By default, the function
waits for five minutes. Specify 0 or None to skip waiting.
"""
return Model.log(
artifact_path=artifact_path,
flavor=mlflow.tensorflow,
tf_saved_model_dir=tf_saved_model_dir,
tf_meta_graph_tags=tf_meta_graph_tags,
tf_signature_def_key=tf_signature_def_key,
conda_env=conda_env,
registered_model_name=registered_model_name,
signature=signature,
input_example=input_example,
await_registration_for=await_registration_for,
)
def save_model(
cb_model,
path,
conda_env=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
**kwargs,
):
"""
Save a CatBoost model to a path on the local file system.
:param cb_model: CatBoost model (an instance of `CatBoost`_, `CatBoostClassifier`_,
or `CatBoostRegressor`_) to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: {{ conda_env }}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
:param kwargs: kwargs to pass to `CatBoost.save_model`_ method.
"""
import catboost as cb
_validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
model_data_path = os.path.join(path, _MODEL_BINARY_FILE_NAME)
cb_model.save_model(model_data_path, **kwargs)
model_bin_kwargs = {_MODEL_BINARY_KEY: _MODEL_BINARY_FILE_NAME}
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.catboost",
env=_CONDA_ENV_FILE_NAME,
**model_bin_kwargs,
)
flavor_conf = {
_MODEL_TYPE_KEY: cb_model.__class__.__name__,
_SAVE_FORMAT_KEY: kwargs.get("format", "cbm"),
**model_bin_kwargs,
}
mlflow_model.add_flavor(FLAVOR_NAME, catboost_version=cb.__version__, **flavor_conf)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
def save_model(
tf_saved_model_dir,
tf_meta_graph_tags,
tf_signature_def_key,
path,
mlflow_model=None,
conda_env=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
):
"""
Save a *serialized* collection of TensorFlow graphs and variables as an MLflow model
to a local path. This method operates on TensorFlow variables and graphs that have been
serialized in TensorFlow's ``SavedModel`` format. For more information about ``SavedModel``
format, see the TensorFlow documentation:
https://www.tensorflow.org/guide/saved_model#save_and_restore_models.
:param tf_saved_model_dir: Path to the directory containing serialized TensorFlow variables and
graphs in ``SavedModel`` format.
:param tf_meta_graph_tags: A list of tags identifying the model's metagraph within the
serialized ``SavedModel`` object. For more information, see the
``tags`` parameter of the
``tf.saved_model.builder.savedmodelbuilder`` method.
:param tf_signature_def_key: A string identifying the input/output signature associated with the
model. This is a key within the serialized ``savedmodel``
signature definition mapping. For more information, see the
``signature_def_map`` parameter of the
``tf.saved_model.builder.savedmodelbuilder`` method.
:param path: Local path where the MLflow model is to be saved.
:param mlflow_model: MLflow model configuration to which to add the ``tensorflow`` flavor.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is added to the model. The
following is an *example* dictionary representation of a Conda environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'tensorflow=1.8.0'
]
}
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
"""
_logger.info(
"Validating the specified TensorFlow model by attempting to load it in a new TensorFlow"
" graph...")
_validate_saved_model(
tf_saved_model_dir=tf_saved_model_dir,
tf_meta_graph_tags=tf_meta_graph_tags,
tf_signature_def_key=tf_signature_def_key,
)
_logger.info("Validation succeeded!")
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path),
DIRECTORY_NOT_EMPTY)
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
root_relative_path = _copy_file_or_tree(src=tf_saved_model_dir,
dst=path,
dst_dir=None)
model_dir_subpath = "tfmodel"
shutil.move(os.path.join(path, root_relative_path),
os.path.join(path, model_dir_subpath))
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env()
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
mlflow_model.add_flavor(
FLAVOR_NAME,
saved_model_dir=model_dir_subpath,
meta_graph_tags=tf_meta_graph_tags,
signature_def_key=tf_signature_def_key,
)
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow.tensorflow",
env=conda_env_subpath)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def _save_model_with_class_artifacts_params(path,
python_model,
artifacts=None,
conda_env=None,
code_paths=None,
mlflow_model=Model()):
"""
:param path: The path to which to save the Python model.
:param python_model: An instance of a subclass of :class:`~PythonModel`. ``python_model``
defines how the model loads artifacts and how it performs inference.
:param artifacts: A dictionary containing ``<name, artifact_uri>`` entries. Remote artifact URIs
will be resolved to absolute filesystem paths, producing a dictionary of
``<name, absolute_path>`` entries. ``python_model`` can reference these
resolved entries as the ``artifacts`` property of the ``context`` attribute.
If *None*, no artifacts will be added to the model.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :data:`mlflow.pyfunc.DEFAULT_CONDA_ENV`. If `None`, the default
:data:`mlflow.pyfunc.DEFAULT_CONDA_ENV` environment will be added to the
model.
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files will be *prepended* to the system
path before the model is loaded.
:param mlflow_model: The model configuration to which to add the ``mlflow.pyfunc`` flavor.
"""
if os.path.exists(path):
raise MlflowException(message="Path '{}' already exists".format(path),
error_code=RESOURCE_ALREADY_EXISTS)
os.makedirs(path)
custom_model_config_kwargs = {
CONFIG_KEY_CLOUDPICKLE_VERSION: cloudpickle.__version__,
}
if isinstance(python_model, PythonModel):
saved_python_model_subpath = "python_model.pkl"
with open(os.path.join(path, saved_python_model_subpath), "wb") as out:
cloudpickle.dump(python_model, out)
custom_model_config_kwargs[
CONFIG_KEY_PYTHON_MODEL] = saved_python_model_subpath
else:
raise MlflowException(message=(
"`python_model` must be a subclass of `PythonModel`. Instead, found an"
" object of type: {python_model_type}".format(
python_model_type=type(python_model))),
error_code=INVALID_PARAMETER_VALUE)
if artifacts:
saved_artifacts_config = {}
with TempDir() as tmp_artifacts_dir:
tmp_artifacts_config = {}
saved_artifacts_dir_subpath = "artifacts"
for artifact_name, artifact_uri in artifacts.items():
tmp_artifact_path = _download_artifact_from_uri(
artifact_uri=artifact_uri,
output_path=tmp_artifacts_dir.path())
tmp_artifacts_config[artifact_name] = tmp_artifact_path
saved_artifact_subpath = os.path.join(
saved_artifacts_dir_subpath,
os.path.relpath(path=tmp_artifact_path,
start=tmp_artifacts_dir.path()))
saved_artifacts_config[artifact_name] = {
CONFIG_KEY_ARTIFACT_RELATIVE_PATH: saved_artifact_subpath,
CONFIG_KEY_ARTIFACT_URI: artifact_uri,
}
shutil.move(tmp_artifacts_dir.path(),
os.path.join(path, saved_artifacts_dir_subpath))
custom_model_config_kwargs[
CONFIG_KEY_ARTIFACTS] = saved_artifacts_config
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = DEFAULT_CONDA_ENV
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
saved_code_subpath = None
if code_paths is not None:
saved_code_subpath = "code"
for code_path in code_paths:
_copy_file_or_tree(src=code_path,
dst=path,
dst_dir=saved_code_subpath)
mlflow.pyfunc.add_to_model(model=mlflow_model,
loader_module=__name__,
code=saved_code_subpath,
env=conda_env_subpath,
**custom_model_config_kwargs)
mlflow_model.save(os.path.join(path, 'MLmodel'))
def log_model(artifact_path, loader_module=None, data_path=None, code_path=None, conda_env=None,
python_model=None, artifacts=None):
"""
Log a Pyfunc model with custom inference logic and optional data dependencies as an MLflow
artifact for the current run.
For information about the workflows that this method supports, see :ref:`Workflows for
creating custom pyfunc models <pyfunc-create-custom-workflows>` and
:ref:`Which workflow is right for my use case? <pyfunc-create-custom-selecting-workflow>`.
You cannot specify the parameters for the second workflow: ``loader_module``, ``data_path``
and the parameters for the first workflow: ``python_model``, ``artifacts`` together.
:param artifact_path: The run-relative artifact path to which to log the Python model.
:param loader_module: The name of the Python module that is used to load the model
from ``data_path``. This module must define a method with the prototype
``_load_pyfunc(data_path)``. If not ``None``, this module and its
dependencies must be included in one of the following locations:
- The MLflow library.
- Package(s) listed in the model's Conda environment, specified by
the ``conda_env`` parameter.
- One or more of the files specified by the ``code_path`` parameter.
:param data_path: Path to a file or directory containing model data.
:param code_path: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path before the model is loaded.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. This decribes the environment this model should
be run in. If ``python_model`` is not ``None``, the Conda environment must
at least specify the dependencies contained in
:func:`get_default_conda_env()`. If `None`, the default
:func:`get_default_conda_env()` environment is added to the
model. The following is an *example* dictionary representation of a Conda
environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'cloudpickle==0.5.8'
]
}
:param python_model: An instance of a subclass of :class:`~PythonModel`. This class is
serialized using the CloudPickle library. Any dependencies of the class
should be included in one of the following locations:
- The MLflow library.
- Package(s) listed in the model's Conda environment, specified by
the ``conda_env`` parameter.
- One or more of the files specified by the ``code_path`` parameter.
Note: If the class is imported from another module, as opposed to being
defined in the ``__main__`` scope, the defining module should also be
included in one of the listed locations.
:param artifacts: A dictionary containing ``<name, artifact_uri>`` entries. Remote artifact URIs
are resolved to absolute filesystem paths, producing a dictionary of
``<name, absolute_path>`` entries. ``python_model`` can reference these
resolved entries as the ``artifacts`` property of the ``context`` parameter
in :func:`PythonModel.load_context() <mlflow.pyfunc.PythonModel.load_context>`
and :func:`PythonModel.predict() <mlflow.pyfunc.PythonModel.predict>`.
For example, consider the following ``artifacts`` dictionary::
{
"my_file": "s3://my-bucket/path/to/my/file"
}
In this case, the ``"my_file"`` artifact is downloaded from S3. The
``python_model`` can then refer to ``"my_file"`` as an absolute filesystem
path via ``context.artifacts["my_file"]``.
If ``None``, no artifacts are added to the model.
"""
return Model.log(artifact_path=artifact_path,
flavor=mlflow.pyfunc,
loader_module=loader_module,
data_path=data_path,
code_path=code_path,
python_model=python_model,
artifacts=artifacts,
conda_env=conda_env)
def deploy(app_name, model_uri, execution_role_arn=None, bucket=None,
image_url=None, region_name="us-west-2", mode=DEPLOYMENT_MODE_CREATE, archive=False,
instance_type=DEFAULT_SAGEMAKER_INSTANCE_TYPE,
instance_count=DEFAULT_SAGEMAKER_INSTANCE_COUNT, vpc_config=None, flavor=None,
synchronous=True, timeout_seconds=1200):
"""
Deploy an MLflow model on AWS SageMaker.
The currently active AWS account must have correct permissions set up.
This function creates a SageMaker endpoint. For more information about the input data
formats accepted by this endpoint, see the
:ref:`MLflow deployment tools documentation <sagemaker_deployment>`.
:param app_name: Name of the deployed application.
:param model_uri: The location, in URI format, of the MLflow model to deploy to SageMaker.
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>`_.
:param execution_role_arn: The name of an IAM role granting the SageMaker service permissions to
access the specified Docker image and S3 bucket containing MLflow
model artifacts. If unspecified, the currently-assumed role will be
used. This execution role is passed to the SageMaker service when
creating a SageMaker model from the specified MLflow model. It is
passed as the ``ExecutionRoleArn`` parameter of the `SageMaker
CreateModel API call <https://docs.aws.amazon.com/sagemaker/latest/
dg/API_CreateModel.html>`_. This role is *not* assumed for any other
call. For more information about SageMaker execution roles for model
creation, see
https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html.
:param bucket: S3 bucket where model artifacts will be stored. Defaults to a
SageMaker-compatible bucket name.
:param image_url: URL of the ECR-hosted Docker image the model should be deployed into, produced
by ``mlflow sagemaker build-and-push-container``. This parameter can also
be specified by the environment variable ``MLFLOW_SAGEMAKER_DEPLOY_IMG_URL``.
:param region_name: Name of the AWS region to which to deploy the application.
:param mode: The mode in which to deploy the application. Must be one of the following:
``mlflow.sagemaker.DEPLOYMENT_MODE_CREATE``
Create an application with the specified name and model. This fails if an
application of the same name already exists.
``mlflow.sagemaker.DEPLOYMENT_MODE_REPLACE``
If an application of the specified name exists, its model(s) is replaced with
the specified model. If no such application exists, it is created with the
specified name and model.
``mlflow.sagemaker.DEPLOYMENT_MODE_ADD``
Add the specified model to a pre-existing application with the specified name,
if one exists. If the application does not exist, a new application is created
with the specified name and model. NOTE: If the application **already exists**,
the specified model is added to the application's corresponding SageMaker
endpoint with an initial weight of zero (0). To route traffic to the model,
update the application's associated endpoint configuration using either the
AWS console or the ``UpdateEndpointWeightsAndCapacities`` function defined in
https://docs.aws.amazon.com/sagemaker/latest/dg/API_UpdateEndpointWeightsAndCapacities.html.
:param archive: If ``True``, any pre-existing SageMaker application resources that become
inactive (i.e. as a result of deploying in
``mlflow.sagemaker.DEPLOYMENT_MODE_REPLACE`` mode) are preserved.
These resources may include unused SageMaker models and endpoint configurations
that were associated with a prior version of the application endpoint. If
``False``, these resources are deleted. In order to use ``archive=False``,
``deploy()`` must be executed synchronously with ``synchronous=True``.
:param instance_type: The type of SageMaker ML instance on which to deploy the model. For a list
of supported instance types, see
https://aws.amazon.com/sagemaker/pricing/instance-types/.
:param instance_count: The number of SageMaker ML instances on which to deploy the model.
:param vpc_config: A dictionary specifying the VPC configuration to use when creating the
new SageMaker model associated with this application. The acceptable values
for this parameter are identical to those of the ``VpcConfig`` parameter in
the `SageMaker boto3 client's create_model method
<https://boto3.readthedocs.io/en/latest/reference/services/sagemaker.html
#SageMaker.Client.create_model>`_. For more information, see
https://docs.aws.amazon.com/sagemaker/latest/dg/API_VpcConfig.html.
.. code-block:: python
:caption: Example
import mlflow.sagemaker as mfs
vpc_config = {
'SecurityGroupIds': [
'sg-123456abc',
],
'Subnets': [
'subnet-123456abc',
]
}
mfs.deploy(..., vpc_config=vpc_config)
:param flavor: The name of the flavor of the model to use for deployment. Must be either
``None`` or one of mlflow.sagemaker.SUPPORTED_DEPLOYMENT_FLAVORS. If ``None``,
a flavor is automatically selected from the model's available flavors. If the
specified flavor is not present or not supported for deployment, an exception
will be thrown.
:param synchronous: If ``True``, this function will block until the deployment process succeeds
or encounters an irrecoverable failure. If ``False``, this function will
return immediately after starting the deployment process. It will not wait
for the deployment process to complete; in this case, the caller is
responsible for monitoring the health and status of the pending deployment
via native SageMaker APIs or the AWS console.
:param timeout_seconds: If ``synchronous`` is ``True``, the deployment process will return after
the specified number of seconds if no definitive result (success or
failure) is achieved. Once the function returns, the caller is
responsible for monitoring the health and status of the pending
deployment using native SageMaker APIs or the AWS console. If
``synchronous`` is ``False``, this parameter is ignored.
"""
import boto3
if (not archive) and (not synchronous):
raise MlflowException(
message=(
"Resources must be archived when `deploy()` is executed in non-synchronous mode."
" Either set `synchronous=True` or `archive=True`."),
error_code=INVALID_PARAMETER_VALUE)
if mode not in DEPLOYMENT_MODES:
raise MlflowException(
message="`mode` must be one of: {deployment_modes}".format(
deployment_modes=",".join(DEPLOYMENT_MODES)),
error_code=INVALID_PARAMETER_VALUE)
model_path = _download_artifact_from_uri(model_uri)
model_config_path = os.path.join(model_path, MLMODEL_FILE_NAME)
if not os.path.exists(model_config_path):
raise MlflowException(
message=(
"Failed to find {} configuration within the specified model's"
" root directory.").format(MLMODEL_FILE_NAME),
error_code=INVALID_PARAMETER_VALUE)
model_config = Model.load(model_config_path)
if flavor is None:
flavor = _get_preferred_deployment_flavor(model_config)
else:
_validate_deployment_flavor(model_config, flavor)
_logger.info("Using the %s flavor for deployment!", flavor)
sage_client = boto3.client('sagemaker', region_name=region_name)
s3_client = boto3.client('s3', region_name=region_name)
endpoint_exists = _find_endpoint(endpoint_name=app_name, sage_client=sage_client) is not None
if endpoint_exists and mode == DEPLOYMENT_MODE_CREATE:
raise MlflowException(
message=(
"You are attempting to deploy an application with name: {application_name} in"
" '{mode_create}' mode. However, an application with the same name already"
" exists. If you want to update this application, deploy in '{mode_add}' or"
" '{mode_replace}' mode.".format(
application_name=app_name,
mode_create=DEPLOYMENT_MODE_CREATE,
mode_add=DEPLOYMENT_MODE_ADD,
mode_replace=DEPLOYMENT_MODE_REPLACE)),
error_code=INVALID_PARAMETER_VALUE)
model_name = _get_sagemaker_model_name(endpoint_name=app_name)
if not image_url:
image_url = _get_default_image_url(region_name=region_name)
if not execution_role_arn:
execution_role_arn = _get_assumed_role_arn()
if not bucket:
_logger.info("No model data bucket specified, using the default bucket")
bucket = _get_default_s3_bucket(region_name)
model_s3_path = _upload_s3(local_model_path=model_path,
bucket=bucket,
prefix=model_name,
region_name=region_name,
s3_client=s3_client)
if endpoint_exists:
deployment_operation = _update_sagemaker_endpoint(
endpoint_name=app_name, model_name=model_name, model_s3_path=model_s3_path,
model_uri=model_uri, image_url=image_url, flavor=flavor,
instance_type=instance_type, instance_count=instance_count, vpc_config=vpc_config,
mode=mode, role=execution_role_arn, sage_client=sage_client, s3_client=s3_client)
else:
deployment_operation = _create_sagemaker_endpoint(
endpoint_name=app_name, model_name=model_name, model_s3_path=model_s3_path,
model_uri=model_uri, image_url=image_url, flavor=flavor,
instance_type=instance_type, instance_count=instance_count, vpc_config=vpc_config,
role=execution_role_arn, sage_client=sage_client)
if synchronous:
_logger.info("Waiting for the deployment operation to complete...")
operation_status = deployment_operation.await_completion(timeout_seconds=timeout_seconds)
if operation_status.state == _SageMakerOperationStatus.STATE_SUCCEEDED:
_logger.info("The deployment operation completed successfully with message: \"%s\"",
operation_status.message)
else:
raise MlflowException(
"The deployment operation failed with the following error message:"
" \"{error_message}\"".format(error_message=operation_status.message))
if not archive:
deployment_operation.clean_up()
def save_model(
lgb_model,
path,
conda_env=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Save a LightGBM model to a path on the local file system.
:param lgb_model: LightGBM model (an instance of `lightgbm.Booster`_) or
models that implement the `scikit-learn API`_ to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: {{ conda_env }}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
"""
import lightgbm as lgb
_validate_env_arguments(conda_env, pip_requirements,
extra_pip_requirements)
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
model_data_subpath = "model.lgb" if isinstance(
lgb_model, lgb.Booster) else "model.pkl"
model_data_path = os.path.join(path, model_data_subpath)
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# Save a LightGBM model
_save_model(lgb_model, model_data_path)
lgb_model_class = _get_fully_qualified_class_name(lgb_model)
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.lightgbm",
data=model_data_subpath,
env=_CONDA_ENV_FILE_NAME,
)
mlflow_model.add_flavor(
FLAVOR_NAME,
lgb_version=lgb.__version__,
data=model_data_subpath,
model_class=lgb_model_class,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements(
include_cloudpickle=not isinstance(lgb_model, lgb.Booster))
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(
conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME),
"\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME),
"\n".join(pip_requirements))
def save_model(spark_model,
path,
mlflow_model=None,
conda_env=None,
dfs_tmpdir=None,
sample_input=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None):
"""
Save a Spark MLlib Model to a local path.
By default, this function saves models using the Spark MLlib persistence mechanism.
Additionally, if a sample input is specified using the ``sample_input`` parameter, the model
is also serialized in MLeap format and the MLeap flavor is added.
:param spark_model: Spark model to be saved - MLflow can only save descendants of
pyspark.ml.Model which implement MLReadable and MLWritable.
:param path: Local path where the model is to be saved.
:param mlflow_model: MLflow model config this flavor is being added to.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If `None`, the default
:func:`get_default_conda_env()` environment is added to the model.
The following is an *example* dictionary representation of a Conda
environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'pyspark=2.3.0'
]
}
:param dfs_tmpdir: Temporary directory path on Distributed (Hadoop) File System (DFS) or local
filesystem if running in local mode. The model is be written in this
destination and then copied to the requested local path. This is necessary
as Spark ML models read from and write to DFS if running on a cluster. All
temporary files created on the DFS are removed if this operation
completes successfully. Defaults to ``/tmp/mlflow``.
:param sample_input: A sample input that is used to add the MLeap flavor to the model.
This must be a PySpark DataFrame that the model can evaluate. If
``sample_input`` is ``None``, the MLeap flavor is not added.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
.. code-block:: python
:caption: Example
from mlflow import spark
from pyspark.ml.pipeline.PipelineModel
# your pyspark.ml.pipeline.PipelineModel type
model = ...
mlflow.spark.save_model(model, "spark-model")
"""
_validate_model(spark_model)
from pyspark.ml import PipelineModel
if not isinstance(spark_model, PipelineModel):
spark_model = PipelineModel([spark_model])
if mlflow_model is None:
mlflow_model = Model()
# Spark ML stores the model on DFS if running on a cluster
# Save it to a DFS temp dir first and copy it to local path
if dfs_tmpdir is None:
dfs_tmpdir = DFS_TMP
tmp_path = _tmp_path(dfs_tmpdir)
spark_model.save(tmp_path)
sparkml_data_path = os.path.abspath(
os.path.join(path, _SPARK_MODEL_PATH_SUB))
_HadoopFileSystem.copy_to_local_file(tmp_path,
sparkml_data_path,
remove_src=True)
_save_model_metadata(dst_dir=path,
spark_model=spark_model,
mlflow_model=mlflow_model,
sample_input=sample_input,
conda_env=conda_env,
signature=signature,
input_example=input_example)
def save_model(h2o_model,
path,
conda_env=None,
mlflow_model=Model(),
settings=None):
"""
Save an H2O model to a path on the local file system.
:param h2o_model: H2O model to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is added to the model.
The following is an *example* dictionary representation of a Conda
environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'pip': [
'h2o==3.20.0.8'
]
]
}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
"""
import h2o
path = os.path.abspath(path)
if os.path.exists(path):
raise Exception("Path '{}' already exists".format(path))
model_data_subpath = "model.h2o"
model_data_path = os.path.join(path, model_data_subpath)
os.makedirs(model_data_path)
# Save h2o-model
h2o_save_location = h2o.save_model(model=h2o_model,
path=model_data_path,
force=True)
model_file = os.path.basename(h2o_save_location)
# Save h2o-settings
if settings is None:
settings = {}
settings['full_file'] = h2o_save_location
settings['model_file'] = model_file
settings['model_dir'] = model_data_path
with open(os.path.join(model_data_path, "h2o.yaml"), 'w') as settings_file:
yaml.safe_dump(settings, stream=settings_file)
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env()
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow.h2o",
data=model_data_subpath,
env=conda_env_subpath)
mlflow_model.add_flavor(FLAVOR_NAME,
h2o_version=h2o.__version__,
data=model_data_subpath)
mlflow_model.save(os.path.join(path, "MLmodel"))
def _save_model_with_class_artifacts_params(
path,
python_model,
artifacts=None,
conda_env=None,
code_paths=None,
mlflow_model=None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
:param path: The path to which to save the Python model.
:param python_model: An instance of a subclass of :class:`~PythonModel`. ``python_model``
defines how the model loads artifacts and how it performs inference.
:param artifacts: A dictionary containing ``<name, artifact_uri>`` entries.
Remote artifact URIs
are resolved to absolute filesystem paths, producing a dictionary of
``<name, absolute_path>`` entries. ``python_model`` can reference these
resolved entries as the ``artifacts`` property of the ``context``
attribute. If ``None``, no artifacts are added to the model.
:param conda_env: Either a dictionary representation of a Conda environment or the
path to a Conda environment yaml file. If provided, this decsribes the
environment this model should be run in. At minimum, it should specify
the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is added to the model.
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path before the model is loaded.
:param mlflow_model: The model configuration to which to add the ``mlflow.pyfunc`` flavor.
"""
if mlflow_model is None:
mlflow_model = Model()
custom_model_config_kwargs = {
CONFIG_KEY_CLOUDPICKLE_VERSION: cloudpickle.__version__,
}
if isinstance(python_model, PythonModel):
saved_python_model_subpath = "python_model.pkl"
with open(os.path.join(path, saved_python_model_subpath), "wb") as out:
cloudpickle.dump(python_model, out)
custom_model_config_kwargs[CONFIG_KEY_PYTHON_MODEL] = saved_python_model_subpath
else:
raise MlflowException(
message=(
"`python_model` must be a subclass of `PythonModel`. Instead, found an"
" object of type: {python_model_type}".format(python_model_type=type(python_model))
),
error_code=INVALID_PARAMETER_VALUE,
)
if artifacts:
saved_artifacts_config = {}
with TempDir() as tmp_artifacts_dir:
tmp_artifacts_config = {}
saved_artifacts_dir_subpath = "artifacts"
for artifact_name, artifact_uri in artifacts.items():
tmp_artifact_path = _download_artifact_from_uri(
artifact_uri=artifact_uri, output_path=tmp_artifacts_dir.path()
)
tmp_artifacts_config[artifact_name] = tmp_artifact_path
saved_artifact_subpath = posixpath.join(
saved_artifacts_dir_subpath,
os.path.relpath(path=tmp_artifact_path, start=tmp_artifacts_dir.path()),
)
saved_artifacts_config[artifact_name] = {
CONFIG_KEY_ARTIFACT_RELATIVE_PATH: saved_artifact_subpath,
CONFIG_KEY_ARTIFACT_URI: artifact_uri,
}
shutil.move(tmp_artifacts_dir.path(), os.path.join(path, saved_artifacts_dir_subpath))
custom_model_config_kwargs[CONFIG_KEY_ARTIFACTS] = saved_artifacts_config
saved_code_subpath = None
if code_paths is not None:
saved_code_subpath = "code"
for code_path in code_paths:
_copy_file_or_tree(src=code_path, dst=path, dst_dir=saved_code_subpath)
mlflow.pyfunc.add_to_model(
model=mlflow_model,
loader_module=__name__,
code=saved_code_subpath,
env=_CONDA_ENV_FILE_NAME,
**custom_model_config_kwargs,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
mlflow.pyfunc.FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
def save_model(pytorch_model,
path,
conda_env=None,
mlflow_model=None,
code_paths=None,
pickle_module=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
requirements_file=None,
extra_files=None,
**kwargs):
"""
Save a PyTorch model to a path on the local file system.
:param pytorch_model: PyTorch model to be saved. Can be either an eager model (subclass of
``torch.nn.Module``) or scripted model prepared via ``torch.jit.script``
or ``torch.jit.trace``.
The model accept a single ``torch.FloatTensor`` as
input and produce a single output tensor.
If saving an eager model, any code dependencies of the
model's class, including the class definition itself, should be
included in one of the following locations:
- The package(s) listed in the model's Conda environment, specified
by the ``conda_env`` parameter.
- One or more of the files specified by the ``code_paths`` parameter.
:param path: Local path where the model is to be saved.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is added to the model. The
following is an *example* dictionary representation of a Conda environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'pytorch=0.4.1',
'torchvision=0.2.1'
]
}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param pickle_module: The module that PyTorch should use to serialize ("pickle") the specified
``pytorch_model``. This is passed as the ``pickle_module`` parameter
to ``torch.save()``. By default, this module is also used to
deserialize ("unpickle") the PyTorch model at load time.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example can be a Pandas DataFrame where the given
example will be serialized to json using the Pandas split-oriented
format, or a numpy array where the example will be serialized to json
by converting it to a list. Bytes are base64-encoded.
:param requirements_file: A string containing the path to requirements file. Remote URIs
are resolved to absolute filesystem paths.
For example, consider the following ``requirements_file`` string -
requirements_file = "s3://my-bucket/path/to/my_file"
In this case, the ``"my_file"`` requirements file is downloaded from S3.
If ``None``, no requirements file is added to the model.
:param extra_files: A list containing the paths to corresponding extra files. Remote URIs
are resolved to absolute filesystem paths.
For example, consider the following ``extra_files`` list -
extra_files = ["s3://my-bucket/path/to/my_file1",
"s3://my-bucket/path/to/my_file2"]
In this case, the ``"my_file1 & my_file2"`` extra file is downloaded from S3.
If ``None``, no extra files are added to the model.
:param kwargs: kwargs to pass to ``torch.save`` method.
.. code-block:: python
:caption: Example
import os
import torch
import mlflow.pytorch
# Class defined here
class LinearNNModel(torch.nn.Module):
...
# Initialize our model, criterion and optimizer
...
# Training loop
...
# Save PyTorch models to current working directory
with mlflow.start_run() as run:
mlflow.pytorch.save_model(model, "model")
# Convert to a scripted model and save it
scripted_pytorch_model = torch.jit.script(model)
mlflow.pytorch.save_model(scripted_pytorch_model, "scripted_model")
# Load each saved model for inference
for model_path in ["model", "scripted_model"]:
model_uri = "{}/{}".format(os.getcwd(), model_path)
loaded_model = mlflow.pytorch.load_model(model_uri)
print("Loaded {}:".format(model_path))
for x in [6.0, 8.0, 12.0, 30.0]:
X = torch.Tensor([[x]])
y_pred = loaded_model(X)
print("predict X: {}, y_pred: {:.2f}".format(x, y_pred.data.item()))
print("--")
.. code-block:: text
:caption: Output
Loaded model:
predict X: 6.0, y_pred: 11.90
predict X: 8.0, y_pred: 15.92
predict X: 12.0, y_pred: 23.96
predict X: 30.0, y_pred: 60.13
--
Loaded scripted_model:
predict X: 6.0, y_pred: 11.90
predict X: 8.0, y_pred: 15.92
predict X: 12.0, y_pred: 23.96
predict X: 30.0, y_pred: 60.13
"""
import torch
pickle_module = pickle_module or mlflow_pytorch_pickle_module
if not isinstance(pytorch_model, torch.nn.Module):
raise TypeError("Argument 'pytorch_model' should be a torch.nn.Module")
if code_paths is not None:
if not isinstance(code_paths, list):
raise TypeError(
"Argument code_paths should be a list, not {}".format(
type(code_paths)))
path = os.path.abspath(path)
if os.path.exists(path):
raise RuntimeError("Path '{}' already exists".format(path))
if mlflow_model is None:
mlflow_model = Model()
os.makedirs(path)
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
model_data_subpath = "data"
model_data_path = os.path.join(path, model_data_subpath)
os.makedirs(model_data_path)
# Persist the pickle module name as a file in the model's `data` directory. This is necessary
# because the `data` directory is the only available parameter to `_load_pyfunc`, and it
# does not contain the MLmodel configuration; therefore, it is not sufficient to place
# the module name in the MLmodel
#
# TODO: Stop persisting this information to the filesystem once we have a mechanism for
# supplying the MLmodel configuration to `mlflow.pytorch._load_pyfunc`
pickle_module_path = os.path.join(model_data_path,
_PICKLE_MODULE_INFO_FILE_NAME)
with open(pickle_module_path, "w") as f:
f.write(pickle_module.__name__)
# Save pytorch model
model_path = os.path.join(model_data_path,
_SERIALIZED_TORCH_MODEL_FILE_NAME)
if isinstance(pytorch_model, torch.jit.ScriptModule):
torch.jit.ScriptModule.save(pytorch_model, model_path)
else:
torch.save(pytorch_model,
model_path,
pickle_module=pickle_module,
**kwargs)
torchserve_artifacts_config = {}
if requirements_file:
if not isinstance(requirements_file, str):
raise TypeError("Path to requirements file should be a string")
with TempDir() as tmp_requirements_dir:
_download_artifact_from_uri(
artifact_uri=requirements_file,
output_path=tmp_requirements_dir.path())
rel_path = os.path.basename(requirements_file)
torchserve_artifacts_config[_REQUIREMENTS_FILE_KEY] = {
"path": rel_path
}
shutil.move(tmp_requirements_dir.path(rel_path), path)
if extra_files:
torchserve_artifacts_config[_EXTRA_FILES_KEY] = []
if not isinstance(extra_files, list):
raise TypeError("Extra files argument should be a list")
with TempDir() as tmp_extra_files_dir:
for extra_file in extra_files:
_download_artifact_from_uri(
artifact_uri=extra_file,
output_path=tmp_extra_files_dir.path())
rel_path = posixpath.join(
_EXTRA_FILES_KEY,
os.path.basename(extra_file),
)
torchserve_artifacts_config[_EXTRA_FILES_KEY].append(
{"path": rel_path})
shutil.move(
tmp_extra_files_dir.path(),
posixpath.join(path, _EXTRA_FILES_KEY),
)
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env()
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
if code_paths is not None:
code_dir_subpath = "code"
for code_path in code_paths:
_copy_file_or_tree(src=code_path,
dst=path,
dst_dir=code_dir_subpath)
else:
code_dir_subpath = None
mlflow_model.add_flavor(
FLAVOR_NAME,
model_data=model_data_subpath,
pytorch_version=torch.__version__,
**torchserve_artifacts_config,
)
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.pytorch",
data=model_data_subpath,
pickle_module_name=pickle_module.__name__,
code=code_dir_subpath,
env=conda_env_subpath,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(
onnx_model,
path,
conda_env=None,
code_paths=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
onnx_execution_providers=None,
):
"""
Save an ONNX model to a path on the local file system.
:param onnx_model: ONNX model to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: {{ conda_env }}
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example can be a Pandas DataFrame where the given
example will be serialized to json using the Pandas split-oriented
format, or a numpy array where the example will be serialized to json
by converting it to a list. Bytes are base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
:param onnx_execution_providers: List of strings defining onnxruntime execution providers.
Defaults to example:
``['CUDAExecutionProvider', 'CPUExecutionProvider']``
This uses GPU preferentially over CPU.
See onnxruntime API for further descriptions:
https://onnxruntime.ai/docs/execution-providers/
"""
import onnx
if onnx_execution_providers is None:
onnx_execution_providers = ONNX_EXECUTION_PROVIDERS
_validate_env_arguments(conda_env, pip_requirements,
extra_pip_requirements)
path = os.path.abspath(path)
_validate_and_prepare_target_save_path(path)
code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
model_data_subpath = "model.onnx"
model_data_path = os.path.join(path, model_data_subpath)
# Save onnx-model
onnx.save_model(onnx_model, model_data_path)
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.onnx",
data=model_data_subpath,
env=_CONDA_ENV_FILE_NAME,
code=code_dir_subpath,
)
mlflow_model.add_flavor(
FLAVOR_NAME,
onnx_version=onnx.__version__,
data=model_data_subpath,
providers=onnx_execution_providers,
code=code_dir_subpath,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(
conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME),
"\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME),
"\n".join(pip_requirements))
def log_model(pytorch_model,
artifact_path,
conda_env=None,
code_paths=None,
pickle_module=None,
registered_model_name=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
requirements_file=None,
extra_files=None,
**kwargs):
"""
Log a PyTorch model as an MLflow artifact for the current run.
:param pytorch_model: PyTorch model to be saved. Can be either an eager model (subclass of
``torch.nn.Module``) or scripted model prepared via ``torch.jit.script``
or ``torch.jit.trace``.
The model accept a single ``torch.FloatTensor`` as
input and produce a single output tensor.
If saving an eager model, any code dependencies of the
model's class, including the class definition itself, should be
included in one of the following locations:
- The package(s) listed in the model's Conda environment, specified
by the ``conda_env`` parameter.
- One or more of the files specified by the ``code_paths`` parameter.
:param artifact_path: Run-relative artifact path.
:param conda_env: Path to a Conda environment file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is added to the model. The
following is an *example* dictionary representation of a Conda environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'pytorch=0.4.1',
'torchvision=0.2.1'
]
}
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param pickle_module: The module that PyTorch should use to serialize ("pickle") the specified
``pytorch_model``. This is passed as the ``pickle_module`` parameter
to ``torch.save()``. By default, this module is also used to
deserialize ("unpickle") the PyTorch model at load time.
:param registered_model_name: (Experimental) If given, create a model version under
``registered_model_name``, also creating a registered model if one
with the given name does not exist.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example can be a Pandas DataFrame where the given
example will be serialized to json using the Pandas split-oriented
format, or a numpy array where the example will be serialized to json
by converting it to a list. Bytes are base64-encoded.
:param await_registration_for: Number of seconds to wait for the model version to finish
being created and is in ``READY`` status. By default, the function
waits for five minutes. Specify 0 or None to skip waiting.
:param requirements_file: A string containing the path to requirements file. Remote URIs
are resolved to absolute filesystem paths.
For example, consider the following ``requirements_file`` string -
requirements_file = "s3://my-bucket/path/to/my_file"
In this case, the ``"my_file"`` requirements file is downloaded from S3.
If ``None``, no requirements file is added to the model.
:param extra_files: A list containing the paths to corresponding extra files. Remote URIs
are resolved to absolute filesystem paths.
For example, consider the following ``extra_files`` list -
extra_files = ["s3://my-bucket/path/to/my_file1",
"s3://my-bucket/path/to/my_file2"]
In this case, the ``"my_file1 & my_file2"`` extra file is downloaded from S3.
If ``None``, no extra files are added to the model.
:param kwargs: kwargs to pass to ``torch.save`` method.
.. code-block:: python
:caption: Example
import numpy as np
import torch
import mlflow.pytorch
class LinearNNModel(torch.nn.Module):
def __init__(self):
super(LinearNNModel, self).__init__()
self.linear = torch.nn.Linear(1, 1) # One in and one out
def forward(self, x):
y_pred = self.linear(x)
return y_pred
def gen_data():
# Example linear model modified to use y = 2x
# from https://github.com/hunkim/PyTorchZeroToAll
# X training data, y labels
X = torch.arange(1.0, 25.0).view(-1, 1)
y = torch.from_numpy(np.array([x * 2 for x in X])).view(-1, 1)
return X, y
# Define model, loss, and optimizer
model = LinearNNModel()
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
# Training loop
epochs = 250
X, y = gen_data()
for epoch in range(epochs):
# Forward pass: Compute predicted y by passing X to the model
y_pred = model(X)
# Compute the loss
loss = criterion(y_pred, y)
# Zero gradients, perform a backward pass, and update the weights.
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Log the model
with mlflow.start_run() as run:
mlflow.pytorch.log_model(model, "model")
# convert to scripted model and log the model
scripted_pytorch_model = torch.jit.script(model)
mlflow.pytorch.log_model(scripted_pytorch_model, "scripted_model")
# Fetch the logged model artifacts
print("run_id: {}".format(run.info.run_id))
for artifact_path in ["model/data", "scripted_model/data"]:
artifacts = [f.path for f in MlflowClient().list_artifacts(run.info.run_id,
artifact_path)]
print("artifacts: {}".format(artifacts))
.. code-block:: text
:caption: Output
run_id: 1a1ec9e413ce48e9abf9aec20efd6f71
artifacts: ['model/data/model.pth',
'model/data/pickle_module_info.txt']
artifacts: ['scripted_model/data/model.pth',
'scripted_model/data/pickle_module_info.txt']
.. figure:: ../_static/images/pytorch_logged_models.png
PyTorch logged models
"""
pickle_module = pickle_module or mlflow_pytorch_pickle_module
Model.log(
artifact_path=artifact_path,
flavor=mlflow.pytorch,
pytorch_model=pytorch_model,
conda_env=conda_env,
code_paths=code_paths,
pickle_module=pickle_module,
registered_model_name=registered_model_name,
signature=signature,
input_example=input_example,
await_registration_for=await_registration_for,
requirements_file=requirements_file,
extra_files=extra_files,
**kwargs,
)
def log_model(
gluon_model,
artifact_path,
conda_env=None,
registered_model_name=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Log a Gluon model as an MLflow artifact for the current run.
:param gluon_model: Gluon model to be saved. Must be already hybridized.
:param artifact_path: Run-relative artifact path.
:param conda_env: Either a dictionary representation of a Conda environment or
the path to a Conda environment yaml file.
If provided, this decribes the environment this model should be
run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`mlflow.gluon.get_default_conda_env()` environment is added to
the model. The following is an *example* dictionary representation of a
Conda environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'mxnet=1.5.0'
]
}
:param registered_model_name: (Experimental) If given, create a model version under
``registered_model_name``, also creating a registered model if one
with the given name does not exist.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example can be a Pandas DataFrame where the given
example will be serialized to json using the Pandas split-oriented
format, or a numpy array where the example will be serialized to json
by converting it to a list. Bytes are base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
.. code-block:: python
:caption: Example
from mxnet.gluon import Trainer
from mxnet.gluon.contrib import estimator
from mxnet.gluon.loss import SoftmaxCrossEntropyLoss
from mxnet.gluon.nn import HybridSequential
from mxnet.metric import Accuracy
import mlflow
# Build, compile, and train your model
net = HybridSequential()
with net.name_scope():
...
net.hybridize()
net.collect_params().initialize()
softmax_loss = SoftmaxCrossEntropyLoss()
trainer = Trainer(net.collect_params())
est = estimator.Estimator(net=net, loss=softmax_loss, metrics=Accuracy(), trainer=trainer)
# Log metrics and log the model
with mlflow.start_run():
est.fit(train_data=train_data, epochs=100, val_data=validation_data)
mlflow.gluon.log_model(net, "model")
"""
Model.log(
artifact_path=artifact_path,
flavor=mlflow.gluon,
gluon_model=gluon_model,
conda_env=conda_env,
registered_model_name=registered_model_name,
signature=signature,
input_example=input_example,
pip_requirements=pip_requirements,
extra_pip_requirements=extra_pip_requirements,
)
def log_explainer(
explainer,
artifact_path,
serialize_model_using_mlflow=True,
conda_env=None,
code_paths=None,
registered_model_name=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Log an SHAP explainer as an MLflow artifact for the current run.
:param explainer: SHAP explainer to be saved.
:param artifact_path: Run-relative artifact path.
:param serialize_model_using_mlflow: When set to True, MLflow will extract the underlying
model and serialize it as an MLmodel, otherwise it
uses SHAP's internal serialization. Defaults to True.
Currently MLflow serialization is only supported for
models of 'sklearn' or 'pytorch' flavors.
:param conda_env: {{ conda_env }}
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param registered_model_name: If given, create a model version under
``registered_model_name``, also creating a registered model if one
with the given name does not exist.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param await_registration_for: Number of seconds to wait for the model version to finish
being created and is in ``READY`` status. By default, the function
waits for five minutes. Specify 0 or None to skip waiting.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
"""
Model.log(
artifact_path=artifact_path,
flavor=mlflow.shap,
explainer=explainer,
conda_env=conda_env,
code_paths=code_paths,
serialize_model_using_mlflow=serialize_model_using_mlflow,
registered_model_name=registered_model_name,
signature=signature,
input_example=input_example,
await_registration_for=await_registration_for,
pip_requirements=pip_requirements,
extra_pip_requirements=extra_pip_requirements,
)
def log_model(spark_model,
artifact_path,
conda_env=None,
dfs_tmpdir=None,
sample_input=None):
"""
Log a Spark MLlib model as an MLflow artifact for the current run. This uses the
MLlib persistence format and produces an MLflow Model with the Spark flavor.
:param spark_model: Spark model to be saved - MLFlow can only save descendants of
pyspark.ml.Model which implement MLReadable and MLWritable.
:param artifact_path: Run relative artifact path.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If `None`, the default
:func:`get_default_conda_env()` environment is added to the model.
The following is an *example* dictionary representation of a Conda
environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'pyspark=2.3.0'
]
}
:param dfs_tmpdir: Temporary directory path on Distributed (Hadoop) File System (DFS) or local
filesystem if running in local mode. The model is written in this
destination and then copied into the model's artifact directory. This is
necessary as Spark ML models read from and write to DFS if running on a
cluster. If this operation completes successfully, all temporary files
created on the DFS are removed. Defaults to ``/tmp/mlflow``.
:param sample_input: A sample input used to add the MLeap flavor to the model.
This must be a PySpark DataFrame that the model can evaluate. If
``sample_input`` is ``None``, the MLeap flavor is not added.
>>> from pyspark.ml import Pipeline
>>> from pyspark.ml.classification import LogisticRegression
>>> from pyspark.ml.feature import HashingTF, Tokenizer
>>> training = spark.createDataFrame([
... (0, "a b c d e spark", 1.0),
... (1, "b d", 0.0),
... (2, "spark f g h", 1.0),
... (3, "hadoop mapreduce", 0.0) ], ["id", "text", "label"])
>>> tokenizer = Tokenizer(inputCol="text", outputCol="words")
>>> hashingTF = HashingTF(inputCol=tokenizer.getOutputCol(), outputCol="features")
>>> lr = LogisticRegression(maxIter=10, regParam=0.001)
>>> pipeline = Pipeline(stages=[tokenizer, hashingTF, lr])
>>> model = pipeline.fit(training)
>>> mlflow.spark.log_model(model, "spark-model")
"""
from py4j.protocol import Py4JJavaError
_validate_model(spark_model)
from pyspark.ml import PipelineModel
if not isinstance(spark_model, PipelineModel):
spark_model = PipelineModel([spark_model])
run_id = mlflow.tracking.fluent._get_or_start_run().info.run_id
run_root_artifact_uri = mlflow.get_artifact_uri()
# If the artifact URI is a local filesystem path, defer to Model.log() to persist the model,
# since Spark may not be able to write directly to the driver's filesystem. For example,
# writing to `file:/uri` will write to the local filesystem from each executor, which will
# be incorrect on multi-node clusters - to avoid such issues we just use the Model.log() path
# here.
if mlflow.tracking.utils._is_local_uri(run_root_artifact_uri):
return Model.log(artifact_path=artifact_path,
flavor=mlflow.spark,
spark_model=spark_model,
conda_env=conda_env,
dfs_tmpdir=dfs_tmpdir,
sample_input=sample_input)
# If Spark cannot write directly to the artifact repo, defer to Model.log() to persist the
# model
model_dir = os.path.join(run_root_artifact_uri, artifact_path)
try:
spark_model.save(os.path.join(model_dir, _SPARK_MODEL_PATH_SUB))
except Py4JJavaError:
return Model.log(artifact_path=artifact_path,
flavor=mlflow.spark,
spark_model=spark_model,
conda_env=conda_env,
dfs_tmpdir=dfs_tmpdir,
sample_input=sample_input)
# Otherwise, override the default model log behavior and save model directly to artifact repo
mlflow_model = Model(artifact_path=artifact_path, run_id=run_id)
with TempDir() as tmp:
tmp_model_metadata_dir = tmp.path()
_save_model_metadata(tmp_model_metadata_dir, spark_model, mlflow_model,
sample_input, conda_env)
mlflow.tracking.fluent.log_artifacts(tmp_model_metadata_dir,
artifact_path)
def save_explainer(
explainer,
path,
serialize_model_using_mlflow=True,
conda_env=None,
code_paths=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Save a SHAP explainer to a path on the local file system. Produces an MLflow Model
containing the following flavors:
- :py:mod:`mlflow.shap`
- :py:mod:`mlflow.pyfunc`
:param explainer: SHAP explainer to be saved.
:param path: Local path where the explainer is to be saved.
:param serialize_model_using_mlflow: When set to True, MLflow will extract the underlying
model and serialize it as an MLmodel, otherwise it
uses SHAP's internal serialization. Defaults to True.
Currently MLflow serialization is only supported for
models of 'sklearn' or 'pytorch' flavors.
:param conda_env: {{ conda_env }}
:param code_paths: A list of local filesystem paths to Python file dependencies (or directories
containing file dependencies). These files are *prepended* to the system
path when the model is loaded.
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
"""
import shap
_validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
_validate_and_prepare_target_save_path(path)
code_dir_subpath = _validate_and_copy_code_paths(code_paths, path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
underlying_model_flavor = None
underlying_model_path = None
serializable_by_mlflow = False
# saving the underlying model if required
if serialize_model_using_mlflow:
underlying_model_flavor = get_underlying_model_flavor(explainer.model)
if underlying_model_flavor != _UNKNOWN_MODEL_FLAVOR:
serializable_by_mlflow = True # prevents SHAP from serializing the underlying model
underlying_model_path = os.path.join(path, _UNDERLYING_MODEL_SUBPATH)
else:
warnings.warn(
"Unable to serialize underlying model using MLflow, will use SHAP serialization"
)
if underlying_model_flavor == mlflow.sklearn.FLAVOR_NAME:
mlflow.sklearn.save_model(explainer.model.inner_model.__self__, underlying_model_path)
elif underlying_model_flavor == mlflow.pytorch.FLAVOR_NAME:
mlflow.pytorch.save_model(explainer.model.inner_model, underlying_model_path)
# saving the explainer object
explainer_data_subpath = "explainer.shap"
explainer_output_path = os.path.join(path, explainer_data_subpath)
with open(explainer_output_path, "wb") as explainer_output_file_handle:
if serialize_model_using_mlflow and serializable_by_mlflow:
explainer.save(explainer_output_file_handle, model_saver=False)
else:
explainer.save(explainer_output_file_handle)
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.shap",
model_path=explainer_data_subpath,
underlying_model_flavor=underlying_model_flavor,
env=_CONDA_ENV_FILE_NAME,
code=code_dir_subpath,
)
mlflow_model.add_flavor(
FLAVOR_NAME,
shap_version=shap.__version__,
serialized_explainer=explainer_data_subpath,
underlying_model_flavor=underlying_model_flavor,
code=code_dir_subpath,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)
if underlying_model_path is not None:
underlying_model_conda_env = _get_conda_env_for_underlying_model(underlying_model_path)
conda_env = _merge_environments(conda_env, underlying_model_conda_env)
pip_requirements = _get_pip_deps(conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
def save_model(
pd_model,
path,
training=False,
conda_env=None,
mlflow_model=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Save a paddle model to a path on the local file system. Produces an MLflow Model
containing the following flavors:
- :py:mod:`mlflow.paddle`
- :py:mod:`mlflow.pyfunc`. NOTE: This flavor is only included for paddle models
that define `predict()`, since `predict()` is required for pyfunc model inference.
:param pd_model: paddle model to be saved.
:param path: Local path where the model is to be saved.
:param training: Only valid when saving a model trained using the PaddlePaddle high level API.
If set to True, the saved model supports both re-training and
inference. If set to False, it only supports inference.
:param conda_env: {{ conda_env }}
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
.. code-block:: python
:caption: Example
import mlflow.paddle
import paddle
from paddle.nn import Linear
import paddle.nn.functional as F
import numpy as np
import os
import random
from sklearn.datasets import load_boston
from sklearn.model_selection import train_test_split
from sklearn import preprocessing
def load_data():
# dataset on boston housing prediction
X, y = load_boston(return_X_y=True)
min_max_scaler = preprocessing.MinMaxScaler()
X_min_max = min_max_scaler.fit_transform(X)
X_normalized = preprocessing.scale(X_min_max, with_std=False)
X_train, X_test, y_train, y_test = train_test_split(
X_normalized, y, test_size=0.2, random_state=42)
y_train = y_train.reshape(-1, 1)
y_test = y_test.reshape(-1, 1)
return np.concatenate(
(X_train, y_train), axis=1),np.concatenate((X_test, y_test), axis=1
)
class Regressor(paddle.nn.Layer):
def __init__(self):
super(Regressor, self).__init__()
self.fc = Linear(in_features=13, out_features=1)
@paddle.jit.to_static
def forward(self, inputs):
x = self.fc(inputs)
return x
model = Regressor()
model.train()
training_data, test_data = load_data()
opt = paddle.optimizer.SGD(learning_rate=0.01, parameters=model.parameters())
EPOCH_NUM = 10
BATCH_SIZE = 10
for epoch_id in range(EPOCH_NUM):
np.random.shuffle(training_data)
mini_batches = [training_data[k : k + BATCH_SIZE]
for k in range(0, len(training_data), BATCH_SIZE)]
for iter_id, mini_batch in enumerate(mini_batches):
x = np.array(mini_batch[:, :-1]).astype('float32')
y = np.array(mini_batch[:, -1:]).astype('float32')
house_features = paddle.to_tensor(x)
prices = paddle.to_tensor(y)
predicts = model(house_features)
loss = F.square_error_cost(predicts, label=prices)
avg_loss = paddle.mean(loss)
if iter_id%20==0:
print("epoch: {}, iter: {}, loss is: {}".format(
epoch_id, iter_id, avg_loss.numpy()))
avg_loss.backward()
opt.step()
opt.clear_grad()
mlflow.log_param('learning_rate', 0.01)
mlflow.paddle.log_model(model, "model")
sk_path_dir = './test-out'
mlflow.paddle.save_model(model, sk_path_dir)
print("Model saved in run %s" % mlflow.active_run().info.run_uuid)
"""
import paddle
_validate_env_arguments(conda_env, pip_requirements,
extra_pip_requirements)
if os.path.exists(path):
raise MlflowException(message="Path '{}' already exists".format(path),
error_code=RESOURCE_ALREADY_EXISTS)
os.makedirs(path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
model_data_subpath = "model"
output_path = os.path.join(path, model_data_subpath)
if isinstance(pd_model, paddle.Model):
pd_model.save(output_path, training=training)
else:
paddle.jit.save(pd_model, output_path)
# `PyFuncModel` only works for paddle models that define `predict()`.
pyfunc.add_to_model(
mlflow_model,
loader_module="mlflow.paddle",
model_path=model_data_subpath,
env=_CONDA_ENV_FILE_NAME,
)
mlflow_model.add_flavor(
FLAVOR_NAME,
pickled_model=model_data_subpath,
paddle_version=paddle.__version__,
)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(
conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME),
"\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME),
"\n".join(pip_requirements))
def save_model(
h2o_model,
path,
conda_env=None,
mlflow_model=None,
settings=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Save an H2O model to a path on the local file system.
:param h2o_model: H2O model to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: {{ conda_env }}
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param mlflow_model: :py:mod:`mlflow.models.Model` this flavor is being added to.
"""
import h2o
_validate_env_arguments(conda_env, pip_requirements,
extra_pip_requirements)
path = os.path.abspath(path)
if os.path.exists(path):
raise Exception("Path '{}' already exists".format(path))
model_data_subpath = "model.h2o"
model_data_path = os.path.join(path, model_data_subpath)
os.makedirs(model_data_path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# Save h2o-model
if hasattr(h2o, "download_model"):
h2o_save_location = h2o.download_model(model=h2o_model,
path=model_data_path)
else:
warnings.warn(
"If your cluster is remote, H2O may not store the model correctly. "
"Please upgrade H2O version to a newer version")
h2o_save_location = h2o.save_model(model=h2o_model,
path=model_data_path,
force=True)
model_file = os.path.basename(h2o_save_location)
# Save h2o-settings
if settings is None:
settings = {}
settings["full_file"] = h2o_save_location
settings["model_file"] = model_file
settings["model_dir"] = model_data_path
with open(os.path.join(model_data_path, "h2o.yaml"), "w") as settings_file:
yaml.safe_dump(settings, stream=settings_file)
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow.h2o",
data=model_data_subpath,
env=_CONDA_ENV_FILE_NAME)
mlflow_model.add_flavor(FLAVOR_NAME,
h2o_version=h2o.__version__,
data=model_data_subpath)
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
if conda_env is None:
if pip_requirements is None:
default_reqs = get_default_pip_requirements()
# To ensure `_load_pyfunc` can successfully load the model during the dependency
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
inferred_reqs = mlflow.models.infer_pip_requirements(
path,
FLAVOR_NAME,
fallback=default_reqs,
)
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
else:
default_reqs = None
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
default_reqs,
pip_requirements,
extra_pip_requirements,
)
else:
conda_env, pip_requirements, pip_constraints = _process_conda_env(
conda_env)
with open(os.path.join(path, _CONDA_ENV_FILE_NAME), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# Save `constraints.txt` if necessary
if pip_constraints:
write_to(os.path.join(path, _CONSTRAINTS_FILE_NAME),
"\n".join(pip_constraints))
# Save `requirements.txt`
write_to(os.path.join(path, _REQUIREMENTS_FILE_NAME),
"\n".join(pip_requirements))
def save_model(keras_model,
path,
conda_env=None,
mlflow_model=None,
custom_objects=None,
keras_module=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
**kwargs):
"""
Save a Keras model to a path on the local file system.
:param keras_model: Keras model to be saved.
:param path: Local path where the model is to be saved.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the
dependencies contained in :func:`get_default_conda_env()`. If
``None``, the default :func:`get_default_conda_env()` environment is
added to the model. The following is an *example* dictionary
representation of a Conda environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'keras=2.2.4',
'tensorflow=1.8.0'
]
}
:param mlflow_model: MLflow model config this flavor is being added to.
:param custom_objects: A Keras ``custom_objects`` dictionary mapping names (strings) to
custom classes or functions associated with the Keras model. MLflow saves
these custom layers using CloudPickle and restores them automatically
when the model is loaded with :py:func:`mlflow.keras.load_model` and
:py:func:`mlflow.pyfunc.load_model`.
:param keras_module: Keras module to be used to save / load the model
(``keras`` or ``tf.keras``). If not provided, MLflow will
attempt to infer the Keras module based on the given model.
:param kwargs: kwargs to pass to ``keras_model.save`` method.
:param signature: (Experimental) :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: (Experimental) Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
.. code-block:: python
:caption: Example
import mlflow
# Build, compile, and train your model
keras_model = ...
keras_model_path = ...
keras_model.compile(optimizer="rmsprop", loss="mse", metrics=["accuracy"])
results = keras_model.fit(
x_train, y_train, epochs=20, batch_size = 128, validation_data=(x_val, y_val))
# Save the model as an MLflow Model
mlflow.keras.save_model(keras_model, keras_model_path)
"""
if keras_module is None:
def _is_plain_keras(model):
try:
import keras
if LooseVersion(keras.__version__) < LooseVersion("2.2.0"):
import keras.engine
return isinstance(model, keras.engine.Model)
else:
# NB: Network is the first parent with save method
import keras.engine.network
return isinstance(model, keras.engine.network.Network)
except ImportError:
return False
def _is_tf_keras(model):
try:
# NB: Network is not exposed in tf.keras, we check for Model instead.
import tensorflow.keras.models
return isinstance(model, tensorflow.keras.models.Model)
except ImportError:
return False
if _is_plain_keras(keras_model):
keras_module = importlib.import_module("keras")
elif _is_tf_keras(keras_model):
keras_module = importlib.import_module("tensorflow.keras")
else:
raise MlflowException(
"Unable to infer keras module from the model, please specify "
"which keras module ('keras' or 'tensorflow.keras') is to be "
"used to save and load the model.")
elif type(keras_module) == str:
keras_module = importlib.import_module(keras_module)
# check if path exists
path = os.path.abspath(path)
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path))
# construct new data folder in existing path
data_subpath = "data"
data_path = os.path.join(path, data_subpath)
os.makedirs(data_path)
if mlflow_model is None:
mlflow_model = Model()
if signature is not None:
mlflow_model.signature = signature
if input_example is not None:
_save_example(mlflow_model, input_example, path)
# save custom objects if there are custom objects
if custom_objects is not None:
_save_custom_objects(data_path, custom_objects)
# save keras module spec to path/data/keras_module.txt
with open(os.path.join(data_path, _KERAS_MODULE_SPEC_PATH), "w") as f:
f.write(keras_module.__name__)
# By default, Keras uses the SavedModel format -- specified by "tf"
# However, we choose to align with prior default of mlflow, HDF5
save_format = kwargs.get("save_format", "h5")
# save keras save_format to path/data/save_format.txt
with open(os.path.join(data_path, _KERAS_SAVE_FORMAT_PATH), "w") as f:
f.write(save_format)
# save keras model
# To maintain prior behavior, when the format is HDF5, we save
# with the h5 file extension. Otherwise, model_path is a directory
# where the saved_model.pb will be stored (for SavedModel format)
file_extension = ".h5" if save_format == "h5" else ""
model_subpath = os.path.join(data_subpath, _MODEL_SAVE_PATH)
model_path = os.path.join(path, model_subpath) + file_extension
if path.startswith("/dbfs/"):
# The Databricks Filesystem uses a FUSE implementation that does not support
# random writes. It causes an error.
with tempfile.NamedTemporaryFile(suffix=".h5") as f:
keras_model.save(f.name, **kwargs)
f.flush() # force flush the data
shutil.copyfile(src=f.name, dst=model_path)
else:
keras_model.save(model_path, **kwargs)
# update flavor info to mlflow_model
mlflow_model.add_flavor(
FLAVOR_NAME,
keras_module=keras_module.__name__,
keras_version=keras_module.__version__,
save_format=save_format,
data=data_subpath,
)
# save conda.yaml info to path/conda.yml
if conda_env is None:
conda_env = get_default_conda_env(include_cloudpickle=custom_objects
is not None,
keras_module=keras_module)
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, _CONDA_ENV_SUBPATH), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
# append loader_module, data and env data to mlflow_model
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow.keras",
data=data_subpath,
env=_CONDA_ENV_SUBPATH)
# save mlflow_model to path/MLmodel
mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
def save_model(tf_saved_model_dir,
tf_meta_graph_tags,
tf_signature_def_key,
path,
mlflow_model=Model(),
conda_env=None):
"""
Save a *serialized* collection of TensorFlow graphs and variables as an MLflow model
to a local path. This method operates on TensorFlow variables and graphs that have been
serialized in TensorFlow's ``SavedModel`` format. For more information about ``SavedModel``
format, see the TensorFlow documentation:
https://www.tensorflow.org/guide/saved_model#save_and_restore_models.
:param tf_saved_model_dir: Path to the directory containing serialized TensorFlow variables and
graphs in ``SavedModel`` format.
:param tf_meta_graph_tags: A list of tags identifying the model's metagraph within the
serialized ``SavedModel`` object. For more information, see the
``tags`` parameter of the
``tf.saved_model.builder.savedmodelbuilder`` method.
:param tf_signature_def_key: A string identifying the input/output signature associated with the
model. This is a key within the serialized ``savedmodel``
signature definition mapping. For more information, see the
``signature_def_map`` parameter of the
``tf.saved_model.builder.savedmodelbuilder`` method.
:param path: Local path where the MLflow model is to be saved.
:param mlflow_model: MLflow model configuration to which to add the ``tensorflow`` flavor.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If ``None``, the default
:func:`get_default_conda_env()` environment is added to the model. The
following is an *example* dictionary representation of a Conda environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'tensorflow=1.8.0'
]
}
"""
_logger.info(
"Validating the specified TensorFlow model by attempting to load it in a new TensorFlow"
" graph...")
_validate_saved_model(tf_saved_model_dir=tf_saved_model_dir,
tf_meta_graph_tags=tf_meta_graph_tags,
tf_signature_def_key=tf_signature_def_key)
_logger.info("Validation succeeded!")
if os.path.exists(path):
raise MlflowException("Path '{}' already exists".format(path),
DIRECTORY_NOT_EMPTY)
os.makedirs(path)
root_relative_path = _copy_file_or_tree(src=tf_saved_model_dir,
dst=path,
dst_dir=None)
model_dir_subpath = "tfmodel"
shutil.move(os.path.join(path, root_relative_path),
os.path.join(path, model_dir_subpath))
conda_env_subpath = "conda.yaml"
if conda_env is None:
conda_env = get_default_conda_env()
elif not isinstance(conda_env, dict):
with open(conda_env, "r") as f:
conda_env = yaml.safe_load(f)
with open(os.path.join(path, conda_env_subpath), "w") as f:
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
mlflow_model.add_flavor(FLAVOR_NAME,
saved_model_dir=model_dir_subpath,
meta_graph_tags=tf_meta_graph_tags,
signature_def_key=tf_signature_def_key)
pyfunc.add_to_model(mlflow_model,
loader_module="mlflow.tensorflow",
env=conda_env_subpath)
mlflow_model.save(os.path.join(path, "MLmodel"))
