Merge pull request #435 from ungarj/future_batches

use future batches when yielding dask futures
ungarj · Mar 1, 2022 · 95bffc6 · 95bffc6
2 parents d787b46 + 1f34a2f
commit 95bffc6
Showing 1 changed file with 13 additions and 12 deletions.
diff --git a/mapchete/_processing.py b/mapchete/_processing.py
@@ -554,23 +554,24 @@ def _compute_task_graph(
     logger.debug("%s tasks sent to scheduler in %s", len(futures), t)
 
     logger.debug("wait for tasks to finish...")
-    for future in as_completed(
+    for batch in as_completed(
         futures,
         with_results=with_results,
         raise_errors=raise_errors,
         loop=executor._executor.loop,
-    ):
-        futures.remove(future)
-        if process.config.output.write_in_parent_process:
-            yield FinishedFuture(
-                result=_write(
-                    process_info=future.result(),
-                    output_writer=process.config.output,
-                    append_output=True,
+    ).batches():
+        for future in batch:
+            futures.remove(future)
+            if process.config.output.write_in_parent_process:
+                yield FinishedFuture(
+                    result=_write(
+                        process_info=future.result(),
+                        output_writer=process.config.output,
+                        append_output=True,
+                    )
                 )
-            )
-        else:
-            yield future
+            else:
+                yield future
 
 
 def _compute_tasks(