adjust for quaver parser

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).")

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