From f17f77902669db4d07f930e70b2e07fd76e9528f Mon Sep 17 00:00:00 2001 From: Rui Wang Date: Wed, 1 Aug 2018 11:27:58 -0700 Subject: [PATCH] Allow SecondaryStatusTransitions to be missing in secondary_training_status_message --- src/sagemaker/utils.py | 5 ++++- tests/unit/test_utils.py | 11 +++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/src/sagemaker/utils.py b/src/sagemaker/utils.py index f91a614d6a..f83c047cf2 100644 --- a/src/sagemaker/utils.py +++ b/src/sagemaker/utils.py @@ -173,7 +173,10 @@ def secondary_training_status_message(job_description, prev_description): or len(job_description.get('SecondaryStatusTransitions')) == 0: return '' - prev_transitions_num = len(prev_description['SecondaryStatusTransitions']) if prev_description is not None else 0 + prev_description_secondary_transitions = prev_description.get('SecondaryStatusTransitions')\ + if prev_description is not None else None + prev_transitions_num = len(prev_description['SecondaryStatusTransitions'])\ + if prev_description_secondary_transitions is not None else 0 current_transitions = job_description['SecondaryStatusTransitions'] if len(current_transitions) == prev_transitions_num: diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a2cf1081f4..d8b0bc67ec 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -155,3 +155,14 @@ 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 + + +def test_secondary_training_status_message_prev_missing(): + now = datetime.now() + TRAINING_JOB_DESCRIPTION_1['SecondaryStatusTransitions'][-1]['StartTime'] = now + 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, {}) == expected