nvalchemi.distributed.distributed_method

nvalchemi.distributed.distributed_method(body)

Decorate a MethodAdapter body that only diverges under domain decomposition.

Removes the boilerplate guard repeated across method adapters so the body holds only the distributed behavior. The wrapped replacement runs the original method verbatim whenever the live context is not distributed (single-process, or any call outside a distributed forward), and otherwise invokes body with the live context already in hand. Gating on is_distributed (not is_halo) keeps the body policy-agnostic: it fires under any strategy (halo, graph-parallel, graph-replicate), and its cross-rank steps are expressed through the policy-dispatched intent verbs.

Parameters:

body (callable) – The halo behavior, called as body(ctx, original, instance, *args, **kwargs) where ctx is the live DistributedContext and original is the unpatched method.

Returns:

An (original, instance, *args, **kwargs) replacement suitable for MethodAdapter.

Return type:

callable

Examples

>>> @distributed_method
... def _refresh_block(ctx, original, block, x, *args, **kwargs):
...     return original(block, refresh_neighbors(x), *args, **kwargs)