Add whole-run zarr outputs across inference segments - #1369
Draft
elynnwu wants to merge 3 commits into
Draft
Conversation
Previously the coupled outer loop dropped each component's stepper state (RNG and corrector state) at every coupled-step boundary: per-step initial conditions were rebuilt without it, the SST prescription discarded it from the atmosphere state, and the terminal StepOutputs were built with stepper_state=None. As a result coupled restart files carried no embedded state and stochastic modules reseeded each coupled step. Now the terminal per-component stepper states are threaded through the per-coupled-step initial conditions, preserved by the SST prescription, and attached to the component StepOutputs, so BatchData serialization embeds them in restart files and restored states continue seamlessly.
Adds run_segmented_inference to fme.coupled.inference and a --segments
CLI flag, mirroring the fme.ace segmented driver: each segment runs
n_coupled_steps coupled steps in its own segment_{n:04d} directory, a
segment is complete when both its ocean and atmosphere restart files
exist, completed segments are skipped on re-invocation, and each
segment after the first initializes from the previous segment's
restart files.
CoupledInitialConditionConfig now reads paired restart files, which
have no sample coordinate to select by: the ocean and atmosphere
datasets are aligned positionally, with validation that sample counts
match and both restarts sit at the same coupled step boundary.
CoupledForcingDataLoaderConfig.build_inference_config now hands out
copies of its dataset configs: the built loader updates the atmosphere
subset in place to align it with the ocean start, which previously
corrupted the user config and broke any second use of it, such as the
second segment of a segmented run within one process.
Zarr outputs configured through DataWriterConfig.files now accumulate into a single whole-run store when running segmented coupled inference, instead of producing per-segment files: - ZarrWriterAdapter and SeparateICZarrWriterAdapter take a start_timestep: the store is sized for the whole run, each segment region-writes its slice of the time axis (idempotent under segment re-runs), and resumed segments open the existing store in append mode. Only the run's first segment may create the store, since its initial condition times determine the whole store's time coordinates. - New MonthlyZarrWriter, built from a files entry with zarr format and monthly_mean time coarsening (previously rejected), with the month axis pre-allocated from the whole run's length. Because monthly accumulation is not idempotent, finalize writes a netCDF snapshot to the segment directory (before the restart files, so restarts remain the completion marker) and the next segment's writer restores from it, making segment re-runs exact. MonthlyDataWriter is untouched. - ZarrWriterConfig gains an optional path for placing stores outside the experiment directory, e.g. directly on a remote bucket. - A SegmentContext (segment index and total, run/segment directories) is threaded at runtime from the coupled segmented driver through the data writer build path; no user-facing config is needed and non-segmented runs are unchanged. Block-mean time coarsening needs no special handling across segments: it is stateless per batch and its validation requires the per-segment step count to be divisible by the coarsen factor, so no coarsening block can straddle a segment boundary; whole-run store offsets are simply expressed in coarsened units.
Contributor
Author
|
The idea for this is that we can write the monthly zarr to gcs like this: data_writer:
ocean:
save_prediction_files: false
save_monthly_files: false
files:
- label: output # -> gs://bucket/output.zarr
save_reference: false # predictions only; no "_predictions" suffix
time_coarsen:
method: monthly_mean
format:
name: zarr
path: gs://bucket # store dir; {label}.zarr is created inside
# names: [sst, o_prog] # optional: subset of variables to save
atmosphere:
save_prediction_files: false
save_monthly_files: false
files:
- label: output # -> gs://bucket/atmosphere/output.zarr
save_reference: false
time_coarsen:
method: monthly_mean
format:
name: zarr
path: gs://bucket/atmosphere |
2 tasks
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
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.
Segmented inference writes each segment's outputs to its own directory, which fragments long-run outputs (e.g. monthly means) into per-segment files. This PR makes zarr outputs configured through
DataWriterConfig.filesaccumulate into a single whole-run store shared by all segments, written to any fsspec-compatible location (including directly to a remote bucket), so a preempted-and-resumed run produces the same single stores as an uninterrupted one. netCDF outputs are unaffected andMonthlyDataWriter/save_monthly_filesare untouched.Block-mean
time_coarsenneeds no special handling across segments: it is stateless per batch and its validation requires the per-segment step count to be divisible by the coarsen factor, so no coarsening block can straddle a segment boundary; whole-run store offsets are simply expressed in coarsened units.Changes:
fme.ace.inference.data_writer.zarr.ZarrWriterAdapterandSeparateICZarrWriterAdaptertake astart_timestep: the store is sized for the whole run, each segment region-writes its slice of the time axis, and resumed segments open the existing store in append mode; only the run's first segment may create the store, since its initial condition times determine the store's time coordinatesfme.ace.inference.data_writer.monthly_zarr.MonthlyZarrWriter, built from afilesentry with zarr format andmonthly_meantime coarsening (previously rejected), with the month axis pre-allocated from the whole run's length; because monthly accumulation is not idempotent,finalizewrites a netCDF snapshot to the segment directory (before the restart files, so restarts remain the completion marker) and the next segment's writer restores from it, making segment re-runs exactfme.ace.inference.data_writer.zarr.ZarrWriterConfiggains an optionalpathfor placing stores outside the experiment directoryfme.ace.inference.data_writer.segment.SegmentContextthreaded at runtime from the coupled segmented driver throughDataWriterConfig/FileWriterConfigbuild paths; no user-facing config changes and non-segmented runs are unchangedTests added