Repository: tomwhite/cubed
Pull request #291 (2023-08-14 09:57:46)

Add support for batching inputs to avoid overwhelming the backend service
See #239 and #276 


Patch:
diff --git a/cubed/runtime/executors/asyncio.py b/cubed/runtime/executors/asyncio.py
index 0beb2a95..0ab20900 100644
--- a/cubed/runtime/executors/asyncio.py
+++ b/cubed/runtime/executors/asyncio.py
@@ -5,6 +5,7 @@ from asyncio import Future
 from typing import Any, AsyncIterator, Callable, Dict, Iterable, List, Optional, Tuple
 
 from cubed.runtime.backup import should_launch_backup
+from cubed.runtime.utils import batched
 
 
 async def async_map_unordered(
@@ -14,6 +15,7 @@ async def async_map_unordered(
     create_backup_futures_func: Optional[
         Callable[..., List[Tuple[Any, Future]]]
     ] = None,
+    batch_size: Optional[int] = None,
     return_stats: bool = False,
     name: Optional[str] = None,
     **kwargs,
@@ -25,8 +27,14 @@ async def async_map_unordered(
     if create_backup_futures_func is None:
         create_backup_futures_func = create_futures_func
 
+    if batch_size is None:
+        inputs = input
+    else:
+        input_batches = batched(input, batch_size)
+        inputs = next(input_batches)
+
     task_create_tstamp = time.time()
-    tasks = {task: i for i, task in create_futures_func(input, **kwargs)}
+    tasks = {task: i for i, task in create_futures_func(inputs, **kwargs)}
     pending = set(tasks.keys())
     t = time.monotonic()
     start_times = {f: t for f in pending}
@@ -81,3 +89,14 @@ async def async_map_unordered(
                     pending.add(new_task)
                     backups[task] = new_task
                     backups[new_task] = task
+
+        if batch_size is not None and len(pending) < batch_size:
+            inputs = next(input_batches, None)  # type: ignore
+            if inputs is not None:
+                new_tasks = {
+                    task: i for i, task in create_futures_func(inputs, **kwargs)
+                }
+                tasks.update(new_tasks)
+                pending.update(new_tasks.keys())
+                t = time.monotonic()
+                start_times = {f: t for f in new_tasks.keys()}
diff --git a/cubed/runtime/executors/dask_distributed_async.py b/cubed/runtime/executors/dask_distributed_async.py
index 5b35673f..5fe6538b 100644
--- a/cubed/runtime/executors/dask_distributed_async.py
+++ b/cubed/runtime/executors/dask_distributed_async.py
@@ -38,6 +38,7 @@ async def map_unordered(
     map_iterdata: Iterable[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]],
     retries: int = 2,
     use_backups: bool = False,
+    batch_size: Optional[int] = None,
     return_stats: bool = False,
     name: Optional[str] = None,
     **kwargs,
@@ -68,6 +69,7 @@ async def map_unordered(
         map_iterdata,
         use_backups=use_backups,
         create_backup_futures_func=create_backup_futures_func,
+        batch_size=batch_size,
         return_stats=return_stats,
         name=name,
         **kwargs,
diff --git a/cubed/runtime/executors/lithops.py b/cubed/runtime/executors/lithops.py
index 19704db6..6a31f7c9 100644
--- a/cubed/runtime/executors/lithops.py
+++ b/cubed/runtime/executors/lithops.py
@@ -24,9 +24,8 @@ from cubed.core.array import Callback, Spec
 from cubed.core.plan import visit_node_generations, visit_nodes
 from cubed.runtime.backup import should_launch_backup
 from cubed.runtime.executors.lithops_retries import (
+    RetryingFunctionExecutor,
     RetryingFuture,
-    map_with_retries,
-    wait_with_retries,
 )
 from cubed.runtime.types import DagExecutor
 from cubed.runtime.utils import handle_callbacks
@@ -40,7 +39,7 @@ def run_func(input, func=None, config=None, name=None):
 
 
 def map_unordered(
-    lithops_function_executor: FunctionExecutor,
+    lithops_function_executor: RetryingFunctionExecutor,
     group_map_functions: Sequence[Callable[..., Any]],
     group_map_iterdata: Sequence[
         Iterable[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]]
@@ -88,8 +87,7 @@ def map_unordered(
         partial_map_function = lambda x: map_function(x, **kwargs)
         group_name_to_function[group_name] = partial_map_function
 
-        futures = map_with_retries(
-            lithops_function_executor,
+        futures = lithops_function_executor.map(
             partial_map_function,
             map_iterdata,
             timeout=timeout,
@@ -101,8 +99,7 @@ def map_unordered(
         pending.extend(futures)
 
     while pending:
-        finished, pending = wait_with_retries(
-            lithops_function_executor,
+        finished, pending = lithops_function_executor.wait(
             pending,
             throw_except=False,
             return_when=return_when,
@@ -142,8 +139,7 @@ def map_unordered(
                 ):
                     input = future.input
                     logger.info("Running backup task for %s", input)
-                    futures = map_with_retries(
-                        lithops_function_executor,
+                    futures = lithops_function_executor.map(
                         group_name_to_function[group_name],
                         [input],
                         timeout=timeout,
@@ -172,16 +168,17 @@ def execute_dag(
 ) -> None:
     use_backups = kwargs.pop("use_backups", False)
     allowed_mem = spec.allowed_mem if spec is not None else None
-    with FunctionExecutor(**kwargs) as executor:
-        runtime_memory_mb = executor.config[executor.backend].get(
-            "runtime_memory", None
-        )
-        if runtime_memory_mb is not None and allowed_mem is not None:
-            runtime_memory = runtime_memory_mb * 1_000_000
-            if runtime_memory < allowed_mem:
-                raise ValueError(
-                    f"Runtime memory ({runtime_memory}) is less than allowed_mem ({allowed_mem})"
-                )
+    function_executor = FunctionExecutor(**kwargs)
+    runtime_memory_mb = function_executor.config[function_executor.backend].get(
+        "runtime_memory", None
+    )
+    if runtime_memory_mb is not None and allowed_mem is not None:
+        runtime_memory = runtime_memory_mb * 1_000_000
+        if runtime_memory < allowed_mem:
+            raise ValueError(
+                f"Runtime memory ({runtime_memory}) is less than allowed_mem ({allowed_mem})"
+            )
+    with RetryingFunctionExecutor(function_executor) as executor:
         if not compute_arrays_in_parallel:
             for name, node in visit_nodes(dag, resume=resume):
                 pipeline = node["pipeline"]
diff --git a/cubed/runtime/executors/lithops_retries.py b/cubed/runtime/executors/lithops_retries.py
index e4cb4903..ba9f29c9 100644
--- a/cubed/runtime/executors/lithops_retries.py
+++ b/cubed/runtime/executors/lithops_retries.py
@@ -2,7 +2,13 @@ from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
 
 from lithops import FunctionExecutor
 from lithops.future import ResponseFuture
-from lithops.wait import ALL_COMPLETED, ALWAYS, ANY_COMPLETED
+from lithops.wait import (
+    ALL_COMPLETED,
+    ALWAYS,
+    ANY_COMPLETED,
+    THREADPOOL_SIZE,
+    WAIT_DUR_SEC,
+)
 from six import reraise
 
 
@@ -86,80 +92,96 @@ class RetryingFuture:
         return res
 
 
-def map_with_retries(
-    function_executor: FunctionExecutor,
-    map_function: Callable[..., Any],
-    map_iterdata: Iterable[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]],
-    timeout: Optional[int] = None,
-    include_modules: Optional[List[str]] = [],
-    retries: Optional[int] = None,
-    group_name: Optional[str] = None,
-) -> List[RetryingFuture]:
+class RetryingFunctionExecutor:
     """
-    A generalisation of Lithops `map`, with retries.
+    A wrapper around Lithops `FunctionExecutor` that supports retries.
     """
 
-    inputs = list(map_iterdata)
-    futures_list = function_executor.map(
-        map_function, inputs, timeout=timeout, include_modules=include_modules
-    )
-    return [
-        RetryingFuture(
-            f,
-            map_function=map_function,
-            input=i,
-            map_kwargs=dict(timeout=timeout, include_modules=include_modules),
-            retries=retries,
-            group_name=group_name,
-        )
-        for i, f in zip(inputs, futures_list)
-    ]
-
-
-def wait_with_retries(
-    function_executor: FunctionExecutor,
-    fs: List[RetryingFuture],
-    throw_except: Optional[bool] = True,
-    return_when: Optional[Any] = ALL_COMPLETED,
-    show_progressbar: Optional[bool] = True,
-) -> Tuple[List[RetryingFuture], List[RetryingFuture]]:
-    """
-    A generalisation of Lithops `wait`, with retries.
-    """
+    def __init__(self, executor: FunctionExecutor):
+        self.executor = executor
 
-    lookup = {f.response_future: f for f in fs}
+    def __enter__(self):
+        self.executor.__enter__()
+        return self
 
-    while True:
-        response_futures = [f.response_future for f in fs]
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.executor.__exit__(exc_type, exc_value, traceback)
 
-        done, pending = function_executor.wait(
-            response_futures,
-            throw_except=throw_except,
-            return_when=return_when,
-            show_progressbar=show_progressbar,
+    def map(
+        self,
+        map_function: Callable[..., Any],
+        map_iterdata: Iterable[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]],
+        *,
+        retries: Optional[int] = None,
+        group_name: Optional[str] = None,
+        **kwargs,
+    ) -> List[RetryingFuture]:
+        inputs = list(map_iterdata)
+        futures_list = self.executor.map(
+            map_function,
+            inputs,
+            **kwargs,
         )
-
-        retrying_done = []
-        retrying_pending = [lookup[response_future] for response_future in pending]
-        for response_future in done:
-            retrying_future = lookup[response_future]
-            if response_future.error:
-                retrying_future._inc_failure_count()
-                if retrying_future._should_retry():
-                    retrying_future._retry(function_executor)
-                    # put back into pending since we are retrying this input
-                    retrying_pending.append(retrying_future)
-                    lookup[retrying_future.response_future] = retrying_future
+        return [
+            RetryingFuture(
+                f,
+                map_function=map_function,
+                input=i,
+                map_kwargs=kwargs,
+                retries=retries,
+                group_name=group_name,
+            )
+            for i, f in zip(inputs, futures_list)
+        ]
+
+    def wait(
+        self,
+        fs: List[RetryingFuture],
+        throw_except: Optional[bool] = True,
+        return_when: Optional[Any] = ALL_COMPLETED,
+        download_results: Optional[bool] = False,
+        timeout: Optional[int] = None,
+        threadpool_size: Optional[int] = THREADPOOL_SIZE,
+        wait_dur_sec: Optional[int] = WAIT_DUR_SEC,
+        show_progressbar: Optional[bool] = True,
+    ) -> Tuple[List[RetryingFuture], List[RetryingFuture]]:
+        lookup = {f.response_future: f for f in fs}
+
+        while True:
+            response_futures = [f.response_future for f in fs]
+
+            done, pending = self.executor.wait(
+                response_futures,
+                throw_except=throw_except,
+                return_when=return_when,
+                download_results=download_results,
+                timeout=timeout,
+                threadpool_size=threadpool_size,
+                wait_dur_sec=wait_dur_sec,
+                show_progressbar=show_progressbar,
+            )
+
+            retrying_done = []
+            retrying_pending = [lookup[response_future] for response_future in pending]
+            for response_future in done:
+                retrying_future = lookup[response_future]
+                if response_future.error:
+                    retrying_future._inc_failure_count()
+                    if retrying_future._should_retry():
+                        retrying_future._retry(self.executor)
+                        # put back into pending since we are retrying this input
+                        retrying_pending.append(retrying_future)
+                        lookup[retrying_future.response_future] = retrying_future
+                    else:
+                        retrying_done.append(retrying_future)
                 else:
                     retrying_done.append(retrying_future)
-            else:
-                retrying_done.append(retrying_future)
 
-        if return_when == ALWAYS:
-            break
-        elif return_when == ANY_COMPLETED and len(retrying_done) > 0:
-            break
-        elif return_when == ALL_COMPLETED and len(retrying_pending) == 0:
-            break
+            if return_when == ALWAYS:
+                break
+            elif return_when == ANY_COMPLETED and len(retrying_done) > 0:
+                break
+            elif return_when == ALL_COMPLETED and len(retrying_pending) == 0:
+                break
 
-    return retrying_done, retrying_pending
+        return retrying_done, retrying_pending
diff --git a/cubed/runtime/executors/modal_async.py b/cubed/runtime/executors/modal_async.py
index 896e8a5d..e02c43c8 100644
--- a/cubed/runtime/executors/modal_async.py
+++ b/cubed/runtime/executors/modal_async.py
@@ -28,6 +28,7 @@ async def map_unordered(
     input: Iterable[Any],
     use_backups: bool = False,
     backup_function: Optional[Function] = None,
+    batch_size: Optional[int] = None,
     return_stats: bool = False,
     name: Optional[str] = None,
     **kwargs,
@@ -44,7 +45,7 @@ async def map_unordered(
     :return: Function values (and optionally stats) as they are completed, not necessarily in the input order.
     """
 
-    if not use_backups:
+    if not use_backups and batch_size is None:
         task_create_tstamp = time.time()
         async for result in app_function.map(input, order_outputs=False, kwargs=kwargs):
             if return_stats:
@@ -76,6 +77,7 @@ async def map_unordered(
         input,
         use_backups=use_backups,
         create_backup_futures_func=create_backup_futures_func,
+        batch_size=batch_size,
         return_stats=return_stats,
         name=name,
         **kwargs,
diff --git a/cubed/runtime/executors/python_async.py b/cubed/runtime/executors/python_async.py
index ff3032f5..588c89a1 100644
--- a/cubed/runtime/executors/python_async.py
+++ b/cubed/runtime/executors/python_async.py
@@ -29,6 +29,7 @@ async def map_unordered(
     input: Iterable[Any],
     retries: int = 2,
     use_backups: bool = False,
+    batch_size: Optional[int] = None,
     return_stats: bool = False,
     name: Optional[str] = None,
     **kwargs,
@@ -54,6 +55,7 @@ async def map_unordered(
         create_futures_func,
         input,
         use_backups=use_backups,
+        batch_size=batch_size,
         return_stats=return_stats,
         name=name,
         **kwargs,
diff --git a/cubed/runtime/utils.py b/cubed/runtime/utils.py
index a8d31856..8ffa32d6 100644
--- a/cubed/runtime/utils.py
+++ b/cubed/runtime/utils.py
@@ -1,5 +1,6 @@
 import time
 from functools import partial
+from itertools import islice
 
 from cubed.utils import peak_measured_mem
 
@@ -52,3 +53,13 @@ def handle_callbacks(callbacks, stats):
         else:
             event = TaskEndEvent(**stats)
         [callback.on_task_end(event) for callback in callbacks]
+
+
+# this will be in Python 3.12 https://docs.python.org/3.12/library/itertools.html#itertools.batched
+def batched(iterable, n):
+    # batched('ABCDEFG', 3) --> ABC DEF G
+    if n < 1:
+        raise ValueError("n must be at least one")
+    it = iter(iterable)
+    while batch := tuple(islice(it, n)):
+        yield batch


Test patch:
diff --git a/cubed/tests/runtime/test_dask_distributed_async.py b/cubed/tests/runtime/test_dask_distributed_async.py
index 74baee84..64b07869 100644
--- a/cubed/tests/runtime/test_dask_distributed_async.py
+++ b/cubed/tests/runtime/test_dask_distributed_async.py
@@ -1,4 +1,5 @@
 import asyncio
+import itertools
 from functools import partial
 
 import pytest
@@ -12,7 +13,7 @@ from dask.distributed import Client
 from cubed.runtime.executors.dask_distributed_async import map_unordered
 
 
-async def run_test(function, input, retries, use_backups=False):
+async def run_test(function, input, retries, use_backups=False, batch_size=None):
     outputs = set()
     async with Client(asynchronous=True) as client:
         async for output in map_unordered(
@@ -21,6 +22,7 @@ async def run_test(function, input, retries, use_backups=False):
             input,
             retries=retries,
             use_backups=use_backups,
+            batch_size=batch_size,
         ):
             outputs.add(output)
     return outputs
@@ -98,3 +100,18 @@ def test_stragglers(tmp_path, timing_map, n_tasks, retries):
     assert outputs == set(range(n_tasks))
 
     check_invocation_counts(tmp_path, timing_map, n_tasks, retries)
+
+
+def test_batch(tmp_path):
+    # input is unbounded, so if entire input were consumed and not read
+    # in batches then it would never return, since it would never
+    # run the first (failing) input
+    with pytest.raises(RuntimeError):
+        asyncio.run(
+            run_test(
+                function=partial(deterministic_failure, tmp_path, {0: [-1]}),
+                input=itertools.count(),
+                retries=0,
+                batch_size=10,
+            )
+        )
diff --git a/cubed/tests/runtime/test_lithops.py b/cubed/tests/runtime/test_lithops.py
index d48b2308..55a04415 100644
--- a/cubed/tests/runtime/test_lithops.py
+++ b/cubed/tests/runtime/test_lithops.py
@@ -10,11 +10,12 @@ pytest.importorskip("lithops")
 from lithops.executors import LocalhostExecutor
 
 from cubed.runtime.executors.lithops import map_unordered
+from cubed.runtime.executors.lithops_retries import RetryingFunctionExecutor
 
 
 def run_test(function, input, retries, timeout=10, use_backups=False):
     outputs = set()
-    with LocalhostExecutor() as executor:
+    with RetryingFunctionExecutor(LocalhostExecutor()) as executor:
         for output in map_unordered(
             executor,
             [function],
diff --git a/cubed/tests/runtime/test_lithops_retries.py b/cubed/tests/runtime/test_lithops_retries.py
index 333d0593..749e7849 100644
--- a/cubed/tests/runtime/test_lithops_retries.py
+++ b/cubed/tests/runtime/test_lithops_retries.py
@@ -4,20 +4,19 @@ pytest.importorskip("lithops")
 
 from lithops.executors import LocalhostExecutor
 
-from cubed.runtime.executors.lithops_retries import map_with_retries, wait_with_retries
+from cubed.runtime.executors.lithops_retries import RetryingFunctionExecutor
 from cubed.tests.runtime.utils import check_invocation_counts, deterministic_failure
 
 
 def run_test(function, input, retries, timeout=10):
-    with LocalhostExecutor() as executor:
-        futures = map_with_retries(
-            executor,
+    with RetryingFunctionExecutor(LocalhostExecutor()) as executor:
+        futures = executor.map(
             function,
             input,
             timeout=timeout,
             retries=retries,
         )
-        done, pending = wait_with_retries(executor, futures, throw_except=False)
+        done, pending = executor.wait(futures, throw_except=False)
         assert len(pending) == 0
     outputs = set(f.result() for f in done)
     return outputs
diff --git a/cubed/tests/runtime/test_modal_async.py b/cubed/tests/runtime/test_modal_async.py
index b9069a61..46a45cdb 100644
--- a/cubed/tests/runtime/test_modal_async.py
+++ b/cubed/tests/runtime/test_modal_async.py
@@ -1,3 +1,5 @@
+import itertools
+
 import pytest
 
 modal = pytest.importorskip("modal")
@@ -48,13 +50,14 @@ def deterministic_failure_modal_long_timeout(i, path=None, timing_map=None):
     return deterministic_failure(path, timing_map, i)
 
 
-async def run_test(app_function, input, use_backups=False, **kwargs):
+async def run_test(app_function, input, use_backups=False, batch_size=None, **kwargs):
     outputs = set()
     async with stub.run():
         async for output in map_unordered(
             app_function,
             input,
             use_backups=use_backups,
+            batch_size=batch_size,
             **kwargs,
         ):
             outputs.add(output)
@@ -181,3 +184,25 @@ def test_stragglers(timing_map, n_tasks, retries, expected_invocation_counts_ove
     finally:
         fs = fsspec.open(tmp_path).fs
         fs.rm(tmp_path, recursive=True)
+
+
+@pytest.mark.cloud
+def test_batch(tmp_path):
+    # input is unbounded, so if entire input were consumed and not read
+    # in batches then it would never return, since it would never
+    # run the first (failing) input
+    try:
+        with pytest.raises(RuntimeError):
+            asyncio.run(
+                run_test(
+                    app_function=deterministic_failure_modal_no_retries,
+                    input=itertools.count(),
+                    path=tmp_path,
+                    timing_map={0: [-1]},
+                    batch_size=10,
+                )
+            )
+
+    finally:
+        fs = fsspec.open(tmp_path).fs
+        fs.rm(tmp_path, recursive=True)
diff --git a/cubed/tests/runtime/test_python_async.py b/cubed/tests/runtime/test_python_async.py
index 3a74ddf1..3e27d797 100644
--- a/cubed/tests/runtime/test_python_async.py
+++ b/cubed/tests/runtime/test_python_async.py
@@ -1,4 +1,5 @@
 import asyncio
+import itertools
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 
@@ -8,7 +9,7 @@ from cubed.runtime.executors.python_async import map_unordered
 from cubed.tests.runtime.utils import check_invocation_counts, deterministic_failure
 
 
-async def run_test(function, input, retries=2, use_backups=False):
+async def run_test(function, input, retries=2, use_backups=False, batch_size=None):
     outputs = set()
     concurrent_executor = ThreadPoolExecutor()
     try:
@@ -18,6 +19,7 @@ async def run_test(function, input, retries=2, use_backups=False):
             input,
             retries=retries,
             use_backups=use_backups,
+            batch_size=batch_size,
         ):
             outputs.add(output)
     finally:
@@ -95,3 +97,18 @@ def test_stragglers(tmp_path, timing_map, n_tasks, retries):
     assert outputs == set(range(n_tasks))
 
     check_invocation_counts(tmp_path, timing_map, n_tasks, retries)
+
+
+def test_batch(tmp_path):
+    # input is unbounded, so if entire input were consumed and not read
+    # in batches then it would never return, since it would never
+    # run the first (failing) input
+    with pytest.raises(RuntimeError):
+        asyncio.run(
+            run_test(
+                function=partial(deterministic_failure, tmp_path, {0: [-1]}),
+                input=itertools.count(),
+                retries=0,
+                batch_size=10,
+            )
+        )