Add variant_name kwarg

Author: hardianlawi

src/sagemaker/huggingface/model.py

@@ -186,6 +186,7 @@ def deploy(
         deserializer=None,
         accelerator_type=None,
         endpoint_name=None,
+        variant_name="AllTraffic",
         tags=None,
         kms_key=None,
         wait=True,
@@ -234,6 +235,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
@@ -281,6 +284,7 @@ def deploy(
             deserializer,
             accelerator_type,
             endpoint_name,
+            variant_name,
             tags,
             kms_key,
             wait,

src/sagemaker/model.py

@@ -956,6 +956,7 @@ def deploy(
         deserializer=None,
         accelerator_type=None,
         endpoint_name=None,
+        variant_name="AllTraffic",
         tags=None,
         kms_key=None,
         wait=True,
@@ -1004,6 +1005,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
@@ -1086,6 +1089,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,
         )

src/sagemaker/multidatamodel.py

@@ -152,6 +152,7 @@ def deploy(
         deserializer=None,
         accelerator_type=None,
         endpoint_name=None,
+        variant_name="AllTraffic",
         tags=None,
         kms_key=None,
         wait=True,
@@ -197,6 +198,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
@@ -244,7 +247,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