File size: 7,838 Bytes
f4a39ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 | """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()
|