Propagate Tags from estimator to model, endpoint, and endpoint config (#699)

* Propagate Tags from estimator to model, endpoint, and endpoint config

Author: mvsusp

src/sagemaker/estimator.py

@@ -350,6 +350,11 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
             update_endpoint (bool): Flag to update the model in an existing Amazon SageMaker endpoint.
                 If True, this will deploy a new EndpointConfig to an already existing endpoint and delete resources
                 corresponding to the previous EndpointConfig. Default: False
+            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.
                 For more, see the implementation docs.
@@ -374,7 +379,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
             initial_instance_count=initial_instance_count,
             accelerator_type=accelerator_type,
             endpoint_name=endpoint_name,
-            update_endpoint=update_endpoint)
+            update_endpoint=update_endpoint,
+            tags=self.tags)
 
     @property
     def model_data(self):

src/sagemaker/model.py

@@ -96,7 +96,7 @@ def enable_network_isolation(self):
         """
         return False
 
-    def _create_sagemaker_model(self, instance_type, accelerator_type=None):
+    def _create_sagemaker_model(self, instance_type, accelerator_type=None, tags=None):
         """Create a SageMaker Model Entity
 
         Args:
@@ -105,6 +105,11 @@ def _create_sagemaker_model(self, instance_type, accelerator_type=None):
             accelerator_type (str): Type of Elastic Inference accelerator to attach to an endpoint for model loading
                 and inference, for example, 'ml.eia1.medium'. If not specified, no Elastic Inference accelerator
                 will be attached to the endpoint.
+            tags(List[dict[str, str]]): Optional. The list of tags to add to the model. 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
+
         """
         container_def = self.prepare_container_def(instance_type, accelerator_type=accelerator_type)
         self.name = self.name or utils.name_from_image(container_def['Image'])

src/sagemaker/session.py

@@ -544,7 +544,8 @@ def transform(self, job_name, model_name, strategy, max_concurrent_transforms, m
         self.sagemaker_client.create_transform_job(**transform_request)
 
     def create_model(self, name, role, container_defs, vpc_config=None,
-                     enable_network_isolation=False, primary_container=None):
+                     enable_network_isolation=False, primary_container=None,
+                     tags=None):
         """Create an Amazon SageMaker ``Model``.
         Specify the S3 location of the model artifacts and Docker image containing
         the inference code. Amazon SageMaker uses this information to deploy the
@@ -570,6 +571,11 @@ def create_model(self, name, role, container_defs, vpc_config=None,
                 You can also specify the return value of ``sagemaker.container_def()``, which is used to create
                 more advanced container configurations, including model containers which need artifacts from S3. This
                 field is deprecated, please use container_defs instead.
+            tags(List[dict[str, str]]): Optional. The list of tags to add to the model. 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
+
 
         Returns:
             str: Name of the Amazon SageMaker ``Model`` created.
@@ -583,12 +589,16 @@ def create_model(self, name, role, container_defs, vpc_config=None,
             container_defs = primary_container
 
         role = self.expand_role(role)
-        create_model_request = {}
+
         if isinstance(container_defs, list):
-            create_model_request = _create_model_request(name=name, role=role, container_def=container_defs)
+            container_definition = container_defs
         else:
-            primary_container = _expand_container_def(container_defs)
-            create_model_request = _create_model_request(name=name, role=role, container_def=primary_container)
+            container_definition = _expand_container_def(container_defs)
+
+        create_model_request = _create_model_request(name=name,
+                                                     role=role,
+                                                     container_def=container_definition,
+                                                     tags=tags)
 
         if vpc_config:
             create_model_request['VpcConfig'] = vpc_config
@@ -702,7 +712,8 @@ def wait_for_model_package(self, model_package_name, poll=5):
                 model_package_name, status, reason))
         return desc
 
-    def create_endpoint_config(self, name, model_name, initial_instance_count, instance_type, accelerator_type=None):
+    def create_endpoint_config(self, name, model_name, initial_instance_count, instance_type,
+                               accelerator_type=None, tags=None):
         """Create an Amazon SageMaker endpoint configuration.
 
         The endpoint configuration identifies the Amazon SageMaker model (created using the
@@ -717,17 +728,24 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta
             instance_type (str): Type of EC2 instance to launch, for example, 'ml.c4.xlarge'.
             accelerator_type (str): Type of Elastic Inference accelerator to attach to the instance. For example,
                 'ml.eia1.medium'. For more information: https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html
+            tags(List[dict[str, str]]): Optional. The list of tags to add to the endpoint config. 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
 
 
         Returns:
             str: Name of the endpoint point configuration created.
         """
         LOGGER.info('Creating endpoint-config with name {}'.format(name))
 
+        tags = tags or []
+
         self.sagemaker_client.create_endpoint_config(
             EndpointConfigName=name,
             ProductionVariants=[production_variant(model_name, instance_type, initial_instance_count,
-                                                   accelerator_type=accelerator_type)]
+                                                   accelerator_type=accelerator_type)],
+            Tags=tags
         )
         return name
 
@@ -1383,19 +1401,18 @@ def __init__(self, model_data, image, env=None):
         self.env = env
 
 
-def _create_model_request(name, role, container_def=None):  # pylint: disable=redefined-outer-name
+def _create_model_request(name, role, container_def=None, tags=None):  # pylint: disable=redefined-outer-name
+    request = {'ModelName': name, 'ExecutionRoleArn': role}
+
     if isinstance(container_def, list):
-        return {
-            'ModelName': name,
-            'Containers': container_def,
-            'ExecutionRoleArn': role
-        }
+        request['Containers'] = container_def
     else:
-        return {
-            'ModelName': name,
-            'PrimaryContainer': container_def,
-            'ExecutionRoleArn': role
-        }
+        request['PrimaryContainer'] = container_def
+
+    if tags:
+        request['Tags'] = tags
+
+    return request
 
 
 def _deployment_entity_exists(describe_fn):

tests/unit/test_create_deploy_entities.py

@@ -70,7 +70,7 @@ def test_create_endpoint_config(sagemaker_session):
                      'InitialVariantWeight': 1,
                      'VariantName': 'AllTraffic'}]
     sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_once_with(
-        EndpointConfigName=ENDPOINT_CONFIG_NAME, ProductionVariants=expected_pvs)
+        EndpointConfigName=ENDPOINT_CONFIG_NAME, ProductionVariants=expected_pvs, Tags=[])
 
 
 def test_create_endpoint_config_with_accelerator(sagemaker_session):
@@ -87,7 +87,7 @@ def test_create_endpoint_config_with_accelerator(sagemaker_session):
                      'VariantName': 'AllTraffic',
                      'AcceleratorType': ACCELERATOR_TYPE}]
     sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_once_with(
-        EndpointConfigName=ENDPOINT_CONFIG_NAME, ProductionVariants=expected_pvs)
+        EndpointConfigName=ENDPOINT_CONFIG_NAME, ProductionVariants=expected_pvs, Tags=[])
 
 
 def test_create_endpoint_no_wait(sagemaker_session):

tests/unit/test_estimator.py

@@ -18,7 +18,7 @@
 from time import sleep
 
 import pytest
-from mock import MagicMock, Mock, patch
+from mock import ANY, MagicMock, Mock, patch
 
 from sagemaker.amazon.amazon_estimator import registry
 from sagemaker.algorithm import AlgorithmEstimator
@@ -882,6 +882,36 @@ def test_unsupported_type_in_dict():
 HP_TRAIN_CALL.update({'hyperparameters': STRINGIFIED_HYPERPARAMS})
 
 
+def test_fit_deploy_keep_tags(sagemaker_session):
+    tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
+    estimator = Estimator(IMAGE_NAME,
+                          ROLE,
+                          INSTANCE_COUNT,
+                          INSTANCE_TYPE,
+                          tags=tags,
+                          sagemaker_session=sagemaker_session)
+
+    estimator.fit()
+
+    estimator.deploy(INSTANCE_COUNT, INSTANCE_TYPE)
+
+    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(job_name,
+                                                                           variant,
+                                                                           tags)
+
+    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)
+
+
 def test_generic_to_fit_no_input(sagemaker_session):
     e = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path=OUTPUT_PATH,
                   sagemaker_session=sagemaker_session)

tests/unit/test_session.py

@@ -19,7 +19,7 @@
 import pytest
 import six
 from botocore.exceptions import ClientError
-from mock import MagicMock, Mock, patch, call, mock_open
+from mock import ANY, MagicMock, Mock, patch, call, mock_open
 
 import sagemaker
 from sagemaker import s3_input, Session, get_execution_role
@@ -758,6 +758,19 @@ def test_create_model(expand_container_def, sagemaker_session):
                                                                        PrimaryContainer=PRIMARY_CONTAINER)
 
 
+@patch('sagemaker.session._expand_container_def', return_value=PRIMARY_CONTAINER)
+def test_create_model_with_tags(expand_container_def, sagemaker_session):
+    tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
+    model = sagemaker_session.create_model(MODEL_NAME, ROLE, PRIMARY_CONTAINER, tags=tags)
+
+    assert model == MODEL_NAME
+    tags = [{'Value': 'TagtestValue', 'Key': 'TagtestKey'}]
+    sagemaker_session.sagemaker_client.create_model.assert_called_with(ExecutionRoleArn=EXPANDED_ROLE,
+                                                                       ModelName=MODEL_NAME,
+                                                                       PrimaryContainer=PRIMARY_CONTAINER,
+                                                                       Tags=tags)
+
+
 @patch('sagemaker.session._expand_container_def', return_value=PRIMARY_CONTAINER)
 def test_create_model_with_primary_container(expand_container_def, sagemaker_session):
     model = sagemaker_session.create_model(MODEL_NAME, ROLE, container_defs=PRIMARY_CONTAINER)
@@ -903,6 +916,17 @@ def test_endpoint_from_production_variants(sagemaker_session):
         ProductionVariants=pvs)
 
 
+def test_create_endpoint_config_with_tags(sagemaker_session):
+    tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
+
+    sagemaker_session.create_endpoint_config('endpoint-test', 'simple-model', 1, 'local', tags=tags)
+
+    sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
+        EndpointConfigName='endpoint-test',
+        ProductionVariants=ANY,
+        Tags=tags)
+
+
 def test_endpoint_from_production_variants_with_tags(sagemaker_session):
     ims = sagemaker_session
     ims.sagemaker_client.describe_endpoint = Mock(return_value={'EndpointStatus': 'InService'})