Repository: mlflow/mlflow

Issue:
[FR] add serverless inference endpoint support for sagemaker deployment
### Willingness to contribute

Yes. I can contribute this feature independently.

### Proposal Summary

Now, we can use the command and deploy model on sagemaker inference.
```
mlflow sagemaker deploy -m <model> -a liang-test-model --execution-role-arn <execution role>
```

Do we have any plan for supporting serverless deployment.

Reference [link](https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints-create.html)
And ref code.
```
response = client.create_endpoint_config(
   EndpointConfigName="<your-endpoint-configuration>",
   KmsKeyId="arn:aws:kms:us-east-1:123456789012:key/143ef68f-76fd-45e3-abba-ed28fc8d3d5e",
   ProductionVariants=[
        {
            "ModelName": "<your-model-name>",
            "VariantName": "AllTraffic",
            "ServerlessConfig": {
                "MemorySizeInMB": 2048,
                "MaxConcurrency": 20
            }
        } 
    ]
)
```

### Motivation

> #### What is the use case for this feature?
We can have near-real-time scenario for inference.
Due to the cost, we would like to use serverless endpoint on sagemaker.

> #### Why is this use case valuable to support for MLflow users in general?
Save cost. Support near-real-time inference.

> #### Why is this use case valuable to support for your project(s) or organization?
Save cost. Support near-real-time inference.

> #### Why is it currently difficult to achieve this use case?
Based on current mlflow code, we need to do two operations.
1. use "mlflow sagemaker deploy" command.
2. use boto3 sagemaker api "create_endpoint_config" and change the endpoint type from real-time to serverless.


### 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
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline 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
- [x] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
@QAM Thank you for the FR! Supporting serverless deployment sounds good to me. Are you interested in creating a PR for this?

@harupy sure. I can. I also copy code from mlflow.sagemaker.__init__.py and modify it locally. 
Now, it works.

I should be able to create the PR.
However i don't have permission to create branch on mlflow for PR.
Could I know any permission I need?
Thanks.
@QAM You can create a fork.

https://docs.github.com/en/get-started/quickstart/fork-a-repo
Will try. Thanks!
<!-- assign-maintainer -->
@BenWilson2 @dbczumar @harupy @WeichenXu123 Please assign a maintainer and start triaging this issue.
Are there any updates on that matter? I would highly appreciate this feature.

Resolving patch:
diff --git a/mlflow/sagemaker/__init__.py b/mlflow/sagemaker/__init__.py
--- a/mlflow/sagemaker/__init__.py
+++ b/mlflow/sagemaker/__init__.py
@@ -177,6 +177,7 @@ def _deploy(
     data_capture_config=None,
     variant_name=None,
     async_inference_config=None,
+    serverless_config=None,
     env=None,
     tags=None,
 ):
@@ -327,7 +328,13 @@ def _deploy(
                                                     "NotificationConfig": {},  # pylint: disable=line-too-long
                                                 },
                                             }
-
+    :param serverless_config: An optional dictionary specifying the serverless configuration
+                                    .. code-block:: python
+                                        :caption: Example
+                                            "ServerlessConfig": {
+                                                "MemorySizeInMB": 2048,
+                                                "MaxConcurrency": 20,
+                                            }
     :param env: An optional dictionary of environment variables to set for the model.
     :param tags: An optional dictionary of tags to apply to the endpoint.
     """
@@ -421,6 +428,7 @@ def _deploy(
             s3_client=s3_client,
             variant_name=variant_name,
             async_inference_config=async_inference_config,
+            serverless_config=serverless_config,
             data_capture_config=data_capture_config,
             env=env,
             tags=tags,
@@ -441,6 +449,7 @@ def _deploy(
             sage_client=sage_client,
             variant_name=variant_name,
             async_inference_config=async_inference_config,
+            serverless_config=serverless_config,
             env=env,
             tags=tags,
         )
@@ -1492,6 +1501,7 @@ def _create_sagemaker_endpoint(
     sage_client,
     variant_name=None,
     async_inference_config=None,
+    serverless_config=None,
     env=None,
     tags=None,
 ):
@@ -1537,10 +1547,15 @@ def _create_sagemaker_endpoint(
     production_variant = {
         "VariantName": variant_name,
         "ModelName": model_name,
-        "InitialInstanceCount": instance_count,
-        "InstanceType": instance_type,
         "InitialVariantWeight": 1,
     }
+
+    if serverless_config:
+        production_variant["ServerlessConfig"] = serverless_config
+    else:
+        production_variant["InstanceType"] = instance_type
+        production_variant["InitialInstanceCount"] = instance_count
+
     config_name = _get_sagemaker_config_name(endpoint_name)
     endpoint_config_kwargs = {
         "EndpointConfigName": config_name,
@@ -1609,6 +1624,7 @@ def _update_sagemaker_endpoint(
     s3_client,
     variant_name=None,
     async_inference_config=None,
+    serverless_config=None,
     data_capture_config=None,
     env=None,
     tags=None,
@@ -1682,10 +1698,15 @@ def _update_sagemaker_endpoint(
     new_production_variant = {
         "VariantName": variant_name,
         "ModelName": model_name,
-        "InitialInstanceCount": instance_count,
-        "InstanceType": instance_type,
         "InitialVariantWeight": new_model_weight,
     }
+
+    if serverless_config:
+        new_production_variant["ServerlessConfig"] = serverless_config
+    else:
+        new_production_variant["InstanceType"] = instance_type
+        new_production_variant["InitialInstanceCount"] = instance_count
+
     production_variants.append(new_production_variant)
 
     # Create the new endpoint configuration and update the endpoint
@@ -1992,6 +2013,7 @@ def _default_deployment_config(self, create_mode=True):
             "env": None,
             "tags": None,
             "async_inference_config": {},
+            "serverless_config": {},
         }
 
         if create_mode:
@@ -2004,7 +2026,14 @@ def _default_deployment_config(self, create_mode=True):
     def _apply_custom_config(self, config, custom_config):
         int_fields = {"instance_count", "timeout_seconds"}
         bool_fields = {"synchronous", "archive"}
-        dict_fields = {"vpc_config", "data_capture_config", "tags", "env", "async_inference_config"}
+        dict_fields = {
+            "vpc_config",
+            "data_capture_config",
+            "tags",
+            "env",
+            "async_inference_config",
+            "serverless_config",
+        }
         for key, value in custom_config.items():
             if key not in config:
                 continue
@@ -2131,6 +2160,8 @@ def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=
                                            production variant.  Defaults to ``None``.
                        - ``async_inference_config``: A dictionary specifying the async_inference_configuration # pylint: disable=line-too-long
 
+                       - ``serverless_config``: A dictionary specifying the serverless_configuration
+
                        - ``env``: A dictionary specifying environment variables as key-value
                          pairs to be set for the deployed model. Defaults to ``None``.
 
@@ -2223,6 +2254,7 @@ def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=
             timeout_seconds=final_config["timeout_seconds"],
             variant_name=final_config["variant_name"],
             async_inference_config=final_config["async_inference_config"],
+            serverless_config=final_config["serverless_config"],
             env=final_config["env"],
             tags=final_config["tags"],
         )
@@ -2483,6 +2515,7 @@ def update_deployment(
             timeout_seconds=final_config["timeout_seconds"],
             variant_name=final_config["variant_name"],
             async_inference_config=final_config["async_inference_config"],
+            serverless_config=final_config["serverless_config"],
             env=final_config["env"],
             tags=final_config["tags"],
         )


Test patch:
diff --git a/tests/sagemaker/test_sagemaker_deployment_client.py b/tests/sagemaker/test_sagemaker_deployment_client.py
--- a/tests/sagemaker/test_sagemaker_deployment_client.py
+++ b/tests/sagemaker/test_sagemaker_deployment_client.py
@@ -364,6 +364,67 @@ def test_update_deployment_without_async_config(
     assert "AsyncInferenceConfig" not in target_config
 
 
+@mock_sagemaker_aws_services
+def test_create_deployment_with_serverless_config(
+    pretrained_model, sagemaker_client, sagemaker_deployment_client
+):
+    app_name = "deploy_with_serverless_config"
+    expected_serverless_config = {
+        "MemorySizeInMB": 2048,
+        "MaxConcurrency": 2,
+    }
+    sagemaker_deployment_client.create_deployment(
+        name=app_name,
+        model_uri=pretrained_model.model_uri,
+        config={"serverless_config": expected_serverless_config},
+    )
+    configs = sagemaker_client.list_endpoint_configs()
+    target_config = None
+    for config in configs["EndpointConfigs"]:
+        if app_name in config["EndpointConfigName"]:
+            target_config = config
+    if target_config is None:
+        raise Exception("Endpoint config not found")
+    endpoint_config = sagemaker_client.describe_endpoint_config(
+        EndpointConfigName=target_config["EndpointConfigName"]
+    )
+
+    for variant in endpoint_config["ProductionVariants"]:
+        assert variant["ServerlessConfig"] == expected_serverless_config
+
+
+@mock_sagemaker_aws_services
+def test_update_deployment_with_serverless_config_when_endpoint_exists(
+    pretrained_model, sagemaker_client, sagemaker_deployment_client
+):
+    app_name = "update_deploy_with_serverless_config"
+    expected_serverless_config = {
+        "MemorySizeInMB": 2048,
+        "MaxConcurrency": 2,
+    }
+    sagemaker_deployment_client.create_deployment(
+        name=app_name, model_uri=pretrained_model.model_uri
+    )
+    sagemaker_deployment_client.update_deployment(
+        name=app_name,
+        model_uri=pretrained_model.model_uri,
+        config={"serverless_config": expected_serverless_config},
+    )
+    configs = sagemaker_client.list_endpoint_configs()
+    target_config = None
+    for config in configs["EndpointConfigs"]:
+        if app_name in config["EndpointConfigName"]:
+            target_config = config
+    if target_config is None:
+        raise Exception("Endpoint config not found")
+    endpoint_config = sagemaker_client.describe_endpoint_config(
+        EndpointConfigName=target_config["EndpointConfigName"]
+    )
+
+    for variant in endpoint_config["ProductionVariants"]:
+        assert variant["ServerlessConfig"] == expected_serverless_config
+
+
 def test_create_deployment_with_unsupported_flavor_raises_exception(
     pretrained_model, sagemaker_deployment_client
 ):