Add segmented inference for coupled runs

[elynnwu]

Coupled inference runs often take hours and lose all progress if preempted, and non-preempted jobs have a limit of 8 hours. This PR adds segmented inference to fme.coupled.inference: the run is split into segments of n_coupled_steps coupled steps each, written to segment_{n:04d} subdirectories. A segment is complete when both its ocean and atmosphere restart files exist, so re-invoking the same command after preemption skips completed segments and resumes at the first incomplete one. Each segment after the first initializes from the previous segment's restart files, which sit at a coupled-step boundary and therefore satisfy the ocean-anchored initial-condition timing. Because StepperState is carried through the coupled rollout (added separately, already on main), a segmented run is bitwise-identical to an unsegmented run of the same length.

Rather than copy the fme.ace segmented driver, this factors the two into one shared driver.

Changes:

• fme.core.generics.inference.run_segments is a single segmented-inference driver, used by both fme.ace and fme.coupled. It owns the parts that were previously duplicated and most likely to drift: the ensemble guard, top-level logging with wandb disabled, the WANDB_NAME capture, and the segment label / skip-if-complete / per-segment wandb run loop. Callers pass the four points that actually differ — steps per segment, restart paths, how the next segment's initial condition is built, and the log description — as explicit values and callbacks, so the driver never reads either config object. get_segment_label moves here from fme.ace.inference.inference
• fme.coupled.inference.inference.run_segmented_inference and a --segments CLI flag on fme.coupled.inference. Each segment gets its own wandb run named <base>-<segment label>
• fme.coupled.inference.inference.CoupledInitialConditionConfig now reads paired restart files, which have no sample coordinate to select by: ocean and atmosphere are aligned positionally, with validation that sample counts match and both restarts have the same times
• fme.coupled.data_loading.inference.CoupledForcingDataLoaderConfig.build_inference_config hands out copies of its dataset configs; the built loader updates the atmosphere subset in place to align with the ocean start, which previously corrupted the user's config and made any second use of it fail (such as the second segment of a segmented run in one process)

Segment-completion caveat, unchanged by this PR and shared by fme.ace segmented runs: the restart files are written when the rollout finishes, but the data writer's final flush and the aggregator diagnostics come after, so a segment interrupted in that window counts as complete despite having incomplete diagnostic output.

• Tests added

[elynnwu]

Pre-review agent: pushed two commits (97ead12 cleanups, d9b653a a bug fix) and ran the audit passes. fme/coupled/ + fme/ace/inference/ are green (436 tests), pre-commit clean.

One real bug, fixed in d9b653a. In the segmented path main() called config.configure_logging(...) with wandb enabled before run_segmented_inference, so the first wandb.init ran against the top-level config. Each segment's run_inference_from_config then called init again with no intervening finish(), and wandb.init returns the active run by default in scripts — so every segment logged into that one run with colliding step counters, and the WANDB_NAME=<base>-<segment_label> assignment was dead code, since wandb reads WANDB_NAME only when the first run starts. fme.ace guards both sides of this (top-level logging with log_to_wandb=False, plus WandB.finish() per segment); the coupled mirror dropped both. No existing test caught it because they all ran log_to_wandb=False. The fix restores ace parity and adds the regression test ace already has for #471 — I confirmed it fails without the fix.

Suggested reviewer focus — the deepcopy fixes a symptom, not the cause. InferenceDataset.__init__ mutates its caller's config in place (fme/coupled/data_loading/inference.py:112), and build_inference_config now deep-copies to defend against that from a distance. Pre-existing behavior, so out of scope here, but worth a decision on whether the root-cause fix (having InferenceDataset subset a local copy instead of its config) belongs in this PR or a follow-up.

Known caveat, not introduced here. A segment counts as complete once both restarts exist, but restarts are written when the rollout ends — writer.finalize() and aggregator.flush_diagnostics() run after. Preemption in that window leaves a "complete" segment with truncated diagnostic output. This comes from the shared run_inference and is identical in fme.ace, so it is not a regression; I corrected the docstring and PR description, which had claimed restarts are written last.

Other cleanups in 97ead12: _get_segment_label → get_segment_label (it was a private name imported across packages once fme.coupled began reusing it), a test for the previously-uncovered mismatched-sample-count validation branch, and an import lifted out of a test body.

Passes with nothing to report: config back-compat (new optional CLI flag, no field removed, checkpoints untouched), behavior change on existing paths (the pre-deepcopy failure mode raised loudly rather than silently drifting), inheritance depth, builder pattern (no new violations beyond the ace-mirrored pattern). No split recommended — 4 files, one feature plus the small fix it depends on.

[mcgibbon]

In the long run it would be nice if we could somehow get both ace and coupled to use the same inference.py but supporting either configuration object, where the configuration type controls which kind of inference gets run. Way out of scope for this PR htough.

[mcgibbon]

Chore: Rewrite lengthy AI prose.

[elynnwu]

Pre-review agent: done in 081ce9d. The lengthy paragraphs here are cut, and the restart-before-final-flush caveat now lives once in the shared run_segments driver rather than being restated in both callers with a "same is true of fme.ace" note.

I also tightened the near-identical comment in the fme.ace copy of _get_initialization_time_and_timestep, slightly outside what you flagged, so the two don't diverge.

One left alone, say the word if you want it gone too: the SEGMENT_LABEL_FORMAT comment's "We can reconsider if this changes".

[mcgibbon]

Suggestion (optional): I asked Claude to look out for cases where we might want to refactor code shared between coupled and ace to use a single shared helper instead of re-implementing, it suggested run_segmented_inference and gave a short sketch of what that could look like. It seems like a good idea, I leave the decision to you though. This could be done in this PR, as a follow-on, or not at all.

Claude output follows.

run_segmented_inference here is a near-verbatim copy of the fme.ace version and could be refactored to a shared helper.

Ignoring docstrings, the two ~80-line functions differ in exactly four places:

• one word in the "Starting segmented inference" log message,
• which config attribute supplies the step count (n_forward_steps vs n_coupled_steps),
• how the restart path(s) are built and checked (one restart.nc vs the ocean+atmosphere pair),
• how the next segment's initial-condition config is constructed from those paths.

Everything else is copied verbatim, including the pieces most likely to drift under maintenance: the n_ensemble_per_ic > 1 guard with its identical error message, the top-level no-wandb logging setup, the config deepcopy + WANDB_NAME capture, and the loop skeleton (segment labeling, skip-if-complete, experiment_dir swap, per-segment WANDB_NAME suffix, GlobalTimer scope, WandB.finish() between segments). A future fix to, e.g., the restart-written-before-final-flush caveat or the per-segment wandb behavior would have to be made twice.

A shared driver would take the invariant loop and parameterize the variation points, roughly:

def run_segmented_inference_loop(
    config, segments, *,
    n_steps: int,
    get_restart_paths: Callable[[str], Sequence[str]],   # segment_dir -> paths
    make_initial_condition: Callable[[Sequence[str]], Any],
    run_segment: Callable[[Any], None],
) -> None

with "segment complete" defined as all returned restart paths existing. The config only needs experiment_dir, logging, n_ensemble_per_ic, and a settable initial_condition, so a small Protocol covers the typing. A natural home is fme/core/generics/inference.py (already holds the shared run_inference/get_record_to_wandb), or it could be exported from fme/ace/inference/inference.py — this PR already established that import direction by making get_segment_label public.

[elynnwu]

Pre-review agent: took the suggestion, done in 081ce9d. The loop now lives in fme.core.generics.inference.run_segments, next to the shared run_inference as you suggested, with get_segment_label moved there too. Each caller keeps a thin run_segmented_inference that owns its own config handling. Net effect: fme/ace/inference/inference.py 577 -> 512 lines, fme/coupled/inference/inference.py 484 -> 444.

Two deliberate deviations from the sketch:

• Explicit values instead of a config Protocol. AGENTS.md treats a free function reading a non-leaf config as a parameter bag as a builder-pattern violation, and dataclasses.asdict on a Protocol-typed value doesn't type-check cleanly. So run_segments takes experiment_dir, logging_config, n_ensemble_per_ic etc. directly, and the three variation points come in as callbacks. Neither config type is visible to the shared driver, which also keeps it importable from fme/core without reaching into fme.ace.
• A fourth variation point, get_initialization. Passing the start time and timestep eagerly would have moved the stepper load ahead of the top-level logging setup, pushing its log lines out of inference_out.log. Taking it as a callable preserves the original ordering.

Both the fme.ace wandb test (#471) and the coupled one now exercise the shared driver from their respective sides, so the behavior is pinned from both callers.

[mcgibbon]

Not ideal that the atmosphere object gets updated in-place, especially when we have code already built to subset the dataset in a new object. But it's a pre-existing issue.

[elynnwu]

Pre-review agent: agreed, and left as-is for this PR. The in-place update at InferenceDataset.__init__ is pre-existing, and the deepcopy here defends against it from a distance rather than fixing it. The root-cause fix would be for InferenceDataset to build a subsetted copy instead of mutating its caller's config, using the existing subsetting code you mention. Happy to open a follow-on issue if you'd like it tracked.

[mcgibbon]

LGTM! Great deduplication!