Async prefetch

The quick start shows the synchronous read pattern:

with batch as t:
    x = torch.from_dlpack(t)
    ...  # train step

x shares the output buffer. Exiting the with block releases the batch reference; the tensor retains its own reference until it is destroyed. Retaining both output buffers prevents the next batch from completing.

For CUDA work that outlives the batch's host scope, pass a consumer event or stream to Batch.release(event=...). Damacy orders subsequent writes behind that event. Release the tensor view as well when it is no longer needed. These CUDA examples use an existing Config-based or injected pipeline; Pipeline composition covers CPU consumers.

Deferred release (preferred)

No D2D copy. The batch handle is held for the duration of the training step; release happens after fwd/bwd has been enqueued on the current stream:

def prefetch(p):
    batch = p.pop()
    tensor = torch.from_dlpack(batch)  # zero-copy view onto damacy's slot
    return tensor, batch

# Main training loop:
tensor, batch = prefetch_future.result()
... # forward / backward on tensor
batch.release(event=torch.cuda.current_stream())  # no host sync
del tensor

event=torch.cuda.current_stream() records a CUevent on the caller's stream and tells damacy to wait on it before re-assembling into the slot's buffer. The host returns immediately.

torch.cuda.current_stream() returns the per-thread currently-active stream on the current device. By default that's the device's default stream; inside a with torch.cuda.stream(s): block it returns s. Call it outside any such block (as above) to capture the stream the training step is enqueued on.

Reach for the dedicated-copy-stream pattern below instead if you need to free damacy's slot before your training step finishes — for example if you're chaining many small steps per batch and want maximum overlap on slot reuse, or if the consumer can't safely hold the slot for the whole step.

Dedicated copy stream

Use a side stream for the D2D copy so the main stream's fwd/bwd overlap with it, and release damacy's slot as soon as the copy is enqueued:

copy_stream = torch.cuda.Stream()

def prefetch(p):
    batch = p.pop()
    view = torch.from_dlpack(batch)
    # Allocate on the DEFAULT stream, not copy_stream — see below.
    tensor = torch.empty_like(view)
    with torch.cuda.stream(copy_stream):
        tensor.copy_(view)
    batch.release(event=copy_stream)  # slot reuse waits on copy_stream
    copy_stream.synchronize()
    return tensor

copy_stream.synchronize() blocks the prefetch thread until the D2D copy is complete; the main thread's fwd/bwd runs in parallel the whole time, so the host-side stall doesn't cost throughput at prefetch depth 1. With deeper prefetch (multiple in-flight futures), drop the synchronize() and have the consumer call torch.cuda.current_stream().wait_stream(copy_stream) before reading tensor — that pushes the wait stream-side and lets the prefetch thread start the next pop immediately.

Allocate on the default stream, copy on the side stream

PyTorch's caching allocator pools per stream. Allocating inside with torch.cuda.stream(copy_stream) (with .clone(), or with torch.empty_like placed inside the block) carves out a separate per-stream pool that fwd/bwd can't reuse — reserved memory grows without bound, and the training step slows down under memory pressure.

The fix is to allocate outside the stream block (so the allocator uses the default pool) and only the .copy_ call inside it. This is the single most common foot-gun with the dedicated-stream pattern.

release(event=...) accepts

Batch.release records or accepts a CUevent and queues a cuStreamWaitEvent against damacy's internal post stream before the slot is reused. Accepted forms:

• None (default) — immediate release, host-syncs the producer stream.
• int — raw CUevent handle.
• torch.cuda.Event — its .cuda_event is read.
• torch.cuda.Stream — record_event() is called on it; the recorded event captures the stream's then-current position.
• cupy.cuda.Event / cupy.cuda.Stream — analogous.

JAX users: JAX does not expose CUevent handles through its public API. Drop to CUDA via DLPack and record your own event with cuda-python or cupy.

Why this matters

damacy's output pool is small (double-buffered by default), so slot reuse is on the critical path of throughput. The default with batch as t: block trades a host sync for simplicity — that sync is fine for synchronous reads but stalls a prefetched loop on every iteration. Deferred release replaces the host sync with a stream-side wait that costs nothing on the host.

Async Metadata I/O

The metadata prefetcher has a separate concern from batch-slot release: hiding latency while resolving zarr metadata, shard indexes, and chunk layouts. prefetch_cache has an async completion boundary: a cache miss creates a pending slot, hands an opaque completion token to an async fetcher, and later completes READY or ERROR without the prefetcher knowing how the work ran.

The production metadata fetchers use a narrow async metadata-store boundary. Array metadata submits an owned whole-file read for zarr.json; shard indexes submit async stat plus an optional footer read; chunk layout submits the 16-byte blosc header probe. The prefetcher observes only cache state and gates:

prefetcher requests metadata key
  -> prefetch cache creates async fetch state
    -> metadata store submits stat/read/open work
      -> completion pump advances cache entry
        -> prefetcher observes ready/error on the next advance

The metadata store is an io_uring driver on Linux. It submits small metadata stat/open/read/close work without blocking the prefetcher and completes cache entries from I/O completion state. metadata_io_concurrency sets the metadata request-concurrency budget; the ring allocates enough entries internally for multi-step requests and driver wakeups. At startup the driver requires kernel support for IORING_OP_STATX, IORING_OP_OPENAT2, IORING_OP_READ, and IORING_OP_CLOSE; there is no thread-pool fallback in the Linux build.

On macOS, CMake selects a POSIX worker-pool backend instead of io_uring. metadata_io_concurrency is the worker count; each worker completes one stat/open/read/close request at a time. Shutdown drains accepted requests. Operation histograms measure syscall duration on macOS, excluding queue wait.

The default metadata concurrency is 32. Treat much deeper values as storage tuning: they are useful when metadata operations have real latency, but each active read can hold an open file descriptor and allocate its destination buffer. On multi-rank jobs, multiply the setting by ranks per node before comparing it with ulimit -n.