feature: Estimator.fit like logs for transformer

Author: imujjwal96

src/sagemaker/session.py

@@ -1164,37 +1164,7 @@ def logs_for_job(self, job_name, wait=False, poll=10):  # noqa: C901 - suppress
         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
-                # may be dynamic until we have a stream for every instance.
-                try:
-                    streams = client.describe_log_streams(logGroupName=log_group, logStreamNamePrefix=job_name + '/',
-                                                          orderBy='LogStreamName', limit=instance_count)
-                    stream_names = [s['logStreamName'] for s in streams['logStreams']]
-                    positions.update([(s, sagemaker.logs.Position(timestamp=0, skip=0))
-                                      for s in stream_names if s not in positions])
-                except ClientError as e:
-                    # On the very first training job run on an account, there's no log group until
-                    # the container starts logging, so ignore any errors thrown about that
-                    err = e.response.get('Error', {})
-                    if err.get('Code', None) != 'ResourceNotFoundException':
-                        raise
-
-            if len(stream_names) > 0:
-                if dot:
-                    print('')
-                    dot = False
-                for idx, event in sagemaker.logs.multi_stream_iter(client, log_group, stream_names, positions):
-                    color_wrap(idx, event['message'])
-                    ts, count = positions[stream_names[idx]]
-                    if event['timestamp'] == ts:
-                        positions[stream_names[idx]] = sagemaker.logs.Position(timestamp=ts, skip=count + 1)
-                    else:
-                        positions[stream_names[idx]] = sagemaker.logs.Position(timestamp=event['timestamp'], skip=1)
-            else:
-                dot = True
-                print('.', end='')
-                sys.stdout.flush()
+            _flush_log_streams(stream_names, instance_count, client, log_group, job_name, positions, dot, color_wrap)
             if state == LogState.COMPLETE:
                 break
 
@@ -1225,6 +1195,87 @@ 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):  # 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.
+
+        Args:
+            job_name (str): Name of the transform job to display the logs for.
+            wait (bool): Whether to keep looking for new log entries until the job completes (default: False).
+            poll (int): The interval in seconds between polling for new log entries and job completion (default: 5).
+
+        Raises:
+            ValueError: If the transform job fails.
+        """
+
+        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
+
+        # Increase retries allowed (from default of 4), as we don't want waiting for a training job
+        # to be interrupted by a transient exception.
+        config = botocore.config.Config(retries={'max_attempts': 15})
+        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 = LogState.TAILING if wait and not job_already_completed else LogState.COMPLETE
+        dot = False
+
+        color_wrap = sagemaker.logs.ColorWrap()
+
+        # The loop below implements a state machine that alternates between checking the job status and
+        # reading whatever is available in the logs at this point. Note, that if we were called with
+        # wait == False, we never check the job status.
+        #
+        # If wait == TRUE and job is not completed, the initial state is TAILING
+        # If wait == FALSE, the initial state is COMPLETE (doesn't matter if the job really is complete).
+        #
+        # The state table:
+        #
+        # STATE               ACTIONS                        CONDITION             NEW STATE
+        # ----------------    ----------------               -----------------     ----------------
+        # TAILING             Read logs, Pause, Get status   Job complete          JOB_COMPLETE
+        #                                                    Else                  TAILING
+        # JOB_COMPLETE        Read logs, Pause               Any                   COMPLETE
+        # COMPLETE            Read logs, Exit                                      N/A
+        #
+        # Notes:
+        # - 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()
+        while True:
+            _flush_log_streams(stream_names, instance_count, client, log_group, job_name, positions, dot, color_wrap)
+            if state == LogState.COMPLETE:
+                break
+
+            time.sleep(poll)
+
+            if state == LogState.JOB_COMPLETE:
+                state = LogState.COMPLETE
+            elif time.time() - last_describe_job_call >= 30:
+                description = self.sagemaker_client.describe_transform_job(TransformJobName=job_name)
+                last_describe_job_call = time.time()
+
+                status = description['TransformJobStatus']
+
+                if status == 'Completed' or status == 'Failed' or status == 'Stopped':
+                    print()
+                    state = LogState.JOB_COMPLETE
+
+        if wait:
+            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):
     """Create a definition for executing a container as part of a SageMaker model.
@@ -1591,3 +1642,37 @@ def _vpc_config_from_training_job(training_job_desc, vpc_config_override=vpc_uti
         return training_job_desc.get(vpc_utils.VPC_CONFIG_KEY)
     else:
         return vpc_utils.sanitize(vpc_config_override)
+
+
+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
+        # may be dynamic until we have a stream for every instance.
+        try:
+            streams = client.describe_log_streams(logGroupName=log_group, logStreamNamePrefix=job_name + '/',
+                                                  orderBy='LogStreamName', limit=instance_count)
+            stream_names = [s['logStreamName'] for s in streams['logStreams']]
+            positions.update([(s, sagemaker.logs.Position(timestamp=0, skip=0))
+                              for s in stream_names if s not in positions])
+        except ClientError as e:
+            # On the very first training job run on an account, there's no log group until
+            # the container starts logging, so ignore any errors thrown about that
+            err = e.response.get('Error', {})
+            if err.get('Code', None) != 'ResourceNotFoundException':
+                raise
+
+    if len(stream_names) > 0:
+        if dot:
+            print('')
+            dot = False
+        for idx, event in sagemaker.logs.multi_stream_iter(client, log_group, stream_names, positions):
+            color_wrap(idx, event['message'])
+            ts, count = positions[stream_names[idx]]
+            if event['timestamp'] == ts:
+                positions[stream_names[idx]] = sagemaker.logs.Position(timestamp=ts, skip=count + 1)
+            else:
+                positions[stream_names[idx]] = sagemaker.logs.Position(timestamp=event['timestamp'], skip=1)
+    else:
+        dot = True
+        print('.', end='')
+        sys.stdout.flush()

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):
+                  job_name=None, wait=True, logs=True):
         """Start a new transform job.
 
         Args:
@@ -97,6 +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).
+            logs (bool): Whether to show the logs produced by the job.
+                Only meaningful when wait is True (default: True).
         """
         local_mode = self.sagemaker_session.local_mode
         if not local_mode and not data.startswith('s3://'):
@@ -113,6 +116,8 @@ def transform(self, data, data_type='S3Prefix', content_type=None, compression_t
 
         self.latest_transform_job = _TransformJob.start_new(self, data, data_type, content_type, compression_type,
                                                             split_type)
+        if wait:
+            self.latest_transform_job.wait(logs=logs)
 
     def delete_model(self):
         """Delete the corresponding SageMaker model for this Transformer.
@@ -130,9 +135,9 @@ def _retrieve_image_name(self):
                              'Local instance types require locally created models.'
                              % self.model_name)
 
-    def wait(self):
+    def wait(self, logs=True):
         self._ensure_last_transform_job()
-        self.latest_transform_job.wait()
+        self.latest_transform_job.wait(logs=logs)
 
     def _ensure_last_transform_job(self):
         if self.latest_transform_job is None:
@@ -205,8 +210,11 @@ def start_new(cls, transformer, data, data_type, content_type, compression_type,
 
         return cls(transformer.sagemaker_session, transformer._current_job_name)
 
-    def wait(self):
-        self.sagemaker_session.wait_for_transform_job(self.job_name)
+    def wait(self, logs=True):
+        if logs:
+            self.sagemaker_session.logs_for_transform_job(self.job_name, wait=True)
+        else:
+            self.sagemaker_session.wait_for_transform_job(self.job_name)
 
     @staticmethod
     def _load_config(data, data_type, content_type, compression_type, split_type, transformer):

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': {},