Repository: mlflow/mlflow
Pull request: Set production variant name
<!-- 🚨 We recommend pull requests be filed from a non-master branch on a repository fork (e.g. <username>:fix-xxx). 🚨 -->

## Related Issues/PRs 

<!--
Please reference any related feature requests, issues, or PRs here. For example, `#123`. To automatically close the referenced items when this PR is merged, please use a closing keyword (close, fix, or resolve). For example, `Close #123`. See https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue for more information.
-->
Close #6178
## What changes are proposed in this pull request?

The option to manually specify a production variant name rather than using the default model name when deploying a model to an endpoint

## How is this patch tested?

<!--
If you're unsure about what to test, where to add tests, or how to run tests, please feel free to ask. We'd be happy to help.
-->
- [x] I have written tests (not required for typo or doc fix) and confirmed the proposed feature/bug-fix/change works.
- I tested variant name creation, as well as the ability to error if using improper names (won't match sagemaker regex), or trying to create a variant with the same name as an existing variant.  The code used for testing requires access to a custom docker image rather than the default image created by mlflow so it is not attached.

<!--
Please describe how you confirmed the proposed feature/bug-fix/change works here. For example, if you fixed an MLflow client API, you could attach the code that didn't work prior to the fix but works now, or if you added a new feature on MLflow UI, you could attach a video that demonstrates the feature.
--> I tested variant name creation, as well as the ability to error if using improper names (won't match sagemaker regex), or trying to create a variant with the same name as an existing variant.  The code used for testing requires access to a custom docker image rather than the default image created by mlflow so it is not attached.

## Does this PR change the documentation?

- [ ] No. You can skip the rest of this section.
- [x ] Yes. Make sure the changed pages / sections render correctly by following the steps below.
If I knew where this was, I'd indicate the pages.

1. Click the `Details` link on the `Preview docs` check.
2. Find the changed pages / sections and make sure they render correctly.

## Release Notes

### Is this a user-facing change?

- [ ] No. You can skip the rest of this section.
- [ x] Yes. Give a description of this change to be included in the release notes for MLflow users.

(Details in 1-2 sentences. You can just refer to another PR with a description if this PR is part of a larger change.)
Allows user to specify a name for a given production variant when deploying mlflow model to Sagemaker endpoint.  If no name is given, the default is used.

### What component(s), interfaces, languages, and integrations does this PR affect?
Components
- [ ] `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

Interface
- [ ] `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

Language
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

<!--
Insert an empty named anchor here to allow jumping to this section with a fragment URL
(e.g. https://github.com/mlflow/mlflow/pull/123#user-content-release-note-category).
Note that GitHub prefixes anchor names in markdown with "user-content-".
-->
<a name="release-note-category"></a>
### How should the PR be classified in the release notes? Choose one:

- [ ] `rn/breaking-change` - The PR will be mentioned in the "Breaking Changes" section
- [ ] `rn/none` - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
- [ x] `rn/feature` - A new user-facing feature worth mentioning in the release notes
- [ ] `rn/bug-fix` - A user-facing bug fix worth mentioning in the release notes
- [ ] `rn/documentation` - A user-facing documentation change worth mentioning in the release notes


Problem statement:
[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

Hints from issue 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!

Commit message: added custom variant name

Signed-off-by: Nicole Farley <nicole.m.h.farley@gmail.com>

Patch to review:
diff --git a/mlflow/sagemaker/__init__.py b/mlflow/sagemaker/__init__.py
index a28a3bd383f0d..2ade6888198b5 100644
--- a/mlflow/sagemaker/__init__.py
+++ b/mlflow/sagemaker/__init__.py
@@ -2,6 +2,7 @@
 The ``mlflow.sagemaker`` module provides an API for deploying MLflow models to Amazon SageMaker.
 """
 import os
+import re
 from subprocess import Popen
 import urllib.parse
 import sys
@@ -163,6 +164,7 @@ def push_image_to_ecr(image=DEFAULT_IMAGE_NAME):
 def deploy(
     app_name,
     model_uri,
+    variant_name=None,
     execution_role_arn=None,
     assume_role_arn=None,
     bucket=None,
@@ -201,6 +203,10 @@ def deploy(
                       `Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
                       artifact-locations>`_.
 
+    :variant_name: (optional) The name to assign to the new production variant.
+                      - If variant_name violates sagemaker constraints on naming an MlflowException 
+                        will be raised.
+
     :param execution_role_arn: The name of an IAM role granting the SageMaker service permissions to
                                access the specified Docker image and S3 bucket containing MLflow
                                model artifacts. If unspecified, the currently-assumed role will be
@@ -388,6 +394,17 @@ def deploy(
         **assume_role_credentials,
     )
 
+    if variant_name:
+        pattern = r"[a-zA-Z0-9](-*[a-zA-Z0-9]){0,62}"
+        if not re.fullmatch(pattern, variant_name):
+            raise MlflowException(
+            message=(
+                "Variant name must match the following pattern" 
+                " {variant_pattern}".format(variant_pattern=pattern)
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
     if endpoint_exists:
         deployment_operation = _update_sagemaker_endpoint(
             endpoint_name=app_name,
@@ -403,6 +420,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 +436,7 @@ def deploy(
             data_capture_config=data_capture_config,
             role=execution_role_arn,
             sage_client=sage_client,
+            variant_name=variant_name,
         )
 
     if synchronous:
@@ -1425,7 +1444,6 @@ def cleanup_fn():
 
     return _SageMakerOperation(status_check_fn=status_check_fn, cleanup_fn=cleanup_fn)
 
-
 def _create_sagemaker_endpoint(
     endpoint_name,
     model_name,
@@ -1439,6 +1457,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 +1475,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.
+    :variant_name: (optional) The name to assign to the new production variant.
     """
     _logger.info("Creating new endpoint with name: %s ...", endpoint_name)
 
@@ -1471,8 +1491,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 +1568,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,7 +1587,8 @@ 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: (optional) The name to assign to the new production variant if it doesn't already exist.
+                     - If variant_name already exists an MlflowException will be raised.    """
     if mode not in [DEPLOYMENT_MODE_ADD, DEPLOYMENT_MODE_REPLACE]:
         msg = "Invalid mode `{md}` for deployment to a pre-existing application".format(md=mode)
         raise ValueError(msg)
@@ -1591,15 +1616,26 @@ 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
+        variant_name_list = [variant["VariantName"] for variant in production_variants]
+        if variant_name and variant_name in variant_name_list:
+            raise MlflowException(
+            message=(
+                "Variant {variant_name} already exists!".format(variant_name=variant_name)
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        )
     elif mode == DEPLOYMENT_MODE_REPLACE:
         new_model_weight = 1
         production_variants = []
 
     new_production_variant = {
-        "VariantName": model_name,
+        "VariantName": variant_name,
         "ModelName": model_name,
         "InitialInstanceCount": instance_count,
         "InstanceType": instance_type,


Code review comments:
diff_hunk: @@ -388,6 +394,17 @@ def deploy(
         **assume_role_credentials,
     )
 
+    if variant_name:
+        pattern = r"[a-zA-Z0-9](-*[a-zA-Z0-9]){0,62}"
+        if not re.fullmatch(pattern, variant_name):
+            raise MlflowException(
+            message=(
+                "Variant name must match the following pattern" 
+                " {variant_pattern}".format(variant_pattern=pattern)
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        )
path: mlflow/sagemaker/__init__.py
text: @user1:
Do we need to validate the variant name in the client, or can SageMaker perform the validation?

@author:
We don't have to.  Sagemaker will throw an error out if it is invalid.  I can remove this if that is preferred 

@user1:
Awesome! Yeah, I think we can remove this :)

diff_hunk: @@ -163,6 +164,7 @@ def push_image_to_ecr(image=DEFAULT_IMAGE_NAME):
 def deploy(
     app_name,
     model_uri,
+    variant_name=None,
path: mlflow/sagemaker/__init__.py
text: @user1:
Nit: For backwards compatibility, can we add `variant_name` to the end of the argument list?

@author:
Can do.

diff_hunk: @@ -1591,15 +1616,26 @@ 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
+        variant_name_list = [variant["VariantName"] for variant in production_variants]
+        if variant_name and variant_name in variant_name_list:
+            raise MlflowException(
+            message=(
+                "Variant {variant_name} already exists!".format(variant_name=variant_name)
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
path: mlflow/sagemaker/__init__.py
text: @user1:
Does SageMaker enforce variant name uniqueness? If not, does the client strictly need to perform this validation?

@author:
I will look into this.

@user1:
Thanks @user2-soaren ! Looking forward to merging once you finish your investigation.

@author:
Thank you @user1.  Looks like Sagemaker will error out if the variant name already exists so I'll remove this part also.