Repository: tomwhite/cubed (python)

Issue:
Lithops executor's maximum retries setting should be per-task not per-array
Currently `LithopsDagExecutor` has a `max_failures` setting that raises an exception if there are more than this number of failures when computing an array (it defaults to 3).

This works but for very large arrays you are more likely to get enough timeout errors to cause the computation to stop.

This should really be a per-task limit, so that it is independent of the size of the array. Modal has this, and the setting is called `retries` - so let's make it the same for Lithops.



Pull request description:
Make Lithops retries apply per-task
Fixes #210 

Patch:
diff --git a/cubed/runtime/executors/lithops.py b/cubed/runtime/executors/lithops.py
index 0d6cde5d..6f1510f1 100644
--- a/cubed/runtime/executors/lithops.py
+++ b/cubed/runtime/executors/lithops.py
@@ -60,7 +60,8 @@ def map_unordered(
     map_function,
     map_iterdata,
     include_modules=[],
-    max_failures=3,
+    timeout=None,
+    retries=2,
     use_backups=False,
     return_stats=False,
     **kwargs,
@@ -75,16 +76,15 @@ def map_unordered(
     :param map_function: The function to map over the data.
     :param map_iterdata: An iterable of input data.
     :param include_modules: Modules to include.
-    :param max_failures: The number of task failures to allow before raising an exception.
+    :param retries: The number of times to retry a failed task before raising an exception.
     :param use_backups: Whether to launch backup tasks to mitigate against slow-running tasks.
     :param return_stats: Whether to return lithops stats.
 
     :return: Function values (and optionally stats) as they are completed, not necessarily in the input order.
     """
-    failures = 0
     return_when = ALWAYS if use_backups else ANY_COMPLETED
 
-    inputs = map_iterdata
+    inputs = list(map_iterdata)
     tasks = {}
     start_times = {}
     end_times = {}
@@ -96,22 +96,27 @@ def map_unordered(
     partial_map_function = lambda x: map_function(x, **kwargs)
 
     futures = lithops_function_executor.map(
-        partial_map_function, inputs, include_modules=include_modules
+        partial_map_function, inputs, timeout=timeout, include_modules=include_modules
     )
     tasks.update({k: v for (k, v) in zip(futures, inputs)})
     start_times.update({k: time.monotonic() for k in futures})
     pending.extend(futures)
 
+    future_to_index = {f: i for i, f in enumerate(futures)}
+    failure_counts = [0] * len(inputs)
+
     while pending:
         finished, pending = lithops_function_executor.wait(
             pending, throw_except=False, return_when=return_when, show_progressbar=False
         )
-
         failed = []
         for future in finished:
             if future.error:
-                failures += 1
-                if failures > max_failures:
+                index = future_to_index[future]
+                failure_counts[index] = failure_counts[index] + 1
+                failure_count = failure_counts[index]
+
+                if failure_count > retries:
                     # re-raise exception
                     # TODO: why does calling status not raise the exception?
                     future.status(throw_except=True)
@@ -124,8 +129,8 @@ def map_unordered(
                 else:
                     yield future.result()
 
+            # remove any backup task
             if use_backups:
-                # remove backups
                 backup = backups.get(future, None)
                 if backup:
                     if backup in pending:
@@ -135,34 +140,45 @@ def map_unordered(
 
         if failed:
             # rerun and add to pending
-            inputs = [v for (fut, v) in tasks.items() if fut in failed]
+            inputs_to_rerun = []
+            input_indexes_to_rerun = []
+            for fut in failed:
+                index = future_to_index[fut]
+                inputs_to_rerun.append(inputs[index])
+                input_indexes_to_rerun.append(index)
+                del future_to_index[fut]
             # TODO: de-duplicate code from above
             futures = lithops_function_executor.map(
                 partial_map_function,
-                inputs,
+                inputs_to_rerun,
+                timeout=timeout,
                 include_modules=include_modules,
             )
-            tasks.update({k: v for (k, v) in zip(futures, inputs)})
+            tasks.update({k: v for (k, v) in zip(futures, inputs_to_rerun)})
             start_times.update({k: time.monotonic() for k in futures})
             pending.extend(futures)
 
+            future_to_index.update(
+                {f: i for i, f in zip(input_indexes_to_rerun, futures)}
+            )
+
         if use_backups:
             now = time.monotonic()
             for future in copy.copy(pending):
                 if future not in backups and should_launch_backup(
                     future, now, start_times, end_times
                 ):
-                    inputs = [v for (fut, v) in tasks.items() if fut == future]
-                    logger.info("Running backup task for %s", inputs)
+                    input = tasks[future]
+                    logger.info("Running backup task for %s", input)
                     futures = lithops_function_executor.map(
                         partial_map_function,
-                        inputs,
+                        [input],
+                        timeout=timeout,
                         include_modules=include_modules,
                     )
-                    tasks.update({k: v for (k, v) in zip(futures, inputs)})
+                    tasks.update({k: v for (k, v) in zip(futures, [input])})
                     start_times.update({k: time.monotonic() for k in futures})
                     pending.extend(futures)
-                    pending.remove(future)  # throw away slow one
                     backup = futures[0]  # TODO: launch multiple backups at once
                     backups[future] = backup
                     backups[backup] = future
diff --git a/cubed/runtime/executors/python_async.py b/cubed/runtime/executors/python_async.py
index 9019f2ee..743945d7 100644
--- a/cubed/runtime/executors/python_async.py
+++ b/cubed/runtime/executors/python_async.py
@@ -33,7 +33,7 @@ async def map_unordered(
     if retries == 0:
         retrying_function = function
     else:
-        retryer = Retrying(stop=stop_after_attempt(retries + 1))
+        retryer = Retrying(reraise=True, stop=stop_after_attempt(retries + 1))
         retrying_function = partial(retryer, function)
 
     task_create_tstamp = time.time()


Test patch:
diff --git a/cubed/tests/runtime/test_lithops.py b/cubed/tests/runtime/test_lithops.py
index 800f5c42..0f0b0f61 100644
--- a/cubed/tests/runtime/test_lithops.py
+++ b/cubed/tests/runtime/test_lithops.py
@@ -1,14 +1,9 @@
+import platform
 from functools import partial
 
 import pytest
 
-from cubed.tests.runtime.utils import (
-    fail_on_first_invocation,
-    never_fail,
-    read_int_from_file,
-    sleep_on_first_invocation,
-)
-from cubed.utils import join_path
+from cubed.tests.runtime.utils import check_invocation_counts, deterministic_failure
 
 pytest.importorskip("lithops")
 
@@ -17,75 +12,90 @@ from lithops.executors import LocalhostExecutor
 from cubed.runtime.executors.lithops import map_unordered
 
 
-def test_map_unordered_no_failures(tmp_path):
-    with LocalhostExecutor() as executor:
-        for _ in map_unordered(
-            executor, partial(never_fail, tmp_path), range(3), max_failures=0
-        ):
-            pass
-
-    assert read_int_from_file(join_path(tmp_path, "0")) == 1
-    assert read_int_from_file(join_path(tmp_path, "1")) == 1
-    assert read_int_from_file(join_path(tmp_path, "2")) == 1
-
-
-def test_map_unordered_recovers_from_failures(tmp_path):
+def run_test(function, input, retries, timeout=10, use_backups=False):
+    outputs = set()
     with LocalhostExecutor() as executor:
-        for _ in map_unordered(
+        for output in map_unordered(
             executor,
-            partial(fail_on_first_invocation, tmp_path),
-            range(3),
-            max_failures=3,
+            function,
+            input,
+            timeout=timeout,
+            retries=retries,
+            use_backups=use_backups,
         ):
-            pass
-
-    assert read_int_from_file(join_path(tmp_path, "0")) == 2
-    assert read_int_from_file(join_path(tmp_path, "1")) == 2
-    assert read_int_from_file(join_path(tmp_path, "2")) == 2
-
-
-def test_map_unordered_too_many_failures(tmp_path):
+            outputs.add(output)
+    return outputs
+
+
+# fmt: off
+@pytest.mark.parametrize(
+    "timing_map, n_tasks, retries",
+    [
+        # no failures
+        ({}, 3, 2),
+        # first invocation fails
+        ({0: [-1], 1: [-1], 2: [-1]}, 3, 2),
+        # first two invocations fail
+        ({0: [-1, -1], 1: [-1, -1], 2: [-1, -1]}, 3, 2),
+        # first input sleeps once
+        ({0: [20]}, 3, 2),
+    ],
+)
+# fmt: on
+def test_success(tmp_path, timing_map, n_tasks, retries):
+    outputs = run_test(
+        function=partial(deterministic_failure, tmp_path, timing_map),
+        input=range(n_tasks),
+        retries=retries,
+    )
+
+    assert outputs == set(range(n_tasks))
+
+    check_invocation_counts(tmp_path, timing_map, n_tasks, retries)
+
+
+# fmt: off
+@pytest.mark.parametrize(
+    "timing_map, n_tasks, retries",
+    [
+        # too many failures
+        ({0: [-1], 1: [-1], 2: [-1, -1, -1]}, 3, 2),
+    ],
+)
+# fmt: on
+def test_failure(tmp_path, timing_map, n_tasks, retries):
     with pytest.raises(RuntimeError):
-        with LocalhostExecutor() as executor:
-            for _ in map_unordered(
-                executor,
-                partial(fail_on_first_invocation, tmp_path),
-                range(3),
-                max_failures=2,
-            ):
-                pass
-
-
-def test_map_unordered_execution_timeout(tmp_path):
-    # set execution timeout to less than sleep value in sleep_on_first_invocation,
-    # to check that task is retried after timeout exception
-    config = {"lithops": {"execution_timeout": 30}}
-    with LocalhostExecutor(config=config) as executor:
-        for _ in map_unordered(
-            executor, partial(sleep_on_first_invocation, tmp_path), range(3)
-        ):
-            pass
+        run_test(
+            function=partial(deterministic_failure, tmp_path, timing_map),
+            input=range(n_tasks),
+            retries=retries,
+        )
 
-    assert read_int_from_file(join_path(tmp_path, "0")) == 2
-    assert read_int_from_file(join_path(tmp_path, "1")) == 1
-    assert read_int_from_file(join_path(tmp_path, "2")) == 1
+    check_invocation_counts(tmp_path, timing_map, n_tasks, retries)
 
 
-def test_map_unordered_stragglers(tmp_path):
+# fmt: off
+@pytest.mark.parametrize(
+    "timing_map, n_tasks, retries",
+    [
+        ({0: [60]}, 10, 2),
+    ],
+)
+# fmt: on
+@pytest.mark.skipif(platform.system() == "Windows", reason="does not run on windows")
+def test_stragglers(tmp_path, timing_map, n_tasks, retries):
     # TODO: run a test like this using a cloud executor
     # Reason: this test passes, but lithops local mode only runs one job at a time,
     # so it actually waits for the first job to finish before running the second one.
 
-    config = {"lithops": {"log_level": "DEBUG"}}
-    with LocalhostExecutor(config=config) as executor:
-        for _ in map_unordered(
-            executor,
-            partial(sleep_on_first_invocation, tmp_path),
-            range(10),
-            use_backups=True,
-        ):
-            pass
+    outputs = run_test(
+        function=partial(deterministic_failure, tmp_path, timing_map),
+        input=range(n_tasks),
+        retries=retries,
+        timeout=500,
+        use_backups=True,
+    )
+
+    assert outputs == set(range(n_tasks))
 
-    assert read_int_from_file(join_path(tmp_path, "0")) == 2
-    for i in range(1, 10):
-        assert read_int_from_file(join_path(tmp_path, f"{i}")) == 1
+    check_invocation_counts(tmp_path, timing_map, n_tasks, retries)
diff --git a/cubed/tests/runtime/test_modal_async.py b/cubed/tests/runtime/test_modal_async.py
index c9fe2ef3..506f965c 100644
--- a/cubed/tests/runtime/test_modal_async.py
+++ b/cubed/tests/runtime/test_modal_async.py
@@ -8,8 +8,7 @@ import fsspec
 import modal.aio
 
 from cubed.runtime.executors.modal_async import map_unordered
-from cubed.tests.runtime.utils import read_int_from_file, write_int_to_file
-from cubed.utils import join_path
+from cubed.tests.runtime.utils import check_invocation_counts, deterministic_failure
 
 tmp_path = "s3://cubed-unittest/map_unordered"
 
@@ -30,162 +29,146 @@ image = modal.Image.debian_slim().pip_install(
 )
 
 
-@stub.function(image=image, secret=modal.Secret.from_name("my-aws-secret"))
-def never_fail(i):
-    invocation_count_file = join_path(tmp_path, f"{i}")
-    write_int_to_file(invocation_count_file, 1)
-    return i
-
-
-@stub.function(image=image, secret=modal.Secret.from_name("my-aws-secret"), retries=2)
-async def fail_on_first_invocation(i):
-    invocation_count_file = join_path(tmp_path, f"{i}")
-    fs = fsspec.open(invocation_count_file).fs
-    if fs.exists(invocation_count_file):
-        count = read_int_from_file(invocation_count_file)
-        write_int_to_file(invocation_count_file, count + 1)
-    else:
-        write_int_to_file(invocation_count_file, 1)
-        raise RuntimeError(f"Deliberately fail on first invocation for input {i}")
-    return i
-
-
-@stub.function(image=image, secret=modal.Secret.from_name("my-aws-secret"))
-async def fail_on_first_invocation_no_retry(i):
-    invocation_count_file = join_path(tmp_path, f"{i}")
-    fs = fsspec.open(invocation_count_file).fs
-    if fs.exists(invocation_count_file):
-        count = read_int_from_file(invocation_count_file)
-        write_int_to_file(invocation_count_file, count + 1)
-    else:
-        write_int_to_file(invocation_count_file, 1)
-        raise RuntimeError(f"Deliberately fail on first invocation for input {i}")
-    return i
-
-
-@stub.function(image=image, secret=modal.Secret.from_name("my-aws-secret"))
-async def sleep_on_first_invocation(i):
-    print(f"sleep_on_first_invocation {i}")
-    invocation_count_file = join_path(tmp_path, f"{i}")
-    fs = fsspec.open(invocation_count_file).fs
-    if fs.exists(invocation_count_file):
-        count = read_int_from_file(invocation_count_file)
-        write_int_to_file(invocation_count_file, count + 1)
-    else:
-        write_int_to_file(invocation_count_file, 1)
-        # only sleep on first invocation of input = 0
-        if i == 0:
-            print("sleeping...")
-            await asyncio.sleep(60)
-            print("... finished sleeping")
-    return i
-
-
 @stub.function(
-    image=image, secret=modal.Secret.from_name("my-aws-secret"), retries=2, timeout=20
+    image=image, secret=modal.Secret.from_name("my-aws-secret"), retries=2, timeout=10
 )
-async def sleep_on_first_invocation_short_timeout(i):
-    print(f"sleep_on_first_invocation {i}")
-    invocation_count_file = join_path(tmp_path, f"{i}")
-    fs = fsspec.open(invocation_count_file).fs
-    if fs.exists(invocation_count_file):
-        count = read_int_from_file(invocation_count_file)
-        write_int_to_file(invocation_count_file, count + 1)
-    else:
-        write_int_to_file(invocation_count_file, 1)
-        # only sleep on first invocation of input = 0
-        if i == 0:
-            print("sleeping...")
-            await asyncio.sleep(60)
-            print("... finished sleeping")
-    return i
-
-
-async def run_test(app_function, input, use_backups=False):
-    async with stub.run():
-        async for _ in map_unordered(app_function, input, use_backups=use_backups):
-            pass
+def deterministic_failure_modal(i, path=None, timing_map=None):
+    return deterministic_failure(path, timing_map, i)
 
 
-@pytest.mark.cloud
-def test_map_unordered_no_failures():
-    try:
-        asyncio.run(run_test(app_function=never_fail, input=range(3)))
+@stub.function(
+    image=image, secret=modal.Secret.from_name("my-aws-secret"), retries=2, timeout=300
+)
+def deterministic_failure_modal_long_timeout(i, path=None, timing_map=None):
+    return deterministic_failure(path, timing_map, i)
 
-        assert read_int_from_file(join_path(tmp_path, "0")) == 1
-        assert read_int_from_file(join_path(tmp_path, "1")) == 1
-        assert read_int_from_file(join_path(tmp_path, "2")) == 1
 
-    finally:
-        fs = fsspec.open(tmp_path).fs
-        fs.rm(tmp_path, recursive=True)
+async def run_test(app_function, input, use_backups=False, **kwargs):
+    outputs = set()
+    async with stub.run():
+        async for output in map_unordered(
+            app_function, input, use_backups=use_backups, **kwargs
+        ):
+            outputs.add(output)
+    return outputs
 
 
+# fmt: off
+@pytest.mark.parametrize(
+    "timing_map, n_tasks, retries",
+    [
+        # no failures
+        ({}, 3, 2),
+        # first invocation fails
+        ({0: [-1], 1: [-1], 2: [-1]}, 3, 2),
+        # first two invocations fail
+        ({0: [-1, -1], 1: [-1, -1], 2: [-1, -1]}, 3, 2),
+        # first input sleeps once
+        ({0: [20]}, 3, 2),
+    ],
+)
+# fmt: on
 @pytest.mark.cloud
-def test_map_unordered_recovers_from_failures():
+def test_success(timing_map, n_tasks, retries):
     try:
-        asyncio.run(run_test(app_function=fail_on_first_invocation, input=range(3)))
+        outputs = asyncio.run(
+            run_test(
+                app_function=deterministic_failure_modal,
+                input=range(n_tasks),
+                path=tmp_path,
+                timing_map=timing_map
+            )
+        )
 
-        assert read_int_from_file(join_path(tmp_path, "0")) == 2
-        assert read_int_from_file(join_path(tmp_path, "1")) == 2
-        assert read_int_from_file(join_path(tmp_path, "2")) == 2
+        assert outputs == set(range(n_tasks))
+        check_invocation_counts(tmp_path, timing_map, n_tasks, retries)
 
     finally:
         fs = fsspec.open(tmp_path).fs
         fs.rm(tmp_path, recursive=True)
 
 
+# fmt: off
+@pytest.mark.parametrize(
+    "timing_map, n_tasks, retries",
+    [
+        # too many failures
+        ({0: [-1], 1: [-1], 2: [-1, -1, -1]}, 3, 2),
+    ],
+)
+# fmt: on
 @pytest.mark.cloud
-def test_map_unordered_too_many_failures():
+def test_failure(timing_map, n_tasks, retries):
     try:
         with pytest.raises(RuntimeError):
             asyncio.run(
                 run_test(
-                    app_function=fail_on_first_invocation_no_retry,
-                    input=range(3),
+                    app_function=deterministic_failure_modal,
+                    input=range(n_tasks),
+                    path=tmp_path,
+                    timing_map=timing_map,
                 )
             )
 
+        check_invocation_counts(tmp_path, timing_map, n_tasks, retries)
+
     finally:
         fs = fsspec.open(tmp_path).fs
         fs.rm(tmp_path, recursive=True)
 
 
+# fmt: off
+@pytest.mark.parametrize(
+    "timing_map, n_tasks, retries",
+    [
+        ({0: [-1], 1: [-1], 2: [-1]}, 1000, 2),
+    ],
+)
+# fmt: on
 @pytest.mark.cloud
-def test_map_unordered_execution_timeout():
-    # sleep_on_first_invocation_short_timeout has a short timeout,
-    # so check that task is retried after timeout exception
-
+def test_large_number_of_tasks(timing_map, n_tasks, retries):
     try:
-        asyncio.run(
+        outputs = asyncio.run(
             run_test(
-                app_function=sleep_on_first_invocation_short_timeout, input=range(3)
+                app_function=deterministic_failure_modal,
+                input=range(n_tasks),
+                path=tmp_path,
+                timing_map=timing_map
             )
         )
 
-        assert read_int_from_file(join_path(tmp_path, "0")) == 2
-        assert read_int_from_file(join_path(tmp_path, "1")) == 1
-        assert read_int_from_file(join_path(tmp_path, "2")) == 1
+        assert outputs == set(range(n_tasks))
+        check_invocation_counts(tmp_path, timing_map, n_tasks, retries)
 
     finally:
         fs = fsspec.open(tmp_path).fs
         fs.rm(tmp_path, recursive=True)
 
 
+# fmt: off
+@pytest.mark.parametrize(
+    "timing_map, n_tasks, retries, expected_invocation_counts_overrides",
+    [
+        # backup succeeds quickly
+        ({0: [60]}, 10, 2, {0: 2}),
+    ],
+)
+# fmt: on
 @pytest.mark.cloud
-def test_map_unordered_stragglers():
+def test_stragglers(timing_map, n_tasks, retries, expected_invocation_counts_overrides):
     try:
-        asyncio.run(
+        outputs = asyncio.run(
             run_test(
-                app_function=sleep_on_first_invocation,
-                input=range(10),
+                app_function=deterministic_failure_modal_long_timeout,
+                input=range(n_tasks),
+                path=tmp_path,
+                timing_map=timing_map,
                 use_backups=True,
             )
         )
 
-        assert read_int_from_file(join_path(tmp_path, "0")) == 2
-        for i in range(1, 10):
-            assert read_int_from_file(join_path(tmp_path, f"{i}")) == 1
+        assert outputs == set(range(n_tasks))
+        check_invocation_counts(tmp_path, timing_map, n_tasks, retries, expected_invocation_counts_overrides)
 
     finally:
         fs = fsspec.open(tmp_path).fs
diff --git a/cubed/tests/runtime/test_python_async.py b/cubed/tests/runtime/test_python_async.py
index a1b7f444..979f86b1 100644
--- a/cubed/tests/runtime/test_python_async.py
+++ b/cubed/tests/runtime/test_python_async.py
@@ -5,46 +5,62 @@ from functools import partial
 import pytest
 
 from cubed.runtime.executors.python_async import map_unordered
-from cubed.tests.runtime.utils import (
-    fail_on_first_invocation,
-    never_fail,
-    read_int_from_file,
-)
-from cubed.utils import join_path
+from cubed.tests.runtime.utils import check_invocation_counts, deterministic_failure
 
 
 async def run_test(function, input, retries=2):
+    outputs = set()
     with ThreadPoolExecutor() as concurrent_executor:
-        async for _ in map_unordered(
+        async for output in map_unordered(
             concurrent_executor, function, input, retries=retries
         ):
-            pass
-
-
-def test_map_unordered_no_failures(tmp_path):
-    asyncio.run(run_test(function=partial(never_fail, tmp_path), input=range(3)))
-
-    assert read_int_from_file(join_path(tmp_path, "0")) == 1
-    assert read_int_from_file(join_path(tmp_path, "1")) == 1
-    assert read_int_from_file(join_path(tmp_path, "2")) == 1
+            outputs.add(output)
+    return outputs
 
 
-def test_map_unordered_recovers_from_failures(tmp_path):
-    asyncio.run(
-        run_test(function=partial(fail_on_first_invocation, tmp_path), input=range(3))
+# fmt: off
+@pytest.mark.parametrize(
+    "timing_map, n_tasks, retries",
+    [
+        # no failures
+        ({}, 3, 2),
+        # first invocation fails
+        ({0: [-1], 1: [-1], 2: [-1]}, 3, 2),
+        # first two invocations fail
+        ({0: [-1, -1], 1: [-1, -1], 2: [-1, -1]}, 3, 2),
+    ],
+)
+# fmt: on
+def test_success(tmp_path, timing_map, n_tasks, retries):
+    outputs = asyncio.run(
+        run_test(
+            function=partial(deterministic_failure, tmp_path, timing_map),
+            input=range(n_tasks),
+            retries=retries,
+        )
     )
 
-    assert read_int_from_file(join_path(tmp_path, "0")) == 2
-    assert read_int_from_file(join_path(tmp_path, "1")) == 2
-    assert read_int_from_file(join_path(tmp_path, "2")) == 2
+    assert outputs == set(range(n_tasks))
+    check_invocation_counts(tmp_path, timing_map, n_tasks)
 
 
-def test_map_unordered_too_many_failures(tmp_path):
+# fmt: off
+@pytest.mark.parametrize(
+    "timing_map, n_tasks, retries",
+    [
+        # too many failures
+        ({0: [-1], 1: [-1], 2: [-1, -1, -1]}, 3, 2),
+    ],
+)
+# fmt: on
+def test_failure(tmp_path, timing_map, n_tasks, retries):
     with pytest.raises(RuntimeError):
         asyncio.run(
             run_test(
-                function=partial(fail_on_first_invocation, tmp_path),
-                input=range(3),
-                retries=0,
+                function=partial(deterministic_failure, tmp_path, timing_map),
+                input=range(n_tasks),
+                retries=retries,
             )
         )
+
+    check_invocation_counts(tmp_path, timing_map, n_tasks, retries)
diff --git a/cubed/tests/runtime/utils.py b/cubed/tests/runtime/utils.py
index b5e13a67..489b2384 100644
--- a/cubed/tests/runtime/utils.py
+++ b/cubed/tests/runtime/utils.py
@@ -1,4 +1,6 @@
+import re
 import time
+from urllib.parse import urlparse
 
 import fsspec
 
@@ -15,33 +17,88 @@ def write_int_to_file(path, i):
         f.write(str(i))
 
 
-def never_fail(path, i):
-    invocation_count_file = join_path(path, f"{i}")
-    write_int_to_file(invocation_count_file, 1)
-    return i
+def deterministic_failure(path, timing_map, i):
+    """A function that can either run normally, run slowly, or raise
+    an exception, depending on input and invocation count.
+
+    The timing_map is a dictionary whose keys are inputs and values
+    are sequences of timing information for each invocation.
+    The maginitude of the value is the time to sleep in seconds, and
+    the sign indicates the input is returned normally (positive, or 0),
+    or an exception is raised (negative).
 
+    If a input is missing then all invocations will run normally.
 
-def fail_on_first_invocation(path, i):
+    If there are subsequent invocations to the ones in the sequence, then
+    they will all run normally.
+    """
+    # increment number of invocations of this function with arg i
     invocation_count_file = join_path(path, f"{i}")
     fs = fsspec.open(invocation_count_file).fs
     if fs.exists(invocation_count_file):
-        count = read_int_from_file(invocation_count_file)
-        write_int_to_file(invocation_count_file, count + 1)
+        invocation_count = read_int_from_file(invocation_count_file)
     else:
-        write_int_to_file(invocation_count_file, 1)
-        raise RuntimeError(f"Deliberately fail on first invocation for input {i}")
-    return i
+        invocation_count = 0
+    write_int_to_file(invocation_count_file, invocation_count + 1)
 
+    timing_code = 0
+    if i in timing_map:
+        timing_codes = timing_map[i]
+        if invocation_count >= len(timing_codes):
+            timing_code = 0
+        else:
+            timing_code = timing_codes[invocation_count]
 
-def sleep_on_first_invocation(path, i):
-    invocation_count_file = join_path(path, f"{i}")
-    fs = fsspec.open(invocation_count_file).fs
-    if fs.exists(invocation_count_file):
-        count = read_int_from_file(invocation_count_file)
-        write_int_to_file(invocation_count_file, count + 1)
+    if timing_code >= 0:
+        time.sleep(timing_code)
+        return i
     else:
-        write_int_to_file(invocation_count_file, 1)
-        # only sleep on first invocation of input = 0
-        if i == 0:
-            time.sleep(60)
-    return i
+        time.sleep(-timing_code)
+        raise RuntimeError(
+            f"Deliberately fail on invocation number {invocation_count+1} for input {i}"
+        )
+
+
+def check_invocation_counts(
+    path, timing_map, n_tasks, retries=None, expected_invocation_counts_overrides=None
+):
+    expected_invocation_counts = {}
+    for i in range(n_tasks):
+        if i not in timing_map:
+            expected_invocation_counts[i] = 1
+        else:
+            timing_codes = timing_map[i]
+            expected_invocation_count = len(timing_codes) + 1
+
+            if retries is not None:
+                # there shouldn't have been more than retries + 1 invocations
+                max_invocations = retries + 1
+                expected_invocation_count = min(
+                    expected_invocation_count, max_invocations
+                )
+
+            expected_invocation_counts[i] = expected_invocation_count
+
+    if expected_invocation_counts_overrides is not None:
+        expected_invocation_counts.update(expected_invocation_counts_overrides)
+
+    # retrieve outputs concurrently, so we can test on large numbers of inputs
+    # see https://filesystem-spec.readthedocs.io/en/latest/async.html#synchronous-api
+    if re.match(r"^[a-zA-Z]:\\", str(path)):  # Windows local file
+        protocol = ""
+    else:
+        protocol = urlparse(str(path)).scheme
+    fs = fsspec.filesystem(protocol)
+    paths = [join_path(path, str(i)) for i in range(n_tasks)]
+    out = fs.cat(paths)
+    path_to_i = lambda p: int(p.rsplit("/", 1)[-1])
+    actual_invocation_counts = {path_to_i(path): int(val) for path, val in out.items()}
+
+    if actual_invocation_counts != expected_invocation_counts:
+        for i, expected_count in expected_invocation_counts.items():
+            actual_count = actual_invocation_counts[i]
+            if actual_count != expected_count:
+                print(
+                    f"Invocation count for {i}, expected: {expected_count}, actual: {actual_count}"
+                )
+    assert actual_invocation_counts == expected_invocation_counts