diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_core.py b/packages/evaluate/src/weathergen/evaluate/export/export_core.py index dfb6c000b..5a6f6e638 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/export_core.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_core.py @@ -1,6 +1,7 @@ import logging from collections import defaultdict from multiprocessing import Pool +from pathlib import Path import numpy as np import xarray as xr @@ -18,17 +19,25 @@ _logger = logging.getLogger(__name__) _logger.setLevel(logging.INFO) -# Module-level cache for the zarr path and open store — resolved once per worker. -_CACHED_FNAME_ZARR: str | None = None -_CACHED_ZIO = None +# Per-worker cache: zarr_path → open zarrio_reader context. +_WORKER_ZIO_CACHE: dict[str, object] = {} -def _init_worker(fname_zarr: str) -> None: - """Pool initializer: open the zarr store once and keep it for the worker's lifetime.""" - global _CACHED_FNAME_ZARR, _CACHED_ZIO - _CACHED_FNAME_ZARR = fname_zarr - _CACHED_ZIO = zarrio_reader(fname_zarr) - _CACHED_ZIO.__enter__() +def _init_worker() -> None: + """Pool initializer: reset the per-worker zarr store cache.""" + global _WORKER_ZIO_CACHE + _WORKER_ZIO_CACHE = {} + + +def _get_or_open_zio(zarr_path: str): + """Return a cached zarrio_reader for *zarr_path*, opening it on first access.""" + global _WORKER_ZIO_CACHE + key = str(zarr_path) + if key not in _WORKER_ZIO_CACHE: + zio = zarrio_reader(Path(zarr_path)) + zio.__enter__() + _WORKER_ZIO_CACHE[key] = zio + return _WORKER_ZIO_CACHE[key] def get_data_worker(args: tuple) -> tuple[int, int, xr.DataArray]: @@ -43,14 +52,16 @@ def get_data_worker(args: tuple) -> tuple[int, int, xr.DataArray]: ------- Tuple of (sample, fstep, xarray.DataArray) with data fully in memory. """ - sample, fstep, stream, dtype = args + global_sample, local_sample, zarr_path, fstep, stream, dtype = args + + zio = _get_or_open_zio(zarr_path) # Navigate directly to the zarr group for this (sample, stream, fstep, dtype). - group_path = f"{sample}/{stream}/{fstep}/{dtype}" - ds_group = _CACHED_ZIO.data_root.get(group_path) + group_path = f"{local_sample}/{stream}/{fstep}/{dtype}" + ds_group = zio.data_root.get(group_path) if ds_group is None: - raise FileNotFoundError(f"Zarr group '{group_path}' not found in {_CACHED_FNAME_ZARR}") + raise FileNotFoundError(f"Zarr group '{group_path}' not found in {zarr_path}") # Read raw arrays as numpy — no dask, no chunking overhead. data_arr = np.asarray(ds_group["data"]) # (npoints, nchannels) or (npoints, nchannels, nens) @@ -83,7 +94,7 @@ def get_data_worker(args: tuple) -> tuple[int, int, xr.DataArray]: da_result = xr.DataArray(data_arr, dims=data_dims, coords=data_coords) - return (sample, fstep, da_result) + return (global_sample, fstep, da_result) def get_fsteps(fsteps, fname_zarr: str): @@ -235,7 +246,8 @@ def get_source_info(fname_zarr, stream, samples) -> tuple[list[np.datetime64], l - ``source_start = min(source_times)`` - ``source_end = max(source_times)`` - The ``source_end`` also serves as the reference (initialisation) time. + The true forecast initialisation (reference) time is either ``source_start`` + or ``source_end``, selected via the ``init_time_reference`` option. Parameters ---------- @@ -286,9 +298,8 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: """ Retrieve data from Zarr store and export to the requested format. - All (sample, fstep) pairs are submitted to the pool at once so that - every worker stays busy. Results are grouped by sample and handed to - the parser in sample order. + Iterates over all rank files. Each rank gets its own parser instance + (and therefore its own output GRIB file pair named with the rank label). Parameters ---------- @@ -302,44 +313,57 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: kwargs = OmegaConf.create(kwargs) run_id = kwargs.run_id - samples = kwargs.samples - fsteps = kwargs.fsteps + samples_cfg = kwargs.samples + fsteps_cfg = kwargs.fsteps stream = kwargs.stream - channels = kwargs.channels + channels_cfg = kwargs.channels n_processes = kwargs.n_processes epoch = kwargs.epoch rank = kwargs.rank - # if not list convert to list - if not isinstance(epoch, ListConfig): - epoch = [epoch] - if not isinstance(rank, ListConfig): - rank = [rank] + init_time_reference = kwargs.get("init_time_reference", "source_start") + if init_time_reference not in ("source_start", "source_end"): + raise ValueError( + f"Invalid init_time_reference: {init_time_reference}. " + "Must be 'source_start' or 'source_end'." + ) if data_type not in ["target", "prediction"]: raise ValueError(f"Invalid type: {data_type}. Must be 'target' or 'prediction'.") - fname_zarr_list = get_model_results(run_id, epoch, rank) - processed_samples = [] - for fname_zarr in fname_zarr_list: - fsteps = get_fsteps(fsteps, fname_zarr) - samples = get_samples(samples, fname_zarr) - streams = get_streams(stream, fname_zarr) - _logger.info(f"Streams to process: {streams}") - for stream in streams: - grid_type = get_grid_type(data_type, stream, fname_zarr) - stream_channels = get_channels(channels, stream, fname_zarr) - source_starts, source_ends = get_source_info(fname_zarr, stream, samples) - kwargs["grid_type"] = grid_type - kwargs["channels"] = stream_channels - kwargs["data_type"] = data_type + # --- Discover rank files --- + # get_model_results accepts lists of epochs and ranks ("all" or list of ints). + rank_arg = ["all"] if rank == "all" else (rank if isinstance(rank, list) else [rank]) + rank_files = get_model_results(run_id, [epoch], rank_arg) + if not rank_files: + raise FileNotFoundError(f"No rank files found for run_id={run_id}, epoch={epoch}, rank={rank}") + _logger.info(f"Discovered {len(rank_files)} rank file(s).") + + first_zarr = rank_files[0] + fsteps = get_fsteps(fsteps_cfg, first_zarr) + streams = get_streams(stream, first_zarr) + + for stream in streams: + grid_type = get_grid_type(data_type, stream, first_zarr) + channels = get_channels(channels_cfg, stream, first_zarr) + kwargs["stream"] = stream + kwargs["grid_type"] = grid_type + kwargs["channels"] = channels + kwargs["data_type"] = data_type + + for rank_file in rank_files: + rank_label = rank_file.stem.split("rank")[-1] # e.g. "0000" + _logger.info( + f"RUN {run_id}: Processing rank {rank_label} ({rank_file.name})" + ) + + samples = get_samples(samples_cfg, rank_file) + source_starts, source_ends = get_source_info(rank_file, stream, samples) + kwargs["rank_label"] = rank_label parser = CfParserFactory.get_parser(config=config, **kwargs) n_fsteps = len(fsteps) total_tasks = len(samples) * n_fsteps - - # Batch size in *samples*. Limits how many samples can be in-flight at once, - # bounding peak memory while still allowing read/write overlap within each batch. batch_size = max(1, n_processes * 2) n_batches = (len(samples) + batch_size - 1) // batch_size @@ -350,12 +374,7 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: f"Reading and writing are interleaved within each batch." ) - # Initialise each worker with the zarr path so it is resolved only once. - with Pool( - processes=n_processes, - initializer=_init_worker, - initargs=(fname_zarr,), - ) as pool: + with Pool(processes=n_processes, initializer=_init_worker) as pool: samples_written = 0 for batch_idx in range(n_batches): @@ -369,8 +388,8 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: sample_to_batch_idx = {s: i for i, s in enumerate(batch_samples)} batch_tasks = [ - (sample, fstep, stream, data_type) - for sample in batch_samples + (s, s, str(rank_file), fstep, stream, data_type) + for s in batch_samples for fstep in fsteps ] @@ -387,33 +406,37 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: pbar = tqdm( total=len(batch_tasks), - desc=f" Batch {batch_idx + 1}/{n_batches}", + desc=f" Rank {rank_label} batch {batch_idx + 1}/{n_batches}", ) - processed_samples = [] - - for sample, _fstep, data in pool.imap_unordered( + for global_s, _fstep, data in pool.imap_unordered( get_data_worker, batch_tasks, chunksize=1 ): - sample_results[sample].append(data) + sample_results[global_s].append(data) pbar.update(1) # Check if this sample is complete (all fsteps received). - if len(sample_results[sample]) == n_fsteps: - b_idx = sample_to_batch_idx[sample] + if len(sample_results[global_s]) == n_fsteps: + b_idx = sample_to_batch_idx[global_s] source_start = batch_source_starts[b_idx] source_end = batch_source_ends[b_idx] - results_iter = iter(sample_results[sample]) - processed = parser.process_sample( - results_iter, - ref_time=source_end, + # The forecast init time is either the start or the end + # of the source (conditioning) window, selected via + # `init_time_reference` (e.g. a 00-05 UTC window has + # init = 00:00 for "source_start" or 05:00 for "source_end"). + init_time = ( + source_start + if init_time_reference == "source_start" + else source_end + ) + parser.process_sample( + iter(sample_results[global_s]), + ref_time=init_time, source_interval_start=source_start, - source_interval_end=source_end, + source_interval_end=init_time, ) - processed_samples.append(processed) - # Free memory immediately. - del sample_results[sample] + del sample_results[global_s] batch_written += 1 pbar.close() @@ -428,8 +451,13 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: # Free any remaining refs before next batch. del sample_results - # Only save here if need to merge samples, otherwise saved in process_sample - if processed_samples[0] is not None: - parser.save(processed_samples) - _logger.info(f"Export complete. Wrote {samples_written}/{len(samples)} samples.") + # Flush and close the parser's file handles for this rank. + if hasattr(parser, "close"): + parser.close() + + _logger.info( + f"Rank {rank_label}: wrote {samples_written}/{len(samples)} samples." + ) + + _logger.info(f"Export complete across {len(rank_files)} rank(s).") diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_inference.py b/packages/evaluate/src/weathergen/evaluate/export/export_inference.py index 0e584ac2e..48d3c08d4 100755 --- a/packages/evaluate/src/weathergen/evaluate/export/export_inference.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_inference.py @@ -152,18 +152,28 @@ def parse_args(args: list) -> argparse.Namespace: parser.add_argument( "--epoch", - nargs="+", - default=None, type=int, + default=0, help="Epoch number to identify the Zarr store", ) parser.add_argument( "--rank", - nargs="+", - default=None, - type=int, - help="Rank number to identify the Zarr store", + type=str, + default="all", + help="Rank(s) to process. Use '0' for a single rank, 'all' for every rank file, " + "or a comma-separated list like '0,1,2'.", + ) + + parser.add_argument( + "--init-time-reference", + type=str, + choices=["source_start", "source_end"], + default="source_start", + help="Which end of the source (conditioning) window to use as the forecast " + "initialisation time written to the output metadata (e.g. GRIB 'date'/'time'). " + "'source_start' (default) uses the beginning of the window " + "(e.g. 00 UTC for a 00-05 UTC window); 'source_end' uses the end of the window.", ) parser.add_argument( @@ -290,6 +300,16 @@ def export_from_args(args: list) -> None: if kwargs.get("expver") == "NEW": kwargs["expver"] = generate_new_expver() + + # Normalise rank + raw_rank = kwargs.get("rank", "all") + if raw_rank == "all": + kwargs["rank"] = "all" + elif "," in str(raw_rank): + kwargs["rank"] = [int(r.strip()) for r in str(raw_rank).split(",")] + else: + kwargs["rank"] = int(raw_rank) + _logger.info(kwargs) # Ensure output directory exists diff --git a/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py b/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py index 29fb5c6c2..895d56625 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py @@ -63,8 +63,15 @@ def __init__(self, config: OmegaConf, **kwargs): self.encoder = ekd.create_encoder("grib") - self.pl_file = ekd.create_target("file", self.get_output_filename("pl")) - self.sf_file = ekd.create_target("file", self.get_output_filename("sfc")) + # Open output files once for the lifetime of this parser. + # Using plain open() in "wb" mode so we own the file handle and + # field.to_target("file", fh) appends directly without earthkit + # ever reopening the path. + pl_path = self.get_output_filename("pl") + sfc_path = self.get_output_filename("sfc") + self.pl_file = open(pl_path, "wb") # noqa: SIM115 + self.sf_file = open(sfc_path, "wb") # noqa: SIM115 + _logger.info(f"Opened output files: {pl_path}, {sfc_path}") self.template_cache = self.cache_templates() @@ -96,10 +103,6 @@ def process_sample( result = result.as_xarray().squeeze() result = result.sel(channel=self.channels) - # Each zarr fstep may contain multiple hourly sub-steps - # concatenated along ipoint. Split by unique valid_time so - # each GRIB message gets exactly one time step worth of grid - # points, with step = valid_time - source_interval_start. unique_times = np.unique(result.valid_time.values) for vt in unique_times: @@ -137,10 +140,12 @@ def process_sample( field_list = pl_fields if level_type == "pl" else sf_fields field_list.append(encoded.to_field()) - self.save(pl_fields, "pl") - self.save(sf_fields, "sfc") + for field in pl_fields: + field.to_target("file", self.pl_file) + for field in sf_fields: + field.to_target("file", self.sf_file) - _logger.info(f"Saved sample data to {self.output_format} in {self.output_dir}.") + _logger.info(f"Saved sample to {self.output_format} in {self.output_dir}.") def extract_var_info(self, var: str) -> tuple[str, str, str]: """ @@ -196,11 +201,9 @@ def cache_templates(self) -> dict[tuple[str, str], object]: def get_output_filename(self, level_type: str) -> Path: """ - Generate output filename. + Generate output filename, including rank label when available. Parameters ---------- - data_type : str - Type of data (e.g., 'prediction' or 'target'). level_type : str Level type (e.g., 'sfc', 'pl', etc.). Returns @@ -208,9 +211,11 @@ def get_output_filename(self, level_type: str) -> Path: Path Output filename as a Path object. """ + rank_label = getattr(self, "rank_label", None) + rank_tag = f"_{rank_label}" if rank_label else "" return ( Path(self.output_dir) - / f"{self.data_type}_{level_type}_{self.run_id}_{self.expver}.{self.file_extension}" + / f"{self.data_type}_{level_type}_{self.run_id}_{self.expver}{rank_tag}.{self.file_extension}" ) def assign_coords(self, data: xr.DataArray) -> xr.DataArray: @@ -244,8 +249,9 @@ def get_metadata( Add metadata to the dataset attributes. The GRIB ``step`` is computed as ``valid_time - source_interval_end`` - (in hours), i.e. the lead time relative to the end of the - conditioning window. + (in hours), i.e. the lead time relative to the forecast init time + (``source_interval_end`` is set to ``source_start`` by the caller, + the beginning of the conditioning window). """ step_hours = int((valid_time - source_interval_end) / np.timedelta64(1, "h")) @@ -259,21 +265,14 @@ def get_metadata( metadata["level"] = level return metadata - def save(self, encoded_fields: list, level_type: str): - """ - Save the dataset to a file. - Parameters - ---------- - encoded_fields : List - List of encoded fields to write. - level_type : str - Level type ('pl' or 'sfc'). - Returns - ------- - None - """ - - file = self.pl_file if level_type == "pl" else self.sf_file - - for field in encoded_fields: - file.write(field) + def close(self): + """Flush and close the output file handles.""" + for fh in (self.pl_file, self.sf_file): + try: + fh.flush() + fh.close() + except Exception: + pass + + def __del__(self): + self.close()