Add tags on endpoints (#253)

Author: Pavlos Mitsoulis-Ntompos

CHANGELOG.rst

@@ -9,6 +9,7 @@ CHANGELOG
 * enhancement: Unify generation of model uploaded code location
 * feature: Change minimum required scipy from 1.0.0 to 0.19.0
 * feature: Allow all Framework Estimators to use a custom docker image.
+* feature: Option to add Tags on SageMaker Endpoints
 
 1.5.0
 =====

src/sagemaker/model.py

@@ -66,7 +66,7 @@ def prepare_container_def(self, instance_type):
         """
         return sagemaker.container_def(self.image, self.model_data, self.env)
 
-    def deploy(self, initial_instance_count, instance_type, endpoint_name=None):
+    def deploy(self, initial_instance_count, instance_type, endpoint_name=None, tags=None):
         """Deploy this ``Model`` to an ``Endpoint`` and optionally return a ``Predictor``.
 
         Create a SageMaker ``Model`` and ``EndpointConfig``, and deploy an ``Endpoint`` from this ``Model``.
@@ -82,6 +82,7 @@ def deploy(self, initial_instance_count, instance_type, endpoint_name=None):
                 ``Endpoint`` created from this ``Model``.
             endpoint_name (str): The name of the endpoint to create (default: None).
                 If not specified, a unique endpoint name will be created.
+            tags (list[dict[str, str]]): A list of key-value pairs for tagging the endpoint (default: None).
 
         Returns:
             callable[string, sagemaker.session.Session] or None: Invocation of ``self.predictor_cls`` on
@@ -98,7 +99,7 @@ def deploy(self, initial_instance_count, instance_type, endpoint_name=None):
         self.sagemaker_session.create_model(model_name, self.role, container_def)
         production_variant = sagemaker.production_variant(model_name, instance_type, initial_instance_count)
         self.endpoint_name = endpoint_name or model_name
-        self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant])
+        self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)
         if self.predictor_cls:
             return self.predictor_cls(self.endpoint_name, self.sagemaker_session)
 

src/sagemaker/session.py

@@ -646,12 +646,13 @@ def endpoint_from_model_data(self, model_s3_location, deployment_image, initial_
         self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
         return name
 
-    def endpoint_from_production_variants(self, name, production_variants, wait=True):
+    def endpoint_from_production_variants(self, name, production_variants, tags=None, wait=True):
         """Create an SageMaker ``Endpoint`` from a list of production variants.
 
         Args:
             name (str): The name of the ``Endpoint`` to create.
             production_variants (list[dict[str, str]]): The list of production variants to deploy.
+            tags (list[dict[str, str]]): A list of key-value pairs for tagging the endpoint (default: None).
             wait (bool): Whether to wait for the endpoint deployment to complete before returning (default: True).
 
         Returns:
@@ -660,8 +661,11 @@ def endpoint_from_production_variants(self, name, production_variants, wait=True
 
         if not _deployment_entity_exists(
                 lambda: self.sagemaker_client.describe_endpoint_config(EndpointConfigName=name)):
-            self.sagemaker_client.create_endpoint_config(
-                EndpointConfigName=name, ProductionVariants=production_variants)
+            config_options = {'EndpointConfigName': name, 'ProductionVariants': production_variants}
+            if tags:
+                config_options['Tags'] = tags
+
+            self.sagemaker_client.create_endpoint_config(**config_options)
         return self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
 
     def expand_role(self, role):

tests/unit/test_model.py

@@ -100,7 +100,8 @@ def test_deploy(tfo, time, sagemaker_session):
           'ModelName': 'mi-2017-10-10-14-14-15',
           'InstanceType': INSTANCE_TYPE,
           'InitialInstanceCount': 1,
-          'VariantName': 'AllTraffic'}])
+          'VariantName': 'AllTraffic'}],
+        None)
 
 
 @patch('tarfile.open')
@@ -114,7 +115,24 @@ def test_deploy_endpoint_name(tfo, time, sagemaker_session):
           'ModelName': 'mi-2017-10-10-14-14-15',
           'InstanceType': INSTANCE_TYPE,
           'InitialInstanceCount': 55,
-          'VariantName': 'AllTraffic'}])
+          'VariantName': 'AllTraffic'}],
+        None)
+
+
+@patch('tarfile.open')
+@patch('time.strftime', return_value=TIMESTAMP)
+def test_deploy_tags(tfo, time, sagemaker_session):
+    model = DummyFrameworkModel(sagemaker_session)
+    tags = [{'ModelName': 'TestModel'}]
+    model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=1, tags=tags)
+    sagemaker_session.endpoint_from_production_variants.assert_called_with(
+        'mi-2017-10-10-14-14-15',
+        [{'InitialVariantWeight': 1,
+          'ModelName': 'mi-2017-10-10-14-14-15',
+          'InstanceType': INSTANCE_TYPE,
+          'InitialInstanceCount': 1,
+          'VariantName': 'AllTraffic'}],
+        tags)
 
 
 @patch('sagemaker.model.Session')

tests/unit/test_session.py

@@ -487,6 +487,35 @@ def test_endpoint_from_production_variants(sagemaker_session):
                 'VariantName': 'AllTraffic'}])
 
 
+def test_endpoint_from_production_variants_with_tags(sagemaker_session):
+    ims = sagemaker_session
+    ims.sagemaker_client.describe_endpoint = Mock(return_value={'EndpointStatus': 'InService'})
+    pvs = [sagemaker.production_variant('A', 'ml.p2.xlarge'), sagemaker.production_variant('B', 'p299.4096xlarge')]
+    ex = ClientError({'Error': {'Code': 'ValidationException', 'Message': 'Could not find your thing'}}, 'b')
+    ims.sagemaker_client.describe_endpoint_config = Mock(side_effect=ex)
+    tags = [{'ModelName': 'TestModel'}]
+    sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
+    sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
+                                                                          EndpointName='some-endpoint')
+    sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
+        EndpointConfigName='some-endpoint',
+        ProductionVariants=[
+            {
+                'InstanceType': 'ml.p2.xlarge',
+                'ModelName': 'A',
+                'InitialVariantWeight': 1,
+                'InitialInstanceCount': 1,
+                'VariantName': 'AllTraffic'
+            },
+            {
+                'InstanceType': 'p299.4096xlarge',
+                'ModelName': 'B',
+                'InitialVariantWeight': 1,
+                'InitialInstanceCount': 1,
+                'VariantName': 'AllTraffic'}],
+        Tags=tags)
+
+
 def test_wait_for_tuning_job(sagemaker_session):
     hyperparameter_tuning_job_desc = {'HyperParameterTuningJobStatus': 'Completed'}
     sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(