feature: Estimator.fit like logs for transformer

imujjwal96 · imujjwal96 · commit 813155eb3dd0 · 2019-06-28T08:32:40.000Z
diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py
@@ -1278,7 +1278,7 @@ def logs_for_job(  # noqa: C901 - suppress complexity warning for this method
             poll (int): The interval in seconds between polling for new log entries and job completion (default: 5).
 
         Raises:
-            ValueError: If waiting and the training job fails.
+            ValueError: If the training job fails.
         """
 
         description = self.sagemaker_client.describe_training_job(TrainingJobName=job_name)
@@ -1326,52 +1326,7 @@ def logs_for_job(  # noqa: C901 - suppress complexity warning for this method
         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
 
@@ -1404,6 +1359,87 @@ def logs_for_job(  # noqa: C901 - suppress complexity warning for this method
             ) * 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.
@@ -1795,3 +1831,37 @@ def _vpc_config_from_training_job(
         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()
diff --git a/src/sagemaker/transformer.py b/src/sagemaker/transformer.py
@@ -104,6 +104,8 @@ def transform(
         input_filter=None,
         output_filter=None,
         join_source=None,
+        wait=False,
+        logs=False
     ):
         """Start a new transform job.
 
@@ -131,6 +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).
+            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://"):
@@ -163,6 +168,9 @@ def transform(
             join_source,
         )
 
+        if wait:
+            self.latest_transform_job.wait(logs=logs)
+
     def delete_model(self):
         """Delete the corresponding SageMaker model for this Transformer.
 
@@ -200,9 +208,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:
@@ -300,8 +308,11 @@ def start_new(
 
         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):
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
@@ -340,6 +340,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():
@@ -787,6 +819,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
 
 
@@ -797,6 +830,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
 
 
@@ -811,6 +845,11 @@ def sagemaker_session_ready_lifecycle():
         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
 
 
@@ -825,6 +864,11 @@ def sagemaker_session_full_lifecycle():
         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
 
 
@@ -892,6 +936,60 @@ def test_logs_for_job_full_lifecycle(time, cw, sagemaker_session_full_lifecycle)
 
 
 MODEL_NAME = "some-model"
+
+
+@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'
+>>>>>>> feature: Estimator.fit like logs for transformer
 PRIMARY_CONTAINER = {
     "Environment": {},
     "Image": IMAGE,
