nvalchemi.distributed.DistributedPipelineModel

class nvalchemi.distributed.DistributedPipelineModel(pipeline, domain_config, *, compile=False, compile_kwargs=None)

Domain decomposition of a composed (pipeline) model.

Parameters:

• pipeline (PipelineModelWrapper) – The composed model. Its groups / steps supply the ordered sub-models; each sub-model’s model_config.neighbor_config.cutoff sets that model’s ghost width.

• domain_config (DomainConfig) – Base config carrying the mesh + skin. Per-model configs are derived from it by overriding cutoff with each sub-model’s cutoff. The caller should build the shared ShardedBatch at (at least) the max sub-model cutoff so the one owned partition’s cells hold every model’s ghost layer.

• compile (bool)

• compile_kwargs (dict | None)

Notes

The composite does not build the ShardedBatch — the caller does (once, at the max cutoff), exactly as for a single DistributedModel. The composite only orchestrates per-model halos over it and sums the results.

close()

Tear down the persistent per-model DistributedModel instances (restoring their adapters / compiled state). Idempotent.

Return type:

None