Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / batching / dynamic_batch.py: 28%
272 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"""Candidate buffering and finalization used by dynamic DataLoaders."""
17import operator
18from abc import ABC, abstractmethod
19from collections.abc import Callable, Mapping, Sequence
20from dataclasses import dataclass, field
21from typing import Any
23from hyper_parallel.data.batching.packing import (
24 FirstFitPackingSelector,
25 PackingSelector,
26 SamplePacker,
27)
30@dataclass
31class BaseCandidateBuffer(ABC):
32 """Store candidates and expose the next selected sample group."""
34 token_budget: int
35 min_buffered_samples: int
36 buffer: list[tuple[Mapping[str, Any], int]] = field(default_factory=list, init=False)
37 buffer_output_indices: list[Any | None] = field(default_factory=list, init=False)
38 buffer_token_count: int = field(default=0, init=False)
40 def __post_init__(self) -> None:
41 """Validate the candidate-buffer limits."""
42 if self.token_budget <= 0:
43 raise ValueError("token_budget must be positive")
45 if self.min_buffered_samples <= 0:
46 raise ValueError("min_buffered_samples must be positive")
48 @abstractmethod
49 def build_buffer_entry(self, model_sample: Mapping[str, Any]) -> tuple[Mapping[str, Any], int]:
50 """Attach the scheduling cost used to select this sample."""
51 raise NotImplementedError
53 def put_item(self, model_sample: Mapping[str, Any], output_index_entry: Any | None = None) -> None:
54 """Append one candidate and its optional checkpoint replay index."""
55 buffered_sample = self.build_buffer_entry(model_sample)
56 self.buffer.append(buffered_sample)
57 self.buffer_output_indices.append(output_index_entry)
58 self.buffer_token_count += buffered_sample[1]
60 def is_ready_for_micro_batch(self) -> bool:
61 """Return whether both configured candidate limits are reached."""
62 enough_samples = len(self.buffer) >= self.min_buffered_samples
63 enough_tokens = self.buffer_token_count >= self.token_budget
64 is_ready = enough_samples and enough_tokens
65 return is_ready
67 @abstractmethod
68 def get_micro_batch(self) -> list[Mapping[str, Any]]:
69 """Select and remove the next sample group."""
70 raise NotImplementedError
72 def empty(self) -> bool:
73 """Return whether no candidates remain."""
74 is_empty = not self.buffer
75 return is_empty
78@dataclass
79class TokenBudgetCandidateBuffer(BaseCandidateBuffer):
80 """Select encoded samples in source order within one token budget."""
82 def build_buffer_entry(self, model_sample: Mapping[str, Any]) -> tuple[Mapping[str, Any], int]:
83 """Attach the encoded token length to one candidate."""
84 sample_length = int(model_sample["input_ids"].shape[-1])
85 if sample_length <= 0:
86 raise ValueError("Dynamic batching samples must contain at least one token")
88 buffered_sample = (model_sample, sample_length)
89 return buffered_sample
91 def get_micro_batch(self) -> list[Mapping[str, Any]]:
92 """Select fitting candidates and retain the remainder."""
93 if not self.buffer:
94 raise ValueError("Dynamic token buffer is empty")
96 selected_samples = []
97 remaining_buffer = []
98 remaining_output_indices = []
99 selected_token_count = 0
100 for (model_sample, sample_length), output_index_entry in zip(
101 self.buffer,
102 self.buffer_output_indices,
103 ):
104 sample_fits = selected_token_count == 0 or selected_token_count + sample_length <= self.token_budget
105 if sample_fits:
106 selected_samples.append(model_sample)
107 selected_token_count += sample_length
108 else:
109 remaining_buffer.append((model_sample, sample_length))
110 remaining_output_indices.append(output_index_entry)
112 self.buffer = remaining_buffer
113 self.buffer_output_indices = remaining_output_indices
114 self.buffer_token_count -= selected_token_count
115 return selected_samples
118@dataclass
119class PackingCandidateBuffer(BaseCandidateBuffer):
120 """Delegate candidate-group selection to a configured packing strategy."""
122 packing_selector: PackingSelector
124 def build_buffer_entry(self, model_sample: Mapping[str, Any]) -> tuple[Mapping[str, Any], int]:
125 """Attach selector-defined scheduling cost to one candidate."""
126 sample_cost = self.packing_selector.get_sample_cost(model_sample)
127 buffered_sample = (model_sample, sample_cost)
128 return buffered_sample
130 def get_micro_batch(self) -> list[Mapping[str, Any]]:
131 """Select and remove the next candidate group."""
132 if not self.buffer:
133 raise ValueError("Packing candidate buffer is empty")
135 candidates = []
136 for model_sample, _ in self.buffer:
137 candidates.append(model_sample)
139 selected_indices = self.packing_selector.select_samples_to_pack(
140 candidates,
141 self.token_budget,
142 )
143 resolved_indices = self._validate_selected_indices(selected_indices)
144 selected_index_set = set(resolved_indices)
146 selected_samples = []
147 selected_cost = 0
148 for selected_index in resolved_indices:
149 model_sample, sample_cost = self.buffer[selected_index]
150 selected_samples.append(model_sample)
151 selected_cost += sample_cost
153 remaining_buffer = []
154 remaining_output_indices = []
155 for sample_index, (buffered_sample, output_index) in enumerate(
156 zip(self.buffer, self.buffer_output_indices)
157 ):
158 if sample_index not in selected_index_set:
159 remaining_buffer.append(buffered_sample)
160 remaining_output_indices.append(output_index)
162 self.buffer = remaining_buffer
163 self.buffer_output_indices = remaining_output_indices
164 self.buffer_token_count -= selected_cost
165 return selected_samples
167 def _validate_selected_indices(self, selected_indices: Sequence[int]) -> list[int]:
168 """Validate selector output before mutating the candidate buffer."""
169 resolved_indices = []
170 for selected_index in selected_indices:
171 if isinstance(selected_index, bool):
172 raise TypeError("PackingSelector indices must be integers")
173 try:
174 resolved_index = operator.index(selected_index)
175 except TypeError as exc:
176 raise TypeError("PackingSelector indices must be integers") from exc
178 if resolved_index < 0 or resolved_index >= len(self.buffer):
179 raise IndexError("PackingSelector returned an out-of-range candidate index")
180 resolved_indices.append(resolved_index)
182 if not resolved_indices:
183 raise ValueError("PackingSelector must select at least one candidate")
184 if len(set(resolved_indices)) != len(resolved_indices):
185 raise ValueError("PackingSelector must not select a candidate more than once")
187 return resolved_indices
190def _normalize_source_samples(source_item: Any) -> list[Mapping[str, Any]]:
191 """Normalize one source output into its ordered model samples."""
192 if isinstance(source_item, Mapping):
193 model_samples = [source_item]
194 return model_samples
196 if isinstance(source_item, Sequence) and not isinstance(source_item, (str, bytes)):
197 model_samples = []
198 for model_sample in source_item:
199 if not isinstance(model_sample, Mapping):
200 raise ValueError("Every dynamic source sample must be a mapping")
202 model_samples.append(model_sample)
203 return model_samples
205 raise ValueError("A dynamic source item must be a mapping or a sequence of mappings")
208def _restore_index_buffer(
209 source_dataset: Any,
210 saved_buffer: Sequence[Any],
211 candidate_buffer: BaseCandidateBuffer,
212) -> list[tuple[Mapping[str, Any], int]]:
213 """Rebuild buffered samples from their output and sample indices."""
214 restored_buffer = []
215 cached_output_index: Any = object()
216 cached_model_samples: list[Mapping[str, Any]] = []
217 for output_index_entry in saved_buffer:
218 output_index, sample_index = output_index_entry
219 resolved_sample_index = operator.index(sample_index)
220 if resolved_sample_index < 0:
221 raise ValueError("Buffered sample index must be a non-negative integer")
223 if output_index != cached_output_index:
224 restored_item = source_dataset.get_item(output_index)
225 cached_model_samples = _normalize_source_samples(restored_item)
226 cached_output_index = output_index
227 try:
228 model_sample = cached_model_samples[resolved_sample_index]
229 except IndexError as exc:
230 raise ValueError(
231 f"Buffered sample index {resolved_sample_index} is out of range "
232 f"for output index {output_index!r}"
233 ) from exc
235 buffer_entry = candidate_buffer.build_buffer_entry(model_sample)
236 restored_buffer.append(buffer_entry)
238 return restored_buffer
241class MappingReplayDataset:
242 """Attach a stable replay index to each mapping Dataset output."""
244 def __init__(self, dataset: Any) -> None:
245 """Store the mapping Dataset used to reconstruct buffered samples."""
246 self.dataset = dataset
248 def __len__(self) -> int:
249 """Return the wrapped Dataset length."""
250 dataset_length = len(self.dataset)
251 return dataset_length
253 def __getitem__(self, index: int) -> tuple[Any, int]:
254 """Return one source item and the index that can reproduce it."""
255 source_item = self.get_item(index)
256 indexed_source_item = (source_item, index)
257 return indexed_source_item
259 def get_item(self, index: Any) -> Any:
260 """Rebuild one source item without advancing the source iterator."""
261 resolved_index = operator.index(index)
262 get_item = getattr(self.dataset, "get_item", None)
263 if callable(get_item):
264 source_item = get_item(resolved_index)
265 else:
266 source_item = self.dataset[resolved_index]
268 return source_item
271class IndexReplayCheckpoint:
272 """Checkpoint candidates as compact mapping Dataset replay indices."""
274 save_by_idx = True
276 def __init__(self, source_dataset: Any) -> None:
277 """Store a Dataset that emits source items with stable replay indices."""
278 self.source_dataset = source_dataset
280 @staticmethod
281 def supports(dataset: Any) -> bool:
282 """Return whether an iterable Dataset supports index replay."""
283 get_item = getattr(dataset, "get_item", None)
284 supports_index_replay = callable(get_item) and hasattr(dataset, "output_index_for_resume")
285 return supports_index_replay
287 @staticmethod
288 def put_source_item(source_item: Any, candidate_buffer: BaseCandidateBuffer) -> None:
289 """Append an indexed source item and retain each sample index."""
290 model_samples_item, output_index = source_item
291 model_samples = _normalize_source_samples(model_samples_item)
292 for sample_index, model_sample in enumerate(model_samples):
293 candidate_buffer.put_item(model_sample, (output_index, sample_index))
295 @staticmethod
296 def get_buffer_state(candidate_buffer: BaseCandidateBuffer) -> list[Any]:
297 """Return replay indices for buffered model samples."""
298 if len(candidate_buffer.buffer) != len(candidate_buffer.buffer_output_indices):
299 raise RuntimeError("Dynamic sample and output-index buffers are inconsistent")
301 if any(index is None for index in candidate_buffer.buffer_output_indices):
302 raise RuntimeError("Index-buffer checkpoint requires a replay key for every buffered sample")
304 return candidate_buffer.buffer_output_indices
306 def restore_buffer(
307 self,
308 saved_buffer: Sequence[Any],
309 saved_by_idx: bool,
310 candidate_buffer: BaseCandidateBuffer,
311 ) -> tuple[list[tuple[Mapping[str, Any], int]], list[Any]]:
312 """Rebuild buffered samples from their replay indices."""
313 if not saved_by_idx:
314 if saved_buffer:
315 raise ValueError("save_by_idx=True cannot restore a checkpoint containing full samples")
317 return [], []
319 restored_buffer = _restore_index_buffer(self.source_dataset, saved_buffer, candidate_buffer)
320 restored_indices = list(saved_buffer)
321 return restored_buffer, restored_indices
324class FullSampleCheckpoint:
325 """Checkpoint full candidate samples alongside the source state."""
327 save_by_idx = False
329 def __init__(self, dataset: Any, replay_dataset: Any | None = None) -> None:
330 """Retain the source Dataset and optional mapping replay Dataset."""
331 self.source_dataset = dataset
332 self.replay_dataset = replay_dataset
334 @staticmethod
335 def put_source_item(source_item: Any, candidate_buffer: BaseCandidateBuffer) -> None:
336 """Append one streaming source item to the candidate buffer."""
337 model_samples = _normalize_source_samples(source_item)
338 for model_sample in model_samples:
339 candidate_buffer.put_item(model_sample)
341 @staticmethod
342 def get_buffer_state(candidate_buffer: BaseCandidateBuffer) -> list[Any]:
343 """Return full samples already pulled beyond the source cursor."""
344 if len(candidate_buffer.buffer) != len(candidate_buffer.buffer_output_indices):
345 raise RuntimeError("Dynamic sample and output-index buffers are inconsistent")
347 return candidate_buffer.buffer
349 def restore_buffer(
350 self,
351 saved_buffer: Sequence[Any],
352 saved_by_idx: bool,
353 candidate_buffer: BaseCandidateBuffer,
354 ) -> tuple[list[tuple[Mapping[str, Any], int]], list[Any]]:
355 """Restore full samples or an earlier index-buffer checkpoint."""
356 if saved_by_idx:
357 if saved_buffer and self.replay_dataset is None:
358 raise ValueError("An index-buffer checkpoint requires a replayable Dataset")
360 restored_buffer = _restore_index_buffer(self.replay_dataset, saved_buffer, candidate_buffer)
361 else:
362 restored_buffer = list(saved_buffer)
363 restored_indices = [None] * len(restored_buffer)
364 return restored_buffer, restored_indices
367class _BatchPipeline(ABC):
368 """Define candidate selection and finalization for one dynamic Loader."""
370 @abstractmethod
371 def build_candidate_buffer(
372 self,
373 *,
374 token_budget: int,
375 min_buffered_samples: int,
376 ) -> BaseCandidateBuffer:
377 """Build the candidate buffer used by this pipeline."""
378 raise NotImplementedError
380 @abstractmethod
381 def finalize_selected_samples(
382 self,
383 selected_samples: Sequence[Mapping[str, Any]],
384 ) -> Sequence[Mapping[str, Any]]:
385 """Finish selected samples before final collation."""
386 raise NotImplementedError
388 def encode_batch(self, batch: Mapping[str, Any]) -> Mapping[str, Any]:
389 """Apply an optional final batch encoding stage."""
390 return batch
393class _TokenBatchPipeline(_BatchPipeline):
394 """Build batches directly from fully encoded token samples."""
396 def build_candidate_buffer(
397 self,
398 *,
399 token_budget: int,
400 min_buffered_samples: int,
401 ) -> BaseCandidateBuffer:
402 """Build the token-budget candidate buffer."""
403 candidate_buffer = TokenBudgetCandidateBuffer(
404 token_budget=token_budget,
405 min_buffered_samples=min_buffered_samples,
406 )
407 return candidate_buffer
409 def finalize_selected_samples(
410 self,
411 selected_samples: Sequence[Mapping[str, Any]],
412 ) -> Sequence[Mapping[str, Any]]:
413 """Return fully encoded samples for final collation."""
414 return selected_samples
417@dataclass(frozen=True)
418class _OmniPackingPipeline(_BatchPipeline):
419 """Compose selection, selected-sample encoding, packing, and batch encoding."""
421 selector: PackingSelector
422 selected_sample_encoder: Callable[[Mapping[str, Any]], Mapping[str, Any]]
423 packer: SamplePacker
424 batch_encoder: Callable[[Mapping[str, Any]], Mapping[str, Any]]
426 @classmethod
427 def from_dataset(cls, dataset: Any) -> "_OmniPackingPipeline":
428 """Build the packing pipeline from one normalized Omni Dataset interface."""
429 selected_sample_encoder = getattr(dataset, "encode_selected_sample", None)
430 batch_encoder = getattr(dataset, "encode_batch", None)
431 if not callable(selected_sample_encoder) or not callable(batch_encoder):
432 raise TypeError("OmniPackingLoader requires encode_selected_sample() and encode_batch()")
434 # Hook—a criterion used to select data objects; this criterion can be extended to support load balancing.
435 selector = getattr(dataset, "packing_selector", None)
436 if selector is None:
437 selector = FirstFitPackingSelector()
438 if not isinstance(selector, PackingSelector):
439 raise TypeError("dataset.packing_selector must implement PackingSelector")
441 # This is a hook function used to package data into different fields;
442 # you need to verify that the keyword mappings match.
443 packer = getattr(dataset, "sample_packer", None)
444 if packer is None:
445 packer = SamplePacker()
446 if not isinstance(packer, SamplePacker):
447 raise TypeError("dataset.sample_packer must implement SamplePacker")
449 packing_pipeline = cls(
450 selector=selector,
451 selected_sample_encoder=selected_sample_encoder,
452 packer=packer,
453 batch_encoder=batch_encoder,
454 )
455 return packing_pipeline
457 def build_candidate_buffer(
458 self,
459 *,
460 token_budget: int,
461 min_buffered_samples: int,
462 ) -> BaseCandidateBuffer:
463 """Build candidate selection around the composed selector."""
464 candidate_buffer = PackingCandidateBuffer(
465 token_budget=token_budget,
466 min_buffered_samples=min_buffered_samples,
467 packing_selector=self.selector,
468 )
469 return candidate_buffer
471 def finalize_selected_samples(
472 self,
473 selected_samples: Sequence[Mapping[str, Any]],
474 ) -> Sequence[Mapping[str, Any]]:
475 """Encode and pack one selected sample group."""
476 encoded_samples = []
477 for selected_sample in selected_samples:
478 encoded_sample = self.selected_sample_encoder(selected_sample)
479 if not isinstance(encoded_sample, Mapping):
480 raise TypeError("postencode_sample must return a mapping")
481 encoded_samples.append(encoded_sample)
483 packed_sample = self.packer.pack_selected_samples(encoded_samples)
484 if not isinstance(packed_sample, Mapping):
485 raise TypeError("SamplePacker must return one mapping")
487 packed_samples = [packed_sample]
488 return packed_samples
490 def encode_batch(self, batch: Mapping[str, Any]) -> Mapping[str, Any]:
491 """Apply the composed final batch encoder."""
492 encoded_batch = self.batch_encoder(batch)
493 return encoded_batch
496__all__ = [
497 "BaseCandidateBuffer",
498 "FullSampleCheckpoint",
499 "IndexReplayCheckpoint",
500 "MappingReplayDataset",
501 "PackingCandidateBuffer",
502 "TokenBudgetCandidateBuffer",
503]