Fix logic around setting job name for tuning jobs (aws#47)

We originally weren't honoring job names passed through fit().
This change fixes that.

Author: laurenyu

src/sagemaker/tuner.py

@@ -81,6 +81,8 @@ def __init__(self, min_value, max_value):
 
 
 class HyperparameterTuner(object):
+    TUNING_JOB_NAME_MAX_LENGTH = 32
+
     SAGEMAKER_ESTIMATOR_MODULE = 'sagemaker_estimator_module'
     SAGEMAKER_ESTIMATOR_CLASS_NAME = 'sagemaker_estimator_class_name'
 
@@ -96,18 +98,25 @@ def __init__(self, estimator, objective_metric_name, hyperparameter_ranges, metr
         self.estimator = estimator
         self.objective_metric_name = objective_metric_name
         self.metric_definitions = metric_definitions
+        self._validate_parameter_ranges()
 
         self.strategy = strategy
         self.objective_type = objective_type
 
         self.max_jobs = max_jobs
         self.max_parallel_jobs = max_parallel_jobs
+
         self.base_tuning_job_name = base_tuning_job_name
-        self.metric_definitions = metric_definitions
+        self._current_job_name = None
         self.latest_tuning_job = None
-        self._validate_parameter_ranges()
 
-    def prepare_for_training(self):
+    def prepare_for_training(self, job_name=None):
+        if job_name is not None:
+            self._current_job_name = job_name
+        else:
+            base_name = self.base_tuning_job_name or base_name_from_image(self.estimator.train_image())
+            self._current_job_name = name_from_base(base_name, max_length=self.TUNING_JOB_NAME_MAX_LENGTH, short=True)
+
         self.static_hyperparameters = {to_str(k): to_str(v) for (k, v) in self.estimator.hyperparameters().items()}
         for hyperparameter_name in self._hyperparameter_ranges.keys():
             self.static_hyperparameters.pop(hyperparameter_name, None)
@@ -133,7 +142,7 @@ def fit(self, inputs, job_name=None, **kwargs):
         else:
             self.estimator.prepare_for_training(**kwargs)
 
-        self.prepare_for_training()
+        self.prepare_for_training(job_name=job_name)
         self.latest_tuning_job = _TuningJob.start_new(self, inputs)
 
     @classmethod
@@ -350,8 +359,6 @@ def _validate_parameter_ranges(self):
 
 
 class _TuningJob(_Job):
-    TUNING_JOB_NAME_MAX_LENGTH = 32
-
     def __init__(self, sagemaker_session, tuning_job_name):
         super(_TuningJob, self).__init__(sagemaker_session, tuning_job_name)
 
@@ -368,10 +375,7 @@ def start_new(cls, tuner, inputs):
         """
         config = _Job._load_config(inputs, tuner.estimator)
 
-        base_name = tuner.base_tuning_job_name or base_name_from_image(tuner.estimator.train_image())
-        tuning_job_name = name_from_base(base_name, max_length=cls.TUNING_JOB_NAME_MAX_LENGTH, short=True)
-
-        tuner.estimator.sagemaker_session.tune(job_name=tuning_job_name, strategy=tuner.strategy,
+        tuner.estimator.sagemaker_session.tune(job_name=tuner._current_job_name, strategy=tuner.strategy,
                                                objective_type=tuner.objective_type,
                                                objective_metric_name=tuner.objective_metric_name,
                                                max_jobs=tuner.max_jobs, max_parallel_jobs=tuner.max_parallel_jobs,
@@ -385,7 +389,7 @@ def start_new(cls, tuner, inputs):
                                                resource_config=(config['resource_config']),
                                                stop_condition=(config['stop_condition']))
 
-        return cls(tuner.sagemaker_session, tuning_job_name)
+        return cls(tuner.sagemaker_session, tuner._current_job_name)
 
     def stop(self):
         self.sagemaker_session.stop_tuning_job(HyperParameterTuningJobName=self.name)

tests/unit/test_tuner.py

@@ -148,6 +148,8 @@ def test_prepare_for_training(tuner):
     tuner.estimator.set_hyperparameters(**static_hyperparameters)
     tuner.prepare_for_training()
 
+    assert tuner._current_job_name.startswith(IMAGE_NAME)
+
     assert len(tuner.static_hyperparameters) == 3
     assert tuner.static_hyperparameters['another_one'] == '0'
 
@@ -157,6 +159,14 @@ def test_prepare_for_training(tuner):
     assert tuner.static_hyperparameters['sagemaker_estimator_module'] == module
 
 
+def test_prepare_for_training_with_job_name(tuner):
+    static_hyperparameters = {'validated': 1, 'another_one': 0}
+    tuner.estimator.set_hyperparameters(**static_hyperparameters)
+
+    tuner.prepare_for_training(job_name='some-other-job-name')
+    assert tuner._current_job_name == 'some-other-job-name'
+
+
 def test_validate_parameter_ranges_number_validation_error(sagemaker_session):
     pca = PCA(ROLE, TRAIN_INSTANCE_COUNT, TRAIN_INSTANCE_TYPE, NUM_COMPONENTS,
               base_job_name='pca', sagemaker_session=sagemaker_session)