mlflow.models

The mlflow.models module provides an API for saving machine learning models in “flavors” that can be understood by different downstream tools.

The built-in flavors are:

• mlflow.pyfunc

• mlflow.h2o

• mlflow.keras

• mlflow.lightgbm

• mlflow.pytorch

• mlflow.sklearn

• mlflow.spark

• mlflow.statsmodels

• mlflow.tensorflow

• mlflow.xgboost

• mlflow.spacy

• mlflow.fastai

• mlflow.paddle

For details, see MLflow Models.

class mlflow.models.FlavorBackend(config, **kwargs)

Bases: object

Abstract class for Flavor Backend. This class defines the API interface for local model deployment of MLflow model flavors.

can_build_image()

Returns

True if this flavor has a build_image method defined for building a docker container capable of serving the model, False otherwise.

abstract can_score_model()

Check whether this flavor backend can be deployed in the current environment.

Returns

True if this flavor backend can be applied in the current environment.

abstract predict(model_uri, input_path, output_path, content_type, json_format)

Generate predictions using a saved MLflow model referenced by the given URI. Input and output are read from and written to a file or stdin / stdout.

Parameters

• model_uri – URI pointing to the MLflow model to be used for scoring.

• input_path – Path to the file with input data. If not specified, data is read from stdin.

• output_path – Path to the file with output predictions. If not specified, data is written to stdout.

• content_type – Specifies the input format. Can be one of {json, csv}

• json_format – Only applies if content_type == json. Specifies how is the input data encoded in json. Can be one of {split, records} mirroring the behavior of Pandas orient attribute. The default is split which expects dict like data: {'index' -> [index], 'columns' -> [columns], 'data' -> [values]}, where index is optional. For more information see https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_json.html

prepare_env(model_uri)

Performs any preparation necessary to predict or serve the model, for example downloading dependencies or initializing a conda environment. After preparation, calling predict or serve should be fast.

abstract serve(model_uri, port, host)

Serve the specified MLflow model locally.

Parameters

• model_uri – URI pointing to the MLflow model to be used for scoring.

• port – Port to use for the model deployment.

• host – Host to use for the model deployment. Defaults to localhost.

class mlflow.models.Model(artifact_path=None, run_id=None, utc_time_created=None, flavors=None, signature=None, saved_input_example_info: Optional[Dict[str, Any]] = None, **kwargs)

Bases: object

An MLflow Model that can support multiple model flavors. Provides APIs for implementing new Model flavors.

add_flavor(name, **params)

Add an entry for how to serve the model in a given format.

classmethod from_dict(model_dict)

Load a model from its YAML representation.

get_input_schema()

get_output_schema()

classmethod load(path)

Load a model from its YAML representation.

classmethod log(artifact_path, flavor, registered_model_name=None, await_registration_for=300, **kwargs)

Log model using supplied flavor module. If no run is active, this method will create a new active run.

Parameters

• artifact_path – Run relative path identifying the model.

• flavor – Flavor module to save the model with. The module must have the save_model function that will persist the model as a valid MLflow model.

• 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.

• signature –

  ModelSignature describes model input and output Schema. The model signature can be inferred from datasets representing valid model input (e.g. the training dataset) and valid model output (e.g. model predictions generated on the training dataset), for example:

  from mlflow.models.signature import infer_signature
  train = df.drop_column("target_label")
  signature = infer_signature(train, model.predict(train))
  

• input_example – Input example provides one or several examples 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.

• 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.

• kwargs – Extra args passed to the model flavor.

save(path)

Write the model as a local YAML file.

property saved_input_example_info

property signature

to_dict()

Serialize the model to a dictionary.

to_json()

Write the model as json.

to_yaml(stream=None)

Write the model as yaml string.

class mlflow.models.ModelSignature(inputs: mlflow.types.schema.Schema, outputs: Optional[mlflow.types.schema.Schema] = None)

Bases: object

ModelSignature specifies schema of model’s inputs and outputs.

ModelSignature can be inferred from training dataset and model predictions using or constructed by hand by passing an input and output Schema.

classmethod from_dict(signature_dict: Dict[str, Any])

Deserialize from dictionary representation.

Parameters

signature_dict – Dictionary representation of model signature. Expected dictionary format: {‘inputs’: <json string>, ‘outputs’: <json string>” }

Returns

ModelSignature populated with the data form the dictionary.

to_dict() → Dict[str, Any]

Serialize into a ‘jsonable’ dictionary.

Input and output schema are represented as json strings. This is so that the representation is compact when embedded in a MLmofel yaml file.

Returns

dictionary representation with input and output shcema represented as json strings.

mlflow.models.infer_pip_requirements(model_uri, flavor, fallback=None)

Infers the pip requirements of the specified model by creating a subprocess and loading the model in it to determine which packages are imported.

Parameters

• model_uri – The URI of the model.

• flavor – The flavor name of the model.

• fallback – If provided, an unexpected error during the inference procedure is swallowed and the value of fallback is returned. Otherwise, the error is raised.

Returns

A list of inferred pip requirements (e.g. ["scikit-learn==0.24.2", ...]).

mlflow.models.infer_signature(model_input: Any, model_output: MlflowInferableDataset = None) → mlflow.models.signature.ModelSignature

Infer an MLflow model signature from the training data (input) and model predictions (output).

The signature represents model input and output as data frames with (optionally) named columns and data type specified as one of types defined in mlflow.types.DataType. This method will raise an exception if the user data contains incompatible types or is not passed in one of the supported formats listed below.

The input should be one of these:

• pandas.DataFrame

• dictionary of { name -> numpy.ndarray}

• numpy.ndarray

• pyspark.sql.DataFrame

The element types should be mappable to one of mlflow.types.DataType.

For pyspark.sql.DataFrame inputs, columns of type DateType and TimestampType are both inferred as type datetime, which is coerced to TimestampType at inference.

NOTE: Multidimensional (>2d) arrays (aka tensors) are not supported at this time.

Parameters

• model_input – Valid input to the model. E.g. (a subset of) the training dataset.

• model_output – Valid model output. E.g. Model predictions for the (subset of) training dataset.

Returns

ModelSignature