nvalchemi.dynamics.base.BufferConfig

pydantic model nvalchemi.dynamics.base.BufferConfig

Pre-allocated send/receive buffer capacities for pipeline communication.

A BufferConfig declares the maximum size of the Batch that a distributed dynamics stage may exchange with a neighbouring rank. The three capacities size the flattened graph tensors independently: num_systems bounds how many graphs fit in the buffer, num_nodes bounds the combined atom count across those graphs, and num_edges bounds the combined edge count. Because the buffers are fixed-size, they must be large enough to hold the biggest batch that will ever cross the rank boundary; batches that exceed any capacity cannot be communicated.

You supply a BufferConfig when a stage participates in inter-rank communication, i.e. it is wired into a DistributedPipeline with a prior_rank and/or next_rank. The buffers themselves are created lazily via Batch.empty() on the first simulation step, once a concrete batch is available to act as a dtype/device template, so only the capacities are needed up front. Set a capacity to 0 for a dimension the batch does not carry (for example num_edges=0 when edges are recomputed downstream rather than communicated).

Examples

Size a buffer for up to four graphs totalling 50 atoms, with no edges sent across the boundary:

from nvalchemi.dynamics.base import BufferConfig

buffer_cfg = BufferConfig(num_systems=4, num_nodes=50, num_edges=0)

A buffer that also carries edge connectivity:

buffer_cfg = BufferConfig(num_systems=10, num_nodes=500, num_edges=2000)

Notes

All three fields are constrained to be >= 0. Choose the capacities from the worst-case batch you expect to communicate: undersizing any dimension fails at runtime, while oversizing wastes pre-allocated memory.

field num_systems: int [Required]

Maximum number of graphs the buffer can hold.

Constraints:

• ge = 0

field num_nodes: int [Required]

Total node (atom) capacity across all graphs.

Constraints:

• ge = 0

field num_edges: int [Required]

Total edge capacity across all graphs.

Constraints:

• ge = 0