Add wait for tuner (aws#37)

Author: ChoiByungWook

src/sagemaker/session.py

@@ -479,21 +479,39 @@ def wait_for_job(self, job, poll=5):
             ValueError: If the training job fails.
         """
         desc = _wait_until(lambda: _train_done(self.sagemaker_client, job), poll)
-        self._check_job_status(job, desc)
+        self._check_job_status(job, desc, 'TrainingJobStatus')
         return desc
 
-    def _check_job_status(self, job, desc):
+    def wait_for_tuning_job(self, job, poll=5):
+        """Wait for an Amazon SageMaker tuning job to complete.
+
+        Args:
+            job (str): Name of the tuning job to wait for.
+            poll (int): Polling interval in seconds (default: 5).
+
+        Returns:
+            (dict): Return value from the ``DescribeHyperParameterTuningJob`` API.
+
+        Raises:
+            ValueError: If the hyperparameter tuning job fails.
+        """
+        desc = _wait_until(lambda: _tuning_job_status(self.sagemaker_client, job), poll)
+        self._check_job_status(job, desc, 'HyperParameterTuningJobStatus')
+        return desc
+
+    def _check_job_status(self, job, desc, status_key_name):
         """Check to see if the job completed successfully and, if not, construct and
         raise a ValueError.
 
         Args:
             job (str): The name of the job to check.
             desc (dict[str, str]): The result of ``describe_training_job()``.
+            status_key_name (str): Status key name to check for.
 
         Raises:
             ValueError: If the training job fails.
         """
-        status = desc['TrainingJobStatus']
+        status = desc[status_key_name]
 
         if status != 'Completed' and status != 'Stopped':
             reason = desc.get('FailureReason', '(No reason provided)')
@@ -757,7 +775,7 @@ def logs_for_job(self, job_name, wait=False, poll=10):  # noqa: C901 - suppress
                     state = LogState.JOB_COMPLETE
 
         if wait:
-            self._check_job_status(job_name, description)
+            self._check_job_status(job_name, description, 'TrainingJobStatus')
             if dot:
                 print()
             print('===== Job Complete =====')
@@ -909,7 +927,7 @@ def _train_done(sagemaker_client, job_name):
     return desc
 
 
-def _tune_done(sagemaker_client, job_name):
+def _tuning_job_status(sagemaker_client, job_name):
     tuning_status_codes = {
         'Completed': '!',
         'InProgress': '.',

src/sagemaker/tuner.py

@@ -153,6 +153,12 @@ def stop_tuning_job(self):
         self._ensure_last_tuning_job()
         self.latest_tuning_job.stop()
 
+    def wait(self):
+        """Wait for latest tuning job to finish.
+        """
+        self._ensure_last_tuning_job()
+        self.latest_tuning_job.wait()
+
     def best_training_job(self):
         """Return name of the best training job for the latest tuning job.
         """
@@ -258,4 +264,4 @@ def stop(self):
         self.sagemaker_session.stop_tuning_job(HyperParameterTuningJobName=self.name)
 
     def wait(self):
-        pass
+        self.sagemaker_session.wait_for_tuning_job(self.name)

tests/unit/test_session.py

@@ -22,6 +22,8 @@
 
 from botocore.exceptions import ClientError
 
+from sagemaker.session import _tuning_job_status
+
 REGION = 'us-west-2'
 
 
@@ -451,3 +453,32 @@ def test_endpoint_from_production_variants(sagemaker_session):
                 'InitialVariantWeight': 1,
                 'InitialInstanceCount': 1,
                 'VariantName': 'AllTraffic'}])
+
+
+def test_wait_for_tuning_job(sagemaker_session):
+    hyperparameter_tuning_job_desc = {'HyperParameterTuningJobStatus': 'Completed'}
+    sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(
+        name='describe_hyper_parameter_tuning_job', return_value=hyperparameter_tuning_job_desc)
+
+    result = sagemaker_session.wait_for_tuning_job(JOB_NAME)
+    assert result['HyperParameterTuningJobStatus'] == 'Completed'
+
+
+def test_tune_job_status(sagemaker_session):
+    hyperparameter_tuning_job_desc = {'HyperParameterTuningJobStatus': 'Completed'}
+    sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(
+        name='describe_hyper_parameter_tuning_job', return_value=hyperparameter_tuning_job_desc)
+
+    result = _tuning_job_status(sagemaker_session.sagemaker_client, JOB_NAME)
+
+    assert result['HyperParameterTuningJobStatus'] == 'Completed'
+
+
+def test_tune_job_status_none(sagemaker_session):
+    hyperparameter_tuning_job_desc = {'HyperParameterTuningJobStatus': 'InProgress'}
+    sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(
+        name='describe_hyper_parameter_tuning_job', return_value=hyperparameter_tuning_job_desc)
+
+    result = _tuning_job_status(sagemaker_session.sagemaker_client, JOB_NAME)
+
+    assert result is None

tests/unit/test_tuner.py

@@ -275,6 +275,15 @@ def test_deploy_default(tuner):
     assert predictor.sagemaker_session == tuner.estimator.sagemaker_session
 
 
+def test_wait(tuner):
+    tuner.latest_tuning_job = _TuningJob(tuner.estimator.sagemaker_session, JOB_NAME)
+    tuner.estimator.sagemaker_session.wait_for_tuning_job = Mock(name='wait_for_tuning_job')
+
+    tuner.wait()
+
+    tuner.estimator.sagemaker_session.wait_for_tuning_job.assert_called_once_with(JOB_NAME)
+
+
 #################################################################################
 # _ParameterRange Tests
 
@@ -353,3 +362,12 @@ def test_stop(sagemaker_session):
     tuning_job.stop()
 
     sagemaker_session.stop_tuning_job.assert_called_once_with(HyperParameterTuningJobName=JOB_NAME)
+
+
+def test_tuning_job_wait(sagemaker_session):
+    sagemaker_session.wait_for_tuning_job = Mock(name='wait_for_tuning_job')
+
+    tuning_job = _TuningJob(sagemaker_session, JOB_NAME)
+    tuning_job.wait()
+
+    sagemaker_session.wait_for_tuning_job.assert_called_once_with(JOB_NAME)