diff --git a/darwin/future/exceptions.py b/darwin/future/exceptions.py index 0a77bada2..df48d8424 100644 --- a/darwin/future/exceptions.py +++ b/darwin/future/exceptions.py @@ -135,3 +135,9 @@ class DatasetNotFound(DarwinException): """Raised when the dataset endpoint returns a malformed response.""" ... + + +class MaxRetriesError(DarwinException): + """Raised when a certain API call is re-tried for {x} number of times.""" + + ... diff --git a/darwin/future/meta/objects/stage.py b/darwin/future/meta/objects/stage.py index 3e9fdbc7f..fd49fec01 100644 --- a/darwin/future/meta/objects/stage.py +++ b/darwin/future/meta/objects/stage.py @@ -7,6 +7,7 @@ from darwin.future.core.items import get_item, move_items_to_stage from darwin.future.core.types.query import QueryFilter from darwin.future.data_objects.workflow import WFEdgeCore, WFStageCore +from darwin.future.exceptions import MaxRetriesError from darwin.future.meta.objects.base import MetaBase from darwin.future.meta.queries.item_id import ItemIDQuery @@ -83,16 +84,28 @@ def move_attached_files_to_stage( ) ids = [x.id for x in self.item_ids.collect_all()] + def all_items_complete(item_ids: list[UUID]) -> bool: + for item_id in item_ids: + if get_item(self.client, slug, item_id).processing_status != "complete": + return False + return True + if wait: - while True: - for _id in ids: - if get_item(self.client, slug, _id).processing_status != "complete": - # wait for 0.5 seconds before checking again - time.sleep(0.5) - break - else: - # All items are complete, break the while loop + max_attempts = 5 + wait_time = 0.5 # seconds + + # Try to check if all items are complete for a maximum number of attempts + for attempt in range(1, max_attempts + 1): + complete = all_items_complete(item_ids=ids) + if complete: break + else: + # The wait time increases with each attempt + time.sleep(wait_time * attempt) + else: + raise MaxRetriesError( + "Max retry attempts to check for Item(s) to complete processing" + ) move_items_to_stage(self.client, slug, w_id, d_id, new_stage_id, ids) return self