Slow performance on Zarr DirectoryStore

[huberb]

I'm looking for ways to improve the performance of my dataloading pipeline and I found Zarr. To get an idea about throughput, I started a small benchmark script in python. To get a baseline I also run tests using numpy memory mapped arrays.

I'm working with 4D arrays which are quite large. One of my criterias is that I need to access them as a key-value store. From each value, I access randomly on the first axis.

I created some dummy arrays to test throughput.

Here is my complete benchmarking code that compares Zarr to accessing raw Numpy arrays on disk:

import numpy as np
from pathlib import Path
import time
from tqdm import tqdm
import os
import random
import zarr


def write_numpy_dataset(data_path, raw_data):
    Path(data_path).mkdir(parents=True, exist_ok=True)
    for sample_name, sample_data in tqdm(raw_data.items()):
        mmap = np.memmap(
            f"{data_path}/{sample_name}.npy",
            dtype=sample_data.dtype,
            mode='w+',
            shape=sample_data.shape
        )
        mmap[:] = sample_data[:]
        mmap.flush()


def write_zarr_dataset(data_path, raw_data):
    zarr_store = zarr.DirectoryStore(f'{data_path}/data.zarr')
    zarr_group = zarr.group(zarr_store, overwrite=True)
    for data_key, sample in tqdm(raw_data.items()):
        zarr_group.create_dataset(
            data_key,
            data=sample,
            chunks=[28, 28, 28],  # I also tried sample.shape[1:] and sample.shape here
            overwrite=True,
            compressor=None  # remove compression because I want as much throughput as possible
        )


def load_samples_numpy(data_path):
    numpy_files = [f"{data_path}/{f}" for f in os.listdir(data_path)]
    while True:
        for numpy_file in numpy_files:
            index = random.randint(0, 499)
            # open a new memory map for every access. This way I will never cache samples in RAM
            volume = np.memmap(
                numpy_file,
                mode='r',
                dtype=np.float32
            ).reshape([500, -1, 28, 28, 28])
            yield volume[index]


def load_samples_zarr(data_path):
    zarr_store = zarr.DirectoryStore(f"{data_path}/data.zarr")
    zarr_group = zarr.open_group(zarr_store, mode="r")
    data_keys = [s for s in zarr_group]
    while True:
        for data_key in data_keys:
            index = random.randint(0, 499)
            yield zarr_group[data_key][index]


def benchmark(dataset_path, load_fn):
    total_size = 0
    start_time = time.time()

    print("\nStarting benchmark")
    print(f"reading from: {dataset_path}")
    sample_sums = []
    for i, sample in tqdm(enumerate(load_fn(dataset_path))):
        if i == 2000:
            break
        total_size += sample.nbytes
        sample_sums.append(sample.sum())  # do some dummy calculation to force the data into RAM

    end_time = time.time()

    throughput = (total_size / (end_time - start_time)) / (1024 * 1024)
    print(f"Throughput: {throughput:.2f} MB/s")


if __name__ == '__main__':
    data_path_numpy = "/mnt/cache/dataset_numpy"
    data_path_zarr = "/mnt/cache/dataset_zarr"

    # the second axis of my data samples is not fixed
    raw_data = {
        'sample0': np.random.normal(size=(500, 14, 28, 28, 28)).astype(np.float32),
        'sample1': np.random.normal(size=(500, 16, 28, 28, 28)).astype(np.float32),
        'sample2': np.random.normal(size=(500, 12, 28, 28, 28)).astype(np.float32),
        'sample3': np.random.normal(size=(500, 11, 28, 28, 28)).astype(np.float32),
        'sample4': np.random.normal(size=(500, 13, 28, 28, 28)).astype(np.float32),
        'sample5': np.random.normal(size=(500, 19, 28, 28, 28)).astype(np.float32),
        'sample6': np.random.normal(size=(500, 20, 28, 28, 28)).astype(np.float32),
        'sample7': np.random.normal(size=(500, 15, 28, 28, 28)).astype(np.float32),
        'sample8': np.random.normal(size=(500, 18, 28, 28, 28)).astype(np.float32),
        'sample9': np.random.normal(size=(500, 17, 28, 28, 28)).astype(np.float32),
    }

    write_numpy_dataset(data_path_numpy, raw_data)
    benchmark(data_path_numpy, load_samples_numpy)

    write_zarr_dataset(data_path_zarr, raw_data)
    benchmark(data_path_zarr, load_samples_zarr)

It turns out that accessing Numpy arrays outperforms Zarr by a factor of ~6-7
My maximum disk speed is 500MB/s and I reach roughly 400MB/s using numpy. With Zarr I see a throughput of ~50-60MB/s

This difference is so big that I feel like I must be missing something. I tried different chunk sizes and disabled compression completely. Still, Zarr never reaches a throughput that comes even close to Numpy's memory mapped arrays.

Does anyone have a hint on what I'm missing? Is Zarr generally slow for my usecase of accessing large 4D arrays?
Appreciate any help

[d-v-b]

It's important to distinguish between zarr the format and zarr-python, the python library. On a local file system, zarr-the-format can be read as quickly as you can read + decompress files stored on disk. Since you are not using any compression, the chunks of your arrays are literally the serialized bytes of the in-memory numpy arrays you are trying to create, so you can skip zarr-python entirely for reading and just memory-map each chunk. Happy to elaborate on how this would work if you are interested.

As for zarr-python, you should know that array indexing in zarr-python versions 2.x.x is serial, which is slow. When performing array indexing like array[:], the chunks for that array are fetched one after another. So if you have any array indexing patterns that read multiple chunks, you will wait O(n_chunks * read_time * decompress_time) time for your entire array indexing operation to complete.

We are trying to make indexing faster in zarr-python 3.x, and we have made good progress in that direction. But I think there a limits to how far we can go w.r.t performance in zarr-python, or any python library. If performance is your main goal, you might consider using tensorstore for reading your data. Or just memory-map your chunks as numpy arrays, as I suggested earlier. The python library dask is also a popular solution for concurrent access to zarr arrays, which has the advantage of generalizing to compute clusters.

[huberb]

Thank you for this quick response.

Indeed, performance is my main criteria here. Previously I've used the TFRecords format which worked great for me so far but only allows sequential read access to my data instead of (random) index based access, which would be really beneficial for my use case.

My datasets exceed 50TB nowadays and since I got some new hardware to play with, data loading became my main bottleneck.

I'm currently benchmarking on a Linux system with a SSD that gives me 500MB/s read access, but later I want to read data from a NFS that has much higher throughput but also adds latency when opening filehandles. This is why I don't think I can afford storing each sample as a single file on disk, so I group them into chunks of 500 samples that have common shapes, as you can see in the code example that I gave above. Still, I think I'm loading far too many filehandles when storing my data like this.

I thought about just grouping even more samples into single files and access those with memory mapping. Can you think of any better approach than this?

My main criterias are:

• index based access
• some form of key-value store (so far this is my filesystem, using filenames as keys)
• try to avoid opening too many filehandles because latency will kill my performance
• compression would be nice to save disk space, but throughput has a higher priority for me

Other stuff like prefetching and batching is an afterthought after I figured out how to max out my throughput on a single process.

I'm willing to redo my whole data pipeline if necessary and I'm happy for any hints. So far I've tried TFRecords, HDF5, DALI, memory mapped Numpy arrays and Zarr. tensorstore seems like a good candidate to try. Can you point me to any examples that might fit my usecase? Would you recommend any specific backend?

Thanks again for the response

[d-v-b]

I think the tensorstore tutorial is a good place to start. You would probably not want to use the format in that example (n5); instead, you can use tensorstore to read and write zarr v2 and v3 arrays (to write zarr groups you can just use zarr-python). Tensorstore also supports other formats but I have much less experience with them. Because tensorstore has a performance focus, you might get good pointers by opening a discussion / issue on their github page.

[huberb]

Do you also have by chance any guidance on how to set up a Linux system to improve reading speed through caching? I'm asking because my current benchmarks seems to improve if they run multiple times and it seems like I'm running into some filesystem caching mechanism on Linux even if my files exceed my available RAM. Running some new benchmarks on my system multiple times shows this:

$ python benchmark.py
Bytes read: 20580.00 MB
Throughput: 530.77 MB/s

$ python benchmark.py
Bytes read: 20580.00 MB
Throughput: 775.01 MB/s

$ python benchmark.py
Bytes read: 20580.00 MB
Throughput: 996.57 MB/s 

I've been reading this StackExchange post:
https://unix.stackexchange.com/questions/30286/can-i-configure-my-linux-system-for-more-aggressive-file-system-caching

but I'm not sure if this applies to my NFS with filesystem type nfs4. Filehandle caching could also be helpful to me

Anyway, thanks a lot for the help! We can also close this discussion since it's moving away from zarr related questions.

[d-v-b]

unfortunately I know basically nothing about tuning local file systems for throughput -- most of the data I work with is on cloud storage, where we have different challenges :) That being said, if you come up with a good story for how to get the most out of zarr on local disk(s), I'd be interested in reading about it (and it could potentially end up in our docs)

[huberb]

So what turned out to work really well for me is saving my data as .HDF5 with LZF compression and an appropriate chunk size. I group a ton of samples into single files to avoid file handle overhead. This way I can read 13.4Gbit/s of data on my test setup, even though my bandwidth would only allow for 10Gbit/s.

With Zarr I can't get to similar performance and it seems to me this is only because of the format that Zarr stores on disk. HDF5 puts everything into single files. Zarr stores chunks as separate files instead.
With strace -T I can see that a lot of time is spend on the openat syscall when reading those .zarr files. This only happens when I read from my NFS but I can't find a way to speed this up yet.

Maybe this is too specific to the setup that I'm using but I think there is a usecase for a storage option that minimizes the amount of files on disk.

[d-v-b]

zarr v3 introduces a compression strategy (sharding) that can pack multiple "read chunks" inside a larger file, called a "shard", to address the problem of too many files. The upcoming zarr-python 3 release will support this, and tensorstore supports it today. Using sharding should give you the same IO benefits of hdf5.