Use extra buffer copies for cloud stores like S3 (#703)

Author: tomwhite

cubed/core/ops.py

@@ -23,6 +23,7 @@
 from cubed.core.plan import Plan, new_temp_path
 from cubed.primitive.blockwise import blockwise as primitive_blockwise
 from cubed.primitive.blockwise import general_blockwise as primitive_general_blockwise
+from cubed.primitive.memory import get_buffer_copies
 from cubed.primitive.rechunk import rechunk as primitive_rechunk
 from cubed.spec import spec_from_config
 from cubed.storage.backend import open_backend_array
@@ -293,6 +294,7 @@ def blockwise(
 
     name = gensym()
     spec = check_array_specs(arrays)
+    buffer_copies = get_buffer_copies(spec)
     if target_store is None:
         target_store = new_temp_path(name=name, spec=spec)
     op = primitive_blockwise(
@@ -312,6 +314,7 @@ def blockwise(
         new_axes=new_axes,
         in_names=in_names,
         out_name=name,
+        buffer_copies=buffer_copies,
         extra_func_kwargs=extra_func_kwargs,
         fusable_with_predecessors=fusable_with_predecessors,
         fusable_with_successors=fusable_with_successors,
@@ -442,6 +445,7 @@ def _general_blockwise(
     op_name = kwargs.pop("op_name", "blockwise")
 
     spec = check_array_specs(arrays)
+    buffer_copies = get_buffer_copies(spec)
 
     if isinstance(target_stores, list):  # multiple outputs
         name = [gensym() for _ in range(len(target_stores))]
@@ -461,6 +465,7 @@ def _general_blockwise(
         allowed_mem=spec.allowed_mem,
         reserved_mem=spec.reserved_mem,
         extra_projected_mem=extra_projected_mem,
+        buffer_copies=buffer_copies,
         target_stores=target_stores,
         target_paths=target_paths,
         storage_options=spec.storage_options,
@@ -1157,10 +1162,13 @@ def _rechunk_plan(x, chunks, *, min_mem=None):
     source_chunks = to_chunksize(normalize_chunks(x.chunks, x.shape, dtype=x.dtype))
 
     # rechunker doesn't take account of uncompressed and compressed copies of the
-    # input and output array chunk/selection, so adjust appropriately
-    # note the factor is 5 (not 4) since there is a extra (unnecessary) copy
-    # made when writing out to Zarr
-    rechunker_max_mem = (spec.allowed_mem - spec.reserved_mem) // 5
+    # input and output array chunk/selection, so adjust appropriately:
+    #  1 input array plus copies to read that array from storage,
+    #  1 array for processing,
+    #  1 output array plus copies to write that array to storage
+    buffer_copies = get_buffer_copies(spec)
+    total_copies = 1 + buffer_copies.read + 1 + 1 + buffer_copies.write
+    rechunker_max_mem = (spec.allowed_mem - spec.reserved_mem) // total_copies
     if min_mem is None:
         min_mem = min(rechunker_max_mem // 20, x.nbytes)
     stages = multistage_rechunking_plan(

cubed/primitive/blockwise.py

@@ -162,6 +162,7 @@ def blockwise(
     in_names: Optional[List[str]] = None,
     out_name: Optional[str] = None,
     extra_projected_mem: int = 0,
+    buffer_copies: Optional[BufferCopies] = None,
     extra_func_kwargs: Optional[Dict[str, Any]] = None,
     fusable_with_predecessors: bool = True,
     fusable_with_successors: bool = True,
@@ -199,6 +200,8 @@ def blockwise(
     extra_projected_mem : int
         Extra memory projected to be needed (in bytes) in addition to the memory used reading
         the input arrays and writing the output.
+    buffer_copies: BufferCopies
+        The the number of buffer copies incurred for array storage operations.
     extra_func_kwargs : dict
         Extra keyword arguments to pass to function that can't be passed as regular keyword arguments
         since they clash with other blockwise arguments (such as dtype).
@@ -250,6 +253,7 @@ def blockwise(
         chunkss=[chunks],
         in_names=in_names,
         extra_projected_mem=extra_projected_mem,
+        buffer_copies=buffer_copies,
         extra_func_kwargs=extra_func_kwargs,
         fusable_with_predecessors=fusable_with_predecessors,
         fusable_with_successors=fusable_with_successors,
@@ -274,6 +278,7 @@ def general_blockwise(
     chunkss: List[T_Chunks],
     in_names: Optional[List[str]] = None,
     extra_projected_mem: int = 0,
+    buffer_copies: Optional[BufferCopies] = None,
     extra_func_kwargs: Optional[Dict[str, Any]] = None,
     fusable_with_predecessors: bool = True,
     fusable_with_successors: bool = True,
@@ -312,6 +317,8 @@ def general_blockwise(
     extra_projected_mem : int
         Extra memory projected to be needed (in bytes) in addition to the memory used reading
         the input arrays and writing the output.
+    buffer_copies: BufferCopies
+        The the number of buffer copies incurred for array storage operations.
     extra_func_kwargs : dict
         Extra keyword arguments to pass to function that can't be passed as regular keyword arguments
         since they clash with other blockwise arguments (such as dtype).
@@ -388,8 +395,7 @@ def general_blockwise(
         return_writes_stores,
     )
 
-    # assumes a single buffer copy for reading and writing, compare https://github.com/tomwhite/memray-array
-    buffer_copies = BufferCopies(read=1, write=1)
+    buffer_copies = buffer_copies or BufferCopies(read=1, write=1)
     projected_mem = calculate_projected_mem(
         reserved_mem=reserved_mem,
         inputs=[array_memory(array.dtype, array.chunks) for array in arrays],

cubed/primitive/memory.py

@@ -1,5 +1,8 @@
 from dataclasses import dataclass
-from typing import List
+from typing import List, Optional
+
+from cubed.spec import Spec
+from cubed.utils import is_cloud_storage_path
 
 
 @dataclass
@@ -13,6 +16,25 @@ class BufferCopies:
     """The number of copies made when writing an array to storage."""
 
 
+def get_buffer_copies(spec: Optional[Spec]):
+    """Return the number of buffer copies to use, based on the spec.
+
+    Using cloud storage will result in more buffer copies being accounted for.
+    """
+
+    # See https://github.com/tomwhite/memray-array
+    # More factors (e.g. compression) could be taken into account in the future.
+
+    if (
+        spec is not None
+        and spec.work_dir is not None
+        and is_cloud_storage_path(spec.work_dir)
+    ):
+        return BufferCopies(read=2, write=2)
+
+    return BufferCopies(read=1, write=1)
+
+
 def calculate_projected_mem(
     reserved_mem: int,
     inputs: List[int],

cubed/tests/test_rechunk.py

@@ -54,14 +54,20 @@ def test_rechunk_era5(
     assert max_output_blocks == expected_max_output_blocks
 
 
-def test_rechunk_era5_chunk_sizes():
+@pytest.mark.parametrize(
+    "spec",
+    [
+        cubed.Spec(allowed_mem="2.5GB"),
+        # cloud stores use extra buffer copies, so need more memory for same rechunk plan
+        cubed.Spec("s3://cubed-unittest/rechunk-era5", allowed_mem="3.5GB"),
+    ],
+)
+def test_rechunk_era5_chunk_sizes(spec):
     # from https://github.com/pangeo-data/rechunker/pull/89
     shape = (350640, 721, 1440)
     source_chunks = (31, 721, 1440)
     target_chunks = (350640, 10, 10)
 
-    spec = cubed.Spec(allowed_mem="2.5GB")
-
     a = xp.empty(shape, dtype=xp.float32, chunks=source_chunks, spec=spec)
 
     from cubed.core.ops import _rechunk_plan