Add integrations tests and refactor

Author: imujjwal96

src/sagemaker/session.py

@@ -1124,7 +1124,6 @@ def logs_for_job(self, job_name, wait=False, poll=10):  # noqa: C901 - suppress
         description = self.sagemaker_client.describe_training_job(TrainingJobName=job_name)
         print(secondary_training_status_message(description, None), end='')
         instance_count = description['ResourceConfig']['InstanceCount']
-        status = description['TrainingJobStatus']
 
         stream_names = []  # The list of log streams
         positions = {}     # The current position in each stream, map of stream name -> position
@@ -1135,9 +1134,8 @@ def logs_for_job(self, job_name, wait=False, poll=10):  # noqa: C901 - suppress
         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()
@@ -1211,7 +1209,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
@@ -1222,9 +1219,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()
@@ -1272,9 +1268,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):
@@ -1644,6 +1637,12 @@ def _vpc_config_from_training_job(training_job_desc, vpc_config_override=vpc_uti
         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

@@ -79,7 +79,7 @@ def __init__(self, model_name, instance_count, instance_type, strategy=None, ass
         self.sagemaker_session = sagemaker_session or Session()
 
     def transform(self, data, data_type='S3Prefix', content_type=None, compression_type=None, split_type=None,
-                  job_name=None, wait=True, logs=True):
+                  job_name=None, wait=False, logs=False):
         """Start a new transform job.
 
         Args:
@@ -97,9 +97,9 @@ def transform(self, data, data_type='S3Prefix', content_type=None, compression_t
             split_type (str): The record delimiter for the input object (default: 'None').
                 Valid values: 'None', 'Line', 'RecordIO', and 'TFRecord'.
             job_name (str): job name (default: None). If not specified, one will be generated.
-            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

@@ -148,7 +148,37 @@ def test_transform_mxnet_vpc(sagemaker_session, mxnet_full_version):
         assert [security_group_id] == model_desc['VpcConfig']['SecurityGroupIds']
 
 
-def _create_transformer_and_transform_job(estimator, transform_input, volume_kms_key=None):
+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, wait=False, logs=False):
     transformer = estimator.transformer(1, 'ml.m4.xlarge', volume_kms_key=volume_kms_key)
-    transformer.transform(transform_input, content_type='text/csv')
+    transformer.transform(transform_input, content_type='text/csv', wait=wait, logs=logs)
     return transformer