Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / online / mapping.py: 80%
350 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"""Build finite integer-indexed Online sources."""
17from __future__ import annotations
19import glob
20import json
21import os
22from collections.abc import Callable, Mapping, Sequence
23from contextlib import ExitStack
24from typing import Any
26import numpy as np
28from hyper_parallel.data.constants import ONLINE_SOURCE_PATH_KEY, ONLINE_SPLIT_COUNT, ONLINE_SPLIT_NAMES
29from hyper_parallel.data.dataset_logging import get_dataset_logger
30from hyper_parallel.data.online.source import (
31 OnlineDataPath,
32 OnlinePathSpec,
33 _OnlineSourceBuilder,
34)
36logger = get_dataset_logger(__name__)
38SampleTransform = Callable[[Any], list[Any]]
41def build_online_mapping_source(
42 *,
43 data_config: Mapping[str, Any],
44 data_path: OnlineDataPath | None = None,
45 sample_filter: Callable[[Mapping[str, Any]], bool] | None = None,
46 train_valid_test_num_samples: Sequence[int] | None = None,
47 sampler_type: str = "single",
48) -> Any:
49 """Build a finite integer-indexed Online source.
51 Args:
52 data_config: Online source options.
53 data_path: Local path, Hub dataset ID, or pre-split path mapping. Use
54 ``data_config.sources`` instead for multiple sources.
55 sampler_type: Cyclic leaves shuffle and repetition across epochs to the sampler.
56 sample_filter: Optional raw-sample predicate applied before indexing.
57 train_valid_test_num_samples: Optional target sample count for each
58 train, validation, and test split.
60 Returns:
61 One Mapping Dataset or a train-valid-test tuple.
62 """
63 dataset_builder = _OnlineMappingSourceBuilder(data_path, data_config, train_valid_test_num_samples, sampler_type)
64 mapping_source = dataset_builder.build_source(sample_filter=sample_filter)
65 return mapping_source
68class MappingTransformDataset:
69 """Apply a sample transform lazily to an integer-indexed source Dataset."""
71 def __init__(self, source_dataset: Any, transform: SampleTransform) -> None:
72 """Store the source Dataset and its lazy sample transform."""
73 self.source_dataset = source_dataset
74 self.transform = transform
76 @classmethod
77 def apply(cls, source_dataset: Any, transform: SampleTransform) -> Any:
78 """Transform one Mapping Dataset or each available split.
80 Args:
81 source_dataset: One Mapping Dataset or a train-valid-test tuple.
82 transform: Transform from one RawSample to zero, one, or many ModelSamples.
84 Returns:
85 A transformed Mapping Dataset or transformed split tuple.
86 """
87 if isinstance(source_dataset, tuple) and len(source_dataset) == ONLINE_SPLIT_COUNT:
88 transformed_splits = []
89 for split_dataset in source_dataset:
90 transformed_dataset = None
91 if split_dataset is not None:
92 transformed_dataset = cls(split_dataset, transform)
93 transformed_splits.append(transformed_dataset)
95 mapping_splits = (transformed_splits[0], transformed_splits[1], transformed_splits[2])
96 return mapping_splits
98 if source_dataset is None:
99 return None
101 mapping_dataset = cls(source_dataset, transform)
102 return mapping_dataset
104 def __len__(self) -> int:
105 """Return the source Dataset length."""
106 dataset_length = len(self.source_dataset)
107 return dataset_length
109 def get_item(self, index: int) -> list[Any]:
110 """Transform one source index into zero, one, or many output samples."""
111 dataset_length = len(self.source_dataset)
112 if index < 0:
113 index += dataset_length
114 if index < 0 or index >= dataset_length:
115 raise IndexError("Online Mapping Dataset index out of range")
117 raw_sample = self.source_dataset[index]
118 transformed_samples = self.transform(raw_sample)
119 logger.debug(
120 "Online sample source_index=%s sample_id=%s", index, raw_sample.get("id"),
121 enabled=True,
122 )
123 if not isinstance(transformed_samples, list):
124 raise TypeError("Online Mapping sample transform must return a list")
126 return transformed_samples
128 def __getitem__(self, index: int) -> Any:
129 """Return exactly one sample for fixed-size DataLoader consumers."""
130 transformed_samples = self.get_item(index)
131 if not transformed_samples:
132 raise ValueError(
133 "An empty transform result requires source filtering or TokenBatchLoader"
134 )
135 if len(transformed_samples) != 1:
136 raise ValueError("A multi-sample transform result requires TokenBatchLoader")
138 transformed_sample = transformed_samples[0]
139 return transformed_sample
142class _JsonlMappingSource:
143 """Read heterogeneous JSONL records without imposing an Arrow schema."""
145 def __init__(
146 self,
147 data_files: Sequence[str],
148 sample_filter: Callable[[Mapping[str, Any]], bool] | None,
149 ) -> None:
150 """Index non-empty JSONL records by file path and byte offset."""
151 self._record_locations: list[tuple[str, int]] = []
152 self._record_start = 0
153 self._record_end = 0
154 self._file_handles: dict[str, Any] = {}
155 self._file_stack = ExitStack()
156 self._reader_process_id = os.getpid()
158 for data_file in data_files:
159 self._index_file(data_file, sample_filter)
160 self._record_end = len(self._record_locations)
162 def __len__(self) -> int:
163 """Return the number of indexed JSONL records."""
164 record_count = self._record_end - self._record_start
165 return record_count
167 def __getitem__(self, index: int) -> dict[str, Any]:
168 """Deserialize one record and attach its source file context."""
169 record_count = len(self)
170 if index < 0:
171 index += record_count
172 if index < 0 or index >= record_count:
173 raise IndexError("Online JSONL source index out of range")
175 record_index = self._record_start + index
176 data_file, byte_offset = self._record_locations[record_index]
177 file_handle = self._get_file_handle(data_file)
178 file_handle.seek(byte_offset)
179 raw_record = file_handle.readline()
180 source_record = self._parse_record(raw_record, data_file)
181 source_record[ONLINE_SOURCE_PATH_KEY] = data_file
182 return source_record
184 def split_by_percentages(
185 self,
186 percentages: Sequence[int],
187 ) -> tuple[_JsonlMappingSource | None, _JsonlMappingSource | None, _JsonlMappingSource | None]:
188 """Create train, validation, and test views over the filtered index."""
189 record_count = len(self)
190 train_end = int(round(percentages[0] * record_count / 100.0))
191 valid_end = int(round((percentages[0] + percentages[1]) * record_count / 100.0))
192 split_ranges = ((0, train_end), (train_end, valid_end), (valid_end, record_count))
194 split_sources = []
195 for percentage, (split_start, split_end) in zip(percentages, split_ranges):
196 split_source = None
197 if percentage > 0:
198 split_source = self._create_view(
199 self._record_locations,
200 self._record_start + split_start,
201 self._record_start + split_end,
202 )
203 split_sources.append(split_source)
204 jsonl_splits = (split_sources[0], split_sources[1], split_sources[2])
205 return jsonl_splits
207 def __getstate__(self) -> dict[str, Any]:
208 """Exclude process-local file handles when DataLoader workers pickle the source."""
209 source_state = dict(self.__dict__)
210 source_state["_file_handles"] = {}
211 source_state["_file_stack"] = ExitStack()
212 source_state["_reader_process_id"] = None
213 return source_state
215 def _index_file(
216 self,
217 data_file: str,
218 sample_filter: Callable[[Mapping[str, Any]], bool] | None,
219 ) -> None:
220 """Record byte offsets and evaluate an optional raw-sample filter."""
221 with open(data_file, "rb") as file_handle:
222 while True:
223 byte_offset = file_handle.tell()
224 raw_record = file_handle.readline()
225 if not raw_record:
226 break
227 if not raw_record.strip():
228 continue
229 if sample_filter is not None:
230 source_record = self._parse_record(raw_record, data_file)
231 if not sample_filter(source_record):
232 continue
233 self._record_locations.append((data_file, byte_offset))
235 @classmethod
236 def _create_view(
237 cls, record_locations: list[tuple[str, int]], start: int, end: int,
238 ) -> _JsonlMappingSource:
239 """Create a zero-copy view over indexed records with independent readers."""
240 split_source = cls.__new__(cls)
241 split_source._record_locations = record_locations
242 split_source._record_start = start
243 split_source._record_end = end
244 split_source._file_handles = {}
245 split_source._file_stack = ExitStack()
246 split_source._reader_process_id = os.getpid()
247 return split_source
249 def _get_file_handle(self, data_file: str) -> Any:
250 """Return a binary reader owned by the current DataLoader process."""
251 process_id = os.getpid()
252 if self._reader_process_id != process_id:
253 self._close_file_handles()
254 self._reader_process_id = process_id
256 file_handle = self._file_handles.get(data_file)
257 if file_handle is None:
258 file_handle = self._file_stack.enter_context(open(data_file, "rb"))
259 self._file_handles[data_file] = file_handle
260 return file_handle
262 def _close_file_handles(self) -> None:
263 """Close readers inherited from another process or opened locally."""
264 self._file_stack.close()
265 self._file_stack = ExitStack()
266 self._file_handles = {}
268 @staticmethod
269 def _parse_record(raw_record: bytes, data_file: str) -> dict[str, Any]:
270 """Decode one JSON object while preserving heterogeneous nested values."""
271 try:
272 source_record = json.loads(raw_record)
273 except (UnicodeDecodeError, json.JSONDecodeError) as error:
274 raise ValueError(f"Invalid JSONL record in {data_file}") from error
275 if not isinstance(source_record, dict):
276 raise ValueError(f"Online JSONL records must be objects: {data_file}")
277 return source_record
280class _OnlineBlendedMappingDataset:
281 """Expose Mapping sources through Megatron-style deterministic indices."""
283 def __init__(
284 self,
285 source_datasets: Sequence[Any],
286 weights: Sequence[float],
287 size: int,
288 random_seed: int,
289 shuffle: bool = True,
290 ) -> None:
291 """Build the source schedule and source-local shuffle indices."""
292 if not source_datasets or len(source_datasets) != len(weights):
293 raise ValueError("Mapping blend sources and weights must be non-empty and aligned")
294 if len(source_datasets) >= np.iinfo(np.int16).max:
295 raise ValueError("Mapping blend supports fewer than 32767 sources")
296 if size <= 0:
297 raise ValueError("Mapping blend size must be positive")
299 source_lengths = []
300 for source_dataset in source_datasets:
301 source_length = len(source_dataset)
302 if source_length <= 0:
303 raise ValueError("Mapping blend sources must not be empty")
304 source_lengths.append(source_length)
306 self.source_datasets = tuple(source_datasets)
307 self.weights = self._normalize_weights(weights)
308 self.size = size
309 self.source_index, self.source_sample_index = self._build_blend_indices()
310 self.source_shuffle_indices = self._build_source_shuffle_indices(
311 source_lengths,
312 random_seed,
313 shuffle,
314 )
316 def __len__(self) -> int:
317 """Return the requested global sample count."""
318 return self.size
320 def __getitem__(self, index: int) -> Any:
321 """Resolve one global index to its source and shuffled source sample."""
322 source_id = int(self.source_index[index])
323 source_sample_id = int(self.source_sample_index[index])
324 shuffled_sample_id = int(self.source_shuffle_indices[source_id][source_sample_id])
325 source_sample = self.source_datasets[source_id][shuffled_sample_id]
326 return source_sample
328 @staticmethod
329 def _normalize_weights(weights: Sequence[float]) -> tuple[float, ...]:
330 """Validate and normalize source weights."""
331 if not weights or any(weight <= 0.0 for weight in weights):
332 raise ValueError("Mapping blend weights must be positive")
334 weight_sum = sum(weights)
335 normalized_weights = []
336 for weight in weights:
337 normalized_weights.append(float(weight) / weight_sum)
338 return tuple(normalized_weights)
340 def _build_blend_indices(self) -> tuple[np.ndarray, np.ndarray]:
341 """Build Megatron's deterministic largest-deficit source schedule."""
342 source_index = np.empty(self.size, dtype=np.int16)
343 source_logical_index = np.empty(self.size, dtype=np.int64)
344 source_sample_counts = [0] * len(self.source_datasets)
346 for global_sample_index in range(self.size):
347 scheduled_samples = max(global_sample_index, 1)
348 selected_source = 0
349 largest_deficit = self.weights[0] * scheduled_samples - source_sample_counts[0]
350 for source_id in range(1, len(self.source_datasets)):
351 source_deficit = (
352 self.weights[source_id] * scheduled_samples
353 - source_sample_counts[source_id]
354 )
355 if source_deficit > largest_deficit:
356 selected_source = source_id
357 largest_deficit = source_deficit
359 source_index[global_sample_index] = selected_source
360 source_logical_index[global_sample_index] = source_sample_counts[selected_source]
361 source_sample_counts[selected_source] += 1
363 return source_index, source_logical_index
365 def _build_source_shuffle_indices(
366 self,
367 source_lengths: Sequence[int],
368 random_seed: int,
369 shuffle: bool,
370 ) -> tuple[np.ndarray, ...]:
371 """Map source-local counters to repeated deterministic sample permutations."""
372 source_sample_counts = np.bincount(
373 self.source_index,
374 minlength=len(self.source_datasets),
375 )
376 source_random = np.random.RandomState(random_seed)
377 source_shuffle_indices = []
378 for source_length, requested_samples in zip(source_lengths, source_sample_counts):
379 repeated_epochs = int(np.ceil(int(requested_samples) / source_length))
380 source_sample_index = np.tile(
381 np.arange(source_length, dtype=np.int64),
382 repeated_epochs,
383 )
384 if shuffle:
385 source_random.shuffle(source_sample_index)
386 source_shuffle_indices.append(source_sample_index[:requested_samples])
387 return tuple(source_shuffle_indices)
390class _OnlineMappingSourceBuilder(_OnlineSourceBuilder):
391 """Build split-local, DP-global Mapping sample indices."""
393 def __init__(
394 self,
395 data_path: OnlineDataPath | None,
396 data_config: Mapping[str, Any],
397 train_valid_test_num_samples: Sequence[int] | None,
398 sampler_type: str = "single",
399 ) -> None:
400 """Store source configuration and optional training sample targets."""
401 super().__init__(data_path, data_config)
402 if sampler_type not in ("single", "cyclic"):
403 raise ValueError(f"Unsupported Mapping sampler_type: {sampler_type!r}")
404 self.sampler_type = sampler_type
405 self.train_valid_test_num_samples = self._validate_split_sizes(train_valid_test_num_samples)
407 def build_source(
408 self,
409 *,
410 sample_filter: Callable[[Mapping[str, Any]], bool] | None = None,
411 ) -> Any:
412 """Load Mapping sources and compose their global split graphs."""
413 source_loaders, weights = self._create_source_loaders()
414 source_results = []
415 for source_loader in source_loaders:
416 source_result = self._build_leaf_source(source_loader, sample_filter)
417 source_results.append(source_result)
419 # Cyclic sampling owns epoch shuffle and repetition, including explicit single sources.
420 if self.sampler_type == "cyclic" and len(source_results) == 1:
421 source = source_results[0]
422 if isinstance(source, tuple):
423 return tuple(split if split is not None and len(split) else None for split in source)
424 return source if len(source) else None
426 # Example: source A returns Dataset, while source B returns (train, valid, test).
427 # The combined result follows B's split structure, and A contributes only to train.
428 has_split_result = False
429 for source_result in source_results:
430 if isinstance(source_result, tuple):
431 has_split_result = True
432 break
434 if not has_split_result:
435 target_size = self._get_target_size(0, source_results)
436 mapping_dataset = self._build_blended_split(source_results, weights, target_size)
437 return mapping_dataset
439 # source of train, valid, test => each source blend
440 blended_splits = []
441 for split_index in range(ONLINE_SPLIT_COUNT):
442 split_datasets = []
443 split_weights = []
444 for source_result, weight in zip(source_results, weights):
445 if isinstance(source_result, tuple):
446 split_dataset = source_result[split_index]
447 elif split_index == 0:
448 split_dataset = source_result
449 else:
450 split_dataset = None
452 if split_dataset is not None:
453 split_datasets.append(split_dataset)
454 split_weights.append(weight)
456 target_size = self._get_target_size(split_index, split_datasets)
457 blended_split = self._build_blended_split(
458 split_datasets,
459 split_weights,
460 target_size,
461 )
462 blended_splits.append(blended_split)
464 # train, valid, test
465 mapping_splits = (blended_splits[0], blended_splits[1], blended_splits[2])
466 return mapping_splits
468 @classmethod
469 def _build_leaf_source(
470 cls,
471 source_loader: Any,
472 sample_filter: Callable[[Mapping[str, Any]], bool] | None,
473 ) -> Any:
474 """Build one leaf through native JSONL indexing or the general loader."""
475 data_path = source_loader.data_path
476 split_config = source_loader.data_config.get("split")
477 if isinstance(data_path, Mapping) and split_config is None:
478 jsonl_split_sources = cls._try_build_presplit_jsonl_sources(
479 data_path,
480 sample_filter,
481 )
482 if jsonl_split_sources is not None:
483 return jsonl_split_sources
485 if not isinstance(data_path, Mapping):
486 jsonl_files = cls._resolve_native_jsonl_files(data_path)
487 logger.debug("Resolved native JSONL files: %s", jsonl_files)
488 if jsonl_files is not None:
489 jsonl_source = _JsonlMappingSource(jsonl_files, sample_filter)
490 if split_config is not None:
491 split_percentages = source_loader.parse_split_percentages(split_config)
492 jsonl_splits = jsonl_source.split_by_percentages(split_percentages)
493 return jsonl_splits
494 return jsonl_source
496 mapping_source = source_loader.load(
497 streaming=False,
498 sample_filter=sample_filter,
499 )
500 return mapping_source
502 @classmethod
503 def _try_build_presplit_jsonl_sources(
504 cls,
505 data_path: Mapping[str, OnlinePathSpec],
506 sample_filter: Callable[[Mapping[str, Any]], bool] | None,
507 ) -> tuple[Any, Any, Any] | None:
508 """Try to preserve configured splits with native JSONL sources."""
509 if not data_path:
510 return None
512 resolved_split_files = {}
513 for split_name, split_path in data_path.items():
514 if split_name not in ONLINE_SPLIT_NAMES:
515 return None
516 jsonl_files = cls._resolve_native_jsonl_files(split_path)
517 if jsonl_files is None:
518 return None
519 resolved_split_files[split_name] = jsonl_files
521 split_sources = []
522 for split_name in ONLINE_SPLIT_NAMES:
523 split_source = None
524 jsonl_files = resolved_split_files.get(split_name)
525 if jsonl_files is not None:
526 split_source = _JsonlMappingSource(jsonl_files, sample_filter)
527 split_sources.append(split_source)
528 jsonl_split_sources = (split_sources[0], split_sources[1], split_sources[2])
529 return jsonl_split_sources
531 @staticmethod
532 def _resolve_native_jsonl_files(data_path: OnlinePathSpec) -> list[str] | None:
533 """Resolve a local JSONL file, glob, directory, or ordered file list."""
534 configured_paths = [data_path] if isinstance(data_path, str) else list(data_path)
535 data_files = []
536 for configured_path in configured_paths:
537 matched_paths = sorted(glob.glob(configured_path))
538 if os.path.isdir(configured_path):
539 matched_paths = sorted(glob.glob(os.path.join(configured_path, "*.jsonl")))
540 if not matched_paths:
541 return None
542 data_files.extend(matched_paths)
544 if not data_files:
545 return None
546 for data_file in data_files:
547 if not os.path.isfile(data_file) or os.path.splitext(data_file)[1].lower() != ".jsonl":
548 return None
549 return data_files
551 def _get_target_size(self, split_index: int, source_datasets: Sequence[Any]) -> int:
552 """Return the configured sample target or one natural blended epoch."""
553 if self.sampler_type != "cyclic" and self.train_valid_test_num_samples is not None:
554 return self.train_valid_test_num_samples[split_index]
556 natural_size = 0
557 for source_dataset in source_datasets:
558 natural_size += len(source_dataset)
559 return natural_size
561 def _build_blended_split(
562 self,
563 source_datasets: Sequence[Any],
564 weights: Sequence[float],
565 target_size: int,
566 ) -> Any:
567 """Build one deterministic source and sample index graph."""
568 nonempty_datasets = []
569 nonempty_weights = []
570 for source_dataset, weight in zip(source_datasets, weights):
571 if len(source_dataset) == 0:
572 continue
573 nonempty_datasets.append(source_dataset)
574 nonempty_weights.append(weight)
576 if not nonempty_datasets or target_size == 0:
577 return None
579 blended_dataset = _OnlineBlendedMappingDataset(
580 nonempty_datasets,
581 nonempty_weights,
582 target_size,
583 self._random_seed,
584 shuffle=self.sampler_type != "cyclic",
585 )
586 return blended_dataset
588 @staticmethod
589 def _validate_split_sizes(split_sizes: Sequence[int] | None) -> tuple[int, int, int] | None:
590 """Validate optional train, validation, and test sample targets."""
591 if split_sizes is None:
592 return None
593 if len(split_sizes) != ONLINE_SPLIT_COUNT:
594 raise ValueError("Online Mapping sample targets must contain train, validation, and test sizes")
596 normalized_sizes = (int(split_sizes[0]), int(split_sizes[1]), int(split_sizes[2]))
597 if any(size < 0 for size in normalized_sizes):
598 raise ValueError("Online Mapping sample targets must be non-negative")
599 return normalized_sizes
602__all__ = [
603 "build_online_mapping_source",
604 "MappingTransformDataset",
605]