Add integrations tests and refactor

Author: imujjwal96

src/sagemaker/session.py

@@ -1295,11 +1295,8 @@ def logs_for_job(  # noqa: C901 - suppress complexity warning for this method
         client = self.boto_session.client("logs", config=config)
         log_group = "/aws/sagemaker/TrainingJobs"
 
-        job_already_completed = (
-            True if status == "Completed" or status == "Failed" or status == "Stopped" else False
-        )
+        state = _get_initial_job_state(description, 'TrainingJobStatus', wait)
 
-        state = LogState.TAILING if wait and not job_already_completed else LogState.COMPLETE
         dot = False
 
         color_wrap = sagemaker.logs.ColorWrap()
@@ -1375,7 +1372,6 @@ def logs_for_transform_job(self, job_name, wait=False, poll=10):  # noqa: C901 -
 
         description = self.sagemaker_client.describe_transform_job(TransformJobName=job_name)
         instance_count = description['TransformResources']['InstanceCount']
-        status = description['TransformJobStatus']
 
         stream_names = []  # The list of log streams
         positions = {}     # The current position in each stream, map of stream name -> position
@@ -1386,9 +1382,8 @@ def logs_for_transform_job(self, job_name, wait=False, poll=10):  # noqa: C901 -
         client = self.boto_session.client('logs', config=config)
         log_group = '/aws/sagemaker/TransformJobs'
 
-        job_already_completed = True if status == 'Completed' or status == 'Failed' or status == 'Stopped' else False
+        state = _get_initial_job_state(description, 'TransformJobStatus', wait)
 
-        state = LogState.TAILING if wait and not job_already_completed else LogState.COMPLETE
         dot = False
 
         color_wrap = sagemaker.logs.ColorWrap()
@@ -1436,9 +1431,6 @@ def logs_for_transform_job(self, job_name, wait=False, poll=10):  # noqa: C901 -
             self._check_job_status(job_name, description, 'TransformJobStatus')
             if dot:
                 print()
-            # Customers are not billed for hardware provisioning, so billable time is less than total time
-            billable_time = (description['TransformEndTime'] - description['TransformStartTime']) * instance_count
-            print('Billable seconds:', int(billable_time.total_seconds()) + 1)
 
 
 def container_def(image, model_data_url=None, env=None):
@@ -1833,6 +1825,12 @@ def _vpc_config_from_training_job(
         return vpc_utils.sanitize(vpc_config_override)
 
 
+def _get_initial_job_state(description, status_key, wait):
+    status = description[status_key]
+    job_already_completed = True if status == 'Completed' or status == 'Failed' or status == 'Stopped' else False
+    return LogState.TAILING if wait and not job_already_completed else LogState.COMPLETE
+
+
 def _flush_log_streams(stream_names, instance_count, client, log_group, job_name, positions, dot, color_wrap):
     if len(stream_names) < instance_count:
         # Log streams are created whenever a container starts writing to stdout/err, so this list

src/sagemaker/transformer.py

@@ -133,9 +133,9 @@ def transform(
                 meaning the entire input record will be joined to the inference result.
                 You can use OutputFilter to select the useful portion before uploading to S3. (default: None).
                 Valid values: Input, None.
-            wait (bool): Whether the call should wait until the job completes (default: True).
+            wait (bool): Whether the call should wait until the job completes (default: False).
             logs (bool): Whether to show the logs produced by the job.
-                Only meaningful when wait is True (default: True).
+                Only meaningful when wait is True (default: False).
         """
         local_mode = self.sagemaker_session.local_mode
         if not local_mode and not data.startswith("s3://"):

tests/integ/test_transformer.py

@@ -301,13 +301,45 @@ def test_transform_byo_estimator(sagemaker_session):
         assert tags == model_tags
 
 
+def test_transform_mxnet_logs(sagemaker_session, mxnet_full_version):
+    data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
+    script_path = os.path.join(data_path, 'mnist.py')
+
+    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)
+
+    with timeout(minutes=45):
+        transformer = _create_transformer_and_transform_job(mx, transform_input, wait=True, logs=True)
+
+    with timeout_and_delete_model_with_transformer(transformer, sagemaker_session,
+                                                   minutes=TRANSFORM_DEFAULT_TIMEOUT_MINUTES):
+        transformer.wait()
+
+
 def _create_transformer_and_transform_job(
     estimator,
     transform_input,
     volume_kms_key=None,
     input_filter=None,
     output_filter=None,
     join_source=None,
+    wait=False,
+    logs=False,
 ):
     transformer = estimator.transformer(1, "ml.m4.xlarge", volume_kms_key=volume_kms_key)
     transformer.transform(
@@ -316,5 +348,8 @@ def _create_transformer_and_transform_job(
         input_filter=input_filter,
         output_filter=output_filter,
         join_source=join_source,
+        wait=wait,
+        logs=logs,
     )
     return transformer
+