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