From 46cd4e596035d787feaa8a61b17dff7f0a1bf09d Mon Sep 17 00:00:00 2001 From: Ranbir Aulakh Date: Tue, 24 Dec 2024 15:58:36 +0000 Subject: [PATCH] fix: check if processing_duration is None --- src/aws/osml/utils/integ_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/aws/osml/utils/integ_utils.py b/src/aws/osml/utils/integ_utils.py index 793c338..88726eb 100644 --- a/src/aws/osml/utils/integ_utils.py +++ b/src/aws/osml/utils/integ_utils.py @@ -117,7 +117,10 @@ def monitor_job_status(sqs_client: boto3.resource, image_id: str) -> None: logging.info("\tIN_PROGRESS message found! Waiting for SUCCESS message...") elif message_image_status == "SUCCESS" and message_image_id == image_id: processing_duration = message_attributes.get("processing_duration", {}).get("Value") - assert float(processing_duration) > 0 + if processing_duration is not None: + assert float(processing_duration) > 0 + else: + logging.warning("Processing duration is None for SUCCESS message") done = True logging.info(f"\tSUCCESS message found! Image took {processing_duration} seconds to process") elif (