[ MLflow Roadmap - FR - Design Required ] Auth Support for MLflow AI Gateway

[BenWilson2]

Willingness to contribute

No. I cannot contribute this feature at this time.

Proposal Summary

A common ask that has come up for adopters of the MLflow AI Gateway is adding integrated support for authorized access to the server.

We'd like to help a contributor work on designing and implementing a solution for a plugin-based system for auth support. As part of this implementation effort, the initial solution should include support for basic auth (username & password), configured by the server admin.

If you are interested in contributing to this, please respond to this issue and we can walk you through the process of submitting a design, refining it, and helping you through the process of implementation.

Motivation

What is the use case for this feature?

To provide authorized user support to the MLflow AI Gateway, providing an integrated means for setting route access permissions.

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

Currently, users are forced to implement their own access controls when using the MLflow AI Gateway.

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

For cost control and to lessen the burden on an MLOps team that is managing an AI Gateway deployment.

Why is it currently difficult to achieve this use case?

Currently, a 3rd party ACL solution needs to be used to control access to the running Gunicorn server, which puts an undue burden on the MLOps team managing such a deployment.

Details

No response

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
• [X] 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

[danilopeixoto]

@BenWilson2 Could we consider defining MLFLOW_GATEWAY_TOKEN and inject it to gateway proxy handler used by Prompt Engineering UI?

[github-actions[bot]]

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

[TomeHirata]

Hi, @BenWilson2. Would you mind my contributing to this ticket? cc: @harupy

[BenWilson2]

@TomeHirata it's all yours! Created an entry for the 2.12 release in the project display here: https://github.com/orgs/mlflow/projects/4/views/1

[TomeHirata]

Hi, @BenWilson2. Let me ask some questions/asks:

1. Do I need to write up the proper design document for this task? If so, can I know the design doc format you are using for mlflow?
2. Do you want to make it a plugin instead of incorporating the functionality in the configuration yaml of mlflow deployment? It will be like the yaml below and I think this direction is worth exploring

endpoints:
  - name: completions
    endpoint_type: llm/v1/completions
    model:
      provider: openai
      name: gpt-3.5-turbo
      config:
        openai_api_key: $OPENAI_API_KEY
    limit:
      renewal_period: minute
      calls: 10
auth:
  basic:
    username: tomehirata
    password: mypassword

3. Would you mind putting my name in the 2.12 release in the project display here?

[TomeHirata]

Hi, @BenWilson2 @harupy @B-Step62. Sorry for the broad mention here, could any of you take a look at my previous comment? Thank you.

[harupy]

@TomeHirata Sorry for the delay.

1. Yes, I'll send you the template.
2. What if you have multiple users? The YAML file would look like this?

auth:
  basic:
    - username: tomehirata
      password: mypassword
    - username: haru
      password: haru

3. Sorry, we can only assign collaborators or mlflow org members.

[TomeHirata]

What if you have multiple users? The YAML file would look like this?

Yes, if we want to support multiple user cases, then the yaml will look like it.

[harupy]

I don't think we should expose credentials like that. It's inconvenient to edit the config file every time a new user needs to be added.

[harupy]

@TomeHirata Sorry for being late to send this, but this is what our design doc template looks like: Design doc template (2).pdf

[TomeHirata]

I don't think we should expose credentials like that.

For the actual production usage, we can allow users to use templating and read the credentials from the env variables as the Spring property does.

auth:
  basic:
    - username: ${USER_NAME_TOMU}
      password: ${USER_PASS_TOMU}
    - username: ${USER_NAME_HARU}
      password: ${USER_PASS_HARU}

It's inconvenient to edit the config file every time a new user needs to be added.

That's a good point. I didn't expect the arbitrary number of users would be added to the Basic Authentication user list. In this case, it would be better to persist the users in a DB and provide an API for the server admin to add a new user. Since there is a similar feature for tracking, can I ask if we want to share the server admin and user list between "tracking" and "deployment"?

[harupy]

@TomeHirata Sorry for the late reply.

can I ask if we want to share the server admin and user list between "tracking" and "deployment"?

Good question. Yes, sharing admin and user sounds good to me.