diff --git a/virtualizarr/parallel.py b/virtualizarr/parallel.py index f067189d..b8774ccb 100644 --- a/virtualizarr/parallel.py +++ b/virtualizarr/parallel.py @@ -302,12 +302,16 @@ def map( ------- Generator of results """ - import lithops # type: ignore[import-untyped] + from lithops.retries import ( + RetryingFunctionExecutor, # type: ignore[import-untyped] + ) - fexec = lithops.FunctionExecutor() + with RetryingFunctionExecutor(self.lithops_client) as fexec: + # TODO retries should be exposed as a configuration arg in lithops, see https://github.com/lithops-cloud/lithops/issues/1412 + futures = fexec.map(fn, *iterables, retries=2) + finished_futures, pending = fexec.wait(futures) - futures = fexec.map(fn, *iterables) - results = fexec.get_result(futures) + results = [future.result() for future in finished_futures] return results