[FR] Migrate MLflow SageMaker deployment functionality to deployments interface

MLflow Roadmap Item

This is an MLflow Roadmap item that has been prioritized by the MLflow maintainers. We’re seeking help with the implementation of roadmap items tagged with the help wanted label.

For requirements clarifications and implementation questions, or to request a PR review, please tag @dbczumar in your communications related to this issue.

Proposal Summary

MLflow provides a deployments API / CLI for CRUDing real-time scoring deployments from MLflow models. This is a pluggable / extendable interface, and plugins currently exist for deployment to Azure ML, TorchServe, Redis, and more. MLflow’s tool for deploying models to SageMaker isn’t currently exposed through the deployments API / CLI; rather, it has a separate mlflow sagemaker CLI / mlflow.sagemaker API for legacy reasons. For consistency, we propose to move SageMaker deployment functionality into the mlflow deployments CLI / API.

Motivation

• What is the use case for this feature? This feature unifies MLflow’s real-time scoring deployment integrations behind a single API.
• Why is this use case valuable to support for MLflow users in general? It’s difficult for users to reason about separate APIs for different deployment environments. A unified API will reduce friction.
• Why is this use case valuable to support for your project(s) or organization? Databricks customers deploy models to a variety of locations. It’s difficult for them to reason about separate APIs for SageMaker deployment.
• Why is it currently difficult to achieve this use case? There is no SageMaker deployment functionality conforming to the MLflow plugin specification.

What component(s), interfaces, languages, and integrations does this feature affect?

Components

• area/artifacts: Artifact stores and artifact logging
• area/build: Build and test infrastructure for MLflow
• area/docs: MLflow documentation pages
• area/examples: Example code
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/projects: MLproject format, project running backends
• area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• area/server-infra: MLflow Tracking server backend
• area/tracking: Tracking Service, tracking client APIs, autologging

Interfaces

• area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• area/docker: Docker use across MLflow’s components, such as MLflow Projects and MLflow Models
• area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• area/windows: Windows support

Languages

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

Integrations

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

Details

We propose to migrate the mlflow.sagemaker module to a plugin specification that still lives within the main MLflow repository and implements the deployment plugin interface: https://github.com/mlflow/mlflow/blob/master/mlflow/deployments/base.py.

[dbczumar]

Closing this out now that https://github.com/mlflow/mlflow/pull/4971 is merged. Looking forward to following on with an implementation for predict()! Thanks @jamestran201 !

[jamestran201]

@dbczumar No worries at all! That PR is quite long 😅. Both of your suggestions look very interesting. I’ll try them out, thanks!