Skip to content

Commit f4021f9

Browse files
authored
Allow SecondaryStatusTransitions to be missing in secondary_training_status_message (#331)
1 parent f650e11 commit f4021f9

File tree

2 files changed

+15
-1
lines changed

2 files changed

+15
-1
lines changed

src/sagemaker/utils.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,10 @@ def secondary_training_status_message(job_description, prev_description):
173173
or len(job_description.get('SecondaryStatusTransitions')) == 0:
174174
return ''
175175

176-
prev_transitions_num = len(prev_description['SecondaryStatusTransitions']) if prev_description is not None else 0
176+
prev_description_secondary_transitions = prev_description.get('SecondaryStatusTransitions')\
177+
if prev_description is not None else None
178+
prev_transitions_num = len(prev_description['SecondaryStatusTransitions'])\
179+
if prev_description_secondary_transitions is not None else 0
177180
current_transitions = job_description['SecondaryStatusTransitions']
178181

179182
if len(current_transitions) == prev_transitions_num:

tests/unit/test_utils.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,3 +155,14 @@ def test_secondary_training_status_message_status_not_changed():
155155
now = datetime.now()
156156
TRAINING_JOB_DESCRIPTION_1['SecondaryStatusTransitions'][-1]['StartTime'] = now
157157
assert secondary_training_status_message(TRAINING_JOB_DESCRIPTION_1, TRAINING_JOB_DESCRIPTION_2) == MESSAGE
158+
159+
160+
def test_secondary_training_status_message_prev_missing():
161+
now = datetime.now()
162+
TRAINING_JOB_DESCRIPTION_1['SecondaryStatusTransitions'][-1]['StartTime'] = now
163+
expected = '{} {} - {}'.format(
164+
datetime.utcfromtimestamp(time.mktime(now.timetuple())).strftime('%Y-%m-%d %H:%M:%S'),
165+
STATUS,
166+
MESSAGE
167+
)
168+
assert secondary_training_status_message(TRAINING_JOB_DESCRIPTION_1, {}) == expected

0 commit comments

Comments
 (0)