Repository: tomwhite/cubed
Pull request #259 (2023-07-13 16:25:12)

Lithops compute arrays in parallel
This is #80 for Lithops.

I'm still experimenting with this, but you can see the effect here, where the first two arrays (from `lithops-add-random.py`) are computed at the same time, rather than sequentially:

![1689264590_timeline](https://github.com/tomwhite/cubed/assets/85085/9a040382-2e86-4cef-87cb-5dad2992e300)



Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 18d756f2..c4c339d2 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -12,7 +12,7 @@ repos:
       - id: isort
         args: ["--profile", "black"]
   - repo: https://github.com/psf/black
-    rev: 23.3.0
+    rev: 23.7.0
     hooks:
       - id: black
         exclude: vendor/
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index 6a64558c..1068358e 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -649,6 +649,34 @@ def rechunk(x, chunks, target_store=None):
         return Array(name, pipeline2.target_array, spec, plan2)
 
 
+def merge_chunks(x, chunks):
+    target_chunks = normalize_chunks(chunks, x.shape, dtype=x.dtype)
+
+    target_chunksize = to_chunksize(target_chunks)
+    if len(target_chunksize) != x.ndim:
+        raise ValueError(
+            f"Chunks {target_chunksize} must have same number of dimensions as array ({x.ndim})"
+        )
+    if not all(c1 % c0 == 0 for c0, c1 in zip(x.chunksize, target_chunksize)):
+        raise ValueError(
+            f"Chunks {target_chunksize} must be a multiple of array's chunks {x.chunksize}"
+        )
+
+    return map_direct(
+        _copy_chunk,
+        x,
+        shape=x.shape,
+        dtype=x.dtype,
+        chunks=target_chunks,
+        extra_projected_mem=0,
+        target_chunks=target_chunks,
+    )
+
+
+def _copy_chunk(e, x, target_chunks=None, block_id=None):
+    return x.zarray[get_item(target_chunks, block_id)]
+
+
 def reduction(
     x: "Array",
     func,
@@ -694,9 +722,9 @@ def reduction(
             adjust_chunks=adjust_chunks,
         )
 
-    # rechunk/reduce along axis in multiple rounds until there's a single block in each reduction axis
+    # merge/reduce along axis in multiple rounds until there's a single block in each reduction axis
     while any(n > 1 for i, n in enumerate(result.numblocks) if i in axis):
-        # rechunk along axis
+        # merge along axis
         target_chunks = list(result.chunksize)
         chunk_mem = chunk_memory(intermediate_dtype, result.chunksize)
         for i, s in enumerate(result.shape):
@@ -708,7 +736,7 @@ def reduction(
                     target_chunks[i] = min(s, x.chunksize[i])
                 else:
                     # single axis: see how many result chunks fit in max_mem
-                    # factor of 4 is memory for {compressed, uncompressed} x {input, output} (see rechunk.py)
+                    # factor of 4 is memory for {compressed, uncompressed} x {input, output}
                     target_chunk_size = (max_mem - chunk_mem) // (chunk_mem * 4)
                     if target_chunk_size <= 1:
                         raise ValueError(
@@ -716,7 +744,7 @@ def reduction(
                         )
                     target_chunks[i] = min(s, target_chunk_size)
         _target_chunks = tuple(target_chunks)
-        result = rechunk(result, _target_chunks)
+        result = merge_chunks(result, _target_chunks)
 
         # reduce chunks (if any axis chunksize is > 1)
         if any(s > 1 for i, s in enumerate(result.chunksize) if i in axis):
diff --git a/cubed/runtime/executors/lithops.py b/cubed/runtime/executors/lithops.py
index 071607d3..f1cceffb 100644
--- a/cubed/runtime/executors/lithops.py
+++ b/cubed/runtime/executors/lithops.py
@@ -1,6 +1,8 @@
+import collections
 import copy
 import logging
 import time
+from functools import partial
 from typing import (
     Any,
     Callable,
@@ -19,7 +21,7 @@ from lithops.wait import ALWAYS, ANY_COMPLETED
 from networkx import MultiDiGraph
 
 from cubed.core.array import Callback
-from cubed.core.plan import visit_nodes
+from cubed.core.plan import visit_node_generations, visit_nodes
 from cubed.runtime.backup import should_launch_backup
 from cubed.runtime.executors.lithops_retries import (
     RetryingFuture,
@@ -39,8 +41,11 @@ def run_func(input, func=None, config=None, name=None):
 
 def map_unordered(
     lithops_function_executor: FunctionExecutor,
-    map_function: Callable[..., Any],
-    map_iterdata: Iterable[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]],
+    group_map_functions: Sequence[Callable[..., Any]],
+    group_map_iterdata: Sequence[
+        Iterable[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]]
+    ],
+    group_names: Sequence[str],
     include_modules: List[str] = [],
     timeout: Optional[int] = None,
     retries: int = 2,
@@ -55,8 +60,9 @@ def map_unordered(
     A generalisation of Lithops `map`, with retries, and relaxed return ordering.
 
     :param lithops_function_executor: The Lithops function executor to use.
-    :param map_function: The function to map over the data.
-    :param map_iterdata: An iterable of input data.
+    :param group_map_functions: A sequence of functions to map over the data.
+    :param group_map_iterdata: A sequence of iterables of input data.
+    :param group_names: The names of the function/iterable groups.
     :param include_modules: Modules to include.
     :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.
@@ -66,28 +72,33 @@ def map_unordered(
     """
     return_when = ALWAYS if use_backups else ANY_COMPLETED
 
-    inputs = list(map_iterdata)
-    tasks = {}
-    start_times = {}
-    end_times = {}
+    group_name_to_function: Dict[str, Callable[..., Any]] = {}
+    # backups are launched based on task start and end times for the group
+    start_times: Dict[str, Dict[RetryingFuture, float]] = {}
+    end_times: Dict[str, Dict[RetryingFuture, float]] = collections.defaultdict(dict)
     backups: Dict[RetryingFuture, RetryingFuture] = {}
-    pending = []
-
-    # can't use functools.partial here as we get an error in lithops
-    # also, lithops extra_args doesn't work for this case
-    partial_map_function = lambda x: map_function(x, **kwargs)
-
-    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)
+    pending: List[RetryingFuture] = []
+    group_name: str
+
+    for map_function, map_iterdata, group_name in zip(
+        group_map_functions, group_map_iterdata, group_names
+    ):
+        # can't use functools.partial here as we get an error in lithops
+        # also, lithops extra_args doesn't work for this case
+        partial_map_function = lambda x: map_function(x, **kwargs)
+        group_name_to_function[group_name] = partial_map_function
+
+        futures = map_with_retries(
+            lithops_function_executor,
+            partial_map_function,
+            map_iterdata,
+            timeout=timeout,
+            include_modules=include_modules,
+            retries=retries,
+            group_name=group_name,
+        )
+        start_times[group_name] = {k: time.monotonic() for k in futures}
+        pending.extend(futures)
 
     while pending:
         finished, pending = wait_with_retries(
@@ -105,9 +116,10 @@ def map_unordered(
                     if not backup.done or not backup.error:
                         continue
                 future.status(throw_except=True)
-            end_times[future] = time.monotonic()
+            group_name = future.group_name  # type: ignore[assignment]
+            end_times[group_name][future] = time.monotonic()
             if return_stats:
-                yield future.result(), standardise_lithops_stats(future.stats)
+                yield future.result(), standardise_lithops_stats(future)
             else:
                 yield future.result()
 
@@ -124,21 +136,24 @@ def map_unordered(
         if use_backups:
             now = time.monotonic()
             for future in copy.copy(pending):
+                group_name = future.group_name  # type: ignore[assignment]
                 if future not in backups and should_launch_backup(
-                    future, now, start_times, end_times
+                    future, now, start_times[group_name], end_times[group_name]
                 ):
-                    input = tasks[future]
+                    input = future.input
                     logger.info("Running backup task for %s", input)
                     futures = map_with_retries(
                         lithops_function_executor,
-                        partial_map_function,
+                        group_name_to_function[group_name],
                         [input],
                         timeout=timeout,
                         include_modules=include_modules,
                         retries=0,  # don't retry backup tasks
+                        group_name=group_name,
+                    )
+                    start_times[group_name].update(
+                        {k: time.monotonic() for k in futures}
                     )
-                    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]
                     backups[future] = backup
@@ -151,32 +166,57 @@ def execute_dag(
     callbacks: Optional[Sequence[Callback]] = None,
     array_names: Optional[Sequence[str]] = None,
     resume: Optional[bool] = None,
+    compute_arrays_in_parallel: Optional[bool] = None,
     **kwargs,
 ) -> None:
     use_backups = kwargs.pop("use_backups", False)
     with FunctionExecutor(**kwargs) as executor:
-        for name, node in visit_nodes(dag, resume=resume):
-            pipeline = node["pipeline"]
-            for stage in pipeline.stages:
-                if stage.mappable is not None:
-                    for _, stats in map_unordered(
-                        executor,
-                        run_func,
-                        stage.mappable,
-                        func=stage.function,
-                        config=pipeline.config,
-                        name=name,
-                        use_backups=use_backups,
-                        return_stats=True,
-                    ):
-                        stats["array_name"] = name
-                        handle_callbacks(callbacks, stats)
-                else:
-                    raise NotImplementedError()
-
-
-def standardise_lithops_stats(stats):
+        if not compute_arrays_in_parallel:
+            for name, node in visit_nodes(dag, resume=resume):
+                pipeline = node["pipeline"]
+                for stage in pipeline.stages:
+                    if stage.mappable is not None:
+                        for _, stats in map_unordered(
+                            executor,
+                            [run_func],
+                            [stage.mappable],
+                            [name],
+                            func=stage.function,
+                            config=pipeline.config,
+                            name=name,
+                            use_backups=use_backups,
+                            return_stats=True,
+                        ):
+                            handle_callbacks(callbacks, stats)
+                    else:
+                        raise NotImplementedError()
+        else:
+            for gen in visit_node_generations(dag, resume=resume):
+                group_map_functions = []
+                group_map_iterdata = []
+                group_names = []
+                for name, node in gen:
+                    pipeline = node["pipeline"]
+                    stage = pipeline.stages[0]  # assume one
+                    f = partial(run_func, func=stage.function, config=pipeline.config)
+                    group_map_functions.append(f)
+                    group_map_iterdata.append(stage.mappable)
+                    group_names.append(name)
+                for _, stats in map_unordered(
+                    executor,
+                    group_map_functions,
+                    group_map_iterdata,
+                    group_names,
+                    use_backups=use_backups,
+                    return_stats=True,
+                ):
+                    handle_callbacks(callbacks, stats)
+
+
+def standardise_lithops_stats(future: RetryingFuture) -> Dict[str, Any]:
+    stats = future.stats
     return dict(
+        array_name=future.group_name,
         task_create_tstamp=stats["host_job_create_tstamp"],
         function_start_tstamp=stats["worker_func_start_tstamp"],
         function_end_tstamp=stats["worker_func_end_tstamp"],
diff --git a/cubed/runtime/executors/lithops_retries.py b/cubed/runtime/executors/lithops_retries.py
index 9306363d..e4cb4903 100644
--- a/cubed/runtime/executors/lithops_retries.py
+++ b/cubed/runtime/executors/lithops_retries.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
 
 from lithops import FunctionExecutor
 from lithops.future import ResponseFuture
@@ -18,12 +18,14 @@ class RetryingFuture:
         input: Any,
         map_kwargs: Any = None,
         retries: Optional[int] = None,
+        group_name: Optional[str] = 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.group_name = group_name
         self.failure_count = 0
         self.cancelled = False
 
@@ -87,10 +89,11 @@ class RetryingFuture:
 def map_with_retries(
     function_executor: FunctionExecutor,
     map_function: Callable[..., Any],
-    map_iterdata: List[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]],
+    map_iterdata: Iterable[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]],
     timeout: Optional[int] = None,
     include_modules: Optional[List[str]] = [],
     retries: Optional[int] = None,
+    group_name: Optional[str] = None,
 ) -> List[RetryingFuture]:
     """
     A generalisation of Lithops `map`, with retries.
@@ -107,6 +110,7 @@ def map_with_retries(
             input=i,
             map_kwargs=dict(timeout=timeout, include_modules=include_modules),
             retries=retries,
+            group_name=group_name,
         )
         for i, f in zip(inputs, futures_list)
     ]
diff --git a/cubed/storage/virtual.py b/cubed/storage/virtual.py
index 5f585f2a..9203dae1 100644
--- a/cubed/storage/virtual.py
+++ b/cubed/storage/virtual.py
@@ -29,6 +29,8 @@ class VirtualOffsetsArray:
 
 
 def _key_to_index_tuple(selection):
+    if isinstance(selection, slice):
+        selection = (selection,)
     assert all(isinstance(s, (slice, Integral)) for s in selection)
     sel = []
     for s in selection:
diff --git a/cubed/vendor/rechunker/api.py b/cubed/vendor/rechunker/api.py
deleted file mode 100644
index 0cdc92ca..00000000
--- a/cubed/vendor/rechunker/api.py
+++ /dev/null
@@ -1,360 +0,0 @@
-"""User-facing functions."""
-from collections import defaultdict
-
-import zarr
-
-import cubed.vendor.dask
-
-from .algorithm import rechunking_plan
-from .types import ArrayProxy, CopySpec
-
-
-def _shape_dict_to_tuple(dims, shape_dict):
-    # convert a dict of shape
-    shape = [shape_dict[dim] for dim in dims]
-    return tuple(shape)
-
-
-def _get_dims_from_zarr_array(z_array):
-    # use Xarray convention
-    # http://xarray.pydata.org/en/stable/internals.html#zarr-encoding-specification
-    return z_array.attrs["_ARRAY_DIMENSIONS"]
-
-
-def _encode_zarr_attributes(attrs):
-    from xarray.backends.zarr import encode_zarr_attr_value
-
-    return {k: encode_zarr_attr_value(v) for k, v in attrs.items()}
-
-
-def _zarr_empty(shape, store_or_group, chunks, dtype, name=None, **kwargs):
-    # wrapper that maybe creates the array within a group
-    if name is not None:
-        assert isinstance(store_or_group, zarr.hierarchy.Group)
-        return store_or_group.empty(
-            name, shape=shape, chunks=chunks, dtype=dtype, **kwargs
-        )
-    else:
-        return zarr.empty(
-            shape, chunks=chunks, dtype=dtype, store=store_or_group, **kwargs
-        )
-
-
-ZARR_OPTIONS = [
-    "compressor",
-    "filters",
-    "order",
-    "cache_metadata",
-    "cache_attrs",
-    "overwrite",
-    "write_empty_chunks",
-]
-
-
-def _validate_options(options):
-    if not options:
-        return
-    for o in options:
-        if o not in ZARR_OPTIONS:
-            raise ValueError(
-                f"Zarr options must not include {o} (got {o}={options[o]}). "
-                f"Only the following options are supported: {ZARR_OPTIONS}."
-            )
-
-
-def get_dim_chunk(da, dim, target_chunks):
-    if dim in target_chunks.keys():
-        if target_chunks[dim] > len(da[dim]) or target_chunks[dim] < 0:
-            dim_chunk = len(da[dim])
-        else:
-            dim_chunk = target_chunks[dim]
-    else:
-        # if not isinstance(da.data, dask.array.Array):
-        dim_chunk = len(da[dim])
-        # else:
-        #     existing_chunksizes = {k: v for k, v in zip(da.dims, da.data.chunksize)}
-        #     dim_chunk = existing_chunksizes[dim]
-    return dim_chunk
-
-
-def parse_target_chunks_from_dim_chunks(ds, target_chunks):
-    """
-    Calculate ``target_chunks`` suitable for ``rechunker.rechunk()`` using chunks defined for
-    dataset dimensions (similar to xarray's ``.rechunk()``) .
-
-    - If a dimension is missing from ``target_chunks`` then use the full length from ``ds``.
-    - If a chunk in ``target_chunks`` is larger than the full length of the variable in ``ds``,
-      then, again, use the full length from the dataset.
-    - If a dimension chunk is specified as -1, again, use the full length from the dataset.
-
-    """
-    group_chunks = defaultdict(list)
-
-    for var in ds.variables:
-        for dim in ds[var].dims:
-            group_chunks[var].append(get_dim_chunk(ds[var], dim, target_chunks))
-
-    # rechunk() expects chunks values to be a tuple. So let's convert them
-    group_chunks_tuples = {var: tuple(chunks) for (var, chunks) in group_chunks.items()}
-    return group_chunks_tuples
-
-
-def _copy_group_attributes(source, target):
-    """Visit every source group and create it on the target and move any attributes found."""
-
-    def _update_group_attrs(name):
-        if isinstance(source.get(name), zarr.Group):
-            group = target.create_group(name)
-            group.attrs.update(source.get(name).attrs)
-
-    source.visit(_update_group_attrs)
-
-
-def _setup_rechunk(
-    source,
-    target_chunks,
-    max_mem,
-    target_store,
-    target_options=None,
-    temp_store=None,
-    temp_options=None,
-):
-    if temp_options is None:
-        temp_options = target_options
-    target_options = target_options or {}
-    temp_options = temp_options or {}
-
-    # import xarray dynamically since it is not a required dependency
-    try:
-        import xarray
-        from xarray.backends.zarr import (
-            DIMENSION_KEY,
-            encode_zarr_attr_value,
-            encode_zarr_variable,
-            extract_zarr_variable_encoding,
-        )
-        from xarray.conventions import encode_dataset_coordinates
-    except ImportError:
-        xarray = None
-
-    if xarray and isinstance(source, xarray.Dataset):
-        import dask.array
-
-        if not isinstance(target_chunks, dict):
-            raise ValueError(
-                "You must specify ``target-chunks`` as a dict when rechunking a dataset."
-            )
-
-        variables, attrs = encode_dataset_coordinates(source)
-        attrs = _encode_zarr_attributes(attrs)
-
-        if temp_store is not None:
-            temp_group = zarr.group(temp_store)
-        else:
-            temp_group = None
-        target_group = zarr.group(target_store)
-        target_group.attrs.update(attrs)
-
-        # if ``target_chunks`` is specified per dimension (xarray ``.rechunk`` style),
-        # parse chunks for each coordinate/variable
-        if all([k in source.dims for k in target_chunks.keys()]):
-            # ! We can only apply this when all keys are indeed dimension, otherwise it falls back to the old method
-            target_chunks = parse_target_chunks_from_dim_chunks(source, target_chunks)
-
-        copy_specs = []
-        for name, variable in variables.items():
-            # This isn't strictly necessary because a shallow copy
-            # also occurs in `encode_dataset_coordinates` but do it
-            # anyways in case the coord encoding function changes
-            variable = variable.copy()
-
-            # Update the array encoding with provided options and apply it;
-            # note that at this point the `options` may contain any valid property
-            # applicable for the `encoding` parameter in Dataset.to_zarr other than "chunks"
-            options = target_options.get(name, {})
-            if "chunks" in options:
-                raise ValueError(
-                    f"Chunks must be provided in ``target_chunks`` rather than options (variable={name})"
-                )
-            variable.encoding.update(options)
-            variable = encode_zarr_variable(variable)
-
-            # Extract the array encoding to get a default chunking, a step
-            # which will also ensure that the target chunking is compatible
-            # with the current chunking (only necessary for on-disk arrays)
-            variable_encoding = extract_zarr_variable_encoding(
-                variable, raise_on_invalid=False, name=name
-            )
-            variable_chunks = target_chunks.get(name, variable_encoding["chunks"])
-            if isinstance(variable_chunks, dict):
-                variable_chunks = _shape_dict_to_tuple(variable.dims, variable_chunks)
-
-            # Restrict options to only those that are specific to zarr and
-            # not managed internally
-            options = {k: v for k, v in options.items() if k in ZARR_OPTIONS}
-            _validate_options(options)
-
-            # Extract array attributes along with reserved property for
-            # xarray dimension names
-            variable_attrs = _encode_zarr_attributes(variable.attrs)
-            variable_attrs[DIMENSION_KEY] = encode_zarr_attr_value(variable.dims)
-
-            copy_spec = _setup_array_rechunk(
-                dask.array.asarray(variable),
-                variable_chunks,
-                max_mem,
-                target_group,
-                target_options=options,
-                temp_store_or_group=temp_group,
-                temp_options=options,
-                name=name,
-            )
-            copy_spec.write.array.attrs.update(variable_attrs)  # type: ignore
-            copy_specs.append(copy_spec)
-
-        return copy_specs, temp_group, target_group
-
-    elif isinstance(source, zarr.hierarchy.Group):
-        if not isinstance(target_chunks, dict):
-            raise ValueError(
-                "You must specify ``target-chunks`` as a dict when rechunking a group."
-            )
-
-        if temp_store is not None:
-            temp_group = zarr.group(temp_store)
-        else:
-            temp_group = None
-        target_group = zarr.group(target_store)
-        _copy_group_attributes(source, target_group)
-        target_group.attrs.update(source.attrs)
-
-        copy_specs = []
-        for array_name, array_target_chunks in target_chunks.items():
-            copy_spec = _setup_array_rechunk(
-                source[array_name],
-                array_target_chunks,
-                max_mem,
-                target_group,
-                target_options=target_options.get(array_name),
-                temp_store_or_group=temp_group,
-                temp_options=temp_options.get(array_name),
-                name=array_name,
-            )
-            copy_specs.append(copy_spec)
-
-        return copy_specs, temp_group, target_group
-
-    # elif isinstance(source, (zarr.core.Array, dask.array.Array)):
-    elif isinstance(source, zarr.core.Array):
-
-        copy_spec = _setup_array_rechunk(
-            source,
-            target_chunks,
-            max_mem,
-            target_store,
-            target_options=target_options,
-            temp_store_or_group=temp_store,
-            temp_options=temp_options,
-        )
-        intermediate = copy_spec.intermediate.array
-        target = copy_spec.write.array
-        return [copy_spec], intermediate, target
-
-    else:
-        raise ValueError(
-            f"Source must be a Zarr Array, Zarr Group, Dask Array or Xarray Dataset (not {type(source)})."
-        )
-
-
-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,
-) -> CopySpec:
-    _validate_options(target_options)
-    _validate_options(temp_options)
-    shape = source_array.shape
-    # source_chunks = (
-    #     source_array.chunksize
-    #     if isinstance(source_array, dask.array.Array)
-    #     else source_array.chunks
-    # )
-    source_chunks = source_array.chunks
-    dtype = source_array.dtype
-    itemsize = dtype.itemsize
-
-    if target_chunks is None:
-        # this is just a pass-through copy
-        target_chunks = source_chunks
-
-    if isinstance(target_chunks, dict):
-        array_dims = _get_dims_from_zarr_array(source_array)
-        try:
-            target_chunks = _shape_dict_to_tuple(array_dims, target_chunks)
-        except KeyError:
-            raise KeyError(
-                "You must explicitly specify each dimension size in target_chunks. "
-                f"Got array_dims {array_dims}, target_chunks {target_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
-    shape = tuple(int(x) for x in shape)  # ensure python ints for serialization
-    target_chunks = tuple(int(x) for x in target_chunks)
-    int_chunks = tuple(int(x) for x in int_chunks)
-    write_chunks = tuple(int(x) for x in write_chunks)
-
-    target_array = _zarr_empty(
-        shape,
-        target_store_or_group,
-        target_chunks,
-        dtype,
-        name=name,
-        **(target_options or {}),
-    )
-    try:
-        target_array.attrs.update(source_array.attrs)
-    except AttributeError:
-        pass
-
-    if read_chunks == write_chunks:
-        int_array = None
-    else:
-        # do intermediate store
-        if temp_store_or_group is None:
-            raise ValueError(
-                "A temporary store location must be provided{}.".format(
-                    f" (array={name})" if name else ""
-                )
-            )
-        int_array = _zarr_empty(
-            shape,
-            temp_store_or_group,
-            int_chunks,
-            dtype,
-            name=name,
-            **(temp_options or {}),
-        )
-
-    read_proxy = ArrayProxy(source_array, read_chunks)
-    int_proxy = ArrayProxy(int_array, int_chunks)
-    write_proxy = ArrayProxy(target_array, write_chunks)
-    return CopySpec(read_proxy, int_proxy, write_proxy)


Test patch:
diff --git a/cubed/tests/runtime/test_lithops.py b/cubed/tests/runtime/test_lithops.py
index 0f0b0f61..d48b2308 100644
--- a/cubed/tests/runtime/test_lithops.py
+++ b/cubed/tests/runtime/test_lithops.py
@@ -17,8 +17,9 @@ def run_test(function, input, retries, timeout=10, use_backups=False):
     with LocalhostExecutor() as executor:
         for output in map_unordered(
             executor,
-            function,
-            input,
+            [function],
+            [input],
+            ["group0"],
             timeout=timeout,
             retries=retries,
             use_backups=use_backups,
diff --git a/cubed/tests/storage/test_virtual.py b/cubed/tests/storage/test_virtual.py
new file mode 100644
index 00000000..652582bd
--- /dev/null
+++ b/cubed/tests/storage/test_virtual.py
@@ -0,0 +1,31 @@
+from itertools import product
+from math import prod
+
+import numpy as np
+import pytest
+
+from cubed.storage.virtual import virtual_offsets
+
+
+@pytest.mark.parametrize("shape", [(), (3,), (3, 2)])
+def test_virtual_offsets(shape):
+    v_offsets = virtual_offsets(shape)
+    offsets = np.arange(prod(shape)).reshape(shape, order="C")
+    for t in product(*(range(n) for n in shape)):
+        assert v_offsets[t] == offsets[t]
+
+    # test some length 1 slices
+    if len(shape) == 1:
+        assert v_offsets[1:2] == offsets[1:2]
+    elif len(shape) == 2:
+        assert v_offsets[1:2, 0:1] == offsets[1:2, 0:1]
+
+
+def test_virtual_offsets_fails():
+    with pytest.raises(NotImplementedError):
+        v_offsets = virtual_offsets((3,))
+        v_offsets[0:2]
+
+    with pytest.raises(NotImplementedError):
+        v_offsets = virtual_offsets((3, 2))
+        v_offsets[0:2, 1]
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index f4519815..ee252936 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -10,6 +10,7 @@ from numpy.testing import assert_array_equal
 import cubed
 import cubed.array_api as xp
 import cubed.random
+from cubed.core.ops import merge_chunks
 from cubed.extensions.history import HistoryCallback
 from cubed.extensions.timeline import TimelineVisualizationCallback
 from cubed.extensions.tqdm import TqdmProgressBar
@@ -319,12 +320,12 @@ def test_reduction_multiple_rounds(tmp_path, executor):
     spec = cubed.Spec(tmp_path, allowed_mem=1000)
     a = xp.ones((100, 10), dtype=np.uint8, chunks=(1, 10), spec=spec)
     b = xp.sum(a, axis=0, dtype=np.uint8)
-    # check that there is > 1 rechunk step
-    rechunks = [
-        n for (n, d) in b.plan.dag.nodes(data=True) if d["op_name"] == "rechunk"
+    # check that there is > 1 blockwise step (after optimization)
+    blockwises = [
+        n for (n, d) in b.plan.dag.nodes(data=True) if d["op_name"] == "blockwise"
     ]
-    assert len(rechunks) > 1
-    assert b.plan.max_projected_mem() == 1000
+    assert len(blockwises) > 1
+    assert b.plan.max_projected_mem() <= 1000
     assert_array_equal(b.compute(executor=executor), np.ones((100, 10)).sum(axis=0))
 
 
@@ -335,6 +336,23 @@ def test_reduction_not_enough_memory(tmp_path):
         xp.sum(a, axis=0, dtype=np.uint8)
 
 
+@pytest.mark.parametrize("target_chunks", [(2, 3), (4, 3), (2, 6), (4, 6)])
+def test_merge_chunks(spec, target_chunks):
+    a = xp.ones((10, 10), dtype=np.uint8, chunks=(2, 3), spec=spec)
+    b = merge_chunks(a, target_chunks)
+    assert b.chunksize == target_chunks
+    assert_array_equal(b.compute(), np.ones((10, 10)))
+
+
+@pytest.mark.parametrize(
+    "target_chunks", [(2,), (2, 3, 1), (3, 2), (1, 3), (5, 5), (12, 12)]
+)
+def test_merge_chunks_fails(spec, target_chunks):
+    a = xp.ones((10, 10), dtype=np.uint8, chunks=(2, 3), spec=spec)
+    with pytest.raises(ValueError):
+        merge_chunks(a, target_chunks)
+
+
 def test_compute_multiple():
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2))
     b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2))
@@ -557,7 +575,16 @@ def test_plan_scaling(tmp_path, factor):
 def test_compute_arrays_in_parallel(spec, any_executor, compute_arrays_in_parallel):
     from cubed.runtime.executors.python_async import AsyncPythonDagExecutor
 
-    if not isinstance(any_executor, AsyncPythonDagExecutor):
+    supported_executors = [AsyncPythonDagExecutor]
+
+    try:
+        from cubed.runtime.executors.lithops import LithopsDagExecutor
+
+        supported_executors.append(LithopsDagExecutor)
+    except ImportError:
+        pass
+
+    if not isinstance(any_executor, tuple(supported_executors)):
         pytest.skip(f"{type(any_executor)} does not support compute_arrays_in_parallel")
 
     a = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)