Repository: tomwhite/cubed
Pull request #241 (2023-06-28 16:52:51)

Dask delayed executor
Would close #161. Also an example of #224.

The idea of this executor is that instead of guaranteeing that each instance doesn't exceed `allowed_mem` because it only sees one function to run (i.e. serverless), we are relying on dask.distributed's workers to each fully execute one delayed func before starting another (i.e. to not fall foul of [root task overproduction](https://www.coiled.io/blog/reducing-dask-memory-usage)). The embarassingly parallel nature of the list of `dask.Delayed` objects is what should hopefully keep the scheduler behaving predictably. Then the peak memory usage of the dask cluster *should* be just `n_workers` * `projected_mem`.

I haven't properly tried this yet, I just wanted to push this as a WIP, because it seems to work for a simple test case

![Screenshot from 2023-06-28 12-47-32](https://github.com/tomwhite/cubed/assets/35968931/f78ada57-999b-4947-b767-7b0bbd486b6b)


I literally just copied how the sequential python executor works but wrapped `exec_stage_func` in `dask.delayed`. Currently this means that a new list of `dask.Delayed` objects is created then computed for each step in the cubed pipeline. I think the more advanced implementation would look more like [rechunker's dask executor](https://github.com/pangeo-data/rechunker/blob/master/rechunker/executors/dask.py), which builds each layer then creates a HLG before anything is computed.

TomNicholas: This did not behave the way I expected it to when I ran with a `LocalCluster` :confused: 

![Screenshot from 2023-06-28 14-14-08](https://github.com/tomwhite/cubed/assets/35968931/d1a96efb-db2c-4387-92c1-4e9ea568a97f)

TomNicholas: Okay I was doing something monumentally stupid :sweat_smile: - I had rushed to create a notebook to test this idea but forgot to change the chunked array type from dask to cubed! So I was accidentally calling `.compute()` on a dask array. (For some reason dask didn't mind that I was passing an argument called `executor` to `dask.array.Array.compute()`...)
TomNicholas: Here's what actually happens, I get an error from within Zarr

https://gist.github.com/TomNicholas/b1ce906e85fa7245bf85484479f609a0

Patch:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 03cfbe9a..a61e1c16 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -51,6 +51,11 @@ jobs:
     if: github.event_name == 'release'
     runs-on: ubuntu-latest
     steps:
+      - name: Download artifacts
+        uses: actions/download-artifact@v3
+        with:
+          name: releases
+          path: dist
       - name: Publish distribution 📦 to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
diff --git a/README.md b/README.md
index 3a7148b3..b4fd3284 100644
--- a/README.md
+++ b/README.md
@@ -11,5 +11,6 @@ Cubed is a distributed N-dimensional array library implemented in Python using b
 - Follows [Dask Array](https://docs.dask.org/en/stable/array.html)'s chunked array API (`map_blocks`, `rechunk`, `apply_gufunc`, etc)
 - [Zarr](https://zarr.readthedocs.io/en/stable/) for persistent and intermediate storage
 - Multiple serverless runtimes: Python (in-process), [Lithops](https://lithops-cloud.github.io/), [Modal](https://modal.com/), [Apache Beam](https://beam.apache.org/)
+- Integration with [Xarray](https://xarray.dev/) via [cubed-xarray](https://github.com/xarray-contrib/cubed-xarray)
 
 [Documentation](https://tomwhite.github.io/cubed)
diff --git a/cubed/array_api/creation_functions.py b/cubed/array_api/creation_functions.py
index 53cec6e1..8c8b6266 100644
--- a/cubed/array_api/creation_functions.py
+++ b/cubed/array_api/creation_functions.py
@@ -6,6 +6,7 @@ from zarr.util import normalize_shape
 from cubed.core import Plan, gensym, map_blocks
 from cubed.core.ops import map_direct
 from cubed.core.plan import new_temp_path
+from cubed.storage.virtual import virtual_offsets
 from cubed.storage.zarr import lazy_from_array, lazy_full
 from cubed.utils import to_chunksize
 from cubed.vendor.dask.array.core import normalize_chunks
@@ -153,6 +154,16 @@ def full(
     return Array(name, target, spec, plan)
 
 
+def offsets_array(shape, spec=None) -> "Array":
+    name = gensym()
+    target = virtual_offsets(shape)
+
+    from .array_object import Array
+
+    plan = Plan._new(name, "block_ids", target)
+    return Array(name, target, spec, plan)
+
+
 def full_like(
     x, /, fill_value, *, dtype=None, device=None, chunks=None, spec=None
 ) -> "Array":
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index 74dcb6cf..ee92a09e 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -423,20 +423,16 @@ def map_blocks(
 ) -> "Array":
     """Apply a function to corresponding blocks from multiple input arrays."""
     if has_keyword(func, "block_id"):
-        from cubed.array_api import asarray
+        from cubed.array_api.creation_functions import offsets_array
 
         # Create an array of index offsets with the same chunk structure as the args,
         # which we convert to block ids (chunk coordinates) later.
         a = args[0]
-        _offsets = np.arange(a.npartitions, dtype=np.int32).reshape(a.numblocks)
-        offsets = asarray(_offsets, spec=a.spec)
-        # rechunk in a separate operation to avoid potentially writing lots of zarr chunks from the client
-        offsets_chunks = (1,) * len(a.numblocks)
-        offsets = rechunk(offsets, offsets_chunks)
+        offsets = offsets_array(a.numblocks, a.spec)
         new_args = args + (offsets,)
 
         def offset_to_block_id(offset):
-            return list(np.ndindex(*(a.numblocks)))[offset]
+            return np.unravel_index(offset, a.numblocks)
 
         def func_with_block_id(func):
             def wrap(*a, **kw):
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
index f003cffa..a5cdbadb 100644
--- a/cubed/core/plan.py
+++ b/cubed/core/plan.py
@@ -137,16 +137,21 @@ class Plan:
         # find all lazy zarr arrays in dag
         all_array_nodes = []
         lazy_zarr_arrays = []
+        reserved_mem_values = []
         for n, d in dag.nodes(data=True):
+            if "reserved_mem" in d and d["reserved_mem"] is not None:
+                reserved_mem_values.append(d["reserved_mem"])
             if isinstance(d["target"], LazyZarrArray):
                 all_array_nodes.append(n)
                 lazy_zarr_arrays.append(d["target"])
 
+        reserved_mem = max(reserved_mem_values, default=0)
+
         if len(lazy_zarr_arrays) > 0:
             # add new node and edges
             name = "create-arrays"
             op_name = name
-            pipeline = create_zarr_arrays(lazy_zarr_arrays)
+            pipeline = create_zarr_arrays(lazy_zarr_arrays, reserved_mem)
             dag.add_node(
                 name,
                 name=name,
@@ -154,7 +159,7 @@ class Plan:
                 target=None,
                 pipeline=pipeline,
                 projected_mem=pipeline.projected_mem,
-                reserved_mem=0,  # TODO: get from Spec
+                reserved_mem=reserved_mem,
                 num_tasks=pipeline.num_tasks,
             )
             # make create arrays node a dependency of all lazy array nodes
@@ -216,11 +221,7 @@ class Plan:
         tasks = 0
         for _, node in visit_nodes(dag, resume=resume):
             pipeline = node["pipeline"]
-            for stage in pipeline.stages:
-                if stage.mappable is not None:
-                    tasks += len(list(stage.mappable))
-                else:
-                    tasks += 1
+            tasks += pipeline.num_tasks
         return tasks
 
     def num_arrays(self, optimize_graph: bool = True) -> int:
@@ -413,7 +414,7 @@ def create_zarr_array(lazy_zarr_array, *, config=None):
     lazy_zarr_array.create(mode="a")
 
 
-def create_zarr_arrays(lazy_zarr_arrays):
+def create_zarr_arrays(lazy_zarr_arrays, reserved_mem):
     stages = [
         Stage(
             create_zarr_array,
@@ -422,13 +423,18 @@ def create_zarr_arrays(lazy_zarr_arrays):
         )
     ]
 
-    # projected memory is size of largest initial values, or 0 if there aren't any
-    projected_mem = max(
-        [
-            lza.initial_values.nbytes if lza.initial_values is not None else 0
-            for lza in lazy_zarr_arrays
-        ],
-        default=0,
+    # projected memory is size of largest initial values, or dtype size if there aren't any
+    projected_mem = (
+        max(
+            [
+                lza.initial_values.nbytes
+                if lza.initial_values is not None
+                else lza.dtype.itemsize
+                for lza in lazy_zarr_arrays
+            ],
+            default=0,
+        )
+        + reserved_mem
     )
     num_tasks = len(lazy_zarr_arrays)
 
diff --git a/cubed/primitive/blockwise.py b/cubed/primitive/blockwise.py
index ada0b2cb..2be0af83 100644
--- a/cubed/primitive/blockwise.py
+++ b/cubed/primitive/blockwise.py
@@ -1,13 +1,15 @@
 import itertools
+import math
 from dataclasses import dataclass
 from functools import partial
-from typing import Callable, Dict
+from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union
 
 import toolz
 import zarr
 from toolz import map
 
-from cubed.storage.zarr import lazy_empty
+from cubed.storage.zarr import T_ZarrArray, lazy_empty
+from cubed.types import T_Chunks, T_DType, T_Shape, T_Store
 from cubed.utils import chunk_memory, get_item, to_chunksize
 from cubed.vendor.dask.array.core import normalize_chunks
 from cubed.vendor.dask.blockwise import _get_coord_mapping, _make_dims, lol_product
@@ -19,7 +21,7 @@ from .types import CubedArrayProxy, CubedPipeline
 sym_counter = 0
 
 
-def gensym(name):
+def gensym(name: str) -> str:
     global sym_counter
     sym_counter += 1
     return f"{name}-{sym_counter:03}"
@@ -43,19 +45,21 @@ class BlockwiseSpec:
         Write proxy with an ``array`` attribute that supports ``__setitem__``.
     """
 
-    block_function: Callable
-    function: Callable
+    block_function: Callable[..., Any]
+    function: Callable[..., Any]
     reads_map: Dict[str, CubedArrayProxy]
     write: CubedArrayProxy
 
 
-def apply_blockwise(out_key, *, config=BlockwiseSpec):
+def apply_blockwise(out_key: List[int], *, config: BlockwiseSpec) -> None:
     """Stage function for blockwise."""
     # lithops needs params to be lists not tuples, so convert back
-    out_key = tuple(out_key)
-    out_chunk_key = key_to_slices(out_key, config.write.array, config.write.chunks)
+    out_key_tuple = tuple(out_key)
+    out_chunk_key = key_to_slices(
+        out_key_tuple, config.write.array, config.write.chunks
+    )
     args = []
-    name_chunk_inds = config.block_function(("out",) + out_key)
+    name_chunk_inds = config.block_function(("out",) + out_key_tuple)
     for name_chunk_ind in name_chunk_inds:
         name = name_chunk_ind[0]
         chunk_ind = name_chunk_ind[1:]
@@ -72,25 +76,27 @@ def apply_blockwise(out_key, *, config=BlockwiseSpec):
         config.write.open()[out_chunk_key] = result
 
 
-def key_to_slices(key, arr, chunks=None):
+def key_to_slices(
+    key: Tuple[int, ...], arr: T_ZarrArray, chunks: Optional[T_Chunks] = None
+) -> Tuple[slice, ...]:
     """Convert a chunk index key to a tuple of slices"""
     chunks = normalize_chunks(chunks or arr.chunks, shape=arr.shape, dtype=arr.dtype)
     return get_item(chunks, key)
 
 
 def blockwise(
-    func,
-    out_ind,
-    *args,
-    allowed_mem,
-    reserved_mem,
-    target_store,
-    shape,
-    dtype,
-    chunks,
-    new_axes=None,
-    in_names=None,
-    out_name=None,
+    func: Callable[..., Any],
+    out_ind: Sequence[Union[str, int]],
+    *args: Any,
+    allowed_mem: int,
+    reserved_mem: int,
+    target_store: T_Store,
+    shape: T_Shape,
+    dtype: T_DType,
+    chunks: T_Chunks,
+    new_axes: Optional[Dict[int, int]] = None,
+    in_names: Optional[List[str]] = None,
+    out_name: Optional[str] = None,
     **kwargs,
 ):
     """Apply a function across blocks from multiple source Zarr arrays.
@@ -126,20 +132,20 @@ def blockwise(
     """
 
     # Use dask's make_blockwise_graph
-    arrays = args[::2]
+    arrays: Sequence[T_ZarrArray] = args[::2]
     array_names = in_names or [f"in_{i}" for i in range(len(arrays))]
     array_map = {name: array for name, array in zip(array_names, arrays)}
 
-    inds = args[1::2]
+    inds: Sequence[Union[str, int]] = args[1::2]
 
-    numblocks = {}
+    numblocks: Dict[str, Tuple[int, ...]] = {}
     for name, array in zip(array_names, arrays):
         input_chunks = normalize_chunks(
             array.chunks, shape=array.shape, dtype=array.dtype
         )
         numblocks[name] = tuple(map(len, input_chunks))
 
-    argindsstr = []
+    argindsstr: List[Any] = []
     for name, ind in zip(array_names, inds):
         argindsstr.extend((name, ind))
 
@@ -157,7 +163,18 @@ def blockwise(
         new_axes=new_axes,
     )
 
-    output_blocks = get_output_blocks(
+    output_blocks_generator_fn = partial(
+        get_output_blocks,
+        func,
+        out_name or "out",
+        out_ind,
+        *argindsstr,
+        numblocks=numblocks,
+        new_axes=new_axes,
+    )
+    output_blocks = IterableFromGenerator(output_blocks_generator_fn)
+
+    num_tasks = num_output_blocks(
         func,
         out_name or "out",
         out_ind,
@@ -211,15 +228,13 @@ def blockwise(
             f"Projected blockwise memory ({projected_mem}) exceeds allowed_mem ({allowed_mem}), including reserved_mem ({reserved_mem})"
         )
 
-    num_tasks = len(output_blocks)
-
     return CubedPipeline(stages, spec, target_array, None, projected_mem, num_tasks)
 
 
 # Code for fusing pipelines
 
 
-def is_fuse_candidate(pipeline):
+def is_fuse_candidate(pipeline: CubedPipeline) -> bool:
     """
     Return True if a pipeline is a candidate for blockwise fusion.
     """
@@ -227,13 +242,13 @@ def is_fuse_candidate(pipeline):
     return len(stages) == 1 and stages[0].function == apply_blockwise
 
 
-def can_fuse_pipelines(pipeline1, pipeline2):
+def can_fuse_pipelines(pipeline1: CubedPipeline, pipeline2: CubedPipeline) -> bool:
     if is_fuse_candidate(pipeline1) and is_fuse_candidate(pipeline2):
         return pipeline1.num_tasks == pipeline2.num_tasks
     return False
 
 
-def fuse(pipeline1, pipeline2):
+def fuse(pipeline1: CubedPipeline, pipeline2: CubedPipeline) -> CubedPipeline:
     """
     Fuse two blockwise pipelines into a single pipeline, avoiding writing to (or reading from) the target of the first pipeline.
     """
@@ -273,8 +288,13 @@ def fuse(pipeline1, pipeline2):
 
 
 def make_blockwise_function(
-    func, output, out_indices, *arrind_pairs, numblocks=None, new_axes=None
-):
+    func: Callable[..., Any],
+    output: str,
+    out_indices: Sequence[Union[str, int]],
+    *arrind_pairs: Any,
+    numblocks: Optional[Dict[str, Tuple[int, ...]]] = None,
+    new_axes: Optional[Dict[int, int]] = None,
+) -> Callable[[List[int]], Any]:
     """Make a function that is the equivalent of make_blockwise_graph."""
 
     if numblocks is None:
@@ -326,8 +346,13 @@ def make_blockwise_function(
 
 
 def make_blockwise_function_flattened(
-    func, output, out_indices, *arrind_pairs, numblocks=None, new_axes=None
-):
+    func: Callable[..., Any],
+    output: str,
+    out_indices: Sequence[Union[str, int]],
+    *arrind_pairs: Any,
+    numblocks: Optional[Dict[str, Tuple[int, ...]]] = None,
+    new_axes: Optional[Dict[int, int]] = None,
+) -> Callable[[List[int]], Any]:
     # TODO: make this a part of make_blockwise_function?
     blockwise_fn = make_blockwise_function(
         func, output, out_indices, *arrind_pairs, numblocks=numblocks, new_axes=new_axes
@@ -344,8 +369,13 @@ def make_blockwise_function_flattened(
 
 
 def get_output_blocks(
-    func, output, out_indices, *arrind_pairs, numblocks=None, new_axes=None
-):
+    func: Callable[..., Any],
+    output: str,
+    out_indices: Sequence[Union[str, int]],
+    *arrind_pairs: Any,
+    numblocks: Optional[Dict[str, Tuple[int, ...]]] = None,
+    new_axes: Optional[Dict[int, int]] = None,
+) -> Iterator[List[int]]:
     if numblocks is None:
         raise ValueError("Missing required numblocks argument.")
     new_axes = new_axes or {}
@@ -355,7 +385,31 @@ def get_output_blocks(
     dims = _make_dims(argpairs, numblocks, new_axes)
 
     # return a list of lists, not of tuples, otherwise lithops breaks
-    output_blocks_lists = [
-        list(tup) for tup in itertools.product(*[range(dims[i]) for i in out_indices])
-    ]
-    return output_blocks_lists
+    for tup in itertools.product(*[range(dims[i]) for i in out_indices]):
+        yield list(tup)
+
+
+class IterableFromGenerator:
+    def __init__(self, generator_fn: Callable[[], Iterator[List[int]]]):
+        self.generator_fn = generator_fn
+
+    def __iter__(self):
+        return self.generator_fn()
+
+
+def num_output_blocks(
+    func: Callable[..., Any],
+    output: str,
+    out_indices: Sequence[Union[str, int]],
+    *arrind_pairs: Any,
+    numblocks: Optional[Dict[str, Tuple[int, ...]]] = None,
+    new_axes: Optional[Dict[int, int]] = None,
+) -> int:
+    if numblocks is None:
+        raise ValueError("Missing required numblocks argument.")
+    new_axes = new_axes or {}
+    argpairs = list(toolz.partition(2, arrind_pairs))
+
+    # Dictionary mapping {i: 3, j: 4, ...} for i, j, ... the dimensions
+    dims = _make_dims(argpairs, numblocks, new_axes)
+    return math.prod(dims[i] for i in out_indices)
diff --git a/cubed/primitive/rechunk.py b/cubed/primitive/rechunk.py
index 58bd3942..b8d4de38 100644
--- a/cubed/primitive/rechunk.py
+++ b/cubed/primitive/rechunk.py
@@ -1,18 +1,22 @@
 from math import ceil, prod
+from typing import Any, Dict, List, Optional, Tuple
 
-import zarr
-
-import cubed
-from cubed.primitive.types import CubedArrayProxy, CubedCopySpec
+from cubed.primitive.types import CubedArrayProxy, CubedCopySpec, CubedPipeline
 from cubed.runtime.pipeline import spec_to_pipeline
-from cubed.storage.zarr import lazy_empty
+from cubed.storage.zarr import T_ZarrArray, lazy_empty
+from cubed.types import T_RegularChunks, T_Shape, T_Store
 from cubed.vendor.rechunker.algorithm import rechunking_plan
 from cubed.vendor.rechunker.api import _validate_options
 
 
 def rechunk(
-    source, target_chunks, allowed_mem, reserved_mem, target_store, temp_store=None
-):
+    source: T_ZarrArray,
+    target_chunks: T_RegularChunks,
+    allowed_mem: int,
+    reserved_mem: int,
+    target_store: T_Store,
+    temp_store: Optional[T_Store] = None,
+) -> CubedPipeline:
     """Rechunk a Zarr array to have target_chunks.
 
     Parameters
@@ -36,7 +40,7 @@ def rechunk(
 
     # rechunker doesn't take account of uncompressed and compressed copies of the
     # input and output array chunk/selection, so adjust appropriately
-    rechunker_max_mem = (allowed_mem - reserved_mem) / 4
+    rechunker_max_mem = (allowed_mem - reserved_mem) // 4
 
     copy_specs, intermediate, target = _setup_rechunk(
         source=source,
@@ -62,14 +66,14 @@ def rechunk(
 
 # from rechunker, but simpler since it only has to handle Zarr arrays
 def _setup_rechunk(
-    source,
-    target_chunks,
-    max_mem,
-    target_store,
-    target_options=None,
-    temp_store=None,
-    temp_options=None,
-):
+    source: T_ZarrArray,
+    target_chunks: T_RegularChunks,
+    max_mem: int,
+    target_store: T_Store,
+    target_options: Optional[Dict[Any, Any]] = None,
+    temp_store: Optional[T_Store] = None,
+    temp_options: Optional[Dict[Any, Any]] = None,
+) -> Tuple[List[CubedCopySpec], T_ZarrArray, T_ZarrArray]:
     if temp_options is None:
         temp_options = target_options
     target_options = target_options or {}
@@ -90,14 +94,14 @@ def _setup_rechunk(
 
 
 def _setup_array_rechunk(
-    source_array,
-    target_chunks,
-    max_mem,
-    target_store_or_group,
-    target_options=None,
-    temp_store_or_group=None,
-    temp_options=None,
-    name=None,
+    source_array: T_ZarrArray,
+    target_chunks: T_RegularChunks,
+    max_mem: int,
+    target_store_or_group: T_Store,
+    target_options: Optional[Dict[Any, Any]] = None,
+    temp_store_or_group: Optional[T_Store] = None,
+    temp_options: Optional[Dict[Any, Any]] = None,
+    name: Optional[str] = None,
 ) -> CubedCopySpec:
     _validate_options(target_options)
     _validate_options(temp_options)
@@ -115,18 +119,12 @@ def _setup_array_rechunk(
         # this is just a pass-through copy
         target_chunks = source_chunks
 
-    # TODO: rewrite to avoid the hard dependency on dask
-    max_mem = cubed.vendor.dask.utils.parse_bytes(max_mem)
-
-    # don't consolidate reads for Dask arrays
-    consolidate_reads = isinstance(source_array, zarr.core.Array)
     read_chunks, int_chunks, write_chunks = rechunking_plan(
         shape,
         source_chunks,
         target_chunks,
         itemsize,
         max_mem,
-        consolidate_reads=consolidate_reads,
     )
 
     # create target
@@ -143,11 +141,6 @@ def _setup_array_rechunk(
         **(target_options or {}),
     )
 
-    try:
-        target_array.attrs.update(source_array.attrs)
-    except AttributeError:
-        pass
-
     if read_chunks == write_chunks:
         int_array = None
     else:
@@ -172,6 +165,6 @@ def _setup_array_rechunk(
     return CubedCopySpec(read_proxy, int_proxy, write_proxy)
 
 
-def total_chunks(shape, chunks):
+def total_chunks(shape: T_Shape, chunks: T_RegularChunks) -> int:
     # cf rechunker's chunk_keys
     return prod(ceil(s / c) for s, c in zip(shape, chunks))
diff --git a/cubed/primitive/types.py b/cubed/primitive/types.py
index accb4ef2..3abab9d7 100644
--- a/cubed/primitive/types.py
+++ b/cubed/primitive/types.py
@@ -1,7 +1,10 @@
 from dataclasses import dataclass
-from typing import Any, Iterable, Optional
+from typing import Any, Optional, Sequence
 
-from cubed.storage.zarr import open_if_lazy_zarr_array
+import zarr
+
+from cubed.storage.zarr import T_ZarrArray, open_if_lazy_zarr_array
+from cubed.types import T_RegularChunks
 from cubed.vendor.rechunker.types import Config, Stage
 
 
@@ -9,7 +12,7 @@ from cubed.vendor.rechunker.types import Config, Stage
 class CubedPipeline:
     """Generalisation of rechunker ``Pipeline`` with extra attributes."""
 
-    stages: Iterable[Stage]
+    stages: Sequence[Stage]
     config: Config
     target_array: Any
     intermediate_array: Optional[Any]
@@ -20,11 +23,11 @@ class CubedPipeline:
 class CubedArrayProxy:
     """Generalisation of rechunker ``ArrayProxy`` with support for ``LazyZarrArray``."""
 
-    def __init__(self, array, chunks):
+    def __init__(self, array: T_ZarrArray, chunks: T_RegularChunks):
         self.array = array
         self.chunks = chunks
 
-    def open(self):
+    def open(self) -> zarr.Array:
         return open_if_lazy_zarr_array(self.array)
 
 
diff --git a/cubed/runtime/executors/dask.py b/cubed/runtime/executors/dask.py
new file mode 100644
index 00000000..24c10f1c
--- /dev/null
+++ b/cubed/runtime/executors/dask.py
@@ -0,0 +1,40 @@
+import dask
+
+from cubed.core.plan import visit_nodes
+from cubed.runtime.types import DagExecutor
+
+
+def exec_stage_func(func, *args, **kwargs):
+    # TODO would be good to give the dask tasks useful names
+    return dask.delayed(func)(*args, **kwargs)  # should we add pure=True?
+
+
+class DaskDelayedExecutor(DagExecutor):
+    """Executes each stage using dask.Delayed functions."""
+
+    def execute_dag(
+        self,
+        dag,
+        callbacks=None,
+        array_names=None,
+        resume=None,
+        **compute_kwargs,
+    ):
+        # Note this currently only builds the task graph for each stage once it gets to that stage in computation
+        for name, node in visit_nodes(dag, resume=resume):
+            pipeline = node["pipeline"]
+            for stage in pipeline.stages:
+                if stage.mappable is not None:
+                    stage_delayed_funcs = []
+                    for m in stage.mappable:
+                        delayed_func = exec_stage_func(
+                            stage.function, m, config=pipeline.config
+                        )
+                        stage_delayed_funcs.append(delayed_func)
+                else:
+                    delayed_func = exec_stage_func(
+                        stage.function, config=pipeline.config
+                    )
+                    stage_delayed_funcs = [delayed_func]
+
+                dask.compute(*stage_delayed_funcs, **compute_kwargs)
diff --git a/cubed/runtime/executors/lithops.py b/cubed/runtime/executors/lithops.py
index 6f1510f1..34d2ec9c 100644
--- a/cubed/runtime/executors/lithops.py
+++ b/cubed/runtime/executors/lithops.py
@@ -6,10 +6,10 @@ from typing import Callable, Iterable
 
 from lithops.executors import FunctionExecutor
 from lithops.wait import ALWAYS, ANY_COMPLETED
-from six import reraise
 
 from cubed.core.plan import visit_nodes
 from cubed.runtime.backup import should_launch_backup
+from cubed.runtime.executors.lithops_retries import map_with_retries, wait_with_retries
 from cubed.runtime.types import DagExecutor
 from cubed.runtime.utils import handle_callbacks
 from cubed.vendor.rechunker.types import ParallelPipelines, PipelineExecutor
@@ -95,39 +95,39 @@ def map_unordered(
     # also, lithops extra_args doesn't work for this case
     partial_map_function = lambda x: map_function(x, **kwargs)
 
-    futures = lithops_function_executor.map(
-        partial_map_function, inputs, timeout=timeout, include_modules=include_modules
+    futures = map_with_retries(
+        lithops_function_executor,
+        partial_map_function,
+        inputs,
+        timeout=timeout,
+        include_modules=include_modules,
+        retries=retries,
     )
     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
+        finished, pending = wait_with_retries(
+            lithops_function_executor,
+            pending,
+            throw_except=False,
+            return_when=return_when,
+            show_progressbar=False,
         )
-        failed = []
         for future in finished:
             if future.error:
-                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)
-                    reraise(*future._exception)
-                failed.append(future)
+                # if the task has a backup that is not done, or is done with no exception, then don't raise this exception
+                backup = backups.get(future, None)
+                if backup:
+                    if not backup.done or not backup.error:
+                        continue
+                future.status(throw_except=True)
+            end_times[future] = time.monotonic()
+            if return_stats:
+                yield future.result(), standardise_lithops_stats(future.stats)
             else:
-                end_times[future] = time.monotonic()
-                if return_stats:
-                    yield future.result(), standardise_lithops_stats(future.stats)
-                else:
-                    yield future.result()
+                yield future.result()
 
             # remove any backup task
             if use_backups:
@@ -137,30 +137,7 @@ def map_unordered(
                         pending.remove(backup)
                     del backups[future]
                     del backups[backup]
-
-        if failed:
-            # rerun and add to pending
-            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_to_rerun,
-                timeout=timeout,
-                include_modules=include_modules,
-            )
-            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)}
-            )
+                    backup.cancel()
 
         if use_backups:
             now = time.monotonic()
@@ -170,16 +147,18 @@ def map_unordered(
                 ):
                     input = tasks[future]
                     logger.info("Running backup task for %s", input)
-                    futures = lithops_function_executor.map(
+                    futures = map_with_retries(
+                        lithops_function_executor,
                         partial_map_function,
                         [input],
                         timeout=timeout,
                         include_modules=include_modules,
+                        retries=0,  # don't retry backup tasks
                     )
                     tasks.update({k: v for (k, v) in zip(futures, [input])})
                     start_times.update({k: time.monotonic() for k in futures})
                     pending.extend(futures)
-                    backup = futures[0]  # TODO: launch multiple backups at once
+                    backup = futures[0]
                     backups[future] = backup
                     backups[backup] = future
             time.sleep(1)
@@ -195,7 +174,7 @@ def execute_dag(dag, callbacks=None, array_names=None, resume=None, **kwargs):
                     for _, stats in map_unordered(
                         executor,
                         run_func,
-                        list(stage.mappable),
+                        stage.mappable,
                         func=stage.function,
                         config=pipeline.config,
                         name=name,
@@ -229,7 +208,7 @@ def build_stage_mappable_func(
         for _, stats in map_unordered(
             lithops_function_executor,
             sf,
-            list(stage.mappable),
+            stage.mappable,
             use_backups=use_backups,
             return_stats=True,
         ):
diff --git a/cubed/runtime/executors/lithops_retries.py b/cubed/runtime/executors/lithops_retries.py
new file mode 100644
index 00000000..570610a1
--- /dev/null
+++ b/cubed/runtime/executors/lithops_retries.py
@@ -0,0 +1,161 @@
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+from lithops import FunctionExecutor
+from lithops.future import ResponseFuture
+from lithops.wait import ALL_COMPLETED, ALWAYS, ANY_COMPLETED
+from six import reraise
+
+
+class RetryingFuture:
+    """
+    A wrapper around Lithops `ResponseFuture` that takes care of retries.
+    """
+
+    def __init__(
+        self,
+        response_future: ResponseFuture,
+        map_function: Callable,
+        input: Any,
+        map_kwargs: Any = None,
+        retries: Optional[int] = None,
+    ):
+        self.response_future = response_future
+        self.map_function = map_function
+        self.input = input
+        self.map_kwargs = map_kwargs or {}
+        self.retries = retries or 0
+        self.failure_count = 0
+        self.cancelled = False
+
+    def _inc_failure_count(self):
+        self.failure_count += 1
+
+    def _should_retry(self):
+        return not self.cancelled and self.failure_count <= self.retries
+
+    def _retry(self, function_executor: FunctionExecutor):
+        inputs = [self.input]
+        futures_list = function_executor.map(
+            self.map_function, inputs, **self.map_kwargs
+        )
+        self.response_future = futures_list[0]
+
+    def cancel(self):
+        # cancelling will prevent any further retries, but won't affect any running tasks
+        self.cancelled = True
+
+    @property
+    def done(self):
+        return self.response_future.done
+
+    @property
+    def error(self):
+        return self.response_future.error
+
+    @property
+    def _exception(self):
+        return self.response_future._exception
+
+    @property
+    def stats(self):
+        return self.response_future.stats
+
+    def status(
+        self,
+        throw_except: bool = True,
+        internal_storage: Any = None,
+        check_only: bool = False,
+    ):
+        stat = self.response_future.status(
+            throw_except=throw_except,
+            internal_storage=internal_storage,
+            check_only=check_only,
+        )
+        if self.response_future.error:
+            reraise(*self.response_future._exception)
+        return stat
+
+    def result(self, throw_except: bool = True, internal_storage: Any = None):
+        res = self.response_future.result(
+            throw_except=throw_except, internal_storage=internal_storage
+        )
+        if self.response_future.error:
+            reraise(*self.response_future._exception)
+        return res
+
+
+def map_with_retries(
+    function_executor: FunctionExecutor,
+    map_function: Callable,
+    map_iterdata: List[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]],
+    timeout: Optional[int] = None,
+    include_modules: Optional[List[str]] = [],
+    retries: Optional[int] = None,
+) -> List[RetryingFuture]:
+    """
+    A generalisation of Lithops `map`, with 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,
+        )
+        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.
+    """
+
+    lookup = {f.response_future: f for f in fs}
+
+    while True:
+        response_futures = [f.response_future for f in fs]
+
+        done, pending = function_executor.wait(
+            response_futures,
+            throw_except=throw_except,
+            return_when=return_when,
+            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(function_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)
+
+        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
diff --git a/cubed/runtime/executors/modal.py b/cubed/runtime/executors/modal.py
index e694aa2a..8c489ec9 100644
--- a/cubed/runtime/executors/modal.py
+++ b/cubed/runtime/executors/modal.py
@@ -107,7 +107,7 @@ def execute_dag(
                 if stage.mappable is not None:
                     task_create_tstamp = time.time()
                     for _, stats in app_function.map(
-                        list(stage.mappable),
+                        stage.mappable,
                         order_outputs=False,
                         kwargs=dict(func=stage.function, config=pipeline.config),
                     ):
diff --git a/cubed/runtime/executors/modal_async.py b/cubed/runtime/executors/modal_async.py
index 8ee293cf..e7fc74b3 100644
--- a/cubed/runtime/executors/modal_async.py
+++ b/cubed/runtime/executors/modal_async.py
@@ -15,7 +15,13 @@ from cubed.runtime.utils import handle_callbacks
 
 # We need map_unordered for the use_backups implementation
 async def map_unordered(
-    app_function, input, use_backups=False, return_stats=False, name=None, **kwargs
+    app_function,
+    input,
+    use_backups=False,
+    backup_function=None,
+    return_stats=False,
+    name=None,
+    **kwargs,
 ):
     """
     Apply a function to items of an input list, yielding results as they are completed
@@ -42,6 +48,8 @@ async def map_unordered(
                 yield result
         return
 
+    backup_function = backup_function or app_function
+
     tasks = {
         asyncio.ensure_future(app_function.call.aio(i, **kwargs)): i for i in input
     }
@@ -55,12 +63,14 @@ async def map_unordered(
         finished, pending = await asyncio.wait(
             pending, return_when=asyncio.FIRST_COMPLETED, timeout=2
         )
-        # print("finished", finished)
-        # print("pending", pending)
-
         for task in finished:
             # TODO: use exception groups in Python 3.11 to handle case of multiple task exceptions
             if task.exception():
+                # if the task has a backup that is not done, or is done with no exception, then don't raise this exception
+                backup = backups.get(task, None)
+                if backup:
+                    if not backup.done() or not backup.exception():
+                        continue
                 raise task.exception()
             end_times[task] = time.monotonic()
             if return_stats:
@@ -76,9 +86,11 @@ async def map_unordered(
             if use_backups:
                 backup = backups.get(task, None)
                 if backup:
-                    pending.remove(backup)
+                    if backup in pending:
+                        pending.remove(backup)
                     del backups[task]
                     del backups[backup]
+                    backup.cancel()
 
         if use_backups:
             now = time.monotonic()
@@ -87,8 +99,11 @@ async def map_unordered(
                     task, now, start_times, end_times
                 ):
                     # launch backup task
+                    print("Launching backup task")
                     i = tasks[task]
-                    new_task = asyncio.ensure_future(app_function.call.aio(i, **kwargs))
+                    new_task = asyncio.ensure_future(
+                        backup_function.call.aio(i, **kwargs)
+                    )
                     tasks[new_task] = i
                     start_times[new_task] = time.monotonic()
                     pending.add(new_task)
@@ -120,7 +135,7 @@ async def async_execute_dag(
                 if stage.mappable is not None:
                     async for _, stats in map_unordered(
                         app_function,
-                        list(stage.mappable),
+                        stage.mappable,
                         return_stats=True,
                         name=name,
                         func=stage.function,
diff --git a/cubed/runtime/executors/python_async.py b/cubed/runtime/executors/python_async.py
index 743945d7..ce9a43e9 100644
--- a/cubed/runtime/executors/python_async.py
+++ b/cubed/runtime/executors/python_async.py
@@ -72,7 +72,7 @@ def pipeline_to_stream(concurrent_executor, name, pipeline, **kwargs):
                 map_unordered,
                 concurrent_executor,
                 run_func,
-                list(stage.mappable),
+                stage.mappable,
                 return_stats=True,
                 name=name,
                 func=stage.function,
diff --git a/cubed/storage/virtual.py b/cubed/storage/virtual.py
new file mode 100644
index 00000000..5f585f2a
--- /dev/null
+++ b/cubed/storage/virtual.py
@@ -0,0 +1,45 @@
+from numbers import Integral
+
+import numpy as np
+import zarr
+from zarr.indexing import is_slice
+
+from cubed.types import T_Shape
+
+
+class VirtualOffsetsArray:
+    """An array that is never materialized (in memory or on disk) and contains sequentially incrementing integers."""
+
+    def __init__(self, shape: T_Shape):
+        dtype = np.int32
+        chunks = (1,) * len(shape)
+        # use an empty in-memory Zarr array as a template since it normalizes its properties
+        template = zarr.empty(
+            shape, dtype=dtype, chunks=chunks, store=zarr.storage.MemoryStore()
+        )
+        self.shape = template.shape
+        self.dtype = template.dtype
+        self.chunks = template.chunks
+        self.ndim = template.ndim
+
+    def __getitem__(self, key):
+        if key == () and self.shape == ():
+            return np.array(0, dtype=self.dtype)
+        return np.ravel_multi_index(_key_to_index_tuple(key), self.shape)
+
+
+def _key_to_index_tuple(selection):
+    assert all(isinstance(s, (slice, Integral)) for s in selection)
+    sel = []
+    for s in selection:
+        if isinstance(s, Integral):
+            sel.append(s)
+        elif is_slice(s) and s.stop == s.start + 1 and (s.step is None or s.step == 1):
+            sel.append(s.start)
+        else:
+            raise NotImplementedError(f"Offset selection not supported: {selection}")
+    return tuple(sel)
+
+
+def virtual_offsets(shape: T_Shape) -> VirtualOffsetsArray:
+    return VirtualOffsetsArray(shape)
diff --git a/cubed/storage/zarr.py b/cubed/storage/zarr.py
index a99e28f4..cdbd09f7 100644
--- a/cubed/storage/zarr.py
+++ b/cubed/storage/zarr.py
@@ -1,4 +1,9 @@
+from typing import Any, Optional, Union
+
 import zarr
+from numpy import ndarray
+
+from cubed.types import T_DType, T_RegularChunks, T_Shape, T_Store
 
 
 class LazyZarrArray:
@@ -11,12 +16,12 @@ class LazyZarrArray:
 
     def __init__(
         self,
-        shape,
-        dtype,
-        chunks,
-        store,
-        initial_values=None,
-        fill_value=None,
+        shape: T_Shape,
+        dtype: T_DType,
+        chunks: T_RegularChunks,
+        store: T_Store,
+        initial_values: Optional[ndarray] = None,
+        fill_value: Any = None,
         **kwargs,
     ):
         """Create a Zarr array lazily in memory."""
@@ -33,7 +38,7 @@ class LazyZarrArray:
         self.fill_value = fill_value
         self.kwargs = kwargs
 
-    def create(self, mode="w-"):
+    def create(self, mode: str = "w-") -> zarr.Array:
         """Create the Zarr array in storage.
 
         The Zarr array's metadata is initialized in the backing store, and any
@@ -45,7 +50,7 @@ class LazyZarrArray:
             The mode to open the Zarr array with using ``zarr.open``.
             Default is 'w-', which means create, fail it already exists.
         """
-        target = zarr.open(
+        target = zarr.open_array(
             self.store,
             mode=mode,
             shape=self.shape,
@@ -58,13 +63,13 @@ class LazyZarrArray:
             target[...] = self.initial_values
         return target
 
-    def open(self):
+    def open(self) -> zarr.Array:
         """Open the Zarr array for reading or writing and return it.
 
         Note that the Zarr array must have been created or this method will raise an exception.
         """
         # r+ means read/write, fail if it doesn't exist
-        return zarr.open(
+        return zarr.open_array(
             self.store,
             mode="r+",
             shape=self.shape,
@@ -72,24 +77,39 @@ class LazyZarrArray:
             chunks=self.chunks,
         )
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"cubed.storage.zarr.LazyZarrArray<shape={self.shape}, dtype={self.dtype}, chunks={self.chunks}>"
 
 
-def lazy_empty(shape, *, dtype, chunks, store, **kwargs):
+T_ZarrArray = Union[zarr.Array, LazyZarrArray]
+
+
+def lazy_empty(
+    shape: T_Shape, *, dtype: T_DType, chunks: T_RegularChunks, store: T_Store, **kwargs
+) -> LazyZarrArray:
     return LazyZarrArray(shape, dtype, chunks, store, **kwargs)
 
 
-def lazy_from_array(array, *, dtype, chunks, store, **kwargs):
+def lazy_from_array(
+    array: ndarray, *, dtype: T_DType, chunks: T_RegularChunks, store: T_Store, **kwargs
+) -> LazyZarrArray:
     return LazyZarrArray(
         array.shape, dtype, chunks, store, initial_values=array, **kwargs
     )
 
 
-def lazy_full(shape, fill_value, *, dtype, chunks, store, **kwargs):
+def lazy_full(
+    shape: T_Shape,
+    fill_value: Any,
+    *,
+    dtype: T_DType,
+    chunks: T_RegularChunks,
+    store: T_Store,
+    **kwargs,
+) -> LazyZarrArray:
     return LazyZarrArray(shape, dtype, chunks, store, fill_value=fill_value, **kwargs)
 
 
-def open_if_lazy_zarr_array(array):
+def open_if_lazy_zarr_array(array: T_ZarrArray) -> zarr.Array:
     """If array is a LazyZarrArray then open it, leaving other arrays unchanged."""
     return array.open() if isinstance(array, LazyZarrArray) else array
diff --git a/cubed/types.py b/cubed/types.py
new file mode 100644
index 00000000..6de07c1a
--- /dev/null
+++ b/cubed/types.py
@@ -0,0 +1,21 @@
+from typing import Any, Dict, Literal, Tuple, Union
+
+T_Shape = Tuple[int, ...]
+
+T_DType = Any  # TODO: improve this
+
+# Regular chunks are where all chunks must be the same size, whereas
+# rectangular chunks may be different sizes, and contain the size of
+# each chunk along each axis.
+# Zarr currently supports only regular chunks, although rectangular chunks
+# are proposed in https://zarr.dev/zeps/draft/ZEP0003.html.
+# Dask supports rectangular chunks.
+# The general T_Chunks type is for specifying chunks, it is generally
+# converted to a T_RectangularChunks or T_RegularChunks internally.
+T_RegularChunks = Tuple[int, ...]
+T_RectangularChunks = Tuple[Tuple[int, ...], ...]
+T_Chunks = Union[
+    int, T_RegularChunks, T_RectangularChunks, Dict[Any, Any], Literal["auto"]
+]
+
+T_Store = str  # TODO: expand this
diff --git a/cubed/utils.py b/cubed/utils.py
index 96bc55a5..d608ca6e 100644
--- a/cubed/utils.py
+++ b/cubed/utils.py
@@ -9,23 +9,24 @@ from math import prod
 from operator import add
 from pathlib import Path
 from posixpath import join
-from typing import Union
+from typing import Tuple, Union
 from urllib.parse import quote, unquote, urlsplit, urlunsplit
 
 import numpy as np
 import tlz as toolz
 
+from cubed.types import T_DType, T_RectangularChunks, T_RegularChunks
 from cubed.vendor.dask.array.core import _check_regular_chunks
 
 PathType = Union[str, Path]
 
 
-def chunk_memory(dtype, chunksize):
+def chunk_memory(dtype: T_DType, chunksize: T_RegularChunks) -> int:
     """Calculate the amount of memory in bytes that a single chunk uses."""
     return np.dtype(dtype).itemsize * prod(chunksize)
 
 
-def get_item(chunks, idx):
+def get_item(chunks: T_RectangularChunks, idx: Tuple[int, ...]) -> Tuple[slice, ...]:
     """Convert a chunk index to a tuple of slices."""
     # could use Dask's cached_cumsum here if it improves performance
     starts = tuple(_cumsum(c, initial_zero=True) for c in chunks)
@@ -48,7 +49,7 @@ def join_path(dir_url: PathType, child_path: str) -> str:
     return urlunsplit(split_parts)
 
 
-def memory_repr(num):
+def memory_repr(num: int) -> str:
     """Convert bytes to a human-readable string in decimal form.
     1 KB is 1,000 bytes, 1 MB is 1,000,000 bytes, and so on.
 
@@ -61,16 +62,20 @@ def memory_repr(num):
     -------
     str
     """
+    if num < 0:
+        raise ValueError(f"Invalid value: {num}. Expected a positive integer.")
     if num < 1000.0:
         return f"{num} bytes"
-    num /= 1000.0
+    val = num / 1000.0
     for x in ["KB", "MB", "GB", "TB", "PB"]:
-        if num < 1000.0:
-            return f"{num:3.1f} {x}"
-        num /= 1000.0
+        if val < 1000.0:
+            return f"{val:3.1f} {x}"
+        val /= 1000.0
+    # fall back to scientific notation
+    return f"{num:.1e} bytes"
 
 
-def peak_measured_mem():
+def peak_measured_mem() -> int:
     """Measures the peak memory usage in bytes.
 
     Note: this function currently doesn't work on Windows.
@@ -87,7 +92,7 @@ def peak_measured_mem():
     return ru_maxrss * 1024 if platform.system() == "Linux" else ru_maxrss
 
 
-def to_chunksize(chunkset):
+def to_chunksize(chunkset: T_RectangularChunks) -> T_RegularChunks:
     """Convert a chunkset to a chunk size for Zarr.
 
     Parameters
diff --git a/cubed/vendor/dask/array/core.pyi b/cubed/vendor/dask/array/core.pyi
new file mode 100644
index 00000000..c9dfcaf4
--- /dev/null
+++ b/cubed/vendor/dask/array/core.pyi
@@ -0,0 +1,20 @@
+from typing import List, Optional, Tuple
+
+from cubed.types import T_Chunks, T_DType, T_RectangularChunks, T_Shape
+
+def _check_regular_chunks(chunkset: T_RectangularChunks) -> bool: ...
+
+def broadcast_chunks(
+    *chunkss
+) -> Tuple[Tuple[int, ...]]: ...
+
+def common_blockdim(blockdims: List[Tuple[int, ...]]) -> Tuple[int, ...]: ...
+
+def normalize_chunks(
+    chunks: T_Chunks,
+    shape: Optional[T_Shape] = None,
+    limit: Optional[int] = None,
+    dtype: Optional[T_DType] = None,
+    previous_chunks: Optional[T_RectangularChunks] = None,
+) -> T_RectangularChunks:
+    ...
diff --git a/docs/index.md b/docs/index.md
index e7214230..55899f21 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -9,6 +9,7 @@ Cubed is a distributed N-dimensional array library implemented in Python using b
 - Follows [Dask Array](https://docs.dask.org/en/stable/array.html)'s chunked array API (`map_blocks`, `rechunk`, etc)
 - [Zarr](https://zarr.readthedocs.io/en/stable/) for storage
 - Multiple serverless runtimes: Python (in-process), [Lithops](https://lithops-cloud.github.io/), [Modal](https://modal.com/), [Apache Beam](https://beam.apache.org/)
+- Integration with [Xarray](https://xarray.dev/) via [cubed-xarray](https://github.com/xarray-contrib/cubed-xarray)
 
 ## Documentation
 
diff --git a/docs/operations.md b/docs/operations.md
index 751dc97d..0f134b37 100644
--- a/docs/operations.md
+++ b/docs/operations.md
@@ -10,6 +10,7 @@ The following diagram shows the dependencies between operations. Array API funct
 
 Note how fundamental `blockwise` is - all array API functions depend on it.
 
+(elemwise-operation)=
 ## `elemwise`
 
 The simplest core operation is `elemwise`, which maps input array elements to output array elements, a block at a time.
@@ -62,6 +63,7 @@ This example shows how `outer` is implemented using `blockwise`. Each block from
 
 ![The blockwise primitive operation](images/blockwise.svg)
 
+(rechunk-operation)=
 ## `rechunk`
 
 The `rechunk` operation is a primitive operation for changing the chunking of an array, without changing its shape or dtype.
diff --git a/docs/related-projects.md b/docs/related-projects.md
index 2d393677..fbd39502 100644
--- a/docs/related-projects.md
+++ b/docs/related-projects.md
@@ -10,6 +10,12 @@ Dask only has a single distributed runtime, Dask Distributed, whereas Cubed has
 
 The core operations and array API layers in Cubed are heavily influenced by Dask Array - in both naming, as well as implementation (Cubed uses Dask Array for some chunking utilities).
 
+## Xarray
+
+You can use Cubed with Xarray using the [cubed-xarray](https://github.com/xarray-contrib/cubed-xarray) package.
+
+Read more about the integration in [Cubed: Bounded-memory serverless array processing in xarray](https://xarray.dev/blog/cubed-xarray).
+
 ## Previous work
 
 This project is a continuation of a similar project I worked on, called [Zappy](https://github.com/lasersonlab/zappy/tree/master/zappy). What's changed in the intervening three years? Rechunker was created (I wasn't so concerned with memory when working on Zappy). The Python Array API standard has been created, which makes implementing a new array API less daunting than implementing the NumPy API. And I have a better understanding of Dask, and the fundamental nature of the blockwise operation.
diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md
index 416dd5c5..90e3f90f 100644
--- a/docs/user-guide/index.md
+++ b/docs/user-guide/index.md
@@ -10,4 +10,5 @@ executors
 storage
 memory
 reliability
+scaling
 ```
diff --git a/docs/user-guide/reliability.md b/docs/user-guide/reliability.md
index 8bda0102..0e819d27 100644
--- a/docs/user-guide/reliability.md
+++ b/docs/user-guide/reliability.md
@@ -2,6 +2,13 @@
 
 Large scale distributed systems must be fault tolerant, and Cubed is no exception. This section covers some of the features that allow you to run computations in a reliable fashion.
 
+## Strong consistency
+
+Cubed relies on file storage providing strong global consistency. This means that after writing a file to storage, it is
+immediately available for reading, even by another client. For Cubed, this means that once a Zarr chunk has been written to storage it can be read by another process.
+
+Importantly, both [Amazon S3](https://aws.amazon.com/s3/consistency/) and [Google Cloud Storage](https://cloud.google.com/storage/docs/consistency) provide strong consistency.
+
 ## Retries
 
 If a task fails - with an IO exception when reading or writing to Zarr, for example - it will be retried, up to a total of three attempts. If it fails again then the whole computation will fail with an error message. Note that currently the number of retries is not configurable.
@@ -14,6 +21,6 @@ If a task takes longer than a pre-determined amount of time then it will be cons
 
 A few slow running tasks (called stragglers) can disproportionately slow down the whole computation. To mitigate this, speculative duplicate tasks are launched in certain circumstances, acting as backups that complete more quickly than the straggler, hence bringing down the overall time taken.
 
-When a backup task is launched the original task is not cancelled, so it is to be expected that both tasks will complete and write their (identical) output. This only works since tasks are idempotent and each write a single, whole Zarr chunk in an atomic operation.
+When a backup task is launched the original task is not cancelled, so it is to be expected that both tasks will complete and write their (identical) output. This only works since tasks are idempotent and each write a single, whole Zarr chunk in an atomic operation. (Updates to a single key are atomic in both [Amazon S3](https://docs.aws.amazon.com/AmazonS3/latest/userguide/Welcome.html#ConsistencyModel) and Google Cloud Storage.)
 
 Note that this feature is experimental and disabled by default since it has not been tested at scale yet.
diff --git a/docs/user-guide/scaling.md b/docs/user-guide/scaling.md
new file mode 100644
index 00000000..7d7db6a1
--- /dev/null
+++ b/docs/user-guide/scaling.md
@@ -0,0 +1,105 @@
+# Scaling
+
+Cubed is specifically designed to maintain optimal performance while processing large arrays (of order Terabytes so far).
+This page aims to provide a deeper understanding of how Cubed's design scales in both theoretical and practical scenarios.
+
+## Preface: Types of Scaling
+
+There are different types of scaling to consider in distributed computing:
+
+**Horizontal versus Vertical Scaling**: Horizontal scaling refers to adding more machines to the system to improve its throughput, while vertical scaling means upgrading an existing machine to a larger one with more speed and resources.
+
+**Weak versus Strong Scaling**: Strong scaling is defined as how the solution time varies with the number of processors for a fixed total problem size.
+Weak scaling is defined as how the solution time varies with the number of processors for a fixed problem size per processor.
+In other words, strong scaling measures how much faster you can get a given problem done, whereas weak scaling measures how big a problem you can get done in a reasonable time.
+
+Cubed has been designed to scale horizontally, in such a way as to display good weak scaling even when processing large array datasets.
+
+
+## Theoretical vs Practical Scaling of Cubed
+
+To understand factors affecting the scaling and performance of Cubed, we will start by considering the simplest possible calculation and add complexity.
+
+### Single-step Calculation
+
+The  simplest non-trivial operation in Cubed would map one block from the input array to one block in the output array, with no intermediate persistent stores.
+Changing the sign of every element in an array would be an example of this type of operation, known as an [elementwise](#elemwise-operation) operation.
+
+In an ideal environment where the serverless service provides infinite workers, the limiting factor for scaling would be concurrent writes to Zarr.
+In such a case weak scaling should be linear, i.e. an array with more blocks could be processed in the same amount of time given proportionally more workers to process those blocks.
+
+In practice, this ideal scenario may not be achieved, for a number of reasons.
+Firstly, you need to make sure you're using a parallel executor, not the default single-threaded executor.
+You also need enough parallelism to match the scale of your problem.
+Weak scaling requires more workers than output chunks, so for large problems it might be necessary to adjust the executor's configuration to not restrict the ``max_workers``.
+With fewer workers than chunks we would expect linear strong scaling, as every new worker added has nothing to wait for.
+
+Stragglers are tasks that take much longer than average, who disproportionately hold up the next step of the computation.
+To handle stragglers, you should consider turning on backups (with ``use_backups=True``), as any failures that are restarted essentially become stragglers.
+Worker start-up time is another practical speed consideration, though it would delay computations of all scales equally.
+
+### Multi-step Calculation
+
+A multi-step calculation requires writing one or more intermediate arrays to persistent storage.
+One important example in Cubed is the [rechunk](#rechunk-operation) operation, which guarantees bounded memory usage by writing and reading from one intermediate persistent Zarr store.
+
+In multi-step calculations, the number of steps in the plan sets the minimum total execution time.
+Hence, reducing the number of steps in the plan can lead to significant performance improvements.
+Reductions can be carried out in fewer iterative steps if ``allowed_mem`` is larger.
+Cubed automatically fuses some steps to enhance performance, but others (especially rechunk) cannot be fused without requiring a shuffle, which can potentially violate memory constraints.
+
+```{note} In theory multiple blockwise operations can be fused together, enhancing the performance further. However this has not yet been implemented in Cubed.
+```
+
+In practical scenarios, stragglers can hold up the completion of each step separately, thereby cumulatively affecting the overall performance of the calculation.
+
+### Multi-pipeline Calculation
+
+A "pipeline" refers to an independent branch of the calculation.
+For example, if you have two separate arrays to compute simultaneously, full parallelism requires sufficient workers for both tasks.
+The same logic applies if you have two arrays feeding into a single array or vice versa.
+
+
+```{note} Currently Cubed will not necessarily execute independent pipelines in parallel on all executors.
+```
+
+
+## Other Performance Considerations
+
+### Different Executors
+
+Different executors come with their own set of advantages and disadvantages.
+For instance, some serverless executors have much faster worker startup times (such as Modal), while others may have different limits to the maximum workers.
+Cluster-based executors (e.g. the Dask executor) will have different performance characteristics depending on how resources are provisioned for the cluster.
+Some other executors such as Beam will first convert the Cubed Plan to a different representation before executing, which will affect performance.
+
+### Different Cloud Providers
+
+Different cloud providers' serverless offerings may perform differently. For example, Google Cloud Functions (GCF) has been observed to have more stragglers than AWS Lambda.
+
+
+## Diagnosing Performance
+
+To understand how your computation could perform better you first need to diagnose the source of any problems.
+
+### Optimized Plan
+
+Use {py:meth}`Plan.visualize() <cubed.Plan.visualize()>` to view the optimized plan. This allows you to see the number of steps involved in your calculation, the number of tasks in each step, and overall.
+
+### History Callback
+
+The history callback function can help determine how much time was spent in worker startup, as well as how much stragglers affected the overall speed.
+
+### Timeline Visualization Callback
+
+A timeline visualization callback can provide a visual representation of the above points. Ideally, we want vertical lines on this plot, which would represent perfect horizontal scaling.
+
+
+## Tips
+
+In Cubed, there are very few "magic numbers", meaning calculations generally take as long as they need to, with few other parameters to tune. Here are a few suggestions:
+
+* Use {py:func}`measure_reserved_mem <cubed.measure_reserved_mem>`.
+* Stick to ~100MB chunks.
+* Set ``allowed_mem`` to around ~2GB (or larger if necessary).
+* Use Cubed only for the part of the calculation where it's needed.
diff --git a/examples/lithops/gcf/lithops-matmul-ones.py b/examples/lithops/gcf/lithops-matmul-ones.py
new file mode 100644
index 00000000..5ab7d5cb
--- /dev/null
+++ b/examples/lithops/gcf/lithops-matmul-ones.py
@@ -0,0 +1,39 @@
+import logging
+import sys
+
+from tqdm.contrib.logging import logging_redirect_tqdm
+
+import cubed
+import cubed.array_api as xp
+import cubed.random
+from cubed.extensions.history import HistoryCallback
+from cubed.extensions.timeline import TimelineVisualizationCallback
+from cubed.extensions.tqdm import TqdmProgressBar
+from cubed.runtime.executors.lithops import LithopsDagExecutor
+
+logging.basicConfig(level=logging.INFO)
+# suppress harmless connection pool warnings
+logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR)
+
+if __name__ == "__main__":
+    tmp_path = sys.argv[1]
+    runtime = sys.argv[2]
+    spec = cubed.Spec(tmp_path, allowed_mem="2GB")
+    executor = LithopsDagExecutor()
+
+    # Note we use default float dtype, since np.matmul is not optimized for ints
+    a = xp.ones((50000, 50000), chunks=(5000, 5000), spec=spec)
+    b = xp.ones((50000, 50000), chunks=(5000, 5000), spec=spec)
+    c = xp.matmul(a, b)
+    d = xp.all(c == 50000)
+    with logging_redirect_tqdm():
+        progress = TqdmProgressBar()
+        hist = HistoryCallback()
+        timeline_viz = TimelineVisualizationCallback()
+        res = d.compute(
+            executor=executor,
+            callbacks=[progress, hist, timeline_viz],
+            runtime=runtime,
+            runtime_memory=2048,  # Note that Lithops/Google Cloud Functions only accepts powers of 2 for this argument.
+        )
+        assert res, "Validation failed"
diff --git a/examples/modal/aws/README.md b/examples/modal/aws/README.md
index 5d5c3d7e..ae114833 100644
--- a/examples/modal/aws/README.md
+++ b/examples/modal/aws/README.md
@@ -1,5 +1,7 @@
 # Examples running Cubed on Modal
 
+**Warning: Modal does not guarantee that functions run in any particular cloud region, so it is not currently recommended that you run large computations since excessive data transfer fees are possible.**
+
 ## Pre-requisites
 
 1. A [Modal account](https://modal.com/)
diff --git a/examples/modal/gcp/README.md b/examples/modal/gcp/README.md
index 31f8e235..63dad766 100644
--- a/examples/modal/gcp/README.md
+++ b/examples/modal/gcp/README.md
@@ -1,5 +1,7 @@
 # Examples running Cubed on Modal
 
+**Warning: Modal does not guarantee that functions run in any particular cloud region, so it is not currently recommended that you run large computations since excessive data transfer fees are possible.**
+
 ## Pre-requisites
 
 1. A [Modal account](https://modal.com/)
diff --git a/pyproject.toml b/pyproject.toml
index a617f077..e57aaabc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "cubed"
-version = "0.8.0"
+version = "0.9.0"
 authors = [
     {name = "Tom White", email = "tom.e.white@gmail.com"},
 ]


Test patch:
diff --git a/cubed/tests/runtime/test_lithops_retries.py b/cubed/tests/runtime/test_lithops_retries.py
new file mode 100644
index 00000000..333d0593
--- /dev/null
+++ b/cubed/tests/runtime/test_lithops_retries.py
@@ -0,0 +1,72 @@
+import pytest
+
+pytest.importorskip("lithops")
+
+from lithops.executors import LocalhostExecutor
+
+from cubed.runtime.executors.lithops_retries import map_with_retries, wait_with_retries
+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,
+            function,
+            input,
+            timeout=timeout,
+            retries=retries,
+        )
+        done, pending = wait_with_retries(executor, futures, throw_except=False)
+        assert len(pending) == 0
+    outputs = set(f.result() for f in done)
+    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):
+    partial_map_function = lambda x: deterministic_failure(tmp_path, timing_map, x)
+    outputs = run_test(
+        function=partial_map_function,
+        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):
+    partial_map_function = lambda x: deterministic_failure(tmp_path, timing_map, x)
+    with pytest.raises(RuntimeError):
+        run_test(
+            function=partial_map_function,
+            input=range(n_tasks),
+            retries=retries,
+        )
+
+    check_invocation_counts(tmp_path, timing_map, n_tasks, retries)
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index 4f79883e..99ea8138 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -9,6 +9,7 @@ from numpy.testing import assert_array_equal
 
 import cubed
 import cubed.array_api as xp
+import cubed.random
 from cubed.extensions.history import HistoryCallback
 from cubed.extensions.timeline import TimelineVisualizationCallback
 from cubed.extensions.tqdm import TqdmProgressBar
@@ -512,3 +513,20 @@ def test_measure_reserved_mem(executor):
 
     reserved_memory = cubed.measure_reserved_mem(executor=executor)
     assert reserved_memory > 1_000_000  # over 1MB
+
+
+# Test we can create a plan for arrays of up to 5PB, and 100s of billions of tasks
+@pytest.mark.parametrize("factor", [10, 20, 50, 100, 200, 500, 1000, 2000, 5000])
+def test_plan_scaling(tmp_path, factor):
+    spec = cubed.Spec(tmp_path, allowed_mem="2GB")
+    chunksize = 5000
+    a = cubed.random.random(
+        (factor * chunksize, factor * chunksize), chunks=chunksize, spec=spec
+    )
+    b = cubed.random.random(
+        (factor * chunksize, factor * chunksize), chunks=chunksize, spec=spec
+    )
+    c = xp.matmul(a, b)
+
+    assert c.plan.num_tasks() > 0
+    c.visualize(filename=tmp_path / "c")
diff --git a/cubed/tests/test_mem_utilization.py b/cubed/tests/test_mem_utilization.py
index f1be412c..28936631 100644
--- a/cubed/tests/test_mem_utilization.py
+++ b/cubed/tests/test_mem_utilization.py
@@ -20,7 +20,7 @@ def spec(tmp_path, reserved_mem):
 @pytest.fixture(scope="module")
 def reserved_mem():
     executor = LithopsDagExecutor(config=LITHOPS_LOCAL_CONFIG)
-    res = cubed.measure_reserved_mem(executor) * 1.05  # add some wiggle room
+    res = cubed.measure_reserved_mem(executor) * 1.1  # add some wiggle room
     return round_up_to_multiple(res, 10_000_000)  # round up to nearest multiple of 10MB
 
 
diff --git a/cubed/tests/test_utils.py b/cubed/tests/test_utils.py
index 30782341..fd1ebdaa 100644
--- a/cubed/tests/test_utils.py
+++ b/cubed/tests/test_utils.py
@@ -54,6 +54,9 @@ def test_memory_repr():
     assert memory_repr(9_999) == "10.0 KB"
     assert memory_repr(1_000_000) == "1.0 MB"
     assert memory_repr(1_000_000_000_000_000) == "1.0 PB"
+    assert memory_repr(int(1e18)) == "1.0e+18 bytes"
+    with pytest.raises(ValueError):
+        memory_repr(-1)
 
 
 @pytest.mark.skipif(platform.system() == "Windows", reason="does not run on windows")