Add tags for hyperparameter tuning jobs

Author: laurenyu

src/sagemaker/session.py

@@ -241,7 +241,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,
@@ -259,7 +258,7 @@ def train(self, image, input_mode, input_config, role, job_name, output_config,
     def tune(self, job_name, strategy, objective_type, objective_metric_name,
              max_jobs, max_parallel_jobs, parameter_ranges,
              static_hyperparameters, image, input_mode, metric_definitions,
-             role, input_config, output_config, resource_config, stop_condition):
+             role, input_config, output_config, resource_config, stop_condition, tags):
         """Create an Amazon SageMaker hyperparameter tuning job
 
         Args:
@@ -292,6 +291,7 @@ def tune(self, job_name, strategy, objective_type, objective_metric_name,
             instance_type (str): Type of EC2 instance to use for training, for example, 'ml.c4.xlarge'.
             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 the tuning job.
         """
         tune_request = {
             'HyperParameterTuningJobName': job_name,
@@ -324,6 +324,9 @@ def tune(self, job_name, strategy, objective_type, objective_metric_name,
         if metric_definitions is not None:
             tune_request['TrainingJobDefinition']['AlgorithmSpecification']['MetricDefinitions'] = metric_definitions
 
+        if tags is not None:
+            tune_request['Tags'] = tags
+
         LOGGER.info('Creating hyperparameter tuning job with name: {}'.format(job_name))
         LOGGER.debug('tune request: {}'.format(json.dumps(tune_request, indent=4)))
         self.sagemaker_client.create_hyper_parameter_tuning_job(**tune_request)

src/sagemaker/tuner.py

@@ -91,7 +91,7 @@ class HyperparameterTuner(object):
 
     def __init__(self, estimator, objective_metric_name, hyperparameter_ranges, metric_definitions=None,
                  strategy='Bayesian', objective_type='Maximize', max_jobs=1, max_parallel_jobs=1,
-                 base_tuning_job_name=None):
+                 tags=None, base_tuning_job_name=None):
         self._hyperparameter_ranges = hyperparameter_ranges
         if self._hyperparameter_ranges is None or len(self._hyperparameter_ranges) == 0:
             raise ValueError('Need to specify hyperparameter ranges')
@@ -106,6 +106,7 @@ def __init__(self, estimator, objective_metric_name, hyperparameter_ranges, metr
         self.max_jobs = max_jobs
         self.max_parallel_jobs = max_parallel_jobs
 
+        self.tags = tags
         self.base_tuning_job_name = base_tuning_job_name
         self._current_job_name = None
         self.latest_tuning_job = None
@@ -387,7 +388,7 @@ def start_new(cls, tuner, inputs):
                                                role=(config['role']), input_config=(config['input_config']),
                                                output_config=(config['output_config']),
                                                resource_config=(config['resource_config']),
-                                               stop_condition=(config['stop_condition']))
+                                               stop_condition=(config['stop_condition']), tags=tuner.tags)
 
         return cls(tuner.sagemaker_session, tuner._current_job_name)
 

tests/unit/test_tuner.py

@@ -201,9 +201,13 @@ def test_fit_pca(sagemaker_session, tuner):
     pca.subtract_mean = True
     pca.extra_components = 5
 
+    tuner.estimator = pca
+
+    tags = [{'Name': 'some-tag-without-a-value'}]
+    tuner.tags = tags
+
     hyperparameter_ranges = {'num_components': IntegerParameter(2, 4),
                              'algorithm_mode': CategoricalParameter(['regular', 'randomized'])}
-    tuner.estimator = pca
     tuner._hyperparameter_ranges = hyperparameter_ranges
 
     records = RecordSet(s3_data=INPUTS, num_records=1, feature_dim=1)
@@ -215,6 +219,7 @@ def test_fit_pca(sagemaker_session, tuner):
     assert tune_kwargs['static_hyperparameters']['extra_components'] == '5'
     assert len(tune_kwargs['parameter_ranges']['IntegerParameterRanges']) == 1
     assert tune_kwargs['job_name'].startswith('pca')
+    assert tune_kwargs['tags'] == tags
     assert tuner.estimator.mini_batch_size == 9999