Repository: mlflow/mlflow

Issue:
[FR] Currently unable to specify variant names when deploying models to Sagemaker
### Willingness to contribute

Yes. I would be willing to contribute this feature with guidance from the MLflow community.

### Proposal Summary

I would like to add the ability to pass in a specific name for the model variant when deploying to Sagmaker as an additional argument to the deploy method in sagemaker/__init__.py.  Currently, variant names are autogenerated by _get_sagemaker_model_name using the endpoint name, model name and uid.  

### Motivation

> #### What is the use case for this feature?
To supply a specific name to each model variant hosted on Sagemaker.

> #### Why is this use case valuable to support for MLflow users in general?
The name would be more meaningful than the current default.

> #### Why is this use case valuable to support for your project(s) or organization?
This would be especially important for endpoints that host multiple models or different version of the same model (ex: for testing purposes).  It would make it much easier for the user to better understand the purpose of a given model based on its name and distinguish between models.  

> #### Why is it currently difficult to achieve this use case?
Currently, there is not way to specify the variant name when deploying, they are just generated automatically


### Details

I propose adding an argument to the deploy method where a specific name can be entered.  This will throw an error if that name currently exists or if it violates any naming conventions.  Use the current default method if no name is specified when calling this method.

### 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
- [X] `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:
To make the model name unique, shall we define this additional argument as a suffix  ? so the final model_name will be {generated_model_name}-{suffix}
Adding a suffix would allow for some model name customization.  Is the {generated_model_name} you mention in your comment above the full name (endpoint_name-model-uid) that is currently being generated by _get_sagemaker_model_name or do you propose substituting the suffix for the uid?   Is there a particular reason why you would not want to allow the names to be fully customizable?
<!-- assign-maintainer -->
@BenWilson2 @dbczumar @harupy @WeichenXu123 Please assign a maintainer and start triaging this issue.
@nfarley-soaren Your proposal sounds great. If you're still interested in contributing this feature, we'd be happy to review a pull request that extends the `mlflow.sagemaker` module with the custom model naming capability. Please let me know if you have any questions!

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(
     synchronous=True,
     timeout_seconds=1200,
     data_capture_config=None,
+    variant_name=None,
 ):
     """
     Deploy an MLflow model on AWS SageMaker.
@@ -261,19 +262,19 @@ def deploy(
                        #SageMaker.Client.create_model>`_. For more information, see
                        https://docs.aws.amazon.com/sagemaker/latest/dg/API_VpcConfig.html.
 
-    .. code-block:: python
-        :caption: Example
+                       .. code-block:: python
+                           :caption: Example
 
-        import mlflow.sagemaker as mfs
-        vpc_config = {
-                        'SecurityGroupIds': [
-                            'sg-123456abc',
-                        ],
-                        'Subnets': [
-                            'subnet-123456abc',
-                        ]
-                     }
-        mfs.deploy(..., vpc_config=vpc_config)
+                           import mlflow.sagemaker as mfs
+                           vpc_config = {
+                               'SecurityGroupIds': [
+                                   'sg-123456abc',
+                               ],
+                               'Subnets': [
+                                   'subnet-123456abc',
+                               ]
+                            }
+                            mfs.deploy(..., vpc_config=vpc_config)
 
     :param flavor: The name of the flavor of the model to use for deployment. Must be either
                    ``None`` or one of mlflow.sagemaker.SUPPORTED_DEPLOYMENT_FLAVORS. If ``None``,
@@ -293,24 +294,25 @@ def deploy(
                             deployment using native SageMaker APIs or the AWS console. If
                             ``synchronous`` is ``False``, this parameter is ignored.
     :param data_capture_config: A dictionary specifying the data capture configuration to use when
-                       creating the new SageMaker model associated with this application. For more
-                       information, see
-                       https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DataCaptureConfig.html.
-
-    .. code-block:: python
-        :caption: Example
-
-        import mlflow.sagemaker as mfs
-        data_capture_config = {
-                        'EnableCapture': True,
-                        'InitalSamplingPercentage': 100,
-                        'DestinationS3Uri": 's3://my-bucket/path',
-                        'CaptureOptions': [
-                            {'CaptureMode': 'Output'}
-                        ],
-                     }
-        mfs.deploy(..., data_capture_config=data_capture_config)
-
+                                creating the new SageMaker model associated with this application.
+                                For more information, see
+                                https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DataCaptureConfig.html.
+
+                                .. code-block:: python
+                                    :caption: Example
+
+                                    import mlflow.sagemaker as mfs
+                                    data_capture_config = {
+                                                    'EnableCapture': True,
+                                                    'InitalSamplingPercentage': 100,
+                                                    'DestinationS3Uri": 's3://my-bucket/path',
+                                                    'CaptureOptions': [
+                                                        {'CaptureMode': 'Output'}
+                                                    ],
+                                                 }
+                                    mfs.deploy(..., data_capture_config=data_capture_config)
+
+    :param variant_name: The name to assign to the new production variant.
     """
     import boto3
 
@@ -403,6 +405,7 @@ def deploy(
             role=execution_role_arn,
             sage_client=sage_client,
             s3_client=s3_client,
+            variant_name=variant_name,
         )
     else:
         deployment_operation = _create_sagemaker_endpoint(
@@ -418,6 +421,7 @@ def deploy(
             data_capture_config=data_capture_config,
             role=execution_role_arn,
             sage_client=sage_client,
+            variant_name=variant_name,
         )
 
     if synchronous:
@@ -633,19 +637,19 @@ def deploy_transform_job(
                        #SageMaker.Client.create_model>`_. For more information, see
                        https://docs.aws.amazon.com/sagemaker/latest/dg/API_VpcConfig.html.
 
-    .. code-block:: python
-        :caption: Example
+                       .. code-block:: python
+                           :caption: Example
 
-        import mlflow.sagemaker as mfs
-        vpc_config = {
-                        'SecurityGroupIds': [
-                            'sg-123456abc',
-                        ],
-                        'Subnets': [
-                            'subnet-123456abc',
-                        ]
-                     }
-        mfs.deploy_transform_job(..., vpc_config=vpc_config)
+                           import mlflow.sagemaker as mfs
+                           vpc_config = {
+                               'SecurityGroupIds': [
+                                   'sg-123456abc',
+                               ],
+                               'Subnets': [
+                                   'subnet-123456abc',
+                               ]
+                            }
+                            mfs.deploy_transform_job(..., vpc_config=vpc_config)
 
     :param flavor: The name of the flavor of the model to use for deployment. Must be either
                    ``None`` or one of mlflow.sagemaker.SUPPORTED_DEPLOYMENT_FLAVORS. If ``None``,
@@ -935,19 +939,19 @@ def push_model_to_sagemaker(
                        #SageMaker.Client.create_model>`_. For more information, see
                        https://docs.aws.amazon.com/sagemaker/latest/dg/API_VpcConfig.html.
 
-    .. code-block:: python
-        :caption: Example
+                       .. code-block:: python
+                           :caption: Example
 
-        import mlflow.sagemaker as mfs
-        vpc_config = {
-                        'SecurityGroupIds': [
-                            'sg-123456abc',
-                        ],
-                        'Subnets': [
-                            'subnet-123456abc',
-                        ]
-                     }
-        mfs.push_model_to_sagemaker(..., vpc_config=vpc_config)
+                           import mlflow.sagemaker as mfs
+                           vpc_config = {
+                               'SecurityGroupIds': [
+                                   'sg-123456abc',
+                               ],
+                               'Subnets': [
+                                   'subnet-123456abc',
+                               ]
+                            }
+                            mfs.push_model_to_sagemaker(..., vpc_config=vpc_config)
 
     :param flavor: The name of the flavor of the model to use for deployment. Must be either
                    ``None`` or one of mlflow.sagemaker.SUPPORTED_DEPLOYMENT_FLAVORS. If ``None``,
@@ -1077,9 +1081,6 @@ def _sigterm_handler(*_):
 def target_help():
     """
     Provide help information for the SageMaker deployment client.
-
-    :return:
-    :rtype: str
     """
     help_str = """\
     For detailed documentation on the SageMaker deployment client, please visit
@@ -1439,6 +1440,7 @@ def _create_sagemaker_endpoint(
     instance_count,
     role,
     sage_client,
+    variant_name=None,
 ):
     """
     :param endpoint_name: The name of the SageMaker endpoint to create.
@@ -1456,6 +1458,7 @@ def _create_sagemaker_endpoint(
                        creating the new SageMaker model associated with this application.
     :param role: SageMaker execution ARN role.
     :param sage_client: A boto3 client for SageMaker.
+    :param variant_name: The name to assign to the new production variant.
     """
     _logger.info("Creating new endpoint with name: %s ...", endpoint_name)
 
@@ -1471,8 +1474,11 @@ def _create_sagemaker_endpoint(
     )
     _logger.info("Created model with arn: %s", model_response["ModelArn"])
 
+    if not variant_name:
+        variant_name = model_name
+
     production_variant = {
-        "VariantName": model_name,
+        "VariantName": variant_name,
         "ModelName": model_name,
         "InitialInstanceCount": instance_count,
         "InstanceType": instance_type,
@@ -1545,6 +1551,7 @@ def _update_sagemaker_endpoint(
     role,
     sage_client,
     s3_client,
+    variant_name=None,
 ):
     """
     :param endpoint_name: The name of the SageMaker endpoint to update.
@@ -1563,6 +1570,7 @@ def _update_sagemaker_endpoint(
     :param role: SageMaker execution ARN role.
     :param sage_client: A boto3 client for SageMaker.
     :param s3_client: A boto3 client for S3.
+    :variant_name: The name to assign to the new production variant if it doesn't already exist.
     """
     if mode not in [DEPLOYMENT_MODE_ADD, DEPLOYMENT_MODE_REPLACE]:
         msg = "Invalid mode `{md}` for deployment to a pre-existing application".format(md=mode)
@@ -1591,6 +1599,9 @@ def _update_sagemaker_endpoint(
     )
     _logger.info("Created new model with arn: %s", new_model_response["ModelArn"])
 
+    if not variant_name:
+        variant_name = model_name
+
     if mode == DEPLOYMENT_MODE_ADD:
         new_model_weight = 0
         production_variants = deployed_production_variants
@@ -1599,7 +1610,7 @@ def _update_sagemaker_endpoint(
         production_variants = []
 
     new_production_variant = {
-        "VariantName": model_name,
+        "VariantName": variant_name,
         "ModelName": model_name,
         "InitialInstanceCount": instance_count,
         "InstanceType": instance_type,
@@ -1891,6 +1902,7 @@ def _default_deployment_config(self, create_mode=True):
             data_capture_config=None,
             synchronous=True,
             timeout_seconds=1200,
+            variant_name=None,
         )
 
         if create_mode:
@@ -2051,6 +2063,10 @@ def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=
                          For more information, see
                          https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DataCaptureConfig.html.
                          Defaults to ``None``.
+
+                       - ``variant_name``: A string specifying the desired name when creating a
+                                           production variant.  Defaults to ``None``.
+
         :param endpoint: (optional) Endpoint to create the deployment under. Currently unsupported
 
         .. code-block:: python
@@ -2078,7 +2094,8 @@ def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=
                 instance_count=1,
                 synchronous=True,
                 timeout_seconds=300,
-                vpc_config=vpc_config
+                vpc_config=vpc_config,
+                variant_name="prod-variant-1",
             )
             client = SageMakerDeploymentClient("sagemaker")
             client.create_deployment(
@@ -2104,6 +2121,7 @@ def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=
                     -C instance_count=1 \\
                     -C synchronous=True \\
                     -C timeout_seconds=300 \\
+                    -C variant_name=prod-variant-1 \\
                     -C vpc_config='{"SecurityGroupIds": ["sg-123456abc"], \\
                     "Subnets": ["subnet-123456abc"]}' \\
                     -C data_capture_config='{"EnableCapture": True, \\
@@ -2131,6 +2149,7 @@ def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=
             data_capture_config=final_config["data_capture_config"],
             synchronous=final_config["synchronous"],
             timeout_seconds=final_config["timeout_seconds"],
+            variant_name=final_config["variant_name"],
         )
 
         return dict(name=app_name, flavor=flavor)
@@ -2248,6 +2267,9 @@ def update_deployment(self, name, model_uri=None, flavor=None, config=None, endp
                          ``synchronous`` is ``False``, this parameter is ignored.
                          Defaults to ``300``.
 
+                       - ``variant_name``: A string specifying the desired name when creating a
+                                           production variant.  Defaults to ``None``.
+
                        - ``vpc_config``: A dictionary specifying the VPC configuration to use when
                          creating the new SageMaker model associated with this application.
                          The acceptable values for this parameter are identical to those of the
@@ -2263,6 +2285,10 @@ def update_deployment(self, name, model_uri=None, flavor=None, config=None, endp
                          For more information, see
                          https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DataCaptureConfig.html.
                          Defaults to ``None``.
+
+                       - ``variant_name``: A string specifying the desired name when creating a
+                                           production variant.  Defaults to ``None``.
+
         :param endpoint: (optional) Endpoint containing the deployment to update. Currently
                          unsupported
 
@@ -2299,6 +2325,7 @@ def update_deployment(self, name, model_uri=None, flavor=None, config=None, endp
                 instance_count=1,
                 synchronous=True,
                 timeout_seconds=300,
+                variant_name="prod-variant-1",
                 vpc_config=vpc_config
                 data_capture_config=data_capture_config
             )
@@ -2326,6 +2353,7 @@ def update_deployment(self, name, model_uri=None, flavor=None, config=None, endp
                     -C instance_count=1 \\
                     -C synchronous=True \\
                     -C timeout_seconds=300 \\
+                    -C variant_name=prod-variant-1 \\
                     -C vpc_config='{"SecurityGroupIds": ["sg-123456abc"], \\
                     "Subnets": ["subnet-123456abc"]}' \\
                     -C data_capture_config='{"EnableCapture": True, \\
@@ -2366,6 +2394,7 @@ def update_deployment(self, name, model_uri=None, flavor=None, config=None, endp
             data_capture_config=final_config["data_capture_config"],
             synchronous=final_config["synchronous"],
             timeout_seconds=final_config["timeout_seconds"],
+            variant_name=final_config["variant_name"],
         )
 
         return dict(name=app_name, flavor=flavor)


Test patch: