Skip to content

Add whole-run zarr outputs across inference segments - #1369

Draft
elynnwu wants to merge 3 commits into
mainfrom
feature/segmented-zarr-outputs
Draft

Add whole-run zarr outputs across inference segments#1369
elynnwu wants to merge 3 commits into
mainfrom
feature/segmented-zarr-outputs

Conversation

@elynnwu

@elynnwu elynnwu commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

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.files accumulate 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 and MonthlyDataWriter / save_monthly_files are untouched.

Block-mean time_coarsen 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.

Changes:

  • fme.ace.inference.data_writer.zarr.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, 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 coordinates

  • fme.ace.inference.data_writer.monthly_zarr.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

  • fme.ace.inference.data_writer.zarr.ZarrWriterConfig gains an optional path for placing stores outside the experiment directory

  • fme.ace.inference.data_writer.segment.SegmentContext threaded at runtime from the coupled segmented driver through DataWriterConfig / FileWriterConfig build paths; no user-facing config changes and non-segmented runs are unchanged

  • Tests added

elynnwu added 3 commits July 16, 2026 10:05
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.
@elynnwu

elynnwu commented Jul 16, 2026

Copy link
Copy Markdown
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

Base automatically changed from feature/coupled-segmented-inference to main August 25, 2026 20:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant