"""Streaming time-window indices for OneScience ERA5Dataset training. The official NeuralGCM reader continuously samples shuffled temporal windows. This lightweight adapter keeps that behavior while leaving file decoding to OneScience's ``ERA5Dataset``. It never materializes the complete training set; only the indices for the current global batch are returned. """ from __future__ import annotations from dataclasses import dataclass from concurrent.futures import Future, ThreadPoolExecutor from collections import deque import numpy as np @dataclass class WindowBatchStream: """Infinite shuffled stream of dataset indices. ``global_batch`` is the number of distinct windows consumed by one optimizer step across all devices. With ``drop_last=True`` (the default), every pass has a fixed number of complete batches, matching the stable batch contract expected by JAX ``jit``/``pmap``. Dataset passes are an internal reader detail; training progress is measured only in steps, as in the official NeuralGCM ``Experiment``. """ size: int global_batch: int seed: int = 0 shuffle: bool = True drop_last: bool = True def __post_init__(self): self.size = int(self.size) self.global_batch = int(self.global_batch) if self.size <= 0: raise ValueError("stream size must be positive") if self.global_batch <= 0: raise ValueError("global_batch must be positive") if self.drop_last and self.size < self.global_batch: raise ValueError( f"dataset size {self.size} is smaller than global_batch " f"{self.global_batch}" ) self.steps_per_pass = ( self.size // self.global_batch if self.drop_last else (self.size + self.global_batch - 1) // self.global_batch ) if self.steps_per_pass <= 0: raise ValueError("stream has no complete batches") self.samples_seen = 0 self._order = np.arange(self.size, dtype=np.int64) self._rng = np.random.default_rng(self.seed) self._reset_pass() def _reset_pass(self) -> None: if self.shuffle: # The persistent RNG yields a deterministic new permutation on # every pass without introducing an epoch-level training concept. self._order = self._rng.permutation(self.size).astype(np.int64) self._cursor = 0 def next_indices(self) -> np.ndarray: """Returns the next global batch from the repeating data stream.""" if self._cursor + self.global_batch > self.size: if self.drop_last: self._reset_pass() else: # Keep a static batch shape by wrapping into the next pass. remainder = self._order[self._cursor :] self._reset_pass() needed = self.global_batch - len(remainder) indices = np.concatenate((remainder, self._order[:needed])) self._cursor = needed self.samples_seen += self.global_batch return indices indices = self._order[self._cursor : self._cursor + self.global_batch] self._cursor += self.global_batch self.samples_seen += self.global_batch return indices.copy() def state_dict(self) -> dict: """Return enough state to reproduce the next emitted batch exactly.""" return { "format_version": 1, "size": self.size, "global_batch": self.global_batch, "seed": self.seed, "shuffle": self.shuffle, "drop_last": self.drop_last, "samples_seen": self.samples_seen, "order": self._order.copy(), "cursor": self._cursor, "rng_state": self._rng.bit_generator.state, } def load_state_dict(self, state: dict) -> None: """Restore a state produced by :meth:`state_dict`.""" expected = { "size": self.size, "global_batch": self.global_batch, "seed": self.seed, "shuffle": self.shuffle, "drop_last": self.drop_last, } mismatches = { key: (state.get(key), value) for key, value in expected.items() if state.get(key) != value } if mismatches: raise ValueError( "Training data stream settings differ from the resume " f"checkpoint: {mismatches}" ) order = np.asarray(state["order"], dtype=np.int64) if order.shape != (self.size,): raise ValueError( f"resume stream order has shape {order.shape}, expected {(self.size,)}" ) cursor = int(state["cursor"]) if not 0 <= cursor <= self.size: raise ValueError(f"invalid resume stream cursor {cursor}") self._order = order.copy() self._cursor = cursor self.samples_seen = int(state.get("samples_seen", cursor)) self._rng.bit_generator.state = state["rng_state"] class PrefetchedWindowBatches: """Read OneScience ERA5 windows concurrently ahead of the train step. ERA5Dataset opens an independent HDF5 handle in every ``__getitem__`` call, so concurrent reads do not share h5py state. JAX/xarray conversion stays on the consumer thread; only file IO and OneScience sample construction happen in workers. """ def __init__( self, dataset, index_stream: WindowBatchStream, *, num_workers: int = 2, prefetch_batches: int = 1, ): self.dataset = dataset self.index_stream = index_stream self.num_workers = int(num_workers) self.prefetch_batches = int(prefetch_batches) if self.num_workers <= 0: raise ValueError("data_num_workers must be positive") if self.prefetch_batches <= 0: raise ValueError("prefetch_batches must be positive") self._executor = ThreadPoolExecutor( max_workers=self.num_workers, thread_name_prefix="era5-reader", ) self._queue: deque[tuple[dict, np.ndarray, list[Future]]] = deque() self._samples_consumed = 0 for _ in range(self.prefetch_batches): self._submit_batch() @property def samples_seen(self) -> int: return self._samples_consumed def _submit_batch(self) -> None: # Capture the logical reader position before prefetch advances it. A # training checkpoint can then resume at the first unconsumed batch, # including batches already queued by worker threads. stream_state = self.index_stream.state_dict() indices = self.index_stream.next_indices() futures = [ self._executor.submit(self.dataset.__getitem__, int(index)) for index in indices ] self._queue.append((stream_state, indices, futures)) def next_batch(self): """Return ``(indices, samples)`` and immediately enqueue another batch.""" _, indices, futures = self._queue.popleft() self._submit_batch() samples = [future.result() for future in futures] self._samples_consumed += len(indices) return indices, samples def resume_state(self) -> dict: """Return the stream state for the next unconsumed prefetched batch.""" if not self._queue: return self.index_stream.state_dict() state, _, _ = self._queue[0] return state def close(self) -> None: self._executor.shutdown(wait=True, cancel_futures=True) def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): self.close()