wait for status update on partial retries (#360)

Co-authored-by: Matic Lubej <[email protected]>

Author: mlubej

eogrow/pipelines/download_batch.py

@@ -205,6 +205,7 @@ def run_procedure(self) -> tuple[list[str], list[str]]:
                 break
             LOGGER.info("Retrying due to PARTIAL status.")
             self.batch_client.restart_job(batch_request)
+            batch_request = self._wait_for_partial_status_update(batch_request)
             results = self._monitor_job(batch_request)
 
         processed = self._get_tile_names_from_results(results, BatchTileStatus.PROCESSED)
@@ -214,6 +215,17 @@ def run_procedure(self) -> tuple[list[str], list[str]]:
         LOGGER.info(log_msg)
         return processed, failed
 
+    def _wait_for_partial_status_update(self, batch_request: BatchRequest) -> BatchRequest:
+        """Wait for the batch job to update the status to PARTIAL with an exponential backoff."""
+        wait_time, max_wait_time = 1, 5 * 60
+        while wait_time < max_wait_time:
+            time.sleep(wait_time)
+            wait_time *= 2
+            if self.batch_client.get_request(batch_request).status != BatchRequestStatus.PARTIAL:
+                break
+
+        return self.batch_client.get_request(batch_request)
+
     def _create_or_collect_batch_request(self) -> BatchRequest:
         """Either creates a new batch request or collects information about an existing one."""
         if not self.config.batch_id: