Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / data_constructor.py: 88%

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"""Target-rank Data Constructor for plan-ordered packing and collation.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Callable, Mapping, Sequence 

20from typing import Any 

21 

22from hyper_parallel.distributed_data.schema import PackingBinPlan, SampleKey 

23 

24 

25def default_pack_fn(samples: Sequence[Any], seq_len: int) -> tuple[Any, ...]: 

26 """Preserve one planned packing bin without assuming a sample schema. 

27 

28 Args: 

29 samples: Raw samples assigned to one Planner-validated bin, including 

30 an explicitly permitted singleton overflow when configured. 

31 seq_len: Configured normal sequence capacity. 

32 

33 Returns: 

34 Immutable raw samples in their deterministic planned order. 

35 """ 

36 if not isinstance(seq_len, int) or isinstance(seq_len, bool) or seq_len < 1: 

37 raise ValueError(f"seq_len must be a positive integer, but got {seq_len!r}.") 

38 return tuple(samples) 

39 

40 

41def default_collate_fn(packed_sequences: Sequence[Any]) -> tuple[Any, ...]: 

42 """Preserve the planned bins as one immutable rank-local batch. 

43 

44 Args: 

45 packed_sequences: Outputs produced for this rank's packing bins. 

46 

47 Returns: 

48 Immutable bins in contiguous ``pack_index`` order. 

49 """ 

50 return tuple(packed_sequences) 

51 

52 

53class PackingDataConstructor: 

54 """Apply configured packing per planned bin and collate one local batch. 

55 

56 ``pack_fn`` never runs in a Dataset Reader. It receives raw samples only 

57 after sample-level planning and CPU redistribution have completed. 

58 """ 

59 

60 def __init__( 

61 self, 

62 pack_fn: Callable[[Sequence[Any], int], Any], 

63 collate_fn: Callable[[Sequence[Any]], Any], 

64 *, 

65 seq_len: int, 

66 ) -> None: 

67 """Initialize user callbacks. 

68 

69 Args: 

70 pack_fn: Called once per planned sequence bin as 

71 ``pack_fn(raw_samples, seq_len)``. 

72 collate_fn: Called once per local batch with packed sequences. 

73 seq_len: Maximum non-oversized sequence length. 

74 """ 

75 if not callable(pack_fn) or not callable(collate_fn): 

76 raise ValueError("pack_fn and collate_fn must be callable.") 

77 if not isinstance(seq_len, int) or isinstance(seq_len, bool) or seq_len < 1: 

78 raise ValueError(f"seq_len must be a positive integer, but got {seq_len!r}.") 

79 self._pack_fn = pack_fn 

80 self._collate_fn = collate_fn 

81 self._seq_len = seq_len 

82 

83 def construct( 

84 self, 

85 plan: Sequence[PackingBinPlan], 

86 payloads: Mapping[SampleKey, Any], 

87 ) -> Any: 

88 """Construct one target-rank local batch. 

89 

90 Args: 

91 plan: Ordered packing bins assigned to this Data Constructor. 

92 payloads: Received raw samples keyed by stable Dataset Reader and 

93 Dataset-index identity. 

94 

95 Returns: 

96 User-collated rank-local batch. 

97 """ 

98 expected_keys = tuple( 

99 key 

100 for packing_bin in plan 

101 for key in packing_bin.sample_keys 

102 ) 

103 expected_set = set(expected_keys) 

104 if len(payloads) != len(expected_keys) or payloads.keys() != expected_set: 

105 missing = expected_set - payloads.keys() 

106 unexpected = payloads.keys() - expected_set 

107 raise ValueError( 

108 f"Data Constructor payload keys do not match the plan; " 

109 f"missing={sorted(missing)}, unexpected={sorted(unexpected)}." 

110 ) 

111 

112 packed_sequences = [] 

113 for pack_index, packing_bin in enumerate(plan): 

114 if packing_bin.pack_tokens > self._seq_len and not packing_bin.oversized: 

115 raise ValueError( 

116 f"Packing bin {pack_index} has {packing_bin.pack_tokens} tokens, " 

117 f"exceeding seq_len={self._seq_len}." 

118 ) 

119 raw_samples = [payloads[key] for key in packing_bin.sample_keys] 

120 packed_sequences.append(self._pack_fn(raw_samples, self._seq_len)) 

121 return self._collate_fn(packed_sequences) 

122 

123 

124__all__ = ["PackingDataConstructor", "default_collate_fn", "default_pack_fn"]