nvalchemi.distributed.DistributedModel

class nvalchemi.distributed.DistributedModel(wrapper, domain_config, *, spec=None, compile=False, compile_kwargs=None)

Wrap an atomic single-process model wrapper for domain-decomposed inference.

Parameters:

• wrapper (BaseModelMixin) – Atomic BaseModelMixin. Its distribution_spec must be non-None. Composite wrappers (PipelineModelWrapper) are rejected — use DistributedPipelineModel for composition.

• domain_config (DomainConfig) – Shared simulation config carrying the cutoff, skin, mesh, and optional grid_dims. The partitioner and halo config are built lazily from the first ShardedBatch’s geometry.

• spec (MLIPSpec, optional, keyword-only) – Explicit distribution spec (the joint model x strategy product). When None (default), it is obtained from wrapper.distribution_spec(domain_config.strategy); a wrapper with distribution_spec=None requires this argument.

• compile (bool, optional, keyword-only) – When True, compile the energy-autograd forward path (fixed-shape padded, per-rank). The spec carries only the compile contract; this switch enables it. Default False.

• compile_kwargs (dict, optional, keyword-only) – Extra keyword arguments forwarded to the compile of the distributed forward, merged with the spec’s compile contract. Only consulted when compile=True. Default None.

Notes

Construction is side-effect-free. The first call to __call__ initializes the partitioner / halo config / world size from the supplied ShardedBatch and invokes wrapper.distributed_setup.

close() — or __exit__ / __del__ — calls wrapper.distributed_teardown to restore any module-level state. Use as a context manager for scoped lifecycle:

with DistributedModel(wrapper, config) as dist_model:
    out = dist_model(sharded)

close()

Release resources and restore any state setup mutated. Safe to call multiple times.

Restores all adapters installed by adapter_registry (custom_ops + third_party_helpers), then defers to the wrapper’s optional distributed_teardown hook for any wrapper-side runtime state.

Return type:

None

property config: DomainConfig

The DomainConfig held by this adapter.

from_batch(batch, *, src=0)

One-call distributed inference from a full Batch.

The convenience entry for one-off inference: shards batch across the scope’s mesh (via ShardedBatch.from_batch(), using the held DomainConfig) and runs the distributed forward — so a caller never constructs a ShardedBatch by hand. Collective: every rank calls it, with the full system on rank src and None elsewhere; every rank gets the consolidated output dict back.

Parameters:

• batch (Batch | None) – The full-system Batch on rank src; None on the other ranks.

• src (int) – The rank holding the full batch (default 0).

Returns:

The consolidated outputs (owned-shape per-atom + replicated per-system), identical to calling __call__() on a hand-built ShardedBatch.

Return type:

dict[str, Any]

property wrapper: BaseModelMixin

The underlying single-process model wrapper.