Parallel ensemble I/O for State and Increment - #1243
Open
travissluka wants to merge 9 commits into
Open
Conversation
Rebased onto develop after the direct-netCDF soca_io_mod PR (#1242) landed. End-to-end ensemble-parallel I/O on top of the stateless direct-netCDF reader. Opt-in via the new oops StateSet/IncrementSet bulk-readEnsemble/writeEnsemble dispatch path; soca provides the model-level implementations. State/Increment: - State::{readEnsemble,writeEnsemble} and Increment::{readEnsemble, writeEnsemble} (Fields + interface Fortran bindings to match) - writeEnsemble: gather each member to a strided writer PE (soca_io_ensemble_root_pe assignment), then phase-2 per-member netCDF writes happen concurrently across writer PEs - readEnsemble: per-member loop today, honoring single-state read mode (strided vs scatter); parallel-across-members reads via soca_io_readers_commit_ensemble are wired for a follow-up soca_io_mod: - writer staging split into define/gather/write/close phases plus soca_io_writers_commit_ensemble; reader staging mirrored with read/distribute/close and soca_io_readers_commit_ensemble. Reader opens+closes ncid inside stage_read (no global file-handle cache, matching develop's stateless model post-2d7ab3cf) - new public knobs: soca_io_ensemble_write_parallel, soca_io_ensemble_read_scatter, soca_io_single_state_read_scatter, soca_io_async_mpi, soca_io_ensemble_root_pe - mpi_pelist_and_comm helper for the writer/reader PE subsets Geometry: - soca_io_config_from_yaml called from soca_geom_init resolves the parallel/sequential, scatter/strided, async-mpi knobs from geometry.io once; values persist module-level for the run Companion to JCSDA-internal/oops setID fix in StateSet::buildFromConfigs (needed so PseudoModel and HTLM dispatch see correct per-member IDs when StateSet takes the HasReadEnsemble path).
travissluka
force-pushed
the
feature/ensemble-parallel-io
branch
from
May 23, 2026 00:34
0c72b73 to
3b0674f
Compare
…arallel-io-rebased
The companion oops branch now captures the posterior stddev before the in-place ana_pert mutation but prints it in its original position, so the variance-before-mean testref updates are no longer needed. The PR now touches no testrefs (per-member output is byte-identical to the legacy path).
Correctness: - async gather/scatter: give each in-flight MPI_Igatherv/Iscatterv its own rc/disp column (the shared rc(:)/disp(:) were recomputed each loop iteration while a prior nonblocking collective was still using them) - writer async gather: allocate a 1-element dummy recvbuf on non-root PEs instead of passing an unallocated actual argument to MPI_Igatherv Cleanup: - strip profiling instrumentation (OOPS_PERF prints, system_clock blocks, measurement-only mpp_sync barriers between phases) from the bulk read/write orchestrators - remove dead code: 3 unused public getters, the unused pelist out-arg of the comm helper (now current_mpi_comm), unreachable root_pe/reader_pe sentinel loops - assert all writers/readers in a batch share one FMS domain - correct the stale 'follow-up' doc comment on State::readEnsemble
- factor the per-domain atlas->compute-domain copy + writer enqueue out of soca_fields_write_rst and soca_fields_write_ensemble into a shared soca_fields_enqueue_domain helper, so the single-state and bulk paths no longer carry two copies of the mask/fillvalue and enqueue logic - replace the four separate MPI_Allgather calls of the static compute-domain bounds (copy-pasted across the three bulk I/O routines) with one packed Allgather in a gather_compute_domains helper: 12 collectives per call drop to 3, with the rank-ordered layout the Igatherv/Iscatterv paths rely on preserved exactly
geometry.io.'ensemble batch size' caps how many members are built, committed, and freed per bulk read/write pass. Default 0 keeps the single-batch behavior (all members at once, maximum overlap) and is byte-identical to the prior path; a positive N processes the ensemble in waves of at most N members so peak memory is bounded by one wave's writers/readers plus scratch rather than the whole ensemble. The writer/reader root PE is rotated by position within the batch so each wave still spreads its members one-per-PE across the comm. Validated on the 1deg 7-member LETKF: batch 1/2/none give bit-identical analysis (nccmp -d); batch 1 cuts total memory ~27% (21.4 -> 15.6 GB) at the cost of wall time, matching the serial memory floor.
Both async paths posted every (member, var) collective up front, then did one MPI_Waitall, unpacked all, and freed all buffers at the end. That held every request's transient buffer live simultaneously: on the write side the per-root recv buffers (each a full global field), on the read side the per-root send buffers (full global, packed from gbuf). With the whole ensemble in one batch that doubled the global-field memory at peak. Replace the Waitall + unpack-all + free-all with an MPI_Waitany loop: as each gather/scatter completes, free its send buffer, unpack (on the target root) into the gbuf / caller slice, and free its recv buffer immediately. All requests are still posted up front, so the runtime overlap is unchanged; only the transient high-water drops from all requests to the in-flight set. 1deg 7-member LETKF, single batch: total memory 21.38 -> 19.08 GB (-10.8%), max RSS -10.7%, wall unchanged, analysis bit-identical (nccmp -d) and ctests green.
The ensemble write path carried two ways to serialize: the 'ensemble write: sequential' knob (commit each writer fully, one at a time) and, since the batch-size knob landed, batch 1 (commit one member's writers per pass through the orchestrator). They serve the same purpose -- bound peak memory by not holding the whole ensemble's buffers at once -- and batch 1 already matches a serial write's memory profile, so the separate sequential path was redundant A/B scaffolding. Remove ensemble_write_parallel_, its 'ensemble write' YAML key, the public getter, and the fallback branch; soca_fields_write_ensemble now always commits through the orchestrator and relies on 'ensemble batch size' for serialization. The single-writer commit() method stays (still used by the single-state write, geometry, and balance paths). Also drop two dead nx_g/ny_g locals in the write gather. ctests green.
Collaborator
|
I've just tested it with our 0.25 deg run (but with 30 members), nice!: previous run: The memory seems quite a bit larger, I'll look into it later unless @travissluka you have an idea why. |
Contributor
Author
Yeah, some memory increase is inevitable. Each ensemble member requires the extra buffer for the full state on read/write. It's a trade-off. You can set That said, I know there are still areas for memory optimization (the async MPI calls have duplicate read/write buffers that I would like to get rid of). But, good to see it's on the right track, thanks for testing ! |
travissluka
marked this pull request as ready for review
July 21, 2026 17:17
shlyaeva
approved these changes
Jul 22, 2026
shlyaeva
left a comment
shlyaeva
left a comment
Collaborator
There was a problem hiding this comment.
Review solely based on testing and great timing results. I'm not an expert in new or old i/o, and can't review the code changes well. 🎉
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Warning
Make sure you check out the matching oops branch (JCSDA-internal/oops#3316)! This soca branch will still compile and run, but it won't go any faster without the oops dispatch hooks.
Replaces the per-member-serial ensemble read/write path with bulk routines that drive concurrent I/O across rotated reader/writer PEs, with optional MPI-async batching for the gather/scatter collectives. Previously the LETKF was doing all ensemble file writes on PE0... ouch. Ensemble member reads were using "strided" read mode (whereby each PE reads its section of a netcdf file), not the worst, but probably not the best.
This PR uses a
oops::StateSet::readEnsemble/writeEnsembleset of methods that let us bypass the per-member loop, doing ensemble read/write in true parallel.YAML configuration
The default settings now should do parallel ensemble I/O in what is hopefully the most efficient way on an HPC. But there are some knobs to turn if you need. All knobs live under
geometry.io.ensemble readscatter|stridedscatterMPI_Scattervto compute groups; strided = each PE reads its compute-domain tile directly.async mpitrue|falsetrueIgatherv/Iscattervacross all (reader/writer, var) pairs so collectives with different roots overlap. Applies to both read and write.ensemble batch sizeN >= 000= process all members in one batch (maximum I/O overlap, peak memory ~ all members resident at once).N > 0commits the ensemble in waves of at mostNmembers to cap peak memory (setN~ number of PEs); batch1matches a serial write's memory footprint.single state readstrided|scatterstridedMemory
Bulk parallel I/O necessarily keeps more in flight than the old serial path: each rotated writer PE stages a member's whole global field (and on read, each reader PE does the same). Two things keep that bounded:
MPI_Waitany) as each gather/scatter finishes, rather than holding every member's buffer until one bigWaitall. On a 1deg 7-member LETKF this cut peak memory ~11% at no batching, for free (I/O overlap unchanged).ensemble batch sizecaps how many members are resident per pass.ensemble batch size: 1brings the footprint down to the serial path's level; the default single batch trades that memory for maximum overlap. Output is bit-identical across all batch sizes.Expected performance impact
Based on tests with a synthetic memory/file-system delay emulating HPC performance... my assistant forecasts an improvement of 1.5–3× on reads, 10–30× on writes... let's see how good that forecast is!
ensemble read: scatteris faster when the number of PEs/nodes is large. But if read times are slower, it might be worth testing putting this back tostrided, which is the currently implemented methodasync mpi: truein theory should go faster when doing ensemble I/O, but the flag is here in case we want to verifysingle state readi'm honestly not sure which would be faster, in theorystridedshould be faster if number of PEs >> the number of lustre OSTs.ensemble batch size: 0should be the fastest ensemble writes, but also uses the most memory. The knob is there if we need to reduce memory.Issue(s) addressed
Dependencies
This PR depends on:
State/IncrementreadEnsemble/writeEnsembledispatch hooks this PR implements)Merge order: oops#3316 first, then this PR.
build-group=https://github.com/JCSDA-internal/oops/pull/3316
Impact
Behavior is identical to the legacy path. Expected wall-clock speedup on HPC; no expected impact on downstream science code.
Checklist
Note
Validated bit-identical to the legacy path on the soca EnKF/LETKF ctests. Still need to stress test on an HPC.
Did I review my code? Of course not, I figure it's correct judging by the cpu/memory usage, runtime, and the fact the output files are bit identical. Who cares about looking at the code??