Distributed output: zarr and parallel netCDF backends, gather/rank-block write modes

[kotsaloscv]

What

Output now works in distributed (MPI) runs. Until now the driver disabled IO with a
warning as soon as comm_size > 1.

Output distributions (common/io/distributed.py)

New strategy layer between the field-group monitor (scheduling) and the writers (format):

• GatherDistribution (default): halo entries dropped via the owner masks, owned
  entries Gatherv-ed to rank 0 and placed at their global indices — output identical to
  a single-rank run, works with both backends, but root-memory-bound.
• RankBlockDistribution: no data communication; every rank writes its own owned
  entries into a rank-contiguous block of a shared store. Horizontal axes are padded
  to a uniform block size; global_index_<dim> coordinates recover the global order.
• SingleNodeDistribution: passthrough for single-rank runs.

Zarr writer (common/io/zarr_writers.py)

ZarrWriter (zarr format 3) next to NETCDFWriter, both behind a FieldWriter protocol
and sharing all coordinate/CF attribute definitions, so a store reads like the equivalent
netCDF file. Rank-block stores are chunked so that no chunk (or shard) crosses a
rank-block boundary — one chunk per rank block by default — so concurrent writes never
touch the same chunk file; the root rank owns all store-metadata operations, ordered by
one barrier per append.

Parallel netCDF (common/io/netcdf_writers.py)

NETCDFWriter supports the rank-block mode too: on an MPI-parallel netCDF4
installation, distributed + netcdf writes one shared file (parallel=True on the
monitor's communicator) with the exact layout of the rank-block zarr store — padded
per-rank blocks, HDF5 chunks never crossing block boundaries, global_index_<dim>
coordinates — so reassembly and any downstream consumer work on either format.
Metadata operations run collectively on every rank; variables touching the unlimited
time dimension use collective access, with every rank writing its full block (padding
included) so ranks owning no entries still participate in every collective write.
Vertical chunk sizes shrink automatically to respect HDF5's 4 GiB chunk limit.

The PyPI wheels of netCDF4 are serial builds (__has_parallel4_support__ == 0), so
the combination is rejected at configuration time on such installations — the error and
the common.io module docs spell out the exact steps to an MPI-parallel build
(MPI-enabled netCDF-C/HDF5, then pip install --no-binary netcdf4 --no-build-isolation --force-reinstall netcdf4 with mpi4py preinstalled, then verify the support flag).
The installation is re-verified at writer construction and at file open, and every
parallel open logs the netCDF4/netcdf-c/HDF5 versions in use.

Chunking and sharding (per field group)

• horizontal_chunk_size: entries per chunk along the cell/edge/vertex axes, both
  backends. Default: one chunk per rank block (distributed), whole axis / library
  default otherwise.
• horizontal_shard_size (zarr only): groups whole chunks into one storage file each —
  the file-count knob on parallel file systems (e.g. one shard file per rank per time
  slice with many small chunks inside).

In distributed mode the uniform rank-block size is rounded up to a multiple of the
chunk/shard size (RankBlockDistribution block alignment), so chunks and shards never
cross rank-block boundaries and concurrent rank writes stay in disjoint files; the
extra positions are ordinary padding. Misaligned layouts are rejected at writer
construction, on every rank.

Configuration

Per field group: backend (netcdf | zarr), mode (gather | distributed) and
the chunk/shard sizes above, normalized and validated at construction. distributed +
netcdf requires an MPI-parallel netCDF4 installation and is rejected otherwise,
regardless of rank count. Time-delta output intervals must be a multiple of the model
time step (no silent rounding). Driver exposes --output-backend / --output-mode
(chunk/shard sizes are library-level config for now).

Timings

Output overhead is split into distribute and write phases per group and logged as the
maximum over the ranks. The driver adds output_assemble / output_store timers and a
"% of wall time" column to the timer report.

Robustness

Errors raised by the IO layer are raised on all ranks together (overwrite refusal,
partition validation, UGRID export failure, field validation before any file mutation),
so a configuration or data error aborts the job instead of hanging it in the next
collective — and a failed append can never leave a phantom time slice behind.

Tests

• Unit: distributions (incl. block-alignment rounding and the partition-validation
  rejection), zarr writer (rank-block padding, chunk/shard layouts, misalignment
  rejection), netCDF writer rank-block layout and chunking (through a serial handle, so
  it runs on any installation), the parallel-support predicate branches, backend/mode
  and chunk/shard config validation, monitor-to-writer wiring (rank blocks, chunk/shard
  sizes, alignment-aware distribution sharing).
• MPI, data-free: synthetic decomposition (uneven parts, locally shuffled global
  indices, halos, one empty rank) through monitor → distribution → writer with a file
  rollover, verified value by value against an analytic field; a sub-chunked + sharded
  distributed-zarr case additionally pins the on-disk chunk/shard layout. The
  netcdf+distributed case skips with an explicit reason on serial netCDF4 installations.
• MPI, datatest: 4-rank JW driver runs through gathered netCDF, gathered zarr and
  rank-block zarr must reproduce the single-rank reference; rank-block netCDF joins on
  MPI-parallel netCDF4 installations.
• CI: MODEL_MPI_SUBSETS set to all so the data-free MPI tests run.

Not in this PR (follow-ups)

• Parallel netCDF in CI: the CI images install the (serial) PyPI wheel of netCDF4,
  so the netcdf+distributed MPI cases currently skip there; running them end-to-end
  needs an MPI-parallel netCDF4 build in the image.
• Driver knobs for chunk/shard sizes: horizontal_chunk_size /
  horizontal_shard_size are settable on FieldGroupIOConfig but not yet exposed as
  driver CLI options.
• Asynchronous IO: store is still synchronous; the deep copy of the output state is
  a marked TODO that becomes necessary once writes are off the critical path.
• Gridlook export: the icon4py-gridlook viewer-export CLI has been split out of
  this PR (per review) and lives on the gridlook branch; it will come as its own PR.

[msimberg]

Looking pretty nice, thanks @kotsaloscv! Some general comments and questions, not everything is blocking.

[msimberg]

I see the type ignore, but I don't understand what the "on purpose" is meant for. The config should be taking enums, not strings. The yaml config to config conversion should take care of converting strings to enums.

[kotsaloscv]

You are right. Dropping the "on purpose", Claude being peculiar.

[msimberg]

Hmm, dropping the "on purpose" text was not what I meant, sorry for ambiguity. I meant what is the "on purpose" part even referring to. As mentioned in other comments, strings should not be accepted here in my opinion. E.g. this test I think should not be a thing, or should not be allowed:

def test_fieldgroup_config_accepts_backend_and_mode_value_strings() -> None:
    config = FieldGroupIOConfig(
        filename="a.nc",
        variables=["air_density"],
        backend="netcdf",  # type: ignore[arg-type]
        mode="gather",  # type: ignore[arg-type]
    )
    assert config.backend is OutputBackend.NETCDF
    assert config.mode is OutputMode.GATHER

[kotsaloscv]

Agreed, and that test is gone. Now that string acceptance is removed from the dataclass (see the __post_init__ thread), the test is inverted: test_fieldgroup_config_rejects_backend_and_mode_strings pins that passing a value string raises InvalidConfigError. String-to-enum conversion is covered where it now lives — the config-file boundary — by the ExperimentConfig yaml roundtrip tests from #1391.

[msimberg]

Commented elsewhere, but I think we should not be accepting plain strings here. The only place where that should be allowed is in the conversion from yaml dicts to a proper config object.

This is probably a good place to sync with @DropD to decide in which order changes should be made: whether it makes sense to change these later or wait for @DropD's changes so that it can be incorporated here.

[kotsaloscv]

Agree. Let me keep it as is for now, and I will sync with Rico for this.

[msimberg]

I think I still see __post_init__ coercing backend strings to enums... Am I looking at the wrong changes are are there still places where the strings are expected? I'd remove all uses of strings for config options except when reading from a config file.

[kotsaloscv]

You were looking at the right place — the coercion was still there, pending the config-file boundary. With #1391 merged that boundary now exists, so this is done: OutputBackend/OutputMode are registered with common.config.config_io (@config_io.register_enum), and value strings are converted exactly once, when reading a config file. __post_init__ no longer coerces anything — it rejects a non-enum value with an InvalidConfigError pointing to that boundary. The remaining entry points already deliver enums: the driver CLI (typer converts the option strings) and create_io_monitor (enum fields of DriverConfig).

[havogt]

Just some high-level comments: Mainly, I'd split the unrelated features from this PR (especially gridlook which is completely orthogonal, but huge on its own).

[havogt]

I'd move this out of this PR, because it's an unrelated feature.

[kotsaloscv]

While I fully agree with your comment, I would insist on keeping it here.
My approach with gridlook is that since we are actively developing the driver, it is very essential and useful to have an idea of the generated output. Gridlook does it very efficiently and minimally.
Also, all gridlook additions are extremely self-contained. If we see that it is not as useful, we drop it anytime simply by deleting the corresponding file.

[havogt]

Not sure I understand why it needs to come in the same PR? Makes history cleaner especially since it's touching dependencies for the different features.

[havogt]

maybe match for symmetry?

[kotsaloscv]

Regarding your comment on the total_wall_time, it is there because I wanted to show the percentage of the output compared to the whole driver execution.

[nfarabullini]

very trivial comments on my side. I see that you already have others. Perhaps once everything is tackled, I can do a second round of review

[nfarabullini]

for mpi-related tests we usually have a separate file, e.g. for the metrics states: model/common/tests/common/metrics/mpi_tests

[kotsaloscv]

These tests don't use MPI — the recorder stub just mirrors IOMonitor's full constructor signature (including the process_props/decomposition_info parameters) to pin the wiring; everything runs single-process on the simple grid. The driver's real MPI tests do live in a separate file, mpi_tests/test_parallel_driver_io.py ;)

[nfarabullini]

maybe place all modes in indented bullet points?

[kotsaloscv]

Done — the mode entry now lists "gather" and "distributed" as indented bullets (and picked up the new rank-block layout-marker documentation while at it).

[nfarabullini]

Suggested change

	def file_horizontal_size(self) -> base.HorizontalGridSize:
	def output_horizontal_size(self) -> base.HorizontalGridSize:

maybe? file is a bit too vague IMO

[kotsaloscv]

Done!

[nfarabullini]

can you make the docstring such that it has bullet points which describe what each param indicates?

[kotsaloscv]

Done — the docstring now has an Attributes: section with one bullet per field (start, count, size, padded_size, global_size, global_index).

[nfarabullini]

can you reduce the amount of comments here? It seems a bit excessive

[kotsaloscv]

Trimmed — the cross-references and tuning commentary are gone; the details stay in the common.io module docs.

[nfarabullini]

should this be an error or at least a warning?

[kotsaloscv]

It effectively is an error: missing_parallel_support is a predicate that only returns the reason, and every caller raises with it — the writer at construction when a multi-rank communicator shares the file, and again at file open. Returning instead of raising keeps it usable for skip conditions too (the MPI test skips with this exact reason).

[nfarabullini]

what if we have offset dimensions, e.g. E2CDim?

[kotsaloscv]

Fields carrying sparse/offset dimensions are not writable: writers.canonicalize_time_slice accepts only fields with one horizontal (cell/edge/vertex) plus one vertical dimension and raises for anything else, on every rank, before the file is touched — so an offset dimension never reaches createDimension. Output of connectivity-shaped fields would be a new feature (none of the current output variables need it).

[nfarabullini]

why do we not create variables for cell, edge, and vertex as well?

[kotsaloscv]

The horizontal coordinates (lon/lat per cell/edge/vertex) live in the UGRID sidecar file (*_ugrid.nc) the monitor writes once at startup; the data files reference it through the UGRID association instead of duplicating the coordinates into every rollover file. Inlining clon/clat into the data files would be a possible follow-up if a consumer needs fully self-contained files.

[nfarabullini]

maybe move these functions into a common.py file in the test folder

[kotsaloscv]

They're only used by this one file so far, so I'd keep them local for now and move them to a shared module the moment a second file needs them

[nfarabullini]

so is this test not supposed to run for distributed?

[kotsaloscv]

It runs for all modes (gather and distributed are both in the parametrization) and on every rank — the writing part above is fully collective. The early return only limits the verification (reading the finished files back and comparing values) to rank 0, after the barrier; the other ranks have nothing left to do at that point.

[jcanton]

overall very nice.

I would prefer to separate gridlook into its own PR though.
I'm also not sure it belongs inside the driver: it's a postprocessing visualization tool, not a model running tool

[jcanton]

I don't like rows: there are no rows in our unstructured triangular grids,
this is more nr_of_local_cell/edge/vertex right?
horizontal_counts / entry_count / dim_count / local_size (since you're using global_size just below) ?
in any case, these shuold also all be available from DecompositionInfo and not need to be re-computed here

or am I totally misunderstanding?

[jcanton]

again, look at parallel_helpers.py

[kotsaloscv]

Agreed on the naming — there are no rows here. Renamed throughout: entry_counts, _gather_entries, entries. On DecompositionInfo: the local owned count is indeed available from it, but these are the per-rank count vectors (needed for the Gatherv receive counts/displacements and the block sizing), and DecompositionInfo is rank-local and holds no communicator — so an exchange is unavoidable. What was re-computed in two places is now consolidated: both distributions share module-level helpers (_allgather_entry_counts, _gather_entries, _check_partition), and RankBlockDistribution now runs the same partition validation as GatherDistribution through them.

[kotsaloscv]

cc: @jcanton @havogt @msimberg
Done: gridlook is out of this PR entirely — moved to the gridlook branch (module, tests and its three dependencies in the driver's pyproject), and the PR title/description are updated accordingly. It will come back as its own PR, and I agree the "does it belong in the driver" question is best discussed there — a separate tool package is a real option.

[jcanton]

remember to check if this needs to be updated with new instructions model/standalone_driver/README.md

[msimberg]

I think these can be generic over the ndarray type, no? I.e. the functions can receive a data_alloc.NDArray and return an array that is consistent with the input arrays that is determined with array_namespace (see other functions for examples). For allgather_entry_counts there's no input array to determine the array_ns, so array_ns should be explicitly passed to the function instead. This is minor, but means we don't have to worry about these not being generic across cpu/gpu arrays in the future.

[kotsaloscv]

Done, exactly as you describe. Indeed, way more generic ;)

[msimberg]

This function feels a bit confused. It's both checking and returning values in different use cases. I see one call to it where the return value is ignored and one where the return value is used. The "checking" part is effectively re-checking that our domain decomposition is doing the right thing, which we probably don't need to do every time in production. If there's a gap in check, our testing should rather be expanded (there are already quite a few tests for domain decomposition, are these checks missing from there?). And for the return value ignored or not ignored: if there are two separate use cases, would it instead make sense to separate the function(ality) into two separate functions, one for actually collecting the indices and one for checking, so that they can be used separately.

One level more abstract: we may already have the decomposition on every single node at the beginning of a run, if we've done the decomposition ourselves. Is there no way to reuse that instead?

I realize these might not be straightforward refactorings, but I want to at least ask the questions.

[kotsaloscv]

I did the split as you suggested. Now it is more clear.

[msimberg]

Probably not for now, but this looks like something we'll have to unify sooner or later. We should have a common way of adding timers across modules, not ad-hoc in each...

Are these timers used for anything right now, or more for debugging?

[kotsaloscv]

They are for performance visibility: each rank logs its distribute/write totals at finalize, which is how we compare gather vs. distributed (rank-block) on real runs. Agreed on unifying timers across modules.

[kotsaloscv]

@jcanton Good catch — it did need updating: the README still claimed output is single-node only and disabled in MPI runs. It now documents distributed output, the --output-backend/--output-mode options, and includes an mpirun example.

[msimberg]

Looking really good, thanks for the changes @kotsaloscv. I have one minor question/suggestion for the filename handling. Not insisting on it if you feel like it doesn't fit.

[msimberg]

Nice, thanks for the changes @kotsaloscv!

[kotsaloscv]

cscs-ci run default;SESSIONS=model:model_mpi;MODEL_SUBPACKAGES=common:standalone_driver;MODEL_MPI_SUBPACKAGES=common:standalone_driver;MODEL_MPI_SUBSETS=all;LEVELS=unit:integration

[kotsaloscv]

cscs-ci run default;SESSIONS=model:model_mpi;MODEL_SUBPACKAGES=common:standalone_driver;MODEL_MPI_SUBPACKAGES=common:standalone_driver;MODEL_MPI_SUBSETS=all;LEVELS=unit:integration

[github-actions]

When developing, you can test your changes on CSCS CI before merge with the default pipeline: cscs-ci run default. This will run a default subset of tests.

You can pass options to override pipeline variables, for example:

• cscs-ci run default;BACKENDS=gtfn_cpu;LEVELS=unit
• cscs-ci run default;MODEL_SUBPACKAGES=common:standalone_driver;SESSIONS=model
  Avoid running the pipeline for all tests when you are developing.

Available options are:

• SESSIONS: model, model_mpi, or tools (correspond to nox sessions)
• MODEL_SUBSETS: datatest, basic, or stencils (correspond to nox session selections)
• MODEL_SUBPACKAGES: subpackages for non-MPI tests (last component, e.g. diffusion, standalone_driver)
• MODEL_MPI_SUBPACKAGES: subpackages for MPI tests (as above)
• BACKENDS: backends
• GRIDS: grids for stencil tests (simple, icon_regional, or icon_global)
• LEVELS: testing level for non-stencil tests (unit or integration)

For each option, all can be used as a shorthand for all possible values of that variable, e.g. LEVELS=all.

See scripts/python/generate_ci_pipeline.py and noxfile.py for available values for each option.

The all pipeline can be run with cscs-ci run all. This will run all icon4py tests in CSCS CI which can be expensive. This pipeline runs on a schedule on main, and can be run when extensive validation is needed (e.g. before releases).

Merging

Once your PR is approved and ready for merging, add it to the merge queue. The merge CSCS CI pipeline will run automatically on the merge-queue branch and must pass before the PR is merged. A dummy merge check will be triggered on the PR itself since it's required to add a PR to the merge queue.

Optional Tests

To run benchmarks you can use:

• cscs-ci run benchmark-bencher

For more detailed information please look at CI in the EXCLAIM universe.

[msimberg]

@kotsaloscv #1426 is now merged, so GHEX issues should be gone on main (no need to merge/rebase for the merge queue).

[jcanton]

most failure paths in the IO layer are designed to fail collectively (broadcast verdicts, validation before any file mutation), so a root-only raise cannot strand the other ranks in the next collective. This is the one path that is not: if the root's zarr.open_group or array creation raises (disk full, permissions), the non-root ranks block forever in this barrier. Consider wrapping the root-side initialization in the same broadcast-failure pattern used by _write_ugrid. Nothing blocking.

🤖 Written by an agent on behalf of @jcanton

[jcanton]

prepare expects the horizontal dimension first (field.dims[0]), but the writers canonicalize to horizontal-last (canonicalize_time_slice). A field arriving in canonical (level, cell) order would fail here with a confusing "leading dimension 'level'" error before ever reaching the canonicalizer. It works today because the driver always emits horizontal-first, but the contract is implicit; a comment on OutputDistribution.prepare stating the expected dimension order would help. Nothing blocking.

🤖 Written by an agent on behalf of @jcanton

[jcanton]

nc.default_fillvals[f"{data.dtype.kind}{data.dtype.itemsize}"] is unguarded. The table covers the usual u/i/f/c variants, but exotic dtypes (bool, float16) raise a KeyError deep inside append. Nothing blocking; a .get() fallback or a small fill-value table would be more robust.

🤖 Written by an agent on behalf of @jcanton

[jcanton]

looks good, sorry it took me so long, just a few minor findings from a friend below, nothing blocking