Compute zarr shards one chunk at a time to reduce peak memory

src/spikeinterface/core/baserecording.py

@@ -876,6 +876,8 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs):
                 - shards: tuple | None, default: None
                     Number of shard size. If None, no sharding is done. Note that shards dimensions need to be larger than
                     chunk dimensions (if chunks is not None) and that sharding is only done on the first dimension.
+                - buffer_megabytes: int, default: 500
+                    Buffer size in megabytes for writing data to zarr.
                 - compressors: BytesBytesCodec | ArrayBytesCodec | list of them, or None, default: None
                     Global compressors. Only one ArrayBytesCodec (serializer) is supported.
                     If None, Blosc-zstd, level 5, with bit shuffle is used.

src/spikeinterface/core/recording_tools.py

@@ -122,6 +122,7 @@ def write_recording_to_zarr(
     dataset_timestamps_paths=None,
     extra_chunks=None,
     dtype=None,
+    buffer_megabytes=500,
     compressor_data=None,
     filters_data=None,
     compressor_times=None,
@@ -147,6 +148,8 @@ def write_recording_to_zarr(
         controlled by the job_kwargs. Useful to chunk by channel, with ``extra_chunks=(channel_chunk_size,)``.
     dtype : dtype, default: None
         Type of the saved data.
+    buffer_megabytes : int, default: 500
+        Buffer size in megabytes for writing data to zarr.
     compressor_data : zarr compressor or None, default: None
         Zarr compressor for data.
     filters_data : list, default: None
@@ -166,6 +169,7 @@ def write_recording_to_zarr(
         dataset_timestamps_paths=dataset_timestamps_paths,
         extra_chunks=extra_chunks,
         dtype=dtype,
+        buffer_megabytes=buffer_megabytes,
         compressor_data=compressor_data,
         filters_data=filters_data,
         compressor_times=compressor_times,

src/spikeinterface/core/time_series_tools.py

@@ -299,6 +299,7 @@ def _write_time_series_to_zarr(
     dtype=None,
     chunks=None,
     shards=None,
+    buffer_megabytes=500,
     compressor_data=None,
     filters_data=None,
     compressor_times=None,
@@ -320,12 +321,14 @@ def _write_time_series_to_zarr(
         List of paths to traces datasets in the zarr group
     dataset_timestamps_paths : list or None, default: None
         List of paths to timestamps datasets in the zarr group. If None, timestamps are not saved.
+    dtype : dtype, default: None
+        Type of the saved data
     chunks : tuple or None, default: None
         Chunking dimensions to use for the zarr dataset.
     shards : tuple or None, default: None
         Sharding configuration for the zarr dataset.
-    dtype : dtype, default: None
-        Type of the saved data
+    buffer_megabytes : int, default: 500
+        Buffer size in megabytes for writing data to zarr
     compressor_data : zarr compressor or None, default: None
         Zarr compressor for data
     filters_data : list, default: None
@@ -403,7 +406,7 @@ def _write_time_series_to_zarr(
     # use executor (loop or workers)
     func = _write_zarr_chunk
     init_func = _init_zarr_worker
-    init_args = (time_series, zarr_datasets, dtype, zarr_timestamps_datasets)
+    init_args = (time_series, zarr_datasets, dtype, buffer_megabytes, zarr_timestamps_datasets)
     executor = TimeSeriesChunkExecutor(
         time_series, func, init_func, init_args, verbose=verbose, job_name="write_zarr", **job_kwargs
     )
@@ -420,7 +423,7 @@ def _write_time_series_to_zarr(
         zarr_group.create_array(name="t_starts", data=t_starts, compressors=None)
 
 
-def _init_zarr_worker(time_series, zarr_datasets, dtype, zarr_timestamps_datasets=None):
+def _init_zarr_worker(time_series, zarr_datasets, dtype, buffer_megabytes, zarr_timestamps_datasets=None):
     import zarr
 
     # create a local dict per worker
@@ -432,6 +435,7 @@ def _init_zarr_worker(time_series, zarr_datasets, dtype, zarr_timestamps_dataset
     else:
         worker_ctx["zarr_timestamps_datasets"] = None
     worker_ctx["dtype"] = np.dtype(dtype)
+    worker_ctx["buffer_megabytes"] = buffer_megabytes
 
     return worker_ctx
 
@@ -447,15 +451,27 @@ def _write_zarr_chunk(segment_index, start_frame, end_frame, worker_ctx):
         zarr_timestamps_dataset = worker_ctx["zarr_timestamps_datasets"][segment_index]
     else:
         zarr_timestamps_dataset = None
-
-    # apply function
-    data = time_series.get_data(
-        start_frame=start_frame,
-        end_frame=end_frame,
-        segment_index=segment_index,
-    )
-    data = data.astype(dtype)
-    zarr_dataset[start_frame:end_frame, :] = data
+    write_buffer_megabytes = worker_ctx["buffer_megabytes"]
+
+    # With shards, a job covers a whole shard so only one worker writes to it. The data is computed one
+    # zarr chunk at a time to keep the memory of the chain small, and written as many whole chunks at a time
+    # as fit in write_buffer_megabytes because zarr reads and rewrites the whole shard on every write.
+    num_samples_per_read = zarr_dataset.chunks[0]
+    num_values_per_sample = int(np.prod(zarr_dataset.shape[1:]))
+    bytes_per_read = num_samples_per_read * num_values_per_sample * dtype.itemsize
+    num_reads_per_write = max(1, write_buffer_megabytes * 1_000_000 // bytes_per_read)
+    num_samples_per_write = num_reads_per_write * num_samples_per_read
+    for write_start in range(start_frame, end_frame, num_samples_per_write):
+        write_end = min(write_start + num_samples_per_write, end_frame)
+        data = np.empty((write_end - write_start,) + zarr_dataset.shape[1:], dtype=dtype)
+        for read_start in range(write_start, write_end, num_samples_per_read):
+            read_end = min(read_start + num_samples_per_read, write_end)
+            data[read_start - write_start : read_end - write_start] = time_series.get_data(
+                start_frame=read_start,
+                end_frame=read_end,
+                segment_index=segment_index,
+            )
+        zarr_dataset[write_start:write_end, :] = data
 
     if zarr_timestamps_dataset is not None:
         timestamps = time_series.get_times(start_frame=start_frame, end_frame=end_frame, segment_index=segment_index)

src/spikeinterface/core/zarrextractors.py

@@ -977,6 +977,7 @@ def add_recording_to_zarr_group(
         job_kwargs=job_kwargs,
         time_series=recording,
     )
+    buffer_megabytes = zarr_kwargs.get("buffer_megabytes", 500)
 
     _write_time_series_to_zarr(
         time_series=recording,
@@ -991,6 +992,7 @@ def add_recording_to_zarr_group(
         compressor_times=compressor_times,
         filters_times=filters_times,
         verbose=verbose,
+        buffer_megabytes=buffer_megabytes,
         **job_kwargs,
     )