Large Blosc chunks decode about 3x slower than in zarr 2 because numcodecs.blosc.use_threads is set to False on import #4531

Zarr version

v3.4.1

Numcodecs version

v0.17.0

Python Version

3.12

Operating System

Linuxx86_64

Installation

using uv to create a venv

Description

Hi, I'm looking at ECMWF's anemoi-datasets (a library for large zarr datasets to train weather models) to smooth out its support for zarr 3, and while checking some performance claims in ecmwf/anemoi-datasets#486 and ecmwf/anemoi-datasets#786 I found a slowdown that seems to belong here. #3524 looks related, so I'm linking it. The reproducer below needs only zarr and numpy.

Summary

With zarr-python 3.4.1 / numcodecs 0.17.0, reading one large Blosc chunk (~200 MB decoded) takes about 3x longer than with zarr-python 2.18.7 / numcodecs 0.15.1, with the data already in memory (warm page cache). In a single process, setting numcodecs.blosc.use_threads = True after import zarr closes more than half of the gap. #3524 was closed because FusedCodecPipeline (3.3.0) was expected to fix this; for one thread reading a single large chunk it does not (finding 3), though with several threads it does help (see the separate issue in point 7). Cold reads are not affected, and with four chunks per read zarr 3 is faster. The numbers are the same on main (bd0dc84).

In a PyTorch DataLoader it depends on how the workers start. With fork (the Linux default) numcodecs ignores use_threads and decodes on one thread, so True changes nothing. With spawn or forkserver it takes effect: 1.25-1.69x the samples/s of zarr 2 on real data (finding 6).

The script is in steps to reproduce. On my 8-core laptop (warm cache, so this is CPU work, not storage):

Main findings

1. zarr 3 turns Blosc threading off at import. blosc.py#L104 sets numcodecs.blosc.use_threads = False (docstring: "to avoid threading issues in asyncio contexts"), so one large chunk decodes on one thread. Decoding one 219 MB raw Blosc buffer takes 28 ms in a zarr 2 process and 103 ms in a zarr 3 process (raw_decode_486.py, end of steps to reproduce). None does not help: zarr decodes on worker threads and numcodecs only enables threads on the main thread.
2. Most of the gap is decode; a copy explains part of the rest. On a 48-core machine, one 219 MB chunk takes 52.5 ms with zarr 2, 165.2 ms with zarr 3 and 102.8 ms with zarr 3 and use_threads=True. A py-spy profile of 60 reads of that array with zarr 3 (all threads, 2,488 samples) puts 49% of the samples in Blosc decode (_decode_sync, codecs/_v2.py:34) and 30% in NDBuffer.__setitem__, reached from scatter_chunk: the decoded chunk is copied into the output array, while zarr 2 decodes into it. With use_threads=True decode drops to 5% and the copy is 64% of the samples. That the remaining ~45 ms is this copy is my reading of the profile. The commands are in steps to reproduce.
3. FusedCodecPipeline makes no difference for one thread reading one chunk: 170.6 vs 170.7 ms per read. It still decodes and scatters in one worker thread (read calls read_sync through asyncio.to_thread).
4. Chunk layout matters. With 4 chunks along the grid zarr 3 is faster than zarr 2 (87 vs 104 ms, 64 ms with FusedCodecPipeline): small chunks parallelize across threads, one large chunk does not. Cold reads are not affected (zarr 3 / zarr 2 = 0.95-1.06).
5. Real data. anemoi-datasets format 2 ERA5 O96 (104 variables x 40320 points, ~17 MB decoded per chunk, one chunk per date) read through a torch DataLoader, warm, one process: zarr 3 reaches 0.43x of zarr 2's samples/s (0.74x with use_threads=True). With 4-8 workers it is 0.83x-0.94x and use_threads changes nothing, because the workers are forked (finding 6). The remaining 6-17% is not use_threads: one 219 MB chunk read in a forked child took about 168 ms with zarr 3 against about 119 ms with zarr 2.
6. fork vs spawn/forkserver. In a process whose pid differs from the one that imported numcodecs (a forked child), numcodecs returns False before it reads use_threads, for zarr 2 and zarr 3 alike. I logged the setting inside the workers. Samples/s as a ratio to zarr 2 on the O96 data (zarr 3 default / zarr 3 with True):

zarr 2 decodes on one thread in spawned workers, so the gain is threaded against single-threaded decoding. With no workers: 0.25 / 0.56. One machine, no real training run.

Questions

1. Is use_threads = False meant to stay? Would it be acceptable to leave Blosc threading on when a read decodes a single large chunk, or to expose the setting through zarr.config so applications can choose without reaching into numcodecs after import? If the concern is the numcodecs lock under concurrent decodes, is there a variant you would consider safe?
2. Could the decoded chunk be written directly into the output array for whole-chunk selections? I found Codec pipeline performance #2904 and its sub-issues (should encode and decode both take an optional reference to a buffer #3060, Add API to Store interface to read directly into an output buffer #3429) which look like the right place. Is this still planned, and would a benchmark case help?
3. Would a benchmark with a large compressed chunk be welcome? The existing ones use in-memory or uncompressed arrays, which is probably why Array indexing with Zarr 3 is noticeably slower than with Zarr 2 #3524 looked resolved. I have a small pytest-benchmark case (32 MiB Blosc-lz4 chunks, one vs four chunks per read, threads off/on) for tests/benchmarks/ that I can open as a PR if useful.

Steps to reproduce

My example to reproduce things with synthetic data, one operation (arr[i]):

# /// script
# requires-python = "==3.12.*"
# dependencies = [
#     "zarr==3.4.1",
#     "numcodecs==0.17.0",
#     "numpy==2.5.4",
# ]
# ///
"""Single-chunk Blosc read: zarr-python 3 vs 2 (anemoi-datasets #486). Synthetic data, one operation: ``arr[i]``.

zarr 3:  uv run repro_486.py
zarr 2:  uv run --python 3.12 --with 'zarr==2.18.7' --with 'numcodecs==0.15.1' --with 'numpy==2.5.4' --no-project repro_486.py   (override the pin above)
The data are written once, the page cache is warm (the files are read once untimed), so this measures CPU work, not storage.
"""
import statistics as st
import tempfile
import time

import numcodecs
import numpy as np
import zarr
from numcodecs import blosc

SHAPE = (4, 101, 1, 542080)  # ERA5 N320 date: 101 variables x 542080 points, float32, ~219 MB decoded per chunk
CHUNKS = (1, 101, 1, 542080)  # one chunk per date, as in anemoi-datasets #486
ZV = int(zarr.__version__.split(".")[0])


def create(path):
    rng = np.random.default_rng(0)
    base = rng.standard_normal(SHAPE[1:]).cumsum(-1).astype("f4")
    if ZV >= 3:
        arr = zarr.create_array(path, shape=SHAPE, chunks=CHUNKS, dtype="f4", compressors=numcodecs.Blosc("lz4"), zarr_format=2, overwrite=True)
    else:
        arr = zarr.open_array(path, mode="w", shape=SHAPE, chunks=CHUNKS, dtype="f4", compressor=numcodecs.Blosc("lz4"))
    for i in range(SHAPE[0]):
        arr[i] = base + np.float32(i)


def timed(arr, label):
    arr[0]  # warm the page cache and any lazy initialization
    t = [(lambda t0: (arr[i % SHAPE[0]], time.perf_counter() - t0)[1])(time.perf_counter()) for i in range(12)]
    print(f"  {label:34s} {st.median(t) * 1e3:7.1f} ms per chunk  (median of 12)")


with tempfile.TemporaryDirectory() as d:
    create(d + "/a.zarr")
    arr = zarr.open_array(d + "/a.zarr", mode="r")
    print(f"zarr {zarr.__version__}, numcodecs {numcodecs.__version__}, numpy {np.__version__}; numcodecs.blosc.use_threads = {blosc.use_threads}")
    timed(arr, "as installed")
    if ZV >= 3:
        blosc.use_threads = True
        timed(arr, "numcodecs.blosc.use_threads = True")

To run it, save it as repro_486.py. With zarr 3: uv run repro_486.py. With zarr 2: uv run --python 3.12 --with 'zarr==2.18.7' --with 'numcodecs==0.15.1' --with 'numpy==2.5.4' --no-project repro_486.py. Versions: Python 3.12, numpy 2.5.4, zarr 3.4.1 + numcodecs 0.17.0 against zarr 2.18.7 + numcodecs 0.15.1; py-spy 0.4.2. For a format 3 store, replace the zarr.create_array(...) call in create() for zarr 3 with arr = zarr.create_array(path, shape=SHAPE, chunks=CHUNKS, dtype="f4", zarr_format=3, overwrite=True) (zarr 3's default zstd compressor) or add compressors=zarr.codecs.BloscCodec(cname="lz4", clevel=5, shuffle="shuffle") for Blosc.

On a 48-core machine: zarr 2.18.7 52.6-53.6 ms, zarr 3.4.1 164.4-165.9 ms, zarr 3.4.1 with use_threads=True 101.0-101.8 ms.

To profile, I used a second script that creates the same array once and then only loops over reads, so that creating the array is not in the profile (repro_486_loop.py, below), and py-spy 0.4.2 installed in the same virtual environment as zarr:

py-spy --version
python repro_486_loop.py create a.zarr          # once, zarr 2 or zarr 3
py-spy record --format raw --threads --rate 250 -o prof.raw -- python repro_486_loop.py read a.zarr default 60
py-spy record --format raw --threads --rate 250 -o prof_threads.raw -- python repro_486_loop.py read a.zarr threads 60

prof.raw is one line per stack with a sample count; I summed the counts by the innermost frame of each stack. py-spy needs permission to attach to a process (on Linux, kernel.yama.ptrace_scope 0, or run it as root). These 60-read timings were 55.4 ms (zarr 2), 162.9 ms (zarr 3) and 107.0 ms (zarr 3 with use_threads).

"""Read loop for py-spy: create <path> once (zarr 2 env), then loop arr[i] on the same array as repro_486.py.
usage: repro_486_loop.py create <path> | read <path> <default|threads> <n_reads>"""
import sys, time, statistics as st
import numcodecs, numpy as np, zarr
from numcodecs import blosc
SHAPE = (4, 101, 1, 542080); CHUNKS = (1, 101, 1, 542080)
cmd, path = sys.argv[1:3]
ZV = int(zarr.__version__.split(".")[0])
if cmd == "create":
    rng = np.random.default_rng(0); base = rng.standard_normal(SHAPE[1:]).cumsum(-1).astype("f4")
    if ZV >= 3:
        arr = zarr.create_array(path, shape=SHAPE, chunks=CHUNKS, dtype="f4", compressors=numcodecs.Blosc("lz4"), zarr_format=2, overwrite=True)
    else:
        arr = zarr.open_array(path, mode="w", shape=SHAPE, chunks=CHUNKS, dtype="f4", compressor=numcodecs.Blosc("lz4"))
    for i in range(SHAPE[0]): arr[i] = base + np.float32(i)
else:
    mode, n = sys.argv[3], int(sys.argv[4])
    arr = zarr.open_array(path, mode="r")
    if mode == "threads": blosc.use_threads = True
    arr[0]
    t = []
    for k in range(n):
        t0 = time.perf_counter(); arr[k % SHAPE[0]]; t.append(time.perf_counter() - t0)
    print(f"zarr {zarr.__version__} numcodecs {numcodecs.__version__} mode={mode} use_threads={blosc.use_threads} median {st.median(t)*1e3:.1f} ms over {n} reads")

# /// script
# requires-python = "==3.12.*"
# dependencies = ["zarr==3.4.1", "numcodecs==0.17.0", "numpy==2.5.4"]
# ///
"""Raw Blosc decode of one 219 MB chunk, no zarr read path involved. Importing zarr 3 turns numcodecs' Blosc threading off.

zarr 3:  uv run raw_decode_486.py
zarr 2:  uv run --python 3.12 --with 'zarr==2.18.7' --with 'numcodecs==0.15.1' --with 'numpy==2.5.4' --no-project raw_decode_486.py
"""
import time, statistics as st, numpy as np, numcodecs, zarr
from numcodecs import blosc
rng = np.random.default_rng(0)
x = rng.standard_normal((101, 542080)).cumsum(-1).astype("f4")
c = numcodecs.Blosc("lz4", clevel=5, shuffle=1)
enc = c.encode(x)
def t():
    r=[]
    for _ in range(7):
        t0=time.perf_counter(); c.decode(enc); r.append(time.perf_counter()-t0)
    return st.median(r)*1e3
print(f"zarr {zarr.__version__} numcodecs {numcodecs.__version__} use_threads={blosc.use_threads}: raw decode of {x.nbytes/1e6:.0f} MB = {t():.1f} ms")

Additional output

No response