Add variant_name kwarg

Author: hardianlawi

src/sagemaker/huggingface/model.py

@@ -200,6 +200,7 @@ def deploy(
         deserializer=None,
         accelerator_type=None,
         endpoint_name=None,
+        variant_name="AllTraffic",
         tags=None,
         kms_key=None,
         wait=True,
@@ -251,6 +252,8 @@ def deploy(
                 https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html
             endpoint_name (str): The name of the endpoint to create (default:
                 None). If not specified, a unique endpoint name will be created.
+            variant_name (string): The ``VariantName`` of this production variant
+                (default: 'AllTraffic').
             tags (List[dict[str, str]]): The list of tags to attach to this
                 specific endpoint.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
@@ -308,6 +311,7 @@ def deploy(
             deserializer,
             accelerator_type,
             endpoint_name,
+            variant_name,
             tags,
             kms_key,
             wait,

src/sagemaker/model.py

@@ -1023,6 +1023,7 @@ def deploy(
         deserializer=None,
         accelerator_type=None,
         endpoint_name=None,
+        variant_name="AllTraffic",
         tags=None,
         kms_key=None,
         wait=True,
@@ -1074,6 +1075,8 @@ def deploy(
                 https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html
             endpoint_name (str): The name of the endpoint to create (default:
                 None). If not specified, a unique endpoint name will be created.
+            variant_name (string): The ``VariantName`` of this production variant
+                (default: 'AllTraffic').
             tags (List[dict[str, str]]): The list of tags to attach to this
                 specific endpoint.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
@@ -1166,6 +1169,7 @@ def deploy(
             self.name,
             instance_type,
             initial_instance_count,
+            variant_name=variant_name,
             accelerator_type=accelerator_type,
             serverless_inference_config=serverless_inference_config_dict,
             volume_size=volume_size,

src/sagemaker/multidatamodel.py

@@ -158,6 +158,7 @@ def deploy(
         deserializer=None,
         accelerator_type=None,
         endpoint_name=None,
+        variant_name="AllTraffic",
         tags=None,
         kms_key=None,
         wait=True,
@@ -203,6 +204,8 @@ def deploy(
                 https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html
             endpoint_name (str): The name of the endpoint to create (default:
                 None). If not specified, a unique endpoint name will be created.
+            variant_name (string): The ``VariantName`` of this production variant
+                (default: 'AllTraffic').
             tags (List[dict[str, str]]): The list of tags to attach to this
                 specific endpoint.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
@@ -250,7 +253,11 @@ def deploy(
         )
 
         production_variant = sagemaker.production_variant(
-            self.name, instance_type, initial_instance_count, accelerator_type=accelerator_type
+            self.name,
+            instance_type,
+            initial_instance_count,
+            variant_name=variant_name,
+            accelerator_type=accelerator_type,
         )
         if endpoint_name:
             self.endpoint_name = endpoint_name