Add segmented inference for coupled runs #1368
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,8 +1,12 @@ | ||
| import copy | ||
| import dataclasses | ||
| import datetime | ||
| import logging | ||
| import os | ||
| from collections.abc import Sequence | ||
| from typing import Literal | ||
|
|
||
| import cftime | ||
| import dacite | ||
| import torch | ||
| import xarray as xr | ||
|
|
@@ -13,21 +17,28 @@ | |
| InferenceInitialConditionIndices, | ||
| TimestampList, | ||
| ) | ||
| from fme.ace.inference.inference import ( | ||
| InitialConditionConfig, | ||
| get_initial_condition, | ||
| get_segment_label, | ||
| ) | ||
| from fme.ace.requirements import InitialConditionRequirements | ||
| from fme.ace.stepper import StepperOverrideConfig | ||
| from fme.core.cli import prepare_config, prepare_directory | ||
| from fme.core.cloud import exists, makedirs | ||
| from fme.core.derived_variables import get_derived_variable_metadata | ||
| from fme.core.generics.inference import get_record_to_wandb, run_inference | ||
| from fme.core.logging_utils import LoggingConfig | ||
| from fme.core.timing import GlobalTimer | ||
| from fme.core.wandb import WandB | ||
| from fme.coupled.aggregator import InferenceAggregatorConfig | ||
| from fme.coupled.data_loading.batch_data import CoupledPrognosticState | ||
| from fme.coupled.data_loading.getters import get_forcing_data | ||
| from fme.coupled.data_loading.gridded_data import InferenceGriddedData | ||
| from fme.coupled.data_loading.inference import CoupledForcingDataLoaderConfig | ||
| from fme.coupled.inference.data_writer import ( | ||
| ATMOSPHERE_OUTPUT_DIR_NAME, | ||
| OCEAN_OUTPUT_DIR_NAME, | ||
| CoupledDataWriterConfig, | ||
| CoupledPairedDataWriter, | ||
| DatasetMetadata, | ||
|
|
@@ -90,9 +101,34 @@ def get_initial_condition( | |
| ocean = self.ocean.get_dataset(self.start_indices) | ||
| # time is a required variable but not necessarily a dimension | ||
| sample_dim_name = ocean.time.dims[0] | ||
| atmos = self.atmosphere.get_dataset() | ||
| if sample_dim_name in ocean.indexes: | ||
| atmos = atmos.sel({sample_dim_name: ocean[sample_dim_name]}) | ||
| else: | ||
| # Datasets without a sample coordinate (e.g. paired restart files | ||
| # written from a single CoupledPrognosticState) are positionally | ||
| # aligned, so validate the alignment instead of selecting. | ||
| if self.start_indices is not None: | ||
| raise ValueError( | ||
| "start_indices cannot be used with an ocean initial " | ||
| f"condition dataset that has no '{sample_dim_name}' " | ||
| "coordinate to select by." | ||
| ) | ||
| if atmos.sizes[sample_dim_name] != ocean.sizes[sample_dim_name]: | ||
| raise ValueError( | ||
| "Ocean and atmosphere initial condition datasets have no " | ||
| f"'{sample_dim_name}' coordinate and different numbers of " | ||
| f"samples: {ocean.sizes[sample_dim_name]} and " | ||
| f"{atmos.sizes[sample_dim_name]}." | ||
| ) | ||
| if not (atmos["time"].values == ocean["time"].values).all(): | ||
| raise ValueError( | ||
| "Ocean and atmosphere initial condition datasets have no " | ||
| f"'{sample_dim_name}' coordinate and different times; both " | ||
| "must be at the same coupled step boundary. Got ocean " | ||
| f"times {ocean['time'].values} and atmosphere times " | ||
| f"{atmos['time'].values}." | ||
| ) | ||
| return CoupledPrognosticState( | ||
| ocean_data=get_initial_condition( | ||
| ds=ocean, | ||
|
|
@@ -237,6 +273,7 @@ def get_data_writer( | |
|
|
||
| def main( | ||
| yaml_config: str, | ||
| segments: int | None = None, | ||
| override_dotlist: Sequence[str] | None = None, | ||
| ): | ||
| config_data = prepare_config(yaml_config, override=override_dotlist) | ||
|
|
@@ -247,8 +284,116 @@ def main( | |
| ) | ||
| prepare_directory(config.experiment_dir, config_data) | ||
| with torch.no_grad(): | ||
| if segments is None: | ||
| with GlobalTimer(): | ||
| return run_inference_from_config(config) | ||
| else: | ||
| run_segmented_inference(config, segments) | ||
|
|
||
|
|
||
| def _get_initialization_time_and_timestep( | ||
| config: InferenceConfig, | ||
| ) -> tuple[cftime.datetime, datetime.timedelta]: | ||
| # Loading the stepper is expensive, but it is necessary to get the coupled | ||
| # timestep and prognostic names. We only call this function once before the | ||
| # segmented loop starts to minimize the overhead. | ||
| stepper = config.load_stepper() | ||
| initial_condition = config.initial_condition.get_initial_condition( | ||
| ocean_prognostic_names=stepper.config.ocean.stepper.prognostic_names, | ||
| atmosphere_prognostic_names=stepper.config.atmosphere.stepper.prognostic_names, | ||
| n_ensemble_per_ic=config.n_ensemble_per_ic, | ||
| ) | ||
| # Coupled steps are anchored to the ocean, so label segments by the ocean | ||
| # initial condition's start time. | ||
| initialization_time = ( | ||
| initial_condition.ocean_data.as_batch_data().time.isel(sample=0).item() | ||
| ) | ||
| # The coupled ("outer") timestep is the ocean timestep. | ||
| return initialization_time, stepper.config.timestep | ||
|
|
||
|
|
||
| def run_segmented_inference(config: InferenceConfig, segments: int): | ||
|
Contributor
There was a problem hiding this comment. 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 Claude output follows.
Ignoring docstrings, the two ~80-line functions differ in exactly four places:
Everything else is copied verbatim, including the pieces most likely to drift under maintenance: the 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],
) -> Nonewith "segment complete" defined as all returned restart paths existing. The config only needs
Contributor
Author
There was a problem hiding this comment. Pre-review agent: took the suggestion, done in Two deliberate deviations from the sketch:
Both the |
||
| """Run coupled inference in multiple segments. | ||
|
|
||
| Each segment runs ``config.n_coupled_steps`` coupled steps, writing its | ||
| outputs to a subdirectory of the experiment directory labeled by the start | ||
| time of its first (or only) ensemble member. A segment is complete when both | ||
| its ocean and atmosphere restart files exist; these are written once the | ||
| rollout finishes, but before the data writer's final flush, so a segment | ||
| interrupted in that window counts as complete despite having incomplete | ||
| diagnostic output (the same caveat applies to ``fme.ace`` segmented runs). | ||
| Completed segments are skipped, so an interrupted run resumes at the first | ||
| incomplete segment when invoked again with the same configuration. 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. | ||
|
Contributor
There was a problem hiding this comment. Chore: Rewrite lengthy AI prose.
Contributor
Author
There was a problem hiding this comment. Pre-review agent: done in I also tightened the near-identical comment in the One left alone, say the word if you want it gone too: the |
||
|
|
||
| Args: | ||
| config: inference configuration to be used for each individual segment. | ||
| The provided initial condition configuration will only be used for | ||
| the first segment. | ||
| segments: total number of segments desired. Only missing segments will | ||
| be run. | ||
| """ | ||
| if config.n_ensemble_per_ic > 1: | ||
| raise ValueError( | ||
| "Ensemble inference (n_ensemble_per_ic > 1) is not supported with " | ||
| "segmented inference. A segment's restart already carries the " | ||
| "broadcasted ensemble as its sample dimension, so later segments " | ||
| "cannot re-broadcast it consistently. Run with n_ensemble_per_ic=1, " | ||
| "or run a single non-segmented inference for ensemble runs." | ||
| ) | ||
| # Configure top-level logging without a wandb run; each segment owns its run. | ||
| top_level_logging = dataclasses.replace(config.logging, log_to_wandb=False) | ||
| top_level_logging.configure_logging( | ||
| config.experiment_dir, | ||
| "inference_out.log", | ||
| config=dataclasses.asdict(config), | ||
| resumable=False, | ||
| ) | ||
| logging.info( | ||
| f"Starting segmented coupled inference with {segments} segments. " | ||
| f"Saving to {config.experiment_dir}." | ||
| ) | ||
| config_copy = copy.deepcopy(config) | ||
| original_wandb_name = os.environ.get("WANDB_NAME") | ||
|
|
||
| initialization_time, timestep = _get_initialization_time_and_timestep(config) | ||
| n_coupled_steps = config.n_coupled_steps | ||
|
|
||
| for segment in range(segments): | ||
| segment_label = get_segment_label( | ||
| initialization_time, | ||
| timestep, | ||
| segment, | ||
| n_coupled_steps, | ||
| ) | ||
| segment_dir = os.path.join(config.experiment_dir, segment_label) | ||
| ocean_restart_path = os.path.join( | ||
| segment_dir, OCEAN_OUTPUT_DIR_NAME, "restart.nc" | ||
| ) | ||
| atmosphere_restart_path = os.path.join( | ||
| segment_dir, ATMOSPHERE_OUTPUT_DIR_NAME, "restart.nc" | ||
| ) | ||
| if exists(ocean_restart_path) and exists(atmosphere_restart_path): | ||
| logging.info(f"Skipping segment {segment} because it has already been run.") | ||
| else: | ||
| logging.info(f"Running segment {segment}.") | ||
| config_copy.experiment_dir = segment_dir | ||
| if original_wandb_name is not None: | ||
| os.environ["WANDB_NAME"] = f"{original_wandb_name}-{segment_label}" | ||
| with GlobalTimer(): | ||
| run_inference_from_config(config_copy) | ||
| # Finish this segment's run so the next segment starts a fresh one. | ||
| WandB.get_instance().finish() | ||
| config_copy.initial_condition = CoupledInitialConditionConfig( | ||
| ocean=ComponentInitialConditionConfig( | ||
| path=ocean_restart_path, engine="netcdf4" | ||
| ), | ||
| atmosphere=ComponentInitialConditionConfig( | ||
| path=atmosphere_restart_path, engine="netcdf4" | ||
| ), | ||
| ) | ||
|
|
||
|
|
||
| def run_inference_from_config(config: InferenceConfig): | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pre-review agent: agreed, and left as-is for this PR. The in-place update at
InferenceDataset.__init__is pre-existing, and thedeepcopyhere defends against it from a distance rather than fixing it. The root-cause fix would be forInferenceDatasetto 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.