Halos and parallel execution

Halo exchange

Halos are extra cells around each grid piece. They hold values from neighboring pieces or from boundary rules. After replacing State fields manually, use fill_state_overlap() to refresh each field’s boundaries. The x direction is periodic: its ends join. The y boundaries follow Configuration.enable_cyclic_y. For sharded arrays, which are split across devices, keep the jax.set_mesh(mesh) context active. A mesh arranges the devices that share the grid. See Model settings for solid-wall behavior and Experiment definitions for packed storage and meshes.

Two-cell halos with periodic x and optional closed global y boundaries.

veris.grid.halos.fill_circular_overlap(A: Array, enable_cyclic_y: bool = True, boundary: Literal['edge', 'zero', 'normal', 'shear'] = 'edge') → Array

Fill periodic x halos and periodic or closed global y halos.

veris.grid.halos.fill_overlap_shard(var: MaskInput, enable_cyclic_y: bool = True, boundary: Literal['edge', 'zero', 'normal', 'shear'] = 'edge') → Array

runs on each shard, must be inside shard_map

veris.grid.halos.fill_overlap(var: MaskInput, conf: Configuration, *, boundary: Literal['edge', 'zero', 'normal', 'shear'] = 'edge') → Array

Fill halos using initialized settings and the caller’s mesh.

Serial inputs store a single interior with two halo cells on each edge. Sharded inputs pack those local halo regions for every mesh partition. Sharded inputs use NamedSharding(mesh, P("x", "y")). Calls must execute inside jax.set_mesh(mesh); mesh execution context is owned by the caller and is never stored among State leaves. Closed y boundaries use nearest-edge extension by default; zero marks dry exterior cells and normal also closes the southern wall face. shear preserves both physical wall tractions, including the north halo face, while zeroing exterior cells beyond the walls. Periodic mode ignores these wall policies.

veris.grid.halos.fill_overlap_uv(u: MaskInput, v: MaskInput, conf: Configuration) → tuple[Array, Array]

Fill both horizontal velocity components with the same initialized settings.

veris.grid.halos.fill_state_overlap(state: T, conf: Configuration) → T

Refresh State or Diagnostics halos with impermeable-wall field semantics.

Dry exterior masks activate the existing coastline/slip discretization. Ice amounts and tangential velocities vanish outside the domain; normal velocities also vanish on the owned southern wall face. Other fields use constant extension, avoiding invalid zero temperatures and grid metrics. Corner shear stress retains no-slip wall traction or vanishes for free slip.

Reductions

Combine already-local diagnostic totals across explicit mesh axes.

Call from inside shard_map with its mesh axis names to reduce across devices and processes. With no axes, the local scalar or vector is returned unchanged. Callers exclude halo duplicates before forming each local total. JAX supplies the collective’s forward and reverse differentiation rules.

veris.grid.reductions.global_sum(value: SumInput, axis_names: tuple[str, ...] = ()) → SumInput | Array

Sum local totals over named mesh axes, preserving component dimensions.

Parallel dynamics helpers

The public runner prepares JAX for a run across processes before accessing devices or creating arrays. The parallel driver selects the dynamics preset and parallel profile, then calls veris.run_experiment(). The shared veris.advance() step function controls the order of its physics calculations. Launch examples appear in Parallel dynamics.

veris.grid.halos.remove_halos keeps its results as JAX arrays and supports JAX transformations. The host distributed_collector() removes halos from packed storage and gathers selected interior fields for the Input and output API. Every process must join calls that exchange data between processes in the same order. Only the writing process creates the output file.

veris.execution.runtime.distributed_options(environment: Mapping[str, str], backend: str, explicit: DistributedExecution | None = None) → dict[str, Any] | None

Resolve validated explicit controls or the existing JAX/Slurm environment.

veris.execution.runtime.create_mesh(partitions: tuple[int, int] | None, backend: str) → Mesh

Create an all-device x/y mesh, including supported multi-slice topologies.

veris.grid.halos.remove_halos(array: Array, mesh: Mesh) → Array

Trim two boundary cells from each shard, retaining global physical order.

Execution resources

The public Python runner prepares resources before initializing State. It checks explicit DistributedExecution controls or reads JAX/Slurm launch variables once. Explicit controls take precedence. Veris closes services it starts after success, a startup failure or an execution failure. Services started by the caller remain active. The caller also keeps ownership of supplied meshes.

Execution() requests serial execution. all_devices=True requests a mesh using every device discovered at launch. The parallel dynamics profile uses that policy. Explicit partitions must use all available devices of the chosen backend, the CPU or GPU system that runs JAX. Multiple equally sized device groups, called slices, use a hybrid mesh that joins them. These layouts are not supported by jax.make_mesh. The lower-level Simulation interface still expects callers to prepare distributed services and meshes themselves.

An explicitly sized mesh owned by Veris must provide at least one device per process. CPU launches owned by Veris use exactly one device per rank, meaning each process. If the rank count is known, it must match the mesh size before startup. Caller-owned services may expose multiple CPU devices per rank; Veris checks them against the actual mesh. If known rank and mesh counts cannot work together, the request fails before applying launch policy or starting distributed JAX.

class veris.model.experiment.DistributedExecution(coordinator_address: str | None = None, num_processes: int | None = None, process_id: int | None = None, local_device_ids: tuple[int, ...] | None = None, cluster_detection_method: str | None = None)

Bases: object

Explicit distributed bootstrap controls, validated without JAX devices.

Supply coordinator/count/rank together, or a supported cluster detector. Local device IDs are immutable nonnegative indices. Execution resources adapt environment defaults only when no explicit bootstrap is supplied.

class veris.execution.runtime.RuntimeResources(experiment: Experiment, mesh: Mesh | None, device: Any, owns_distributed: bool, metadata: Mapping[str, Any])

Bases: object

Resolved numerical definition, borrowed mesh and immutable launch metadata.

veris.execution.runtime.execution_resources(experiment: Experiment, *, mesh: Mesh | None = None, environment: Mapping[str, str] | None = None) → Iterator[RuntimeResources]

Preflight inputs, acquire requested resources and close owned services.

Distributed bootstrap precedes device enumeration. Explicit CPU launches retain one local device per rank and Gloo collectives. Precision/backend policy is applied only during this explicit launch operation. A supplied mesh with distributed launch controls requires a preinitialized service.