nvalchemi.data.InMemoryDataset

class nvalchemi.data.InMemoryDataset(in_memory_batch=None, *, reader=None, chunk_size=4096, device=None, skip_validation=False, batch_transforms=None)

Resident dataset that holds the whole dataset in memory as one Batch.

InMemoryDataset is the memory-resident counterpart to Dataset in the same pipeline (Reader -> Dataset / InMemoryDataset (-> MultiDataset) -> DataLoader). Instead of streaming from a Reader on every access, it materializes the entire dataset once into a single resident Batch and then serves batches by slicing that cache. Because it implements BatchDatasetProtocol, a DataLoader and MultiDataset consume it interchangeably with Dataset.

Prefer it over Dataset when the dataset fits in (CPU or GPU) memory and is iterated many times – small-to-medium training sets over many epochs, or repeated in-loop evaluation – since after the one-time materialization there is no per-access storage I/O. Prefer Dataset when the data is too large to hold resident or is read only once. Construct it from either a pre-built in_memory_batch or a reader that materializes the full dataset (pass exactly one).

Like Dataset, it exposes the DataLoader-facing batch API (load_batches(), prefetch_fused_batches(), get_fused_batches(), has_pending_fused_batches(), cancel_prefetch()) backed by in-memory implementations, so a DataLoader loads whole Batch objects through those methods rather than calling __getitem__() per sample; __getitem__() and __iter__() remain available for direct, dataset-style access.

Parameters:

• in_memory_batch (Batch | None, default=None) – Fully loaded batch containing all graphs in the dataset. Pass this when the batch has already been built.

• reader (ReaderProtocol | None, default=None) – Reader to materialize the full dataset into in_memory_batch. Pass either in_memory_batch or reader, not both.

• chunk_size (int, default=4096) – Number of reader samples to materialize per intermediate batch.

• device (str | torch.device | None, default=None) – Target device for emitted samples and batches. None leaves emitted batches on the resident cache device. "auto" selects CUDA when available, otherwise CPU.

• skip_validation (bool, default=False) – If True, bypass AtomicData construction and Pydantic validation while materializing from a reader, building batches directly from raw tensor dicts via from_raw_dicts(). Enable this for trusted stores that are already validated.

• batch_transforms (Sequence[BatchTransform] | None, default=None) – Optional per-batch transforms applied while building the resident batch. For reader-backed construction they run on each materialized chunk; for a pre-built in_memory_batch they run once on the full batch. This mirrors the DataLoader(batch_transforms=...) API.

Raises:

• ValueError – If neither or both of in_memory_batch and reader are provided.

• TypeError – If batch_transforms is not a Sequence (e.g. a single callable or a generator was passed).

in_memory_batch

Fully loaded batch containing all graphs in the dataset.

Type:

Batch

target_device

Resolved target device for emitted samples and batches.

Type:

torch.device | None

Examples

>>> from nvalchemi.data.datapipes.in_memory_dataset import InMemoryDataset
>>> # Assuming a concrete Reader implementation exists:
>>> # reader = MyReader("dataset.zarr")
>>> # ds = InMemoryDataset(reader=reader, device="cpu")
>>> # batch = ds.load_batches([[0, 1, 2]])[0]
>>> # trusted = InMemoryDataset(reader=reader, device="cuda", skip_validation=True)

With a pre-built in-memory batch:

>>> # batch = Batch.from_raw_dicts(raw_dicts)
>>> # ds = InMemoryDataset(in_memory_batch=batch)

cancel_prefetch(index=None)

Cancel pending prefetch operations.

Parameters:

index (int | None, default=None) – Unused; retained for API compatibility with Dataset. Clears all pending fused batch requests.

Return type:

None

close()

Release resources held by the dataset.

Clears any pending fused batch prefetch requests.

Return type:

None

property field_names: list[str]

Return field names available in in-memory samples.

Returns:

Field names exposed by the resident batch.

Return type:

list[str]

get_fused_batches()

Consume the pending fused prefetch and yield per-batch results.

Blocks until any queued CUDA transfers complete, then yields one Batch per sub-batch from the prepared request.

Yields:

Batch – One batch per sub-batch from the fused prefetch request.

Raises:

RuntimeError – If no fused prefetch is pending.

Return type:

Iterator[Batch]

get_metadata(index)

Return lightweight graph-size metadata for an in-memory sample.

Parameters:

index (int) – Graph index within in_memory_batch.

Returns:

(num_atoms, num_edges) for the sample.

Return type:

tuple[int, int]

has_pending_fused_batches()

Return whether a fused prefetch chunk is waiting to be consumed.

Return type:

bool

load_batches(batch_index_lists, stream=None)

Load several batches immediately.

This is the synchronous counterpart to prefetch_fused_batches()/get_fused_batches(). Each requested index list selects graphs from in_memory_batch and returns one Batch per input list.

Parameters:

• batch_index_lists (Sequence[Sequence[int]] | Sequence[Tensor]) – Per-batch graph indices.

• stream (torch.cuda.Stream | None, default=None) – CUDA stream for target-device transfers when supported.

Returns:

One Batch per input batch-index list.

Return type:

list[Batch]

property pin_memory: bool

Whether the materialized CPU batch is pinned in page-locked memory.

prefetch(index, stream=None)

Satisfy the sample-prefetch API without queuing work.

InMemoryDataset already holds every graph in in_memory_batch, so there is no reader I/O to overlap for an individual sample.

Parameters:

• index (int) – Graph index within in_memory_batch.

• stream (torch.cuda.Stream | None, default=None) – Unused CUDA stream argument retained for compatibility.

Return type:

None

property prefetch_count: int

Return the number of queued fused-batch prefetch requests.

prefetch_fused_batches(batch_index_lists, stream=None)

Submit multiple batches for fused prefetch.

Selected batches are prepared immediately and queued for later consumption via get_fused_batches().

Parameters:

• batch_index_lists (Sequence[Sequence[int]] | Sequence[Tensor]) – Per-batch graph indices.

• stream (torch.cuda.Stream | None, default=None) – CUDA stream for target-device transfers when supported.

Raises:

RuntimeError – If the fused batch prefetch queue is already full.

Return type:

None