Add test cases for transformer logs

Author: imujjwal96

src/sagemaker/session.py

@@ -1225,7 +1225,7 @@ def logs_for_job(self, job_name, wait=False, poll=10):  # noqa: C901 - suppress
             billable_time = (description['TrainingEndTime'] - description['TrainingStartTime']) * instance_count
             print('Billable seconds:', int(billable_time.total_seconds()) + 1)
 
-    def logs_for_transform_job(self, job_name, wait=False, poll=10):
+    def logs_for_transform_job(self, job_name, wait=False, poll=10):  # noqa: C901 - suppress complexity warning
         """Display the logs for a given transform job, optionally tailing them until the
         job is complete. If the output is a tty or a Jupyter cell, it will be color-coded
         based on which instance the log entry is from.
@@ -1277,7 +1277,6 @@ def logs_for_transform_job(self, job_name, wait=False, poll=10):
         # - The JOB_COMPLETE state forces us to do an extra pause and read any items that got to Cloudwatch after
         #   the job was marked complete.
         last_describe_job_call = time.time()
-        last_description = description
         while True:
             if len(stream_names) < instance_count:
                 # Log streams are created whenever a container starts writing to stdout/err, so this list

tests/unit/test_session.py

@@ -308,6 +308,38 @@ def test_s3_input_all_arguments():
 IN_PROGRESS_DESCRIBE_JOB_RESULT = dict(DEFAULT_EXPECTED_TRAIN_JOB_ARGS)
 IN_PROGRESS_DESCRIBE_JOB_RESULT.update({'TrainingJobStatus': 'InProgress'})
 
+COMPLETED_DESCRIBE_TRANSFORM_JOB_RESULT = {
+    'TransformJobStatus': 'Completed',
+    'ModelName': 'some-model',
+    'TransformJobName': JOB_NAME,
+    'TransformResources': {
+        'InstanceCount': INSTANCE_COUNT,
+        'InstanceType': INSTANCE_TYPE
+    },
+    'TransformEndTime': datetime.datetime(2018, 2, 17, 7, 19, 34, 953000),
+    'TransformStartTime': datetime.datetime(2018, 2, 17, 7, 15, 0, 103000),
+    'TransformOutput': {
+        'AssembleWith': 'None',
+        'KmsKeyId': '',
+        'S3OutputPath': S3_OUTPUT
+    },
+    'TransformInput': {
+        'CompressionType': 'None',
+        'ContentType': 'text/csv',
+        'DataSource': {
+            'S3DataType': 'S3Prefix',
+            'S3Uri': S3_INPUT_URI
+        },
+        'SplitType': 'Line'
+    }
+}
+
+STOPPED_DESCRIBE_TRANSFORM_JOB_RESULT = dict(COMPLETED_DESCRIBE_TRANSFORM_JOB_RESULT)
+STOPPED_DESCRIBE_TRANSFORM_JOB_RESULT.update({'TransformJobStatus': 'Stopped'})
+
+IN_PROGRESS_DESCRIBE_TRANSFORM_JOB_RESULT = dict(COMPLETED_DESCRIBE_TRANSFORM_JOB_RESULT)
+IN_PROGRESS_DESCRIBE_TRANSFORM_JOB_RESULT.update({'TransformJobStatus': 'InProgress'})
+
 
 @pytest.fixture()
 def sagemaker_session():
@@ -653,6 +685,7 @@ def sagemaker_session_complete():
     boto_mock.client('logs').get_log_events.side_effect = DEFAULT_LOG_EVENTS
     ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
     ims.sagemaker_client.describe_training_job.return_value = COMPLETED_DESCRIBE_JOB_RESULT
+    ims.sagemaker_client.describe_transform_job.return_value = COMPLETED_DESCRIBE_TRANSFORM_JOB_RESULT
     return ims
 
 
@@ -663,6 +696,7 @@ def sagemaker_session_stopped():
     boto_mock.client('logs').get_log_events.side_effect = DEFAULT_LOG_EVENTS
     ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=Mock())
     ims.sagemaker_client.describe_training_job.return_value = STOPPED_DESCRIBE_JOB_RESULT
+    ims.sagemaker_client.describe_transform_job.return_value = STOPPED_DESCRIBE_TRANSFORM_JOB_RESULT
     return ims
 
 
@@ -675,6 +709,9 @@ def sagemaker_session_ready_lifecycle():
     ims.sagemaker_client.describe_training_job.side_effect = [IN_PROGRESS_DESCRIBE_JOB_RESULT,
                                                               IN_PROGRESS_DESCRIBE_JOB_RESULT,
                                                               COMPLETED_DESCRIBE_JOB_RESULT]
+    ims.sagemaker_client.describe_transform_job.side_effect = [IN_PROGRESS_DESCRIBE_TRANSFORM_JOB_RESULT,
+                                                               IN_PROGRESS_DESCRIBE_TRANSFORM_JOB_RESULT,
+                                                               COMPLETED_DESCRIBE_TRANSFORM_JOB_RESULT]
     return ims
 
 
@@ -687,6 +724,9 @@ def sagemaker_session_full_lifecycle():
     ims.sagemaker_client.describe_training_job.side_effect = [IN_PROGRESS_DESCRIBE_JOB_RESULT,
                                                               IN_PROGRESS_DESCRIBE_JOB_RESULT,
                                                               COMPLETED_DESCRIBE_JOB_RESULT]
+    ims.sagemaker_client.describe_transform_job.side_effect = [IN_PROGRESS_DESCRIBE_TRANSFORM_JOB_RESULT,
+                                                               IN_PROGRESS_DESCRIBE_TRANSFORM_JOB_RESULT,
+                                                               COMPLETED_DESCRIBE_TRANSFORM_JOB_RESULT]
     return ims
 
 
@@ -740,6 +780,56 @@ def test_logs_for_job_full_lifecycle(time, cw, sagemaker_session_full_lifecycle)
                                    call(0, 'hi there #2a'), call(0, 'hi there #3')]
 
 
+@patch('sagemaker.logs.ColorWrap')
+def test_logs_for_transform_job_no_wait(cw, sagemaker_session_complete):
+    ims = sagemaker_session_complete
+    ims.logs_for_transform_job(JOB_NAME)
+    ims.sagemaker_client.describe_transform_job.assert_called_once_with(TransformJobName=JOB_NAME)
+    cw().assert_called_with(0, 'hi there #1')
+
+
+@patch('sagemaker.logs.ColorWrap')
+def test_logs_for_transform_job_no_wait_stopped_job(cw, sagemaker_session_stopped):
+    ims = sagemaker_session_stopped
+    ims.logs_for_transform_job(JOB_NAME)
+    ims.sagemaker_client.describe_transform_job.assert_called_once_with(TransformJobName=JOB_NAME)
+    cw().assert_called_with(0, 'hi there #1')
+
+
+@patch('sagemaker.logs.ColorWrap')
+def test_logs_for_transform_job_wait_on_completed(cw, sagemaker_session_complete):
+    ims = sagemaker_session_complete
+    ims.logs_for_transform_job(JOB_NAME, wait=True, poll=0)
+    assert ims.sagemaker_client.describe_transform_job.call_args_list == [call(TransformJobName=JOB_NAME,)]
+    cw().assert_called_with(0, 'hi there #1')
+
+
+@patch('sagemaker.logs.ColorWrap')
+def test_logs_for_transform_job_wait_on_stopped(cw, sagemaker_session_stopped):
+    ims = sagemaker_session_stopped
+    ims.logs_for_transform_job(JOB_NAME, wait=True, poll=0)
+    assert ims.sagemaker_client.describe_transform_job.call_args_list == [call(TransformJobName=JOB_NAME,)]
+    cw().assert_called_with(0, 'hi there #1')
+
+
+@patch('sagemaker.logs.ColorWrap')
+def test_logs_for_transform_job_no_wait_on_running(cw, sagemaker_session_ready_lifecycle):
+    ims = sagemaker_session_ready_lifecycle
+    ims.logs_for_transform_job(JOB_NAME)
+    assert ims.sagemaker_client.describe_transform_job.call_args_list == [call(TransformJobName=JOB_NAME,)]
+    cw().assert_called_with(0, 'hi there #1')
+
+
+@patch('sagemaker.logs.ColorWrap')
+@patch('time.time', side_effect=[0, 30, 60, 90, 120, 150, 180])
+def test_logs_for_transform_job_full_lifecycle(time, cw, sagemaker_session_full_lifecycle):
+    ims = sagemaker_session_full_lifecycle
+    ims.logs_for_transform_job(JOB_NAME, wait=True, poll=0)
+    assert ims.sagemaker_client.describe_transform_job.call_args_list == [call(TransformJobName=JOB_NAME,)] * 3
+    assert cw().call_args_list == [call(0, 'hi there #1'), call(0, 'hi there #2'),
+                                   call(0, 'hi there #2a'), call(0, 'hi there #3')]
+
+
 MODEL_NAME = 'some-model'
 PRIMARY_CONTAINER = {
     'Environment': {},