Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / activation_memory / pinned_memory_pool.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"""Framework-independent pool for reusable pinned host memory.""" 

16 

17import threading 

18from collections import defaultdict 

19from dataclasses import dataclass 

20from typing import Any, Dict, List, Optional, Tuple 

21 

22import torch 

23 

24from . import _backend 

25 

26_MIN_BUCKET_BYTES = 1024 

27_IN_USE = "in_use" 

28_AVAILABLE = "available" 

29_PENDING = "pending" 

30 

31 

32@dataclass 

33class _Block: 

34 """A complete allocation owned by :class:`PinnedMemoryPool`.""" 

35 

36 buffer: Any 

37 capacity: int 

38 state: str 

39 event: Optional[Any] = None 

40 storage_key: Optional[int] = None 

41 

42 

43def _valid_positive_int(value: int, name: str) -> None: 

44 if isinstance(value, bool) or not isinstance(value, int) or value <= 0: 

45 raise ValueError(f"{name} must be a positive integer, but got {value!r}.") 

46 

47 

48def _storage_key(tensor: Any) -> int: 

49 """Return a stable identity for a tensor's underlying storage.""" 

50 try: 

51 return int(tensor.untyped_storage().data_ptr()) 

52 except (AttributeError, TypeError, ValueError) as exc: 

53 raise ValueError("release() expects a tensor with an identifiable storage.") from exc 

54 

55 

56class PinnedMemoryPool: 

57 """Thread-safe, capacity-bounded pool of pinned CPU byte buffers. 

58 

59 The pool keeps complete one-dimensional ``uint8`` allocations and returns 

60 length-limited views from :meth:`acquire`. A returned view may be cast to 

61 another dtype by the caller, but no view may be used after :meth:`release`. 

62 When no reusable block is available and a new pinned allocation would 

63 exceed ``max_host_bytes``, :meth:`acquire` raises ``RuntimeError``. 

64 

65 Args: 

66 max_host_bytes: Maximum total capacity retained by the pool. 

67 align_limit: Largest size used by the power-of-two alignment policy. 

68 

69 Raises: 

70 ValueError: If a constructor or ``acquire`` argument is invalid. 

71 RuntimeError: If an allocation would exceed ``max_host_bytes`` and no 

72 sufficient reusable block exists. 

73 """ 

74 

75 def __init__(self, max_host_bytes: int, align_limit: int = 2 * 1024**3) -> None: 

76 _valid_positive_int(max_host_bytes, "max_host_bytes") 

77 _valid_positive_int(align_limit, "align_limit") 

78 self._max_host_bytes = max_host_bytes 

79 self._align_limit = align_limit 

80 self._available: Dict[int, List[_Block]] = defaultdict(list) 

81 self._pending: Dict[int, List[_Block]] = defaultdict(list) 

82 self._blocks: Dict[int, _Block] = {} 

83 self._total_allocated = 0 

84 self._lock = threading.Lock() 

85 

86 @property 

87 def total_allocated(self) -> int: 

88 """Return total aligned pinned capacity held by this pool.""" 

89 with self._lock: 

90 return self._total_allocated 

91 

92 @property 

93 def max_host_bytes(self) -> int: 

94 """Return the pool's hard capacity limit in bytes.""" 

95 return self._max_host_bytes 

96 

97 @property 

98 def align_limit(self) -> int: 

99 """Return the largest alignment bucket size in bytes.""" 

100 return self._align_limit 

101 

102 def _aligned_size(self, size: int) -> int: 

103 if size > self._align_limit: 

104 return size 

105 minimum = max(size, _MIN_BUCKET_BYTES) 

106 power = 1 << (minimum - 1).bit_length() 

107 aligned = min(power, self._align_limit) 

108 return max(aligned, size) 

109 

110 def _reclaim_completed_locked(self, minimum_capacity: int) -> None: 

111 """Move completed pending blocks into the available buckets.""" 

112 # A transfer batch may release many blocks with one event. Retain the 

113 # event objects so each runtime event is queried only once per scan. 

114 event_results: Dict[int, Tuple[Any, bool]] = {} 

115 for capacity in list(self._pending): 

116 if capacity < minimum_capacity: 

117 continue 

118 still_pending = [] 

119 for block in self._pending[capacity]: 

120 event = block.event 

121 event_key = id(event) 

122 event_result = event_results.get(event_key) 

123 if event_result is None: 

124 completed = event.query() 

125 event_results[event_key] = (event, completed) 

126 else: 

127 completed = event_result[1] 

128 if completed: 

129 block.event = None 

130 block.state = _AVAILABLE 

131 self._available[capacity].append(block) 

132 else: 

133 still_pending.append(block) 

134 if still_pending: 

135 self._pending[capacity] = still_pending 

136 else: 

137 del self._pending[capacity] 

138 

139 def _register_view_locked(self, block: _Block, view: Any) -> Any: 

140 """Associate a checked-out view with its owning allocation.""" 

141 key = _storage_key(view) 

142 owner = self._blocks.get(key) 

143 if owner is not None and owner is not block: 

144 raise RuntimeError("Two host buffers exposed the same storage identity.") 

145 old_key = block.storage_key 

146 if old_key is not None and old_key != key and self._blocks.get(old_key) is block: 

147 del self._blocks[old_key] 

148 self._blocks[key] = block 

149 block.storage_key = key 

150 return view 

151 

152 def _checkout_locked(self, block: _Block, size: int) -> Any: 

153 """Create and register the exact view returned by :meth:`acquire`.""" 

154 view = block.buffer if size == block.capacity else block.buffer[:size] 

155 return self._register_view_locked(block, view) 

156 

157 def _find_available_locked(self, minimum_capacity: int) -> Optional[_Block]: 

158 for capacity in sorted(self._available): 

159 if capacity >= minimum_capacity and self._available[capacity]: 

160 block = self._available[capacity].pop() 

161 if not self._available[capacity]: 

162 del self._available[capacity] 

163 block.state = _IN_USE 

164 return block 

165 return None 

166 

167 def _find_pending_locked(self, minimum_capacity: int) -> Optional[_Block]: 

168 for capacity in sorted(self._pending): 

169 if capacity >= minimum_capacity and self._pending[capacity]: 

170 block = self._pending[capacity].pop(0) 

171 if not self._pending[capacity]: 

172 del self._pending[capacity] 

173 block.state = _IN_USE 

174 return block 

175 return None 

176 

177 def acquire(self, size: int) -> Any: 

178 """Acquire a pooled pinned CPU ``uint8`` view.""" 

179 _valid_positive_int(size, "size") 

180 aligned_size = self._aligned_size(size) 

181 

182 with self._lock: 

183 self._reclaim_completed_locked(aligned_size) 

184 block = self._find_available_locked(aligned_size) 

185 if block is not None: 

186 try: 

187 return self._checkout_locked(block, size) 

188 except Exception: 

189 block.state = _AVAILABLE 

190 self._available[block.capacity].append(block) 

191 raise 

192 

193 if self._total_allocated + aligned_size <= self._max_host_bytes: 

194 self._total_allocated += aligned_size 

195 reserved = True 

196 else: 

197 reserved = False 

198 block = self._find_pending_locked(aligned_size) 

199 if block is None: 

200 raise RuntimeError( 

201 "PinnedMemoryPool capacity exceeded: " 

202 f"requested_bytes={size}, aligned_bytes={aligned_size}, " 

203 f"pooled_bytes={self._total_allocated}, " 

204 f"max_host_bytes={self._max_host_bytes}." 

205 ) 

206 

207 if reserved: 

208 try: 

209 buffer = _backend.alloc_tensor_buffer( 

210 aligned_size, 

211 torch.uint8, 

212 device="cpu", 

213 pin_memory=True, 

214 ) 

215 except Exception: 

216 with self._lock: 

217 self._total_allocated -= aligned_size 

218 raise 

219 block = _Block(buffer, aligned_size, _IN_USE) 

220 try: 

221 with self._lock: 

222 view = self._checkout_locked(block, size) 

223 except Exception: 

224 with self._lock: 

225 self._total_allocated -= aligned_size 

226 raise 

227 return view 

228 

229 event = block.event 

230 try: 

231 event.synchronize() 

232 except Exception: 

233 with self._lock: 

234 block.state = _PENDING 

235 block.event = event 

236 self._pending[block.capacity].insert(0, block) 

237 raise 

238 block.event = None 

239 try: 

240 with self._lock: 

241 return self._checkout_locked(block, size) 

242 except Exception: 

243 with self._lock: 

244 block.state = _AVAILABLE 

245 self._available[block.capacity].append(block) 

246 raise 

247 

248 def release(self, tensor: Any, event: Optional[Any] = None) -> None: 

249 """Return an acquired view to the pool, optionally after an async event.""" 

250 key = _storage_key(tensor) 

251 with self._lock: 

252 block = self._blocks.get(key) 

253 if block is None: 

254 raise ValueError("The tensor does not belong to this PinnedMemoryPool.") 

255 if block.state != _IN_USE: 

256 raise ValueError("The tensor has already been released to this PinnedMemoryPool.") 

257 if event is None: 

258 block.state = _AVAILABLE 

259 self._available[block.capacity].append(block) 

260 else: 

261 block.state = _PENDING 

262 block.event = event 

263 self._pending[block.capacity].append(block)