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

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"""Pinned Host staging and asynchronous device copies shared by data loaders.""" 

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

17 

18from __future__ import annotations 

19 

20from collections.abc import Mapping 

21from typing import Any 

22 

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

24 

25 

26def _get_device_module(device: torch.device) -> Any: 

27 """Return the PyTorch device module that owns streams and events.""" 

28 module = getattr(torch, device.type, None) 

29 required = ("Event", "Stream", "current_stream", "stream") 

30 if module is None or any(not callable(getattr(module, name, None)) for name in required): 

31 raise ValueError( 

32 f"Device {device} does not expose PyTorch accelerator Stream/Event APIs." 

33 ) 

34 return module 

35 

36 

37def _move_to_device(value: Any, device: torch.device) -> Any: 

38 """Move tensors, batch objects, and standard containers recursively.""" 

39 if isinstance(value, torch.Tensor): 

40 return value.to(device, non_blocking=True) 

41 to_device = getattr(value, "to", None) 

42 if callable(to_device): 

43 return to_device(device, non_blocking=True) 

44 if isinstance(value, Mapping): 

45 return {key: _move_to_device(item, device) for key, item in value.items()} 

46 if isinstance(value, tuple): 

47 moved = tuple(_move_to_device(item, device) for item in value) 

48 if hasattr(value, "_fields"): 

49 return type(value)(*moved) 

50 return moved 

51 if isinstance(value, list): 

52 return [_move_to_device(item, device) for item in value] 

53 return value 

54 

55 

56def _record_stream(value: Any, stream: Any, visited: set[int] | None = None) -> None: 

57 """Associate device batch storage with its consumer stream.""" 

58 if visited is None: 

59 visited = set() 

60 value_id = id(value) 

61 if value_id in visited: 

62 return 

63 visited.add(value_id) 

64 

65 if isinstance(value, torch.Tensor): 

66 if value.device.type != "cpu": 

67 value.record_stream(stream) 

68 return 

69 record_stream = getattr(value, "record_stream", None) 

70 if callable(record_stream): 

71 record_stream(stream) 

72 return 

73 if isinstance(value, Mapping): 

74 for item in value.values(): 

75 _record_stream(item, stream, visited) 

76 return 

77 if isinstance(value, (tuple, list)): 

78 for item in value: 

79 _record_stream(item, stream, visited) 

80 return 

81 

82 attributes = getattr(value, "__dict__", None) 

83 if isinstance(attributes, dict): 

84 for item in attributes.values(): 

85 _record_stream(item, stream, visited) 

86 

87 

88def _pin_memory(value: Any) -> Any: 

89 """Pin final collated CPU tensors without changing metadata or dtype.""" 

90 if isinstance(value, torch.Tensor): 

91 if value.device.type == "cpu" and value.layout == torch.strided and not value.is_pinned(): 

92 return value.pin_memory() 

93 return value 

94 if isinstance(value, Mapping): 

95 return {key: _pin_memory(item) for key, item in value.items()} 

96 if isinstance(value, tuple): 

97 pinned = tuple(_pin_memory(item) for item in value) 

98 return type(value)(*pinned) if hasattr(value, "_fields") else pinned 

99 if isinstance(value, list): 

100 return [_pin_memory(item) for item in value] 

101 pin = getattr(value, "pin_memory", None) 

102 return pin() if callable(pin) else value 

103 

104 

105class DevicePrefetchedStep: 

106 """Hand the complete device batch and its copy dependency to training.""" 

107 

108 def __init__( 

109 self, 

110 device_batch: Any, 

111 ready_event: Any, 

112 device: torch.device, 

113 ) -> None: 

114 """Store device inputs while the prefetcher owns pending Host sources.""" 

115 self.device_batch = device_batch 

116 self.ready_event = ready_event 

117 self.device = device 

118 

119 def take_batch(self) -> Any: 

120 """Transfer the complete batch's lifetime to the current consumer stream. 

121 

122 Returns: 

123 Device batch, with non-tensor metadata left untouched. 

124 """ 

125 device_module = getattr(torch, self.device.type) 

126 stream = device_module.current_stream(self.device) 

127 stream.wait_event(self.ready_event) 

128 batch = self.device_batch 

129 _record_stream(batch, stream) 

130 self.device_batch = None 

131 return batch 

132 

133 

134class DeviceStepPrefetcher: 

135 """Separate Host pinning from caller-owned asynchronous device copies.""" 

136 

137 def __init__(self, device: Any) -> None: 

138 """Configure final input transfer without allocating a device stream. 

139 

140 Args: 

141 device: Target CUDA or NPU device, independent of payload transport. 

142 """ 

143 self.device = torch.device(device) 

144 if self.device.type not in ("cuda", "npu"): 

145 raise ValueError("DeviceStepPrefetcher requires a CUDA or NPU device.") 

146 self._device_module = _get_device_module(self.device) 

147 self._copy_stream = None 

148 self._pending_staging: list[tuple[Any, Any]] = [] 

149 

150 def close(self) -> None: 

151 """Drain pending copies before dropping their pinned source storage.""" 

152 for event, _ in self._pending_staging: 

153 event.synchronize() 

154 self._pending_staging.clear() 

155 

156 def __call__(self, cpu_batch: Any) -> DevicePrefetchedStep: 

157 """Enqueue the final batch and publish its copy-completion event. 

158 

159 Args: 

160 cpu_batch: Final collated batch to pin and copy. 

161 

162 Returns: 

163 Device batch with a copy-completion event. The 

164 caller does not wait for that event; the consumer stream 

165 orders itself with ``wait_event`` in ``take_batch``. 

166 """ 

167 return self.stage_prepared(self.prepare_host(cpu_batch)) 

168 

169 @staticmethod 

170 def prepare_host(cpu_batch: Any) -> Any: 

171 """Pin a constructed batch on its Host producer before device staging. 

172 

173 Args: 

174 cpu_batch: Collated batch whose CPU tensor leaves need pinned storage. 

175 

176 Returns: 

177 The batch structure with pinned CPU tensors. 

178 """ 

179 return _pin_memory(cpu_batch) 

180 

181 def stage_prepared(self, staging: Any) -> DevicePrefetchedStep: 

182 """Enqueue an already-pinned batch without CPU allocation or copying for pinning. 

183 

184 Args: 

185 staging: Pinned batch prepared by the Host producer. 

186 

187 Returns: 

188 Device input and its copy event. This prefetcher retains pending Host sources. 

189 """ 

190 device_module = self._device_module 

191 # The local loader may create a fresh producer thread each step; device 

192 # selection is thread-local even when payload exchange uses only Gloo. 

193 device_module.set_device(self.device) 

194 if self._copy_stream is None: 

195 self._copy_stream = device_module.Stream(device=self.device) 

196 self._pending_staging = [(event, batch) for event, batch in self._pending_staging if not event.query()] 

197 try: 

198 with device_module.stream(self._copy_stream): 

199 device_batch = _move_to_device(staging, self.device) 

200 ready_event = device_module.Event() 

201 ready_event.record(self._copy_stream) 

202 except BaseException: 

203 # Earlier copies may still read pinned staging when a later 

204 # launch fails. Drain only this stream before releasing sources. 

205 self._copy_stream.synchronize() 

206 raise 

207 self._pending_staging.append((ready_event, staging)) 

208 return DevicePrefetchedStep(device_batch, ready_event, self.device) 

209 

210 

211def _resolve_device(device: Any = None, *, communication_backend: str = "hccl") -> torch.device: 

212 """Keep Gloo on CPU by default; otherwise select the current accelerator.""" 

213 if device is None and communication_backend == "gloo": 

214 return torch.device("cpu") 

215 if device is None: 

216 for device_type in ("npu", "cuda"): 

217 device_module = getattr(torch, device_type, None) 

218 if device_module is not None and device_module.is_available(): 

219 device = torch.device(device_type, device_module.current_device()) 

220 break 

221 else: 

222 device = "cpu" 

223 return torch.device(device) 

224 

225 

226def _create_device_prefetcher(device: Any = None) -> DeviceStepPrefetcher | None: 

227 """Create an H2D copy stage only when an accelerator is selected.""" 

228 device = _resolve_device(device) 

229 if device.type == "cpu": 

230 return None 

231 return DeviceStepPrefetcher(device) 

232 

233 

234__all__ = ["DeviceStepPrefetcher"]