Print secondary status with status message (aws#78)

Author: icywang86rui

src/sagemaker/session.py

@@ -796,7 +796,7 @@ 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))
+        print(secondary_training_status_message(description, None), end='')
         instance_count = description['ResourceConfig']['InstanceCount']
         status = description['TrainingJobStatus']
 

src/sagemaker/utils.py

@@ -173,18 +173,17 @@ def secondary_training_status_message(job_description, prev_description):
     current_transitions = job_description['SecondaryStatusTransitions']
 
     if len(current_transitions) == prev_transitions_num:
-        transitions_to_print = current_transitions[-1:]
+        return current_transitions[-1]['StatusMessage']
     else:
         transitions_to_print = current_transitions[prev_transitions_num - len(current_transitions):]
-
-    status_strs = []
-    for transition in transitions_to_print:
-        message = transition['StatusMessage']
-        time_str = datetime.utcfromtimestamp(
-            time.mktime(transition['StartTime'].timetuple())).strftime('%Y-%m-%d %H:%M:%S')
-        status_strs.append('{} {}'.format(time_str, message))
-
-    return '\n'.join(status_strs)
+        status_strs = []
+        for transition in transitions_to_print:
+            message = transition['StatusMessage']
+            time_str = datetime.utcfromtimestamp(
+                time.mktime(transition['StartTime'].timetuple())).strftime('%Y-%m-%d %H:%M:%S')
+            status_strs.append('{} {} - {}'.format(time_str, transition['Status'], message))
+
+        return '\n'.join(status_strs)
 
 
 class DeferredError(object):

tests/unit/test_utils.py

@@ -97,11 +97,12 @@ def test_name_from_training_arn():
 
 
 MESSAGE = 'message'
+STATUS = 'status'
 TRAINING_JOB_DESCRIPTION_1 = {
-    'SecondaryStatusTransitions': [{'StatusMessage': MESSAGE}]
+    'SecondaryStatusTransitions': [{'StatusMessage': MESSAGE, 'Status': STATUS}]
 }
 TRAINING_JOB_DESCRIPTION_2 = {
-    'SecondaryStatusTransitions': [{'StatusMessage': 'different message'}]
+    'SecondaryStatusTransitions': [{'StatusMessage': 'different message', 'Status': STATUS}]
 }
 TRAINING_JOB_DESCRIPTION_EMPTY = {
     'SecondaryStatusTransitions': []
@@ -123,11 +124,18 @@ def test_secondary_training_status_changed_empty():
     assert changed is False
 
 
-def test_secondary_training_status_message():
+def test_secondary_training_status_message_status_changed():
     now = datetime.now()
     TRAINING_JOB_DESCRIPTION_1['SecondaryStatusTransitions'][-1]['StartTime'] = now
-    expected = '{} {}'.format(
+    expected = '{} {} - {}'.format(
         datetime.utcfromtimestamp(time.mktime(now.timetuple())).strftime('%Y-%m-%d %H:%M:%S'),
+        STATUS,
         MESSAGE
     )
-    assert secondary_training_status_message(TRAINING_JOB_DESCRIPTION_1, TRAINING_JOB_DESCRIPTION_2) == expected
+    assert secondary_training_status_message(TRAINING_JOB_DESCRIPTION_1, TRAINING_JOB_DESCRIPTION_EMPTY) == expected
+
+
+def test_secondary_training_status_message_status_not_changed():
+    now = datetime.now()
+    TRAINING_JOB_DESCRIPTION_1['SecondaryStatusTransitions'][-1]['StartTime'] = now
+    assert secondary_training_status_message(TRAINING_JOB_DESCRIPTION_1, TRAINING_JOB_DESCRIPTION_2) == MESSAGE