Deploy MLflow models with ModelBuilder

You can deploy MLflow models to a SageMaker AI endpoint using Amazon SageMaker AI Model Builder. For more information about Amazon SageMaker AI Model Builder, see Create a model in Amazon SageMaker AI with ModelBuilder.

ModelBuilder is a Python class that takes a framework model or a user-specified inference specification and converts it to a deployable model. For more details about the ModelBuilder class, see ModelBuilder.

To deploy your MLflow model using ModelBuilder, provide a path to your MLflow artifacts in the model_metadata["MLFLOW_MODEL_PATH"] attribute. Read on for more information about valid model path input formats:

Model paths that require an ARN in the model_metadata

The following model paths do require that you specify an ARN in the model_metadata for deployment:

Model paths that do not require an ARN in the model_metadata

The following model paths do not require that you specify an ARN in the model_metadata for deployment:

For more information on how MLflow model deployment works with Amazon SageMaker AI, see Deploy MLflow Model to Amazon SageMaker AI in the MLflow documentation.

If using an Amazon S3 path, you can find the path of your registered model with the following commands:

registered_model = client.get_registered_model(name='AutoRegisteredModel')
source_path = registered_model.latest_versions[0].source

The following sample is an overview of how to deploy your MLflow model using ModelBuilder and an MLflow model registry path. Because this sample provides the model artifact path in the form of an MLflow model registry path, the call to ModelBuilder must also specify a tracking server ARN through the model_metadata["MLFLOW_TRACKING_ARN"] attribute.

from sagemaker.serve import ModelBuilder
from sagemaker.serve.mode.function_pointers import Mode
from sagemaker.serve import SchemaBuilder

my_schema = SchemaBuilder(
    sample_input=sample_input, 
    sample_output=sample_output
)

model_builder = ModelBuilder(
    mode=Mode.SAGEMAKER_ENDPOINT,
    schema_builder=my_schema,
    role_arn="Your-service-role-ARN",
    model_metadata={
        # both model path and tracking server ARN are required if you use an mlflow run ID or mlflow model registry path as input
        "MLFLOW_MODEL_PATH": "models:/sklearn-model/1"
        "MLFLOW_TRACKING_ARN": "arn:aws:sagemaker:region:account-id:mlflow-tracking-server/tracking-server-name"
    }
)
model = model_builder.build()
predictor = model.deploy( initial_instance_count=1, instance_type="ml.c6i.xlarge" )

To maintain lineage tracking for MLflow models deployed using ModelBuilder, you must have the following IAM permissions:

For more information and end-to-end examples, see MLflow tutorials using example Jupyter notebooks.