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

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"""Ahead-of-fetch metadata and plan-aware local sample loading.""" 

16# This distributed-data package is intentionally PyTorch-only. 

17 

18from __future__ import annotations 

19 

20from collections.abc import Mapping, Sequence 

21from typing import Any 

22 

23import torch # pylint: disable=forbidden-backend-import 

24from torch.utils.data import DataLoader, Dataset # pylint: disable=forbidden-backend-import 

25 

26from hyper_parallel.distributed_data.schema import SampleKey 

27from hyper_parallel.distributed_data.dataset_reader import ( 

28 _build_worker_options, 

29 _identity, 

30) 

31 

32 

33class PlannedSampleLoader: 

34 """Use DataLoader workers to fetch only samples selected by the plan.""" 

35 

36 def __init__( 

37 self, 

38 dataset: Any, 

39 *, 

40 num_workers: int, 

41 pin_memory: bool, 

42 prefetch_factor: int | None, 

43 persistent_workers: bool, 

44 seed: int, 

45 dataloader_kwargs: Mapping[str, Any] | None = None, 

46 ) -> None: 

47 """Initialize a reusable plan-aware DataLoader. 

48 

49 Args: 

50 dataset: Mapping-style Dataset aligned with the metadata sequence. 

51 num_workers: Native DataLoader worker count. 

52 pin_memory: Whether workers pin returned sample memory. 

53 prefetch_factor: Samples prefetched by each worker. 

54 persistent_workers: Whether workers persist for the loader lifetime. 

55 seed: Base worker seed. 

56 dataloader_kwargs: Additional validated DataLoader execution options. 

57 """ 

58 getitem = getattr(type(dataset), "__getitem__", None) 

59 if not callable(getitem) or getitem is Dataset.__getitem__ or not hasattr(dataset, "__len__"): 

60 raise ValueError("Planned reads require a mapping-style Dataset with __len__ and __getitem__.") 

61 self._dataset = dataset 

62 self._indices: list[int] = [] 

63 self._seed = seed 

64 self._epoch = 0 

65 self._worker_generator = torch.Generator().manual_seed(seed) 

66 worker_options = _build_worker_options( 

67 num_workers=num_workers, 

68 pin_memory=pin_memory, 

69 prefetch_factor=prefetch_factor, 

70 persistent_workers=persistent_workers, 

71 generator=self._worker_generator, 

72 dataloader_kwargs=dataloader_kwargs, 

73 ) 

74 self._data_loader = DataLoader( 

75 dataset, 

76 batch_size=None, 

77 sampler=self._indices, 

78 collate_fn=_identity, 

79 **worker_options, 

80 ) 

81 

82 def fetch_samples(self, sample_keys: Sequence[SampleKey]) -> dict[SampleKey, Any]: 

83 """Fetch Dataset samples for the requested keys, preserving their order. 

84 

85 Args: 

86 sample_keys: Keys whose local ``dataset_index`` values belong to 

87 this loader's Dataset. 

88 

89 Returns: 

90 Mapping from every requested key to its Dataset sample. 

91 """ 

92 if any(not isinstance(key, SampleKey) for key in sample_keys): 

93 raise ValueError("sample_keys must contain SampleKey values.") 

94 if len(sample_keys) != len(set(sample_keys)): 

95 raise ValueError("sample_keys must not contain duplicates.") 

96 dataset_size = len(self._dataset) 

97 invalid_indices = [key.dataset_index for key in sample_keys if key.dataset_index >= dataset_size] 

98 if invalid_indices: 

99 raise ValueError( 

100 f"Sample plan references Dataset indices outside [0, {dataset_size}): {invalid_indices}." 

101 ) 

102 # DataLoader retains this list, so replace its contents instead of rebinding it. 

103 self._indices[:] = (key.dataset_index for key in sample_keys) 

104 iterator = iter(self._data_loader) 

105 samples_by_key = {} 

106 for key in sample_keys: 

107 try: 

108 payload = next(iterator) 

109 except StopIteration as exc: 

110 raise ValueError("Plan-aware DataLoader exhausted before all requested samples were read.") from exc 

111 samples_by_key[key] = payload 

112 return samples_by_key 

113 

114 def state_dict(self) -> dict[str, Any]: 

115 """Return the epoch required to seed future plan-aware workers.""" 

116 return { 

117 "dataset_size": len(self._dataset), 

118 "seed": self._seed, 

119 "epoch": self._epoch, 

120 } 

121 

122 def load_state_dict(self, state_dict: Mapping[str, Any]) -> None: 

123 """Restore plan-aware reader state before its first fetch. 

124 

125 Args: 

126 state_dict: State produced by :meth:`state_dict` on this constructor. 

127 """ 

128 try: 

129 state = dict(state_dict) 

130 except Exception as exc: 

131 raise ValueError(f"Plan-aware DataLoader state is not a mapping: {exc}") from exc 

132 expected = { 

133 "dataset_size": len(self._dataset), 

134 "seed": self._seed, 

135 } 

136 for name, expected_value in expected.items(): 

137 if state.get(name) != expected_value: 

138 raise ValueError( 

139 f"Plan-aware DataLoader checkpoint {name}={state.get(name)!r} " 

140 f"does not match {expected_value!r}." 

141 ) 

142 self.set_epoch(state.get("epoch")) 

143 

144 def set_epoch(self, epoch: int) -> None: 

145 """Record the epoch and reset the future worker-initialization seed. 

146 

147 Args: 

148 epoch: Non-negative Dataset epoch. 

149 """ 

150 if not isinstance(epoch, int) or isinstance(epoch, bool) or epoch < 0: 

151 raise ValueError(f"epoch must be a non-negative integer, but got {epoch!r}.") 

152 self._epoch = epoch 

153 self._worker_generator.manual_seed(self._seed + epoch) 

154 

155 

156__all__ = ["PlannedSampleLoader"]