Extend Prognostic Prescriber to Coupled Modeling

[Arcomano1234]

Coupled models could not use the prognostic prescriber available in uncoupled ACE, so ocean/atmosphere prognostic variables (e.g. SST, thetao_18) could not be overwritten from forcing data during evaluation/inference. This PR extends the prescriber to coupled modeling and lets it be enabled at inference time via overrides, so checkpoints trained without prescribed names can still prescribe them.

Changes:

• fme.coupled.stepper.CoupledStepperConfig: include each component's prescribed_prognostic_names in the ocean and atmosphere forcing windows and forcing-window data requirements, so prescribed variables are loaded and supplied as next-step inputs; added refresh_{ocean,atmosphere}_forcing_window_names to recompute after inference-time overrides.

• fme.coupled.inference.evaluator: thread optional ocean/atmosphere StepperOverrideConfig through the coupled load paths and InferenceEvaluatorConfig, applying overrides to both loaded steppers and the nested StepperConfig used for forcing-window requirements.

• fme.ace.stepper.single_module: extract apply_stepper_override from load_stepper and consolidate override dispatch into a shared helper reused for both a Stepper and a StepperConfig.

• Tests added

[elynnwu]

If I understand correctly, this only works for coupled eval, not inference. But this feature works for both ace eval/inference. Consider adding ocean_stepper_override / atmosphere_stepper_override to InferenceConfig and threading them into its loaders,, so the feature is available in coupled.inference also. I think it could be useful if we have crashes on a long running job where we don't necessary have target data to compare with.

[elynnwu]

lgtm

[jpdunc23]

Correctness of this PR seems to be good.

That said, I think we can get rid of all of the private attribute accesses (step_selector._step_config_instance, the stepper._config patching in sync_coupled_stepper_runtime_stepper_configs / load_coupled_stepper) and config name recomputations with a bit of straightforward refactoring suggested by Claude Fable 5:

1. Add a public prescribed_prognostic_names getter to StepConfigABC (and expose through StepSelector/StepperConfig/Stepper). The ABC already has the mutator replace_prescribed_prognostic_names and pairs replace_ocean with get_ocean — this is just the missing getter half. multi_call delegates to wrapped_step like its other properties, separate_radiation returns [] to match its no-op replace. That deletes _get_prescribed_prognostic_names_from_step and its isinstance recursion.
2. Make ocean_forcing_window_names / atmosphere_forcing_window_names computed @propertys instead of __post_init__-cached lists. The set ops are cheap, and the refresh_* methods (and all their call sites) disappear — mutations are reflected automatically.
3. Establish the config-aliasing invariant in CoupledStepper.__init__/from_state by setting the nested ComponentConfig.stepper from the public ocean.config / atmosphere.config properties, instead of patching it post-hoc in sync_coupled_stepper_runtime_stepper_configs. Then load_coupled_stepper is a plain from_state again.

[Arcomano1234]

Add a public prescribed_prognostic_names getter to StepConfigABC (and expose through StepSelector/StepperConfig/Stepper). The ABC already has the mutator replace_prescribed_prognostic_names and pairs replace_ocean with get_ocean — this is just the missing getter half. multi_call delegates to wrapped_step like its other properties, separate_radiation returns [] to match its no-op replace. That deletes _get_prescribed_prognostic_names_from_step and its isinstance recursion.

Regarding this comment if you look at this commit I did this exact thing, the problem is it touches a lot of code and seems pretty bloaty so in the end of removed it. Can you take a look at this commit to see if you think its worth it? The other 2 suggestions are worth it IMO but I am on the fence about the first one.

[jpdunc23]

See the output from running the /pre-review skill in a clean session below. To address, I suggest:

• Item 1: Please disallow ocean and derived_forcings overrides on CoupledStepper in this PR. If we need them we can figure out how to enable them in a follow-up PR. See https://github.com/ai2cm/ace/pull/1030/changes#r3562194140.
• Item 2: Please add the suggested "validation error when prescribed_prognostic_names intersects the ocean-to-atmosphere-supplied names".
• Item 3: See https://github.com/ai2cm/ace/pull/1030/changes#r3562082953
• Item 4: As stated.
• Item 5: As stated.
• Item 6: I'd suggest adding CoupledStepperConfig.apply_component_overrides() in this PR, but could be done as a follow-up.

Pre-review agent: I did a pass over the full diff ahead of review (merged latest main locally; the new tests in fme/coupled/test_stepper.py, test_evaluator.py, and test_inference.py all pass on the merged tree). Overall the PR is in good shape — the on-demand forcing-window properties plus the config aliasing in CoupledStepper.__init__ is a clean design, and the standalone-checkpoint validation fails loudly rather than silently ignoring overrides. A few things worth attention before/during review:

1. Stale cached name sets for ocean= / derived_forcings= overrides. CoupledStepperConfig.__post_init__ caches _ocean_forcing_exogenous_names, _atmosphere_forcing_exogenous_names, _shared_forcing_exogenous_names, and _atmosphere_ocean_config from the component configs, and also runs cross-component validation there. The new override path (apply_coupled_stepper_config_inference_overrides / apply_stepper_override on loaded component steppers) mutates the nested stepper configs after __post_init__ has run. That's fine for prescribed_prognostic_names (the new window-name properties compute those on demand), but a StepperOverrideConfig(ocean=...) or derived_forcings=... override could change the atmosphere's ocean config or input names while the cached exogenous sets and _atmosphere_ocean_config stay stale — silently. Suggest either (a) rejecting ocean/derived_forcings overrides in the coupled path for now (like multi_call already is), or (b) recomputing/revalidating after mutation.

2. Prescribed atmosphere names can be silently clobbered by ocean-supplied fields. In _get_atmosphere_forcings, forcing_data.update(forcings_from_ocean) runs after prescribed prognostics are pulled from atmos_data. If someone prescribes a name that the coupler also supplies from the ocean (e.g. the surface temperature name), the prescription is silently overwritten. A validation error when prescribed_prognostic_names intersects the ocean-to-atmosphere-supplied names would make this loud.

3. Private cross-module import. evaluator.py imports _apply_stepper_overrides from fme.ace.stepper.single_module. Along with my earlier inline comment about folding _apply_stepper_overrides into apply_stepper_override: one way to resolve both is to make apply_stepper_override the single public entry point accepting Stepper | StepperConfig, with the multi_call-on-config rejection built in, and drop the private import.

4. PR description is stale. It says the PR "added refresh_{ocean,atmosphere}_forcing_window_names to recompute after inference-time overrides," but the final code has no refresh methods — window names are now on-demand properties. Since the description becomes the squash commit message, please update it to describe the final state.

5. Minor cleanups:

   • from unittest.mock import MagicMock appears inside three test function bodies — lift to module top.
   • A couple of comments/docstrings narrate the refactor history rather than the code ("replaces the previous post-hoc config-sync helper" in CoupledStepper.__init__; "the branch order is preserved so behavior matches the previous per-caller implementations" in _apply_stepper_overrides). These won't mean anything after merge — trim to what the code needs.
   • Stray added blank line in CoupledStepperConfig.__post_init__.

6. (Non-blocking design note.) CoupledStepper.__init__ reaching in with config.ocean.stepper = ocean.config, and the evaluator free functions mutating coupled_config.ocean.stepper, both treat another config's fields as public. A sanctioned method on CoupledStepperConfig (e.g. apply_component_overrides(ocean_override, atmosphere_override)) would keep that traversal inside the owning class. Fine to defer.

One intended behavior change worth stating in the description: get_forcing_window_data_requirements now includes prescribed prognostic names, so any checkpoint whose component steppers already carry prescribed_prognostic_names will now require those variables in the forcing dataset (a loud data-loading failure rather than a silent skip — which is the point of the PR, just worth calling out).

— 🤖 Drafted by Claude (claude-fable-5) via Claude Code

[jpdunc23]

Nit: You could just include the contents of _apply_stepper_overrides in the body of apply_stepper_override. I think the diff would also be a bit smaller if you edited load_stepper above and put apply_stepper_override below it.

[jpdunc23]

Given that this method isn't typed in a way that is specific to fme.coupled, I think it should live in fme/ace/stepper/single_module.py (and maybe just call it apply_stepper_override_to_stepper_config).

This would also avoid the import of the private _apply_stepper_overrides method.

If you remove it, please remove mention of CoupledStepperConfig or "coupled" from the docstring, since fme.ace shouldn't refer back to fme.coupled.

[jpdunc23]

I would just handle the StepperOverrideConfig validation to disallow replace_multi_call, replace_ocean, and replace_derived_forcings as a private method in this module that you can apply to ocean_override and atmosphere_override inside of apply_coupled_stepper_config_inference_overrides().

[jpdunc23]

Nit: call this stepper_or_config rather than target.