Check if SecondaryStatusTransitions for previous DescribeTrainingJob call (aws#325)

Author: icywang86rui

src/sagemaker/utils.py

@@ -146,8 +146,12 @@ def secondary_training_status_changed(current_job_description, prev_job_descript
     if current_secondary_status_transitions is None or len(current_secondary_status_transitions) == 0:
         return False
 
-    last_message = prev_job_description['SecondaryStatusTransitions'][-1]['StatusMessage']\
-        if prev_job_description is not None else ''
+    prev_job_secondary_status_transitions = prev_job_description.get('SecondaryStatusTransitions') \
+        if prev_job_description is not None else None
+
+    last_message = prev_job_secondary_status_transitions[-1]['StatusMessage'] \
+        if prev_job_secondary_status_transitions is not None and len(prev_job_secondary_status_transitions) > 0 else ''
+
     message = current_job_description['SecondaryStatusTransitions'][-1]['StatusMessage']
 
     return message != last_message

tests/unit/test_utils.py

@@ -104,6 +104,7 @@ def test_name_from_training_arn():
 TRAINING_JOB_DESCRIPTION_2 = {
     'SecondaryStatusTransitions': [{'StatusMessage': 'different message', 'Status': STATUS}]
 }
+
 TRAINING_JOB_DESCRIPTION_EMPTY = {
     'SecondaryStatusTransitions': []
 }
@@ -119,6 +120,21 @@ def test_secondary_training_status_changed_false():
     assert changed is False
 
 
+def test_secondary_training_status_changed_prev_missing():
+    changed = secondary_training_status_changed(TRAINING_JOB_DESCRIPTION_1, {})
+    assert changed is True
+
+
+def test_secondary_training_status_changed_prev_none():
+    changed = secondary_training_status_changed(TRAINING_JOB_DESCRIPTION_1, None)
+    assert changed is True
+
+
+def test_secondary_training_status_changed_current_missing():
+    changed = secondary_training_status_changed({}, TRAINING_JOB_DESCRIPTION_1)
+    assert changed is False
+
+
 def test_secondary_training_status_changed_empty():
     changed = secondary_training_status_changed(TRAINING_JOB_DESCRIPTION_EMPTY, TRAINING_JOB_DESCRIPTION_1)
     assert changed is False