MLflow

About this page

This is an API reference for using MLflow in BentoML. Please refer to MLflow guide for more information about how to use MLflow in BentoML.

Note

You can find more examples for MLflow in our examples/mlflow directory.

bentoml.mlflow.import_model(name: Tag | str, model_uri: str, *, signatures: dict[str, ModelSignature] | dict[str, ModelSignatureDict] | None = None, labels: dict[str, str] | None = None, custom_objects: dict[str, t.Any] | None = None, external_modules: t.List[ModuleType] | None = None, metadata: dict[str, t.Any] | None = None) → bentoml.Model

Import MLflow model from a artifact URI to the BentoML model store.

Parameters:

• name – The name to give to the model in the BentoML store. This must be a valid Tag name.

• model_uri – The MLflow model to be saved.

• signatures – Signatures of predict methods to be used. If not provided, the signatures default to {“predict”: {“batchable”: False}}. See ModelSignature for more details.

• labels – A default set of management labels to be associated with the model. For example: {"training-set": "data-v1"}.

• custom_objects – Custom objects to be saved with the model. An example is {"my-normalizer": normalizer}. Custom objects are serialized with cloudpickle.

• metadata –

  Metadata to be associated with the model. An example is {"param_a": .2}.

  Metadata is intended for display in a model management UI and therefore all values in metadata dictionary must be a primitive Python type, such as str or int.

Returns:

A Model instance referencing a saved model in the local BentoML model store.

Example:

import bentoml

bentoml.mlflow.import_model(
    'my_mlflow_model',
    model_uri="runs:/<mlflow_run_id>/run-relative/path/to/model",
    signatures={
        "predict": {"batchable": True},
    }
)

bentoml.mlflow.load_model(bento_model: str | Tag | Model) → mlflow.pyfunc.PyFuncModel

Load the MLflow PyFunc model with the given tag from the local BentoML model store.

Parameters:

bento_model – Either the tag of the model to get from the store, or a BentoML ~bentoml.Model instance to load the model from.

Returns:

The MLflow model loaded as PyFuncModel from the BentoML model store.

Example:

import bentoml
pyfunc_model = bentoml.mlflow.load_model('my_model:latest')
pyfunc_model.predict( input_df )

bentoml.mlflow.get(tag_like: str | Tag) → Model

Get the BentoML model with the given tag.

Parameters:

tag_like – The tag of the model to retrieve from the model store.

Returns:

A BentoML Model with the matching tag.

Return type:

Model

Example:

import bentoml
# target model must be from the BentoML model store
model = bentoml.mlflow.get("my_mlflow_model")

bentoml.mlflow.get_mlflow_model(tag_like: str | Tag) → mlflow.models.Model