fix: fix tags in deploy call for generic estimators (#1146)

Author: Dan

src/sagemaker/estimator.py

@@ -606,6 +606,7 @@ def deploy(
         model_name=None,
         kms_key=None,
         data_capture_config=None,
+        tags=None,
         **kwargs
     ):
         """Deploy the trained model to an Amazon SageMaker endpoint and return a
@@ -639,18 +640,18 @@ def deploy(
                 model completes (default: True).
             model_name (str): Name to use for creating an Amazon SageMaker
                 model. If not specified, the name of the training job is used.
-            tags(List[dict[str, str]]): Optional. The list of tags to attach to this specific
-                endpoint. Example:
-                >>> tags = [{'Key': 'tagname', 'Value': 'tagvalue'}]
-                For more information about tags, see
-                https://boto3.amazonaws.com/v1/documentation\
-                /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags
             kms_key (str): The ARN of the KMS key that is used to encrypt the
                 data on the storage volume attached to the instance hosting the
                 endpoint.
             data_capture_config (sagemaker.model_monitor.DataCaptureConfig): Specifies
                 configuration related to Endpoint data capture for use with
                 Amazon SageMaker Model Monitoring. Default: None.
+            tags(List[dict[str, str]]): Optional. The list of tags to attach to this specific
+                endpoint. Example:
+                >>> tags = [{'Key': 'tagname', 'Value': 'tagvalue'}]
+                For more information about tags, see
+                https://boto3.amazonaws.com/v1/documentation\
+                /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags
             **kwargs: Passed to invocation of ``create_model()``.
                 Implementations may customize ``create_model()`` to accept
                 ``**kwargs`` to customize model creation during deploy.
@@ -685,7 +686,7 @@ def deploy(
             accelerator_type=accelerator_type,
             endpoint_name=endpoint_name,
             update_endpoint=update_endpoint,
-            tags=self.tags,
+            tags=tags or self.tags,
             wait=wait,
             kms_key=kms_key,
             data_capture_config=data_capture_config,

tests/unit/test_estimator.py

@@ -1702,7 +1702,7 @@ def test_unsupported_type_in_dict():
 )
 
 
-def test_fit_deploy_keep_tags(sagemaker_session):
+def test_fit_deploy_tags_in_estimator(sagemaker_session):
     tags = [{"Key": "TagtestKey", "Value": "TagtestValue"}]
     estimator = Estimator(
         IMAGE_NAME,
@@ -1747,6 +1747,46 @@ def test_fit_deploy_keep_tags(sagemaker_session):
     )
 
 
+def test_fit_deploy_tags(sagemaker_session):
+    estimator = Estimator(
+        IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, sagemaker_session=sagemaker_session
+    )
+
+    estimator.fit()
+
+    tags = [{"Key": "TagtestKey", "Value": "TagtestValue"}]
+    estimator.deploy(INSTANCE_COUNT, INSTANCE_TYPE, tags=tags)
+
+    variant = [
+        {
+            "InstanceType": "c4.4xlarge",
+            "VariantName": "AllTraffic",
+            "ModelName": ANY,
+            "InitialVariantWeight": 1,
+            "InitialInstanceCount": 1,
+        }
+    ]
+
+    job_name = estimator._current_job_name
+    sagemaker_session.endpoint_from_production_variants.assert_called_with(
+        name=job_name,
+        production_variants=variant,
+        tags=tags,
+        kms_key=None,
+        wait=True,
+        data_capture_config_dict=None,
+    )
+
+    sagemaker_session.create_model.assert_called_with(
+        ANY,
+        "DummyRole",
+        {"ModelDataUrl": "s3://bucket/model.tar.gz", "Environment": {}, "Image": "fakeimage"},
+        enable_network_isolation=False,
+        vpc_config=None,
+        tags=tags,
+    )
+
+
 def test_generic_to_fit_no_input(sagemaker_session):
     e = Estimator(
         IMAGE_NAME,