feature: handler for stopping transform job

Author: imujjwal96

src/sagemaker/session.py

@@ -1006,6 +1006,27 @@ def wait_for_transform_job(self, job, poll=5):
         self._check_job_status(job, desc, "TransformJobStatus")
         return desc
 
+    def stop_transform_job(self, name):
+        """Stop the Amazon SageMaker hyperparameter tuning job with the specified name.
+
+        Args:
+            name (str): Name of the Amazon SageMaker batch transform job.
+
+        Raises:
+            ClientError: If an error occurs while trying to stop the batch transform job.
+        """
+        try:
+            LOGGER.info('Stopping transform job: {}'.format(name))
+            self.sagemaker_client.stop_transform_job(TransformJobName=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('Transform job: {} is already stopped or not running.'.format(name))
+            else:
+                LOGGER.error('Error occurred while attempting to stop transform job: {}.'.format(name))
+                raise
+
     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.

src/sagemaker/transformer.py

@@ -206,6 +206,12 @@ def wait(self):
         self._ensure_last_transform_job()
         self.latest_transform_job.wait()
 
+    def stop_transform_job(self):
+        """Stop latest running batch transform job.
+        """
+        self._ensure_last_transform_job()
+        self.latest_transform_job.stop()
+
     def _ensure_last_transform_job(self):
         if self.latest_transform_job is None:
             raise ValueError("No transform job available")
@@ -305,6 +311,9 @@ def start_new(
     def wait(self):
         self.sagemaker_session.wait_for_transform_job(self.job_name)
 
+    def stop(self):
+        self.sagemaker_session.stop_transform_job(name=self.job_name)
+
     @staticmethod
     def _load_config(data, data_type, content_type, compression_type, split_type, transformer):
         input_config = _TransformJob._format_inputs_to_input_config(

tests/integ/test_transformer.py

@@ -16,6 +16,7 @@
 import os
 import pickle
 import sys
+import time
 
 import pytest
 
@@ -364,4 +365,50 @@ def _create_transformer_and_transform_job(
         output_filter=output_filter,
         join_source=join_source,
     )
+
+
+def test_stop_transform_job(sagemaker_session, mxnet_full_version):
+    data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
+    script_path = os.path.join(data_path, 'mnist.py')
+    tags = [{'Key': 'some-tag', 'Value': 'value-for-tag'}]
+
+    mx = MXNet(entry_point=script_path, role='SageMakerRole', train_instance_count=1,
+               train_instance_type='ml.c4.xlarge', sagemaker_session=sagemaker_session,
+               framework_version=mxnet_full_version)
+
+    train_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'train'),
+                                                   key_prefix='integ-test-data/mxnet_mnist/train')
+    test_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'test'),
+                                                  key_prefix='integ-test-data/mxnet_mnist/test')
+    job_name = unique_name_from_base('test-mxnet-transform')
+
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+        mx.fit({'train': train_input, 'test': test_input}, job_name=job_name)
+
+    transform_input_path = os.path.join(data_path, 'transform', 'data.csv')
+    transform_input_key_prefix = 'integ-test-data/mxnet_mnist/transform'
+    transform_input = mx.sagemaker_session.upload_data(path=transform_input_path,
+                                                       key_prefix=transform_input_key_prefix)
+
+    transformer = mx.transformer(1, 'ml.m4.xlarge', tags=tags)
+    transformer.transform(transform_input, content_type='text/csv')
+
+    time.sleep(15)
+
+    latest_transform_job_name = transformer.latest_transform_job.name
+
+    print('Attempting to stop {}'.format(latest_transform_job_name))
+
+    transformer.stop_transform_job()
+
+    desc = transformer.latest_transform_job.sagemaker_session.sagemaker_client \
+        .describe_transform_job(TransformJobName=latest_transform_job_name)
+    assert desc['TransformJobStatus'] == 'Stopping'
+
+
+def _create_transformer_and_transform_job(estimator, transform_input, volume_kms_key=None,
+                                          input_filter=None, output_filter=None, join_source=None):
+    transformer = estimator.transformer(1, 'ml.m4.xlarge', volume_kms_key=volume_kms_key)
+    transformer.transform(transform_input, content_type='text/csv', input_filter=input_filter,
+                          output_filter=output_filter, join_source=join_source)
     return transformer

tests/unit/test_transformer.py

@@ -449,3 +449,18 @@ def test_restart_output_path(start_new_job, transformer, sagemaker_session):
 
     transformer.transform(DATA, job_name="job-2")
     assert transformer.output_path == "s3://{}/{}".format(S3_BUCKET, "job-2")
+
+
+def test_stop_transform_job(sagemaker_session, transformer):
+    sagemaker_session.stop_transform_job = Mock(name='stop_transform_job')
+    transformer.latest_transform_job = _TransformJob(sagemaker_session, JOB_NAME)
+
+    transformer.stop_transform_job()
+
+    sagemaker_session.stop_transform_job.assert_called_once_with(name=JOB_NAME)
+
+
+def test_stop_transform_job_no_transform_job(transformer):
+    with pytest.raises(ValueError) as e:
+        transformer.stop_transform_job()
+    assert 'No transform job available' in str(e)