Add stop tuning job (aws#29)

Author: ChoiByungWook

src/sagemaker/session.py

@@ -333,6 +333,25 @@ def tune(self, job_name, strategy, objective, metric_name,
         LOGGER.debug('tune request: {}'.format(json.dumps(tune_request, indent=4)))
         self.sagemaker_client.create_hyper_parameter_tuning_job(**tune_request)
 
+    def stop_tuning_job(self, name):
+        """Attempts to stop tuning job on Amazon SageMaker with specified name.
+
+        Args:
+            name: Name of Amazon SageMaker tuning job.
+        """
+        try:
+            LOGGER.info('Stopping tuning job: {}'.format(name))
+            self.sagemaker_client.stop_hyper_parameter_tuning_job(HyperParameterTuningJobName=name)
+        except ClientError as e:
+            error_code = e.response['Error']['Code']
+            # allow to pass if the job already stopped
+            if error_code == 'ValidationException':
+                LOGGER.info('Tuning job: {} is already stopped or not running.'.format(name))
+                pass
+            else:
+                LOGGER.error('Error occurred while attempting to stop tuning job: {}. Please try again.'.format(name))
+                raise
+
     def create_model(self, name, role, primary_container):
         """Create an Amazon SageMaker ``Model``.
 

src/sagemaker/tuner.py

@@ -76,18 +76,28 @@ def __init__(self, estimator, objective_metric_name, hyperparameter_ranges, metr
         self.latest_tuning_job = None
 
     def fit(self, inputs):
-        """Create HPO job
+        """Create tuning job
 
         Args:
             inputs (str): Parameters used when called  :meth:`~sagemaker.estimator.EstimatorBase.fit`.
         """
         self.latest_tuning_job = _TuningJob.start_new(self, inputs)
 
+    def stop_tuning_job(self):
+        """Stop latest running tuning job.
+        """
+        self._ensure_last_tuning_job()
+        self.latest_tuning_job.stop()
+
+    def _ensure_last_tuning_job(self):
+        if 'latest_tuning_job' not in dir(self) or self.latest_tuning_job is None:
+            raise ValueError('No tuning job available')
+
     def hyperparameter_ranges(self):
         """Return collections of ``ParameterRanges``
 
         Returns:
-            dict: ParameterRanges suitable for HPO tuning job.
+            dict: ParameterRanges suitable for tuning job.
         """
         hyperparameter_ranges = dict()
         for range_type in _ParameterRange.__all_types__:
@@ -108,7 +118,7 @@ def __init__(self, sagemaker_session, tuning_job_name):
 
     @classmethod
     def start_new(cls, tuner, inputs):
-        """Create a new Amazon SageMaker HPO tuning job from the HyperparameterTuner.
+        """Create a new Amazon SageMaker tuning job from the HyperparameterTuner.
 
         Args:
             tuner (sagemaker.tuner.HyperparameterTuner): Tuner object created by the user.
@@ -144,5 +154,8 @@ def start_new(cls, tuner, inputs):
 
         return cls(tuner.estimator.sagemaker_session, tuning_job_name)
 
+    def stop(self):
+        self.sagemaker_session.stop_tuning_job(HyperParameterTuningJobName=self.name)
+
     def wait(self):
         pass

tests/unit/test_session.py

@@ -230,6 +230,40 @@ def test_train_pack_to_request(sagemaker_session):
         'create_training_job', (), DEFAULT_EXPECTED_TRAIN_JOB_ARGS)
 
 
+def test_stop_tuning_job(sagemaker_session):
+    sms = sagemaker_session
+    sms.sagemaker_client.stop_hyper_parameter_tuning_job = Mock(name='stop_hyper_parameter_tuning_job')
+
+    sagemaker_session.stop_tuning_job(JOB_NAME)
+    sms.sagemaker_client.stop_hyper_parameter_tuning_job.assert_called_once_with(HyperParameterTuningJobName=JOB_NAME)
+
+
+def test_stop_tuning_job_client_error_already_stopped(sagemaker_session):
+    sms = sagemaker_session
+    exception = ClientError({'Error': {'Code': 'ValidationException'}}, 'Operation')
+    sms.sagemaker_client.stop_hyper_parameter_tuning_job = Mock(name='stop_hyper_parameter_tuning_job',
+                                                                side_effect=exception)
+    sagemaker_session.stop_tuning_job(JOB_NAME)
+
+    sms.sagemaker_client.stop_hyper_parameter_tuning_job.assert_called_once_with(HyperParameterTuningJobName=JOB_NAME)
+
+
+def test_stop_tuning_job_client_error(sagemaker_session):
+    error_response = {'Error': {'Code': 'MockException', 'Message': 'MockMessage'}}
+    operation = 'Operation'
+    exception = ClientError(error_response, operation)
+
+    sms = sagemaker_session
+    sms.sagemaker_client.stop_hyper_parameter_tuning_job = Mock(name='stop_hyper_parameter_tuning_job',
+                                                                side_effect=exception)
+
+    with pytest.raises(ClientError) as e:
+        sagemaker_session.stop_tuning_job(JOB_NAME)
+
+    sms.sagemaker_client.stop_hyper_parameter_tuning_job.assert_called_once_with(HyperParameterTuningJobName=JOB_NAME)
+    assert 'An error occurred (MockException) when calling the Operation operation: MockMessage' in str(e)
+
+
 @patch('sys.stdout', new_callable=io.BytesIO if six.PY2 else io.StringIO)
 def test_color_wrap(bio):
     color_wrap = sagemaker.logs.ColorWrap()

tests/unit/test_tuner.py

@@ -102,6 +102,21 @@ def test_serialize_nonexistent_parameter_ranges(tuner):
     assert not ranges[parameter_type + 'ParameterRanges']
 
 
+def test_stop_tuning_job(sagemaker_session, tuner):
+    sagemaker_session.stop_tuning_job = Mock(name='stop_hyper_parameter_tuning_job')
+    tuner.latest_tuning_job = _TuningJob(sagemaker_session, JOB_NAME)
+
+    tuner.stop_tuning_job()
+
+    sagemaker_session.stop_tuning_job.assert_called_once_with(HyperParameterTuningJobName=JOB_NAME)
+
+
+def test_stop_tuning_job_no_tuning_job(tuner):
+    with pytest.raises(ValueError) as e:
+            tuner.stop_tuning_job()
+    assert 'No HPO job available' in str(e)
+
+
 #################################################################################
 # _ParameterRange Tests
 
@@ -171,3 +186,10 @@ def test_start_new(tuner, sagemaker_session):
 
     assert started_tuning_job.sagemaker_session == sagemaker_session
     sagemaker_session.tune.assert_called_once()
+
+
+def test_stop(sagemaker_session):
+    tuning_job = _TuningJob(sagemaker_session, JOB_NAME)
+    tuning_job.stop()
+
+    sagemaker_session.stop_tuning_job.assert_called_once_with(HyperParameterTuningJobName=JOB_NAME)