[BUG] Misconfiguration of bedrock provider synonym in gateway

[shabie]

Issues Policy acknowledgement

• [X] I have read and agree to submit bug reports in accordance with the issues policy

Where did you encounter this bug?

Local machine

Willingness to contribute

Yes. I can contribute a fix for this bug independently.

MLflow version

• Client: 2.15.0

System information

• OS: MacOS sonoma 14.5
• Python version: 3.10.14

Describe the problem

When trying to start the deployment server locally using:

mlflow deployments start-server --port 7000 --config-path config.yaml

The config.yaml comes from here:

https://github.com/mlflow/mlflow/blob/d1e442fdd00456d749dddbfbc8d3055a084869de/examples/deployments/deployments_server/bedrock/config.yaml#L1-L11

This throws an error:

File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/mlflow/gateway/providers/__init__.py", line 37, in get_provider
raise MlflowException.invalid_parameter_value(f"Provider {provider} not found")
mlflow.exceptions.MlflowException: Provider amazon-bedrock not found

This is because the providers that map to classes are missing the provider amazon-bedrock is missing:

https://github.com/mlflow/mlflow/blob/d1e442fdd00456d749dddbfbc8d3055a084869de/mlflow/gateway/providers/init.py#L21-L33

Even though it has been defined here as a synonym to bedrock:

https://github.com/mlflow/mlflow/blob/d1e442fdd00456d749dddbfbc8d3055a084869de/mlflow/gateway/config.py#L36-L46

We can see that AMAZON_BEDROCK although defined in the Enum is never mapped to the AmazonBedrockProvider class which raises the error.

I am happy to raise a small PR to fix this issue.

Tracking information

No response

Code to reproduce issue

Create a new config.yaml

endpoints:
  - name: completions
    endpoint_type: llm/v1/completions
    model:
      provider: amazon-bedrock
      name: amazon.titan-tg1-large
      config:
        aws_config:
          aws_region: us-east-1
          aws_access_key_id: $AWS_ACCESS_KEY_ID
          aws_secret_access_key: $AWS_SECRET_ACCESS_KEY

and run mlflow deployments start-server --port 7000 --config-path config.yaml

Stack trace

Traceback (most recent call last):
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/gunicorn/arbiter.py", line 609, in spawn_worker
    worker.init_process()
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/uvicorn/workers.py", line 66, in init_process
    super(UvicornWorker, self).init_process()
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/gunicorn/workers/base.py", line 134, in init_process
    self.load_wsgi()
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/gunicorn/workers/base.py", line 146, in load_wsgi
    self.wsgi = self.app.wsgi()
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/gunicorn/app/base.py", line 67, in wsgi
    self.callable = self.load()
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/gunicorn/app/wsgiapp.py", line 58, in load
    return self.load_wsgiapp()
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/gunicorn/app/wsgiapp.py", line 48, in load_wsgiapp
    return util.import_app(self.app_uri)
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/gunicorn/util.py", line 424, in import_app
    app = app(*args, **kwargs)
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/mlflow/deployments/server/app.py", line 411, in create_app_from_env
    return create_app_from_path(config_path)
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/mlflow/deployments/server/app.py", line 403, in create_app_from_path
    return create_app_from_config(config)
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/mlflow/deployments/server/app.py", line 233, in create_app_from_config
    app = GatewayAPI(
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/mlflow/deployments/server/app.py", line 55, in __init__
    self.set_dynamic_routes(config, limiter)
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/mlflow/deployments/server/app.py", line 64, in set_dynamic_routes
    endpoint=_route_type_to_endpoint(route, limiter, "deployments"),
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/mlflow/deployments/server/app.py", line 131, in _route_type_to_endpoint
    handler = factory(config)
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/mlflow/deployments/server/app.py", line 96, in _create_completions_endpoint
    prov = get_provider(config.model.provider)(config)
  File "/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/mlflow/gateway/providers/__init__.py", line 37, in get_provider
    raise MlflowException.invalid_parameter_value(f"Provider {provider} not found")
mlflow.exceptions.MlflowException: Provider amazon-bedrock not found

Other info / logs

The first half of the logs:

> mlflow deployments start-server --port 7000 --config-path config.yaml
/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/pydantic/_internal/_config.py:334: UserWarning: Valid config keys have changed in V2:
* 'schema_extra' has been renamed to 'json_schema_extra'
  warnings.warn(message, UserWarning)
[2024-08-02 14:42:21 +0200] [99448] [INFO] Starting gunicorn 22.0.0
[2024-08-02 14:42:21 +0200] [99448] [INFO] Listening at: http://127.0.0.1:7000 (99448)
[2024-08-02 14:42:21 +0200] [99448] [INFO] Using worker: uvicorn.workers.UvicornWorker
[2024-08-02 14:42:21 +0200] [99449] [INFO] Booting worker with pid: 99449
[2024-08-02 14:42:21 +0200] [99450] [INFO] Booting worker with pid: 99450
/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/pydantic/_internal/_config.py:334: UserWarning: Valid config keys have changed in V2:
* 'schema_extra' has been renamed to 'json_schema_extra'
  warnings.warn(message, UserWarning)
/Users/johndoe/projects/bedrocks/venv/lib/python3.10/site-packages/pydantic/_internal/_config.py:334: UserWarning: Valid config keys have changed in V2:
* 'schema_extra' has been renamed to 'json_schema_extra'
  warnings.warn(message, UserWarning)
[2024-08-02 14:42:22 +0200] [99449] [ERROR] Exception in worker process
Traceback (most recent call last):

What component(s) does this bug affect?

• [ ] area/artifacts: Artifact stores and artifact logging
• [ ] area/build: Build and test infrastructure for MLflow
• [ ] area/deployments: MLflow Deployments client APIs, server, and third-party Deployments integrations
• [ ] 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/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