Storage Format Intelligence & I/O Optimization — dask_setup Recipe

This notebook demonstrates how to choose and benchmark storage formats and I/O strategies:

• Create sample NetCDF climate data files

• Benchmark netcdf4 vs h5netcdf engine performance

• Convert a multi-file NetCDF dataset to Zarr with rechunking

• Compare NetCDF vs Zarr on the same workload

• Observe cold-cache vs warm-cache I/O effects

• Summary of best practices

Setup and Imports

import importlib.util
import time
from contextlib import contextmanager
from pathlib import Path
import tempfile

import numpy as np
import xarray as xr

from dask_setup import setup_dask_client

# Check optional dependencies
try:
    import zarr
    ZARR_AVAILABLE = True
    print(f"zarr {zarr.__version__}")
except ImportError:
    ZARR_AVAILABLE = False
    print("zarr not installed — Zarr sections will be skipped")

# xarray engine name -> the package that provides it
NETCDF_ENGINES = [
    engine
    for engine, module in (("netcdf4", "netCDF4"), ("h5netcdf", "h5netcdf"))
    if importlib.util.find_spec(module) is not None
]


def blosc_encoding(cname="lz4", clevel=5):
    """Per-variable Zarr compressor encoding for whichever zarr major version is installed."""
    if int(zarr.__version__.split(".")[0]) >= 3:
        from zarr.codecs import BloscCodec
        return {"compressors": [BloscCodec(cname=cname, clevel=clevel)]}
    from numcodecs import Blosc  # zarr 2
    return {"compressor": Blosc(cname=cname, clevel=clevel)}


@contextmanager
def timer(label=""):
    t = {"start": time.perf_counter()}
    try:
        yield t
    finally:
        t["total"] = time.perf_counter() - t["start"]
        print(f"  [{label}] {t['total']:.2f}s")


print("Imports OK")
print(f"NetCDF engines available: {NETCDF_ENGINES}")

zarr 3.4.0
Imports OK
NetCDF engines available: ['netcdf4', 'h5netcdf']

1. Create Sample NetCDF Files

Generate synthetic CF-compliant climate data (temperature + precipitation) for benchmarking. We use a small spatial grid and few timesteps so this runs quickly on a laptop.

# Working directory for generated files
work_dir = Path(tempfile.mkdtemp(prefix="dask_recipe_io_"))
netcdf_dir = work_dir / "sample_netcdf"
netcdf_dir.mkdir()

NUM_FILES = 10
TIME_PER_FILE = 20
LAT_SIZE, LON_SIZE = 60, 120   # smaller grid for speed

lat = np.linspace(-90, 90, LAT_SIZE)
lon = np.linspace(-180, 180, LON_SIZE)
file_paths = []

for i in range(NUM_FILES):
    np.random.seed(42 + i)
    time_coord = np.arange(i * TIME_PER_FILE, (i + 1) * TIME_PER_FILE)
    temp = 273.15 + 15 * np.cos(2 * np.pi * time_coord / 365.25)[:, None, None] \
           + 5 * np.random.random((TIME_PER_FILE, LAT_SIZE, LON_SIZE))
    precip = np.abs(np.random.gamma(2, 2, (TIME_PER_FILE, LAT_SIZE, LON_SIZE)))

    ds = xr.Dataset(
        {
            "temperature": (["time", "lat", "lon"], temp.astype("float32"),
                            {"units": "K", "long_name": "Near-surface air temperature"}),
            "precipitation": (["time", "lat", "lon"], precip.astype("float32"),
                              {"units": "mm/day", "long_name": "Daily precipitation"}),
        },
        coords={
            "time": time_coord,
            "lat": ("lat", lat, {"units": "degrees_north"}),
            "lon": ("lon", lon, {"units": "degrees_east"}),
        },
    )
    ds.attrs["conventions"] = "CF-1.8"

    fp = netcdf_dir / f"sample_{i + 1:03d}.nc"
    ds.to_netcdf(fp, encoding={"temperature": {"zlib": True}, "precipitation": {"zlib": True}})
    file_paths.append(fp)

total_mb = sum(f.stat().st_size for f in file_paths) / 1e6
print(f"Created {NUM_FILES} NetCDF files  ({total_mb:.1f} MB total)")

Created 10 NetCDF files  (8.2 MB total)

2. Start a Cluster for NetCDF

For NetCDF the right choice is workload_type="cpu", even though this is file reading. The netCDF4/HDF5 C library serialises every read behind a single process-wide lock, so the "io" topology (one process, many threads) queues its threads on that lock. Several single-threaded processes each have their own lock and really do read in parallel.

"io" is the right choice for Zarr, whose codecs release the GIL. dask_setup warns if you pass a NetCDF-backed dataset with workload_type="io".

client, cluster, tmp = setup_dask_client(
    workload_type="cpu",
    max_workers=4,
    fallback_on_detection_failure=True,
)
print(f"Workers: {len(client.nthreads())}")
print(f"Temp dir: {tmp}")

Dask Dashboard → http://127.0.0.1:37139/status

[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-2704
Workers: 4 | threads/worker: 1 | processes: True
Memory: total ~15.6 GiB | usable ~11.6 GiB | per-worker ~2.9 GiB
Compression: spill=auto | comm=False
Workers: 4
Temp dir: /tmp/dask-2704

# What happens with the wrong topology: dask_setup inspects the dataset and logs a
# warning naming the fix (look for "WARNING [client]" in this cell's output)
ds_check = xr.open_mfdataset(file_paths[:2], chunks={})
c, cl, _, _ = setup_dask_client(workload_type="io", ds=ds_check, dashboard=False)
c.close(); cl.close()
ds_check.close()

WARNING  [client] workload_type='io' is usually the wrong choice for NetCDF/HDF5 input (reason=HDF5 is not thread-safe, so xarray serialises every read through one process-wide lock; the threads queue rather than read in parallel | suggestion=use workload_type='cpu' (one lock per process) for NetCDF | note=also avoid open_mfdataset(..., parallel=True) with 'io' -- concurrent metadata reads can kill the worker and Dask's silent retries look like a hang)

[setup_dask_client] Configuration summary
temp/spill dir: /tmp/dask-2704
Workers: 1 | threads/worker: 4 | processes: False
Memory: total ~15.6 GiB | usable ~11.6 GiB | per-worker ~11.6 GiB
Compression: spill=auto | comm=False

/home/runner/work/dask_setup/dask_setup/src/dask_setup/client.py:140: UserWarning: Dimension 'time' has very small chunks (0.5 MiB). This may create high task overhead.
  raw = recommend_chunks(
/home/runner/work/dask_setup/dask_setup/src/dask_setup/client.py:140: UserWarning: Dimension 'lat' has very small chunks (0.5 MiB). This may create high task overhead.
  raw = recommend_chunks(
/home/runner/work/dask_setup/dask_setup/src/dask_setup/client.py:140: UserWarning: Dimension 'lon' has very small chunks (0.5 MiB). This may create high task overhead.
  raw = recommend_chunks(

3. NetCDF Engine Benchmark

Compare netcdf4 and h5netcdf for single-file reading, multi-file open_mfdataset, and a chunked spatial reduction.

engine_results = {}

for engine in NETCDF_ENGINES:
    print(f"\nEngine: {engine}")
    try:
        # Single file
        with timer(f"{engine} single-file") as t1:
            ds = xr.open_dataset(file_paths[0], engine=engine)
            _ = ds.temperature.values
            ds.close()

        # Multi-file
        with timer(f"{engine} multi-file (5 files)") as t2:
            ds = xr.open_mfdataset(
                file_paths[:5], engine=engine,
                chunks={"time": 20},  # whole lat/lon: matches how the files are stored
                parallel=True, combine="by_coords",
            )
            _ = ds.temperature.mean().compute()
            ds.close()

        engine_results[engine] = {
            "single_file_s": t1["total"],
            "multi_file_s": t2["total"],
        }
    except Exception as e:
        print(f"  ERROR: {e}")
        engine_results[engine] = {"error": str(e)}

print("\nSummary:")
for eng, r in engine_results.items():
    if "error" not in r:
        print(f"  {eng:12s}  single={r['single_file_s']:.2f}s  multi={r['multi_file_s']:.2f}s")

Engine: netcdf4
  [netcdf4 single-file] 0.00s

  [netcdf4 multi-file (5 files)] 0.62s

Engine: h5netcdf
  [h5netcdf single-file] 0.00s
  ERROR: No module named 'h5py', backend not available. Please install 'h5py' into your Python environment.

Summary:
  netcdf4       single=0.00s  multi=0.62s

4. Convert NetCDF to Zarr

Zarr uses chunked, compressed storage that is friendlier to parallel reads and cloud access. Rechunk during conversion to optimise for analysis access patterns.

zarr_store = work_dir / "converted.zarr"

if ZARR_AVAILABLE:
    engine = NETCDF_ENGINES[0] if NETCDF_ENGINES else "netcdf4"

    with timer("NetCDF → Zarr conversion") as t:
        ds = xr.open_mfdataset(
            file_paths, engine=engine,
            chunks={"time": 40},
            parallel=True, combine="by_coords",
        )
        # Rechunk for better Zarr layout
        ds = ds.chunk({"time": 60, "lat": LAT_SIZE, "lon": LON_SIZE})

        # zarr.Blosc was removed in zarr 3; blosc_encoding() handles both versions
        encoding = {"temperature": blosc_encoding(), "precipitation": blosc_encoding()}
        ds.to_zarr(zarr_store, mode="w", consolidated=True, encoding=encoding)
        ds.close()

    zarr_mb = sum(f.stat().st_size for f in zarr_store.rglob("*") if f.is_file()) / 1e6
    netcdf_mb = sum(f.stat().st_size for f in file_paths) / 1e6
    print(f"  NetCDF: {netcdf_mb:.1f} MB → Zarr: {zarr_mb:.1f} MB  "
          f"(ratio {zarr_mb / netcdf_mb:.2f}x)")
else:
    print("Zarr not available — skipping conversion")

/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/zarr/api/asynchronous.py:246: ZarrUserWarning: Consolidated metadata is currently not part in the Zarr format 3 specification. It may not be supported by other zarr implementations and may change in the future.
  warnings.warn(

  [NetCDF → Zarr conversion] 0.36s
  NetCDF: 8.2 MB → Zarr: 8.8 MB  (ratio 1.07x)

5. NetCDF vs Zarr Performance

Run the same analysis (spatial mean, annual resampling, temporal std) on both formats.

fmt_times = {}

# NetCDF
if NETCDF_ENGINES:
    engine = NETCDF_ENGINES[0]
    with timer("NetCDF analysis") as t:
        ds = xr.open_mfdataset(
            file_paths, engine=engine,
            chunks={"time": 40}, parallel=True,
        )
        _ = ds.temperature.mean().compute()
        _ = ds.temperature.std(dim="time").compute()
        ds.close()
    fmt_times["NetCDF"] = t["total"]

# Zarr
if ZARR_AVAILABLE and zarr_store.exists():
    with timer("Zarr analysis") as t:
        ds = xr.open_zarr(zarr_store, consolidated=True)
        _ = ds.temperature.mean().compute()
        _ = ds.temperature.std(dim="time").compute()
        ds.close()
    fmt_times["Zarr"] = t["total"]

if len(fmt_times) == 2:
    speedup = fmt_times["NetCDF"] / fmt_times["Zarr"]
    print(f"\nZarr speedup: {speedup:.2f}x")

  [NetCDF analysis] 0.20s
  [Zarr analysis] 0.08s

Zarr speedup: 2.38x

6. Filesystem Caching Effects

The first read of a file is slower (cold cache) than later reads (warm OS file cache). Account for this when benchmarking I/O: time a second run, or drop caches between runs.

These files were written moments ago, so they are probably already in the page cache and both reads here are “warm”. On a real dataset on Lustre (/g/data, /scratch) the first read is usually far slower.

if NETCDF_ENGINES:
    engine = NETCDF_ENGINES[0]
    subset = file_paths[:3]
    chunk_kw = {"time": 20}

    with timer("Cold cache (1st read)") as t_cold:
        ds = xr.open_mfdataset(subset, engine=engine, chunks=chunk_kw)
        _ = ds.temperature.mean().compute()
        ds.close()

    with timer("Warm cache (2nd read)") as t_warm:
        ds = xr.open_mfdataset(subset, engine=engine, chunks=chunk_kw)
        _ = ds.temperature.mean().compute()
        ds.close()

    speedup = t_cold["total"] / t_warm["total"]
    print(f"Cache speedup: {speedup:.2f}x")

  [Cold cache (1st read)] 0.04s
  [Warm cache (2nd read)] 0.04s
Cache speedup: 0.99x

client.close()
cluster.close()

# Clean up generated files
import shutil
shutil.rmtree(work_dir, ignore_errors=True)
print("Cluster closed, temp files removed.")

Cluster closed, temp files removed.

Storage Format & I/O Best Practices

Choose your format:

Format	Best for	Key advantage
NetCDF	Archival, sharing with collaborators	Self-describing, CF-compliant
Zarr	Analysis pipelines, cloud storage	Chunked, concurrent writes, cloud-native

Engine selection:

• netcdf4 — the default, widely compatible

• h5netcdf — reads through h5py rather than the netCDF-C library. It is not GIL-free: h5py wraps the same HDF5 C library and serialises calls on its own lock. Benchmark both engines on your data rather than assuming either wins.

Whichever engine you use, NetCDF input wants workload_type="cpu" (parallel processes); Zarr input wants "io" (threads).

# Large multi-file NetCDF datasets
client, cluster, dask_tmp = setup_dask_client(workload_type="cpu")
ds = xr.open_mfdataset(files, parallel=True, combine="by_coords")

Zarr write tips:

# Consolidate metadata. zarr 3 takes `compressors` (a list);
# zarr 2 took `compressor`, and zarr.Blosc no longer exists in zarr 3.
from zarr.codecs import BloscCodec
ds.to_zarr("output.zarr", mode="w", consolidated=True,
           encoding={"temp": {"compressors": [BloscCodec(cname="lz4", clevel=5)]}})

Chunk sizing: target 256–512 MiB per chunk for analysis. Use recommend_chunks() if unsure:

from dask_setup import recommend_chunks
chunks = recommend_chunks(ds, client, verbose=True)