Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / batch_sampler.py: 95%
181 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Transactional Dataset reads driven by HP's native DP BatchSampler."""
17from __future__ import annotations
19import copy
20import operator
21from collections.abc import Callable, Mapping, Sequence
22from dataclasses import dataclass, replace
23from typing import Any
25from hyper_parallel.data.parallel.batch_sampler import sampler_optimizer_step_progress
26from hyper_parallel.distributed_data.schema import SampleKey, SampleMetadata
27from hyper_parallel.distributed_data.metadata import PlannedSampleLoader
30def validate_native_sampler(sampler: Any, *, data_rank: int, dp_size: int, local_batch_size: int) -> None:
31 """Validate the native HP sampler interface, DP layout, and cursor alignment.
33 Args:
34 sampler: Native HP batch sampler to validate.
35 data_rank: Expected DP coordinate for this rank.
36 dp_size: Expected number of DP replicas.
37 local_batch_size: Expected samples per local forward/backward round.
39 Raises:
40 ValueError: If the sampler interface, DP settings, or cursor is incompatible.
41 """
42 for name in ("__iter__", "state_dict", "load_state_dict", "set_epoch", "enable_source_batch_resume"):
43 if not callable(getattr(sampler, name, None)):
44 raise ValueError(f"batch_sampler must implement HP's native sampler contract: missing {name}.")
45 expected = {"dp_rank": data_rank, "dp_world_size": dp_size, "micro_batch_size": local_batch_size,
46 "drop_last": True}
47 for name, value in expected.items():
48 if getattr(sampler, name, None) != value:
49 raise ValueError(f"batch_sampler.{name} must equal {value!r}.")
50 if sampler_optimizer_step_progress(sampler.state_dict()):
51 raise ValueError("batch_sampler.consumed_samples must align to a complete optimizer step.")
54def preserve_sample(samples: Sequence[Any], seq_len: int) -> Any:
55 """Pass one native Dataset output unchanged to the original collator.
57 Args:
58 samples: Source samples for one native sequence bin.
59 seq_len: Sequence capacity, unused when preserving a Dataset output.
61 Returns:
62 The sole source sample, unchanged.
63 """
64 del seq_len
65 if len(samples) != 1:
66 raise ValueError("Native BatchSampler mode must keep one Dataset output per sequence bin.")
67 return samples[0]
70@dataclass(frozen=True)
71class _ReaderPlanInput:
72 """Transfer metadata and source keys, grouped by micro-batch then original pack."""
74 metadata: tuple[SampleMetadata, ...] = ()
75 original_batches: tuple[tuple[tuple[SampleKey, ...], ...], ...] = ()
78@dataclass
79class _ReaderWindow:
80 """Keep one window's data and cursor independent of speculative sampling."""
82 metadata: tuple[SampleMetadata, ...]
83 samples: dict[SampleKey, Any]
84 sampler_state: dict[str, Any]
85 original_batches: tuple[tuple[tuple[SampleKey, ...], ...], ...]
86 delivered_rounds: int = 0
89class BatchSamplerReader:
90 """Read a bounded window of native batches within one optimizer step.
92 Only the DP Constructor owns this reader. Payload loading is deferred until
93 after planning when metadata is available; otherwise it precedes metadata
94 extraction. The sampler's speculative cursor is never a checkpoint cursor.
95 """
97 def __init__(
98 self,
99 sampler: Any,
100 *,
101 reader_rank: int,
102 metadata: Sequence[SampleMetadata] | None,
103 metadata_fn: Callable[[Any], SampleMetadata] | None,
104 sample_loader: PlannedSampleLoader | None,
105 split_fn: Callable[[Any], Sequence[Any]] | None = None,
106 ) -> None:
107 """Store the validated native sampler and metadata/payload sources."""
108 if metadata is None and (sample_loader is None or not callable(metadata_fn)):
109 raise ValueError("Online BatchSampler Reader requires payload loading and metadata_fn.")
110 if split_fn is not None and (not callable(split_fn) or metadata is not None):
111 raise ValueError("split_fn requires online metadata and must be callable.")
112 self._sampler = sampler
113 self._reader_rank = reader_rank
114 self._metadata = metadata
115 self._metadata_fn = metadata_fn
116 self._sample_loader = sample_loader
117 self._split_fn = split_fn
118 self._window: _ReaderWindow | None = None
119 self._sampler.enable_source_batch_resume()
120 self._committed_state = copy.deepcopy(sampler.state_dict())
121 self._iterator = None
122 self._exhausted = False
124 @property
125 def exhausted(self) -> bool:
126 """Return whether the native sampler has exhausted its epoch."""
127 return self._exhausted
129 @property
130 def buffer_size(self) -> int:
131 """Return the number of pending native sample occurrences."""
132 return len(self.metadata())
134 @property
135 def batch_position(self) -> int:
136 """Return the native global cursor before this pending round."""
137 return self._committed_state["consumed_samples"]
139 def prepare_next_step(self, window_micro_batch_num: int) -> None:
140 """Read the window requested by the owning loader without source lookahead.
142 Args:
143 window_micro_batch_num: Positive number of source rounds to read,
144 already clipped to the optimizer boundary by the loader.
145 """
146 if isinstance(window_micro_batch_num, bool) or not isinstance(window_micro_batch_num, int):
147 raise ValueError("window_micro_batch_num must be a positive integer.")
148 if window_micro_batch_num < 1:
149 raise ValueError("window_micro_batch_num must be a positive integer.")
150 if self._window is not None or self._exhausted:
151 return
152 self._window = self._read_window(window_micro_batch_num)
153 self._exhausted = self._window is None
155 def _read_window(self, window_micro_batch_num: int) -> _ReaderWindow | None:
156 """Read a bounded native window without changing the active routing window."""
157 # Only one read runs at a time. It must not mutate the active window
158 # while that window's samples are being routed and committed.
159 try:
160 if self._iterator is None:
161 self._iterator = iter(self._sampler)
162 source_batches = self._read_window_keys(window_micro_batch_num)
163 if not source_batches:
164 return None
165 return self._load_window(source_batches)
166 except StopIteration as exc:
167 # A user callback's StopIteration must not masquerade as dataset EOF.
168 raise RuntimeError(f"Native BatchSampler callback raised StopIteration: {exc}") from exc
170 def _load_window(self, source_batches: tuple[tuple[SampleKey, ...], ...]) -> _ReaderWindow:
171 """Materialize the selected occurrences and freeze their original pack boundaries."""
172 keys = tuple(key for batch in source_batches for key in batch)
173 original_batches = ()
174 samples = {}
175 if self._metadata is None:
176 samples = self._sample_loader.fetch_samples(keys)
177 if self._split_fn is not None:
178 samples, original_batches = self._split_payloads(samples, source_batches)
179 keys = tuple(samples)
180 metadata = [self._metadata_fn(samples[key]) for key in keys]
181 else:
182 metadata = [self._metadata[key.dataset_index] for key in keys]
183 if self._split_fn is None:
184 original_batches = tuple(tuple((key,) for key in batch) for batch in source_batches)
185 if any(not isinstance(item, SampleMetadata) for item in metadata):
186 raise ValueError("BatchSampler metadata must contain SampleMetadata entries.")
187 return _ReaderWindow(
188 metadata=tuple(replace(item, key=key) for key, item in zip(keys, metadata)),
189 samples=samples,
190 sampler_state=copy.deepcopy(self._sampler.state_dict()),
191 original_batches=original_batches,
192 )
194 def _read_window_keys(self, window_micro_batch_num: int) -> tuple[tuple[SampleKey, ...], ...]:
195 source_batches = []
196 for _ in range(window_micro_batch_num):
197 position = self._sampler.consumed_samples
198 try:
199 indices = next(self._iterator)
200 except StopIteration:
201 break
202 source_batches.append(self._sample_keys(indices, position))
203 return tuple(source_batches)
205 def _split_payloads(
206 self, samples: Mapping[SampleKey, Any], source_batches: tuple[tuple[SampleKey, ...], ...],
207 ) -> tuple[dict[SampleKey, Any], tuple[tuple[tuple[SampleKey, ...], ...], ...]]:
208 """Split records into keyed sequences while retaining their original per-round bins."""
209 payloads = {}
210 original_batches = []
211 for batch in source_batches:
212 original_bins = []
213 for key in batch:
214 sequences = tuple(self._split_fn(samples[key]))
215 if not sequences:
216 raise ValueError("split_fn must return at least one independent sequence per Dataset output.")
217 sequence_keys = tuple(replace(key, sequence_index=index) for index in range(len(sequences)))
218 payloads.update(zip(sequence_keys, sequences))
219 original_bins.append(sequence_keys)
220 original_batches.append(tuple(original_bins))
221 return payloads, tuple(original_batches)
223 def close(self) -> None:
224 """Release source buffers after the owning loader has drained its producer."""
225 self._window = None
226 self._iterator = None
228 def _sample_keys(self, indices: Sequence[int], batch_position: int) -> tuple[SampleKey, ...]:
229 """Bind local Dataset indices to distinct global sampling occurrences."""
230 if len(indices) != self._sampler.micro_batch_size:
231 raise ValueError("Native BatchSampler must yield exactly local_batch_size indices on every DP rank.")
232 # A native sampler exposes local indices, not its internal permutation
233 # positions. Stable DP-major occurrence slots preserve duplicates.
234 start = batch_position + self._sampler.dp_rank * self._sampler.micro_batch_size
235 keys = []
236 for offset, index in enumerate(indices):
237 if isinstance(index, bool):
238 raise ValueError("Native BatchSampler indices must be non-negative integers.")
239 keys.append(SampleKey(self._reader_rank, operator.index(index), start + offset))
240 return tuple(keys)
242 def metadata(self) -> tuple[SampleMetadata, ...]:
243 """Return metadata for the undelivered occurrences in the frozen window."""
244 return () if self._window is None else self._window.metadata
246 def planning_input(self) -> _ReaderPlanInput:
247 """Return the window's metadata and original layout without payloads."""
248 if self._window is None:
249 return _ReaderPlanInput()
250 return _ReaderPlanInput(self._window.metadata, self._window.original_batches)
252 def selected_payloads(self, selected_keys: set[SampleKey]) -> tuple[tuple[SampleKey, Any], ...]:
253 """Return already-read payloads without rereading their Dataset indices.
255 Args:
256 selected_keys: Sample occurrences selected from this Reader's buffer.
258 Returns:
259 Selected key-payload pairs in buffered sample order.
260 """
261 return tuple(
262 (item.key, self._window.samples[item.key]) for item in self.metadata() if item.key in selected_keys
263 )
265 def commit(self, selected_keys: set[SampleKey]) -> None:
266 """Retire delivered occurrences and commit the cursor at the window end.
268 Args:
269 selected_keys: Source-owned occurrences delivered in this distributed round.
270 """
271 window = self._window
272 remaining = {item.key for item in self.metadata()}
273 if window is None or not selected_keys <= remaining:
274 raise ValueError("Native BatchSampler commit must consume the complete pending local batch.")
275 last_round = window.delivered_rounds + 1 == len(window.original_batches)
276 if last_round and selected_keys != remaining:
277 raise ValueError("Native BatchSampler commit must consume the complete pending local batch.")
278 window.metadata = tuple(item for item in window.metadata if item.key not in selected_keys)
279 for key in selected_keys:
280 window.samples.pop(key, None)
281 window.delivered_rounds += 1
282 if last_round:
283 self._committed_state = window.sampler_state
284 self._window = None
286 def state_dict(self) -> dict[str, Any]:
287 """Snapshot the delivered native cursor, excluding speculative prefetch."""
288 state = {
289 "epoch": self._committed_state["epoch"],
290 "sampler": self._committed_state,
291 "sample_loader": self._sample_loader.state_dict() if self._sample_loader is not None else None,
292 }
293 return copy.deepcopy(state)
295 def load_state_dict(self, state: Mapping[str, Any]) -> None:
296 """Restore the original native cursor before the next sampler yield.
298 Args:
299 state: Reader checkpoint containing the committed native sampler state.
301 Raises:
302 ValueError: If the checkpoint epoch or payload reader is incompatible.
303 """
304 sampler_state = state["sampler"]
305 if sampler_state.get("epoch") != state.get("epoch"):
306 raise ValueError("Native BatchSampler checkpoint epoch is inconsistent.")
307 if (state.get("sample_loader") is None) != (self._sample_loader is None):
308 raise ValueError("Native BatchSampler checkpoint payload reader ownership changed.")
309 if "window" in state or sampler_optimizer_step_progress(sampler_state):
310 raise ValueError("Native BatchSampler checkpoint must align to a complete optimizer step.")
311 self.close()
312 self._sampler.load_state_dict(sampler_state)
313 if self._sample_loader is not None:
314 self._sample_loader.load_state_dict(state["sample_loader"])
315 self._committed_state = copy.deepcopy(self._sampler.state_dict())
316 self._exhausted = False
318 def set_epoch(self, epoch: int) -> None:
319 """Delegate epoch changes to the native sampler without changing its policy.
321 Args:
322 epoch: Epoch to apply to the sampler and sample loader.
323 """
324 self.close()
325 self._sampler.load_state_dict(self._committed_state)
326 self._sampler.set_epoch(epoch)
327 if self._sample_loader is not None:
328 self._sample_loader.set_epoch(epoch)
329 self._committed_state = copy.deepcopy(self._sampler.state_dict())
330 self._exhausted = False