Add tags for training jobs

Author: laurenyu

CHANGELOG.rst

@@ -6,6 +6,7 @@ CHANGELOG
 ========
 
 * bug-fix: Estimators: Change max_iterations hyperparameter key for KMeans
+* feature: Estimators: add support for tagging training jobs
 
 1.3.0
 =====

src/sagemaker/estimator.py

@@ -44,7 +44,7 @@ class EstimatorBase(with_metaclass(ABCMeta, object)):
 
     def __init__(self, role, train_instance_count, train_instance_type,
                  train_volume_size=30, train_max_run=24 * 60 * 60, input_mode='File',
-                 output_path=None, output_kms_key=None, base_job_name=None, sagemaker_session=None):
+                 output_path=None, output_kms_key=None, base_job_name=None, sagemaker_session=None, tags=None):
         """Initialize an ``EstimatorBase`` instance.
 
         Args:
@@ -73,13 +73,16 @@ def __init__(self, role, train_instance_count, train_instance_type,
             sagemaker_session (sagemaker.session.Session): Session object which manages interactions with
                 Amazon SageMaker APIs and any other AWS services needed. If not specified, the estimator creates one
                 using the default AWS configuration chain.
+            tags (list[dict]): List of tags for labeling a training job. For more, see
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
         """
         self.role = role
         self.train_instance_count = train_instance_count
         self.train_instance_type = train_instance_type
         self.train_volume_size = train_volume_size
         self.train_max_run = train_max_run
         self.input_mode = input_mode
+        self.tags = tags
 
         if self.train_instance_type in ('local', 'local_gpu'):
             if self.train_instance_type == 'local_gpu' and self.train_instance_count > 1:
@@ -345,7 +348,8 @@ def start_new(cls, estimator, inputs):
         estimator.sagemaker_session.train(image=estimator.train_image(), input_mode=estimator.input_mode,
                                           input_config=input_config, role=role, job_name=estimator._current_job_name,
                                           output_config=output_config, resource_config=resource_config,
-                                          hyperparameters=hyperparameters, stop_condition=stop_condition)
+                                          hyperparameters=hyperparameters, stop_condition=stop_condition,
+                                          tags=estimator.tags)
 
         return cls(estimator.sagemaker_session, estimator._current_job_name)
 

src/sagemaker/session.py

@@ -203,7 +203,7 @@ def default_bucket(self):
         return self._default_bucket
 
     def train(self, image, input_mode, input_config, role, job_name, output_config,
-              resource_config, hyperparameters, stop_condition):
+              resource_config, hyperparameters, stop_condition, tags):
         """Create an Amazon SageMaker training job.
 
         Args:
@@ -232,6 +232,8 @@ def train(self, image, input_mode, input_config, role, job_name, output_config,
                 keys and values, but ``str()`` will be called to convert them before training.
             stop_condition (dict): Defines when training shall finish. Contains entries that can be understood by the
                 service like ``MaxRuntimeInSeconds``.
+            tags (list[dict]): List of tags for labeling a training job. For more, see
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
 
         Returns:
             str: ARN of the training job, if it is created.
@@ -242,7 +244,6 @@ def train(self, image, input_mode, input_config, role, job_name, output_config,
                 'TrainingImage': image,
                 'TrainingInputMode': input_mode
             },
-            # 'HyperParameters': hyperparameters,
             'InputDataConfig': input_config,
             'OutputDataConfig': output_config,
             'TrainingJobName': job_name,
@@ -253,6 +254,10 @@ def train(self, image, input_mode, input_config, role, job_name, output_config,
 
         if hyperparameters and len(hyperparameters) > 0:
             train_request['HyperParameters'] = hyperparameters
+
+        if tags is not None:
+            train_request['Tags'] = tags
+
         LOGGER.info('Creating training-job with name: {}'.format(job_name))
         LOGGER.debug('train request: {}'.format(json.dumps(train_request, indent=4)))
         self.sagemaker_client.create_training_job(**train_request)

tests/unit/test_estimator.py

@@ -292,16 +292,18 @@ def test_fit_then_fit_again(sagemaker_session):
 
 @patch('time.strftime', return_value=TIMESTAMP)
 def test_fit_verify_job_name(strftime, sagemaker_session):
+    tags = [{'Name': 'some-tag'}]
     fw = DummyFramework(entry_point=SCRIPT_PATH, role='DummyRole', sagemaker_session=sagemaker_session,
                         train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,
-                        enable_cloudwatch_metrics=True)
+                        enable_cloudwatch_metrics=True, tags=tags)
     fw.fit(inputs=s3_input('s3://mybucket/train'))
 
     _, _, train_kwargs = sagemaker_session.train.mock_calls[0]
 
     assert train_kwargs['hyperparameters']['sagemaker_enable_cloudwatch_metrics']
     assert train_kwargs['image'] == IMAGE_NAME
     assert train_kwargs['input_mode'] == 'File'
+    assert train_kwargs['tags'] == tags
     assert train_kwargs['job_name'] == JOB_NAME
     assert fw.latest_training_job.name == JOB_NAME
 
@@ -475,7 +477,8 @@ def test_unsupported_type_in_dict():
         'InstanceType': INSTANCE_TYPE,
         'VolumeSizeInGB': 30
     },
-    'stop_condition': {'MaxRuntimeInSeconds': 86400}
+    'stop_condition': {'MaxRuntimeInSeconds': 86400},
+    'tags': None,
 }