adding easy callback to check number of pending tasks on thread workers; this can be used to guarantee the download gets finished

Author: HansVRP

openeo/extra/job_management/_manager.py

@@ -400,6 +400,14 @@ def stop_job_thread(self, timeout_seconds: Optional[float] = _UNSET):
             self._download_pool.shutdown()
             self._download_pool = None
 
+        if self._download_pool is not None:
+        # Wait for downloads to complete before shutting down
+            _log.info("Waiting for download tasks to complete before stopping...")
+            while  self._download_pool.num_pending_tasks() > 0:
+                time.sleep(0.5)
+            self._download_pool.shutdown()
+            self._download_pool = None
+
         if self._thread is not None:
             self._stop_thread = True
             if timeout_seconds is _UNSET:
@@ -513,6 +521,8 @@ def run_jobs(
                 ).values()
             )
             > 0
+
+            or self._worker_pool.num_pending_tasks() > 0
         ):
             self._job_update_loop(job_db=job_db, start_job=start_job, stats=stats)
             stats["run_jobs loop"] += 1
@@ -523,7 +533,20 @@ def run_jobs(
             stats["sleep"] += 1
 
         # TODO; run post process after shutdown once more to ensure completion?
-        self.stop_job_thread()
+        # Wait for all download tasks to complete
+        if self._download_results and self._download_pool is not None:
+            _log.info("Waiting for download tasks to complete...")
+            while self._download_pool.num_pending_tasks() > 0:
+                self._process_threadworker_updates(
+                    worker_pool=self._download_pool, 
+                    job_db=job_db, 
+                    stats=stats
+                )
+                time.sleep(1)  # Brief pause to avoid busy waiting
+            _log.info("All download tasks completed.")
+
+        self._worker_pool.shutdown()
+        self._download_pool.shutdown()
 
         return stats
 

openeo/extra/job_management/_thread_worker.py

@@ -251,6 +251,10 @@ def process_futures(self, timeout: Union[float, None] = 0) -> Tuple[List[_TaskRe
 
         self._future_task_pairs = to_keep
         return results, len(to_keep)
+    
+    def num_pending_tasks(self) -> int:
+        """Return the number of tasks that are still pending (not completed)."""
+        return len(self._future_task_pairs)
 
     def shutdown(self) -> None:
         """Shuts down the thread pool gracefully."""