Add best tuning job method (aws#34)

Author: ChoiByungWook

src/sagemaker/tuner.py

@@ -83,7 +83,7 @@ def __init__(self, estimator, objective_metric_name, hyperparameter_ranges, metr
 
         self.max_jobs = max_jobs
         self.max_parallel_jobs = max_parallel_jobs
-        self.tuning_job_name = base_tuning_job_name
+        self.base_tuning_job_name = base_tuning_job_name
         self.metric_definitions = metric_definitions
         self.latest_tuning_job = None
         self._validate_parameter_ranges()
@@ -125,8 +125,22 @@ def stop_tuning_job(self):
         self._ensure_last_tuning_job()
         self.latest_tuning_job.stop()
 
+    def best_training_job(self):
+        """Return name of the best training job for the latest tuning job.
+        """
+        self._ensure_last_tuning_job()
+
+        tuning_job_describe_result = \
+            self.estimator.sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job(
+                HyperParameterTuningJobName=self.latest_tuning_job.name)
+
+        try:
+            return tuning_job_describe_result['BestTrainingJob']['TrainingJobName']
+        except KeyError:
+            raise Exception('Best training job not available for tuning job: {}'.format(self.latest_tuning_job.name))
+
     def _ensure_last_tuning_job(self):
-        if 'latest_tuning_job' not in dir(self) or self.latest_tuning_job is None:
+        if self.latest_tuning_job is None:
             raise ValueError('No tuning job available')
 
     def hyperparameter_ranges(self):

tests/unit/test_tuner.py

@@ -174,6 +174,43 @@ def test_stop_tuning_job_no_tuning_job(tuner):
     assert 'No tuning job available' in str(e)
 
 
+def test_best_tuning_job(tuner):
+    tuning_job_description = {'BestTrainingJob': {'TrainingJobName': JOB_NAME}}
+
+    tuner.estimator.sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(
+        name='describe_hyper_parameter_tuning_job', return_value=tuning_job_description)
+
+    tuner.latest_tuning_job = _TuningJob(tuner.estimator.sagemaker_session, JOB_NAME)
+    best_training_job = tuner.best_training_job()
+
+    assert best_training_job == JOB_NAME
+    tuner.estimator.sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job.assert_called_once_with(
+        HyperParameterTuningJobName=JOB_NAME)
+
+
+def test_best_tuning_job_no_latest_job(tuner):
+    with pytest.raises(Exception) as e:
+        tuner.best_training_job()
+
+    assert 'No tuning job available' in str(e)
+
+
+def test_best_tuning_job_no_best_job(tuner):
+    tuning_job_description = {'BestTrainingJob': {'Mock': None}}
+
+    tuner.estimator.sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(
+        name='describe_hyper_parameter_tuning_job', return_value=tuning_job_description)
+
+    tuner.latest_tuning_job = _TuningJob(tuner.estimator.sagemaker_session, JOB_NAME)
+
+    with pytest.raises(Exception) as e:
+        tuner.best_training_job()
+
+    tuner.estimator.sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job.assert_called_once_with(
+        HyperParameterTuningJobName=JOB_NAME)
+    assert 'Best training job not available for tuning job:' in str(e)
+
+
 #################################################################################
 # _ParameterRange Tests