black check

Author: imujjwal96

src/sagemaker/session.py

@@ -1062,15 +1062,17 @@ def stop_transform_job(self, name):
             ClientError: If an error occurs while trying to stop the batch transform job.
         """
         try:
-            LOGGER.info('Stopping transform job: {}'.format(name))
+            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']
+            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))
+            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))
+                LOGGER.error(
+                    "Error occurred while attempting to stop transform job: {}.".format(name)
+                )
                 raise
 
     def _check_job_status(self, job, desc, status_key_name):

src/sagemaker/transformer.py

@@ -228,11 +228,13 @@ def wait(self):
         self._ensure_last_transform_job()
         self.latest_transform_job.wait()
 
-    def stop_transform_job(self):
+    def stop_transform_job(self, wait=True):
         """Stop latest running batch transform job.
         """
         self._ensure_last_transform_job()
         self.latest_transform_job.stop()
+        if wait:
+            self.latest_transform_job.wait()
 
     def _ensure_last_transform_job(self):
         """Placeholder docstring"""
@@ -352,6 +354,7 @@ def wait(self):
         self.sagemaker_session.wait_for_transform_job(self.job_name)
 
     def stop(self):
+        """Placeholder docstring"""
         self.sagemaker_session.stop_transform_job(name=self.job_name)
 
     @staticmethod

tests/integ/test_transformer.py

@@ -351,66 +351,68 @@ def test_single_transformer_multiple_jobs(sagemaker_session, mxnet_full_version)
         )
 
 
-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,
-    )
-
-
 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'}]
+    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)
+    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')
+    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)
+        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)
+    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')
+    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))
+    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'
+    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)
+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

@@ -452,7 +452,7 @@ def test_restart_output_path(start_new_job, transformer, sagemaker_session):
 
 
 def test_stop_transform_job(sagemaker_session, transformer):
-    sagemaker_session.stop_transform_job = Mock(name='stop_transform_job')
+    sagemaker_session.stop_transform_job = Mock(name="stop_transform_job")
     transformer.latest_transform_job = _TransformJob(sagemaker_session, JOB_NAME)
 
     transformer.stop_transform_job()
@@ -463,4 +463,4 @@ def test_stop_transform_job(sagemaker_session, transformer):
 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)
+    assert "No transform job available" in str(e)