scripts/streaming_data.py

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