[FR] Adding SageMaker provider to MLflow Gateway AI

[pdifranc]

Willingness to contribute

Yes. I can contribute this feature independently.

Proposal Summary

SageMaker JumpStart provides an easy way to deploy LLM endpoints on the SageMaker managed infrastructure. However, as of today, the MLflow Gateway AI does not implement a SageMaker provider.

We would like to introduce the SageMaker provider and adapter together with an example script to show how you can add a SageMaker hosted model and make it available to the MLflow Gateway AI

Motivation

What is the use case for this feature?

SageMaker jumpstarts supports many LLMs with a one click deployment. Would be nice to have a way to interface with these models in a centralized way as offered by mlflow Gateway AI

Why is this use case valuable to support for MLflow users in general?

Sagemaker is a popular tool, and some models are readily available directly there. Would open up to a great deal of possibilities

Why is this use case valuable to support for your project(s) or organization?

I want to have access to more providers

Why is it currently difficult to achieve this use case?

The provider/adapters should not be too hard, of course the evil is in the details.

Details

I would like to provide a basic SageMaker provider, which would handle the credentials and provide the basic AWS specific credentials mechanisms (similar as done already for Bedrock).
I would then provide two Adapter for two models (thinking on providing samples for the two LLama2 models, likely the 7B one).
This can be a baseline for more models deployable via sagemaker that can be extended. Users might be able to quickly integrate new adapters, requiring some sort of plugin feature, but this might be for the future.

What component(s) does this bug affect?

• 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/gateway: AI Gateway service, Gateway client APIs, third-party Gateway integrations
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
• 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

What interface(s) does this bug affect?

• 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

What language(s) does this bug affect?

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

What integration(s) does this bug affect?

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

[kirit93]

+1 for this feature! It would be really helpful to have this.

[github-actions]

@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.

[pdifranc]

maybe rather than adding a new model
provider, give the possibility to add new providers via plugins (currently not possible via regular mlflow plugins).