Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / online / iterable.py: 19%
235 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""Build streaming Online sources and their local/global blend policy."""
17from __future__ import annotations
19import random
20from collections.abc import Callable, Mapping, Sequence
21from dataclasses import replace
22from typing import Any
24from torch.utils.data import IterableDataset # pylint: disable=forbidden-backend-import
26from hyper_parallel.data.constants import ONLINE_BLEND_STRATEGIES
27from hyper_parallel.data.online.source import OnlineDataPath, _OnlineSourceBuilder
28from hyper_parallel.data.parallel import (
29 DataLoaderParallelContext,
30 build_dataset_for_dataloader,
31 split_iterable_dataset_by_dp,
32)
34SampleTransform = Callable[[Any], list[Any]]
37def build_online_iterable_source(
38 *,
39 data_config: Mapping[str, Any],
40 data_path: OnlineDataPath | None = None,
41 dataloader_context: DataLoaderParallelContext | None = None,
42 sample_filter: Callable[[Mapping[str, Any]], bool] | None = None,
43) -> Any:
44 """Build a locally shuffled and optionally DP-sharded Online stream.
46 Args:
47 data_config: Online streaming source options.
48 data_path: Local path or Hub dataset ID. Use ``data_config.sources``
49 instead for multiple sources.
50 dataloader_context: Optional DataLoader ownership and DP policy.
51 sample_filter: Optional raw-sample predicate applied before shuffling.
53 Returns:
54 One Iterable Dataset on each DataLoader-owning rank.
55 """
56 normalized_context = dataloader_context or DataLoaderParallelContext()
57 if normalized_context.data_index_cache:
58 normalized_context = replace(normalized_context, data_index_cache=False)
60 dataset_builder = _OnlineIterableSourceBuilder(data_path, data_config, normalized_context)
62 def dataset_factory() -> Any:
63 """Build the configured single-source or multi-source stream."""
64 iterable_dataset = dataset_builder.build_source(sample_filter=sample_filter)
65 return iterable_dataset
67 iterable_source = build_dataset_for_dataloader(
68 dataset_factory,
69 normalized_context,
70 barrier_needed=False,
71 )
72 return iterable_source
75def _supports_output_index_for_resume(dataset: Any) -> bool:
76 """Return whether a Dataset can emit and rebuild stable output indices."""
77 get_item = getattr(dataset, "get_item", None)
78 supports_output_index = callable(get_item) and hasattr(dataset, "output_index_for_resume")
79 return supports_output_index
82class IterableTransformDataset(IterableDataset):
83 """Apply a sample transform lazily while preserving streaming source state."""
85 @classmethod
86 def apply(cls, source_dataset: Any, transform: SampleTransform) -> Any:
87 """Transform one Iterable Dataset built on a DataLoader-owning rank.
89 Args:
90 source_dataset: Iterable Dataset, or ``None`` on a non-owning rank.
91 transform: Transform from one RawSample to zero, one, or many ModelSamples.
93 Returns:
94 A transformed Iterable Dataset, or ``None`` on a non-owning rank.
95 """
96 if source_dataset is None:
97 return None
99 iterable_dataset = cls(source_dataset, transform)
100 return iterable_dataset
102 def __init__(self, source_dataset: Any, transform: SampleTransform) -> None:
103 """Store the streaming source and its lazy sample transform."""
104 self.source_dataset = source_dataset
105 self.transform = transform
107 @property
108 def output_index_for_resume(self) -> bool:
109 """Return whether the upstream source emits stable output indices."""
110 if not _supports_output_index_for_resume(self.source_dataset):
111 raise AttributeError("The upstream iterable does not support output-index resume")
113 output_index_enabled = bool(self.source_dataset.output_index_for_resume)
114 return output_index_enabled
116 @output_index_for_resume.setter
117 def output_index_for_resume(self, value: bool) -> None:
118 """Forward output-index mode to a replayable upstream source."""
119 if not _supports_output_index_for_resume(self.source_dataset):
120 raise ValueError("The upstream iterable does not support output-index resume")
122 self.source_dataset.output_index_for_resume = value
124 def get_item(self, output_index: Any) -> list[Any]:
125 """Rebuild and transform one upstream item from its stable output index."""
126 if not _supports_output_index_for_resume(self.source_dataset):
127 raise AttributeError("The upstream iterable does not support get_item")
129 raw_sample = self.source_dataset.get_item(output_index)
130 transformed_samples = self._transform_sample(raw_sample)
131 return transformed_samples
133 def __iter__(self) -> Any:
134 """Transform source records lazily while preserving source order."""
135 output_index_enabled = (
136 _supports_output_index_for_resume(self.source_dataset)
137 and bool(self.source_dataset.output_index_for_resume)
138 )
139 for source_item in self.source_dataset:
140 if output_index_enabled:
141 raw_sample, output_index = source_item
142 else:
143 raw_sample = source_item
145 transformed_samples = self._transform_sample(raw_sample)
146 if not transformed_samples:
147 continue
148 if output_index_enabled:
149 yield transformed_samples, output_index
150 elif len(transformed_samples) == 1:
151 yield transformed_samples[0]
152 else:
153 yield transformed_samples
155 def state_dict(self) -> dict[str, Any]:
156 """Forward checkpoint state to a stateful upstream stream."""
157 state_builder = getattr(self.source_dataset, "state_dict", None)
158 source_state = {}
159 if callable(state_builder):
160 source_state = state_builder()
161 transform_state = {"source_dataset": source_state}
162 return transform_state
164 def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
165 """Restore checkpoint state through the upstream stream interface."""
166 state_loader = getattr(self.source_dataset, "load_state_dict", None)
167 if not callable(state_loader):
168 raise ValueError("Online iterable source does not support load_state_dict")
170 state_loader(state_dict["source_dataset"])
172 def set_epoch(self, epoch: int) -> None:
173 """Forward deterministic epoch state when supported upstream."""
174 epoch_setter = getattr(self.source_dataset, "set_epoch", None)
175 if callable(epoch_setter):
176 epoch_setter(epoch)
178 def _transform_sample(self, raw_sample: Any) -> list[Any]:
179 """Transform one RawSample and enforce the zero/one/many contract."""
180 transformed_samples = self.transform(raw_sample)
181 if not isinstance(transformed_samples, list):
182 raise TypeError("Online Iterable sample transform must return a list")
184 return transformed_samples
187class _LocalWeightedIterableDataset(IterableDataset):
188 """Choose among already-sharded Iterable sources on the local rank."""
190 def __init__(
191 self,
192 source_datasets: Sequence[Any],
193 weights: Sequence[float],
194 *,
195 random_seed: int,
196 stopping_strategy: str,
197 vary_by_epoch: bool,
198 ) -> None:
199 """Store the rank-local sources and weighted selection policy."""
200 self.source_datasets = tuple(source_datasets)
201 self.weights = tuple(weights)
202 self.random_seed = random_seed
203 self.stopping_strategy = stopping_strategy
204 self.vary_by_epoch = vary_by_epoch
205 self.epoch = 0
206 self._random_state = None
207 self._active_indices = None
208 self._exhausted_once = None
210 def __iter__(self) -> Any:
211 """Yield samples according to the configured local source weights."""
212 source_epoch = getattr(self.source_datasets[0], "epoch", self.epoch)
213 epoch = 0
214 if self.vary_by_epoch:
215 epoch = int(source_epoch)
217 random_generator = random.Random(self.random_seed + epoch)
218 if self._random_state is not None:
219 random_generator.setstate(self._random_state)
221 source_iterators = []
222 for source_dataset in self.source_datasets:
223 source_iterators.append(iter(source_dataset))
225 if self._active_indices is None:
226 active_indices = list(range(len(source_iterators)))
227 else:
228 active_indices = list(self._active_indices)
230 if self._exhausted_once is None:
231 exhausted_once = set()
232 else:
233 exhausted_once = set(self._exhausted_once)
235 while active_indices:
236 active_weights = []
237 for active_index in active_indices:
238 active_weights.append(self.weights[active_index])
240 source_index = random_generator.choices(active_indices, weights=active_weights, k=1)[0]
241 try:
242 source_sample = next(source_iterators[source_index])
243 except StopIteration:
244 if self.stopping_strategy == "first_exhausted":
245 self._clear_iteration_state()
246 return
248 exhausted_once.add(source_index)
249 if self.stopping_strategy == "all_exhausted_without_replacement":
250 active_indices.remove(source_index)
251 self._save_iteration_state(random_generator, active_indices, exhausted_once)
252 continue
254 if len(exhausted_once) == len(source_iterators):
255 self._clear_iteration_state()
256 return
258 source_iterators[source_index] = iter(self.source_datasets[source_index])
259 self._save_iteration_state(random_generator, active_indices, exhausted_once)
260 continue
262 self._save_iteration_state(random_generator, active_indices, exhausted_once)
263 yield source_sample
265 self._clear_iteration_state()
267 def set_epoch(self, epoch: int) -> None:
268 """Set the local blend epoch and forward it to every source."""
269 self.epoch = epoch
270 self._clear_iteration_state()
271 for source_dataset in self.source_datasets:
272 epoch_setter = getattr(source_dataset, "set_epoch", None)
273 if callable(epoch_setter):
274 epoch_setter(epoch)
276 def state_dict(self) -> dict[str, Any]:
277 """Save source cursors and the local weighted-selection cursor."""
278 source_states = []
279 for source_dataset in self.source_datasets:
280 state_builder = getattr(source_dataset, "state_dict", None)
281 source_state = None
282 if callable(state_builder):
283 source_state = state_builder()
284 source_states.append(source_state)
286 source_epoch = getattr(self.source_datasets[0], "epoch", self.epoch)
287 dataset_state = {
288 "epoch": int(source_epoch),
289 "random_state": self._random_state,
290 "active_indices": self._active_indices,
291 "exhausted_once": self._exhausted_once,
292 "sources": source_states,
293 }
294 return dataset_state
296 def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
297 """Restore source cursors and the local weighted-selection cursor."""
298 source_states = state_dict["sources"]
299 if len(source_states) != len(self.source_datasets):
300 raise ValueError("Online source count does not match the saved local blend state")
302 for source_dataset, source_state in zip(self.source_datasets, source_states):
303 if source_state is None:
304 continue
305 state_loader = getattr(source_dataset, "load_state_dict", None)
306 if not callable(state_loader):
307 raise ValueError("Online iterable source does not support load_state_dict")
308 state_loader(source_state)
310 self.epoch = int(state_dict["epoch"])
311 self._random_state = state_dict["random_state"]
312 self._active_indices = state_dict["active_indices"]
313 self._exhausted_once = state_dict["exhausted_once"]
315 def _save_iteration_state(
316 self,
317 random_generator: random.Random,
318 active_indices: Sequence[int],
319 exhausted_once: set[int],
320 ) -> None:
321 """Record the local source-selection cursor after one decision."""
322 self._random_state = random_generator.getstate()
323 self._active_indices = tuple(active_indices)
324 self._exhausted_once = tuple(sorted(exhausted_once))
326 def _clear_iteration_state(self) -> None:
327 """Reset the local source-selection cursor for a new epoch."""
328 self._random_state = None
329 self._active_indices = None
330 self._exhausted_once = None
333class _OnlineIterableSourceBuilder(_OnlineSourceBuilder):
334 """Build the rank-local or globally interleaved streaming source graph."""
336 def __init__(
337 self,
338 data_path: OnlineDataPath | None,
339 data_config: Mapping[str, Any],
340 dataloader_context: DataLoaderParallelContext,
341 ) -> None:
342 """Store source configuration and DataLoader ownership."""
343 super().__init__(data_path, data_config)
344 self.dataloader_context = dataloader_context
346 def build_source(
347 self,
348 *,
349 sample_filter: Callable[[Mapping[str, Any]], bool] | None = None,
350 ) -> Any:
351 """Load Iterable sources and apply their shard, shuffle, and blend policy."""
352 source_loaders, weights = self._create_source_loaders()
353 source_datasets = []
354 for source_loader in source_loaders:
355 source_dataset = source_loader.load(streaming=True, sample_filter=sample_filter)
356 if isinstance(source_dataset, tuple):
357 raise ValueError("Online Iterable Dataset requires one train source per entry")
358 source_datasets.append(source_dataset)
360 if len(source_datasets) == 1:
361 iterable_dataset = self._prepare_local_source(source_datasets[0])
362 return iterable_dataset
364 blend_strategy = str(self.data_config.get("blend_strategy", "local_weighted"))
365 if blend_strategy not in ONLINE_BLEND_STRATEGIES:
366 raise ValueError(
367 f"Unsupported Online blend_strategy {blend_strategy!r}; "
368 f"expected one of {ONLINE_BLEND_STRATEGIES!r}"
369 )
371 if blend_strategy == "local_weighted":
372 local_sources = []
373 for source_dataset in source_datasets:
374 local_source = self._prepare_local_source(source_dataset)
375 local_sources.append(local_source)
377 iterable_dataset = _LocalWeightedIterableDataset(
378 local_sources,
379 weights,
380 random_seed=self._random_seed,
381 stopping_strategy=self._stopping_strategy,
382 vary_by_epoch=self._shuffle_enabled,
383 )
384 return iterable_dataset
386 shuffled_sources = []
387 for source_dataset in source_datasets:
388 shuffled_source = self._shuffle(source_dataset)
389 shuffled_sources.append(shuffled_source)
391 iterable_dataset = self._interleave_sources(shuffled_sources, weights)
392 if self._shuffle_enabled:
393 # buffer_size=1 preserves output order while making set_epoch reseed the full source graph.
394 iterable_dataset = iterable_dataset.shuffle(seed=self._random_seed, buffer_size=1)
395 if self._split_by_data_parallel:
396 iterable_dataset = split_iterable_dataset_by_dp(iterable_dataset, self.dataloader_context)
397 return iterable_dataset
399 @property
400 def _shuffle_enabled(self) -> bool:
401 """Return whether Iterable sources use local buffer shuffle."""
402 return bool(self.data_config.get("shuffle", True))
404 @property
405 def _split_by_data_parallel(self) -> bool:
406 """Return whether Iterable sources own DP sharding."""
407 return bool(self.data_config.get("split_by_data_parallel", True))
409 def _prepare_local_source(self, source_dataset: Any) -> Any:
410 """Apply per-source DP sharding before local buffer shuffle."""
411 if self._split_by_data_parallel:
412 source_dataset = split_iterable_dataset_by_dp(source_dataset, self.dataloader_context)
413 local_dataset = self._shuffle(source_dataset)
414 return local_dataset
416 def _shuffle(self, source_dataset: Any) -> Any:
417 """Apply deterministic buffer shuffle when configured."""
418 if not self._shuffle_enabled:
419 return source_dataset
421 buffer_size = int(self.data_config.get("shuffle_buffer_size", 10_000))
422 if buffer_size <= 0:
423 raise ValueError("Online shuffle_buffer_size must be positive")
425 shuffled_dataset = source_dataset.shuffle(seed=self._random_seed, buffer_size=buffer_size)
426 return shuffled_dataset
429__all__ = [
430 "IterableTransformDataset",
431 "build_online_iterable_source",
432]