Skip to content

Repository files navigation

Damacy

buildtestfuzzcodecovbenchdocs

High-speed streamed assembly of nD tensors from zarr sources to GPU.

Damacy reads sharded NGFFzarr v3 stores directly on the GPU: per-shard chunk indexing, parallel host I/O, in-flight GPU-side decompression (zstd, blosc1-zstd), and assembly of each batch as a DLPack-ready device tensor.

Quick start

importrandomimportdamacyimporttorchcfg=damacy.Config(
samples_per_batch=8,
sample_shape=(64, 256, 256),
# Resource caps are fixed at construction; nothing grows after.max_gpu_memory_bytes=1<<30, # primary GPU budgetdtype="bf16", # source dtype is cast on assemble# One pipeline binds to one GPU. Omit `device=` to capture the# current CUDA context (handy single-GPU; PyTorch sets one up# implicitly). For multi-GPU pass `device=local_rank` — see# https://nclack.github.io/damacy/distributed/
)
# A Sample names an absolute uri and a per-axis half-open AABB into# the stored array (np.s_[...] also accepted). Build them however# suits — your own sampler, a torch Dataset, a curriculum, a fixed# tile grid, ...volumes= { # absolute uri → full ZYX shape"/data/cells/brain-001.zarr": (512, 4096, 4096),
"/data/cells/brain-002.zarr": (768, 4096, 4096),
"/data/cells/kidney-007.zarr": (256, 2048, 2048),
}
defrandom_crop(size=(64, 256, 256)):
uri, full=random.choice(list(volumes.items()))
origin= [random.randint(0, f-s) forf, sinzip(full, size)]
returndamacy.Sample(uri=uri, aabb=[(o, o+s) foro, sinzip(origin, size)])
samples= [random_crop() for_inrange(64)]
withdamacy.Pipeline(cfg) asp:
p.push(samples) # producer sideforbatchinp.batches(len(samples) //cfg.samples_per_batch):
withbatchast: # consumer sidex=torch.from_dlpack(t) # zero-copy + stream-fenced
... # train step

torch.from_dlpack (or any DLPack v1 consumer — cupy, jax, …) handles the stream handoff: damacy hands over BatchInfo.ready_stream, the consumer records a cuStreamWaitEvent against it, and the resulting tensor is fenced for downstream kernels. Damacy's internal streams are non-blocking with respect to the legacy default stream, so don't read BatchInfo.device_ptr directly without a matching cuStreamWaitEvent on ready_stream.

Streaming

push accepts any iterable, including infinite generators — samples are pulled lazily as pop frees space. For unbounded training, hand the pipeline a generator and let it drain:

defcrops():
whileTrue:
yieldrandom_crop() # from the example abovewithdamacy.Pipeline(cfg) asp:
p.push(crops()) # pulled on demandforstepinrange(N_STEPS):
withp.pop() ast:
x=torch.from_dlpack(t)
... # train step

For training loops that prefetch the next batch on a background thread, see Async prefetch — zero-copy with deferred release, plus the dedicated-copy-stream variant.

Zarr support

Damacy reads zarr v3 (sharded and non-sharded). What's recognized today:

supportednotes
Array versionsv3v2 stores are not read
Layout (sharded)sharding_indexedwith index_location either "start" or "end" (default)
Layout (non-sharded)yeseach chunk is a separate file at c/<i>/<j>/...
Inner / chunk codecbytes (passthrough), zstd, blosc (cname=zstd)blosc with cname=lz4/lz4hc is recognized at parse time and rejected at planning
Sharding index codecbytes + crc32cthe shard index itself; the data codec is separate
Missing chunksyes — read as fill_valueper zarr v3 spec; sharded "empty" entries (offset==nbytes==2^64−1) and missing chunk files both route here

Not yet handled — arrays declaring any of these will fail to parse:

  • Non-trivial transposes (transpose codec)
  • Compression codecs other than the list above (gzip, lz4 raw, crc32c as a data codec, future v3 codecs)
  • Complex / fixed-bytes / variable-length dtypes

If you have data that uses one of the unsupported codecs and you'd like it added, please open an issue with a sample zarr.json.

Runtime dependencies

Damacy links the core CUDA and Linux async-I/O pieces directly. Optional features dlopen their backends lazily, so a damacy binary loads on any host with the required core libraries; optional features simply turn off when their library is not present.

LibraryRequired at runtimeWhat you lose if missingHow damacy finds it
libcuda.so.1alwaysnothing — damacy cannot run without itNVIDIA driver install (/run/opengl-driver/lib, /usr/lib/x86_64-linux-gnu, …)
liburing.soalways on Linux buildsnothing — async metadata stat/open/read/close uses io_uringnormal dynamic loader (LD_LIBRARY_PATH, ld.so.cache, rpath, …)
libnuma.so.1optionalNUMA pinning of pinned-host slabs, bulk I/O workers, scheduler thread, and metadata io_uring driver thread (single-socket hosts: no effect)dlopen via dynamic loader (LD_LIBRARY_PATH, ld.so.cache)
libcufile.so.0optionaldamacy_config.enable_gds = DAMACY_GDS_ON — direct read of compressed chunks into device memory via NVIDIA GPUDirect Storagedlopen via dynamic loader; ships with the CUDA toolkit and with nvidia-fs. Requires a build with -DDAMACY_ENABLE_GDS=ON (default OFF)
libmount.so.1, libudev.so.1required if and only if using GDScuFile dlopen's these at driver init even in compat modedynamic loader

Metadata reads require a Linux kernel with the io_uring operations damacy uses: STATX, OPENAT2, READ, and CLOSE. If the kernel does not advertise those operations, damacy_create fails instead of falling back to a thread pool.

GDS notes:

  • Build with cmake -DDAMACY_ENABLE_GDS=ON to link the cuFile backend. The default-OFF build still accepts enable_gds = DAMACY_GDS_ON but damacy_create returns DAMACY_INVAL (no silent fallback).
  • enable_gds = DAMACY_GDS_AUTO (default, the value from designated-init) defers to env DAMACY_GDS_ENABLE=1; explicit DAMACY_GDS_ON / DAMACY_GDS_OFF override the env.
  • On hosts without nvidia-fs, point CUFILE_ENV_PATH_JSON at a JSON with {"properties":{"allow_compat_mode":true}} to enable cuFile compat mode — reads go through cuFile's host-bounce buffer instead of DMA. Useful for correctness testing on consumer GPUs.
  • If libcufile can't be loaded or cuFileDriverOpen fails, damacy_create returns DAMACY_INVAL.

Documentation

Full API reference and guides: https://nclack.github.io/damacy/

Performance dashboards (auto-published from bench.yml):

About

High-speed streamed assembly of nD tensors from zarr sources to GPU

Topics

Resources

Stars

8 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages