Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / hyper_offload / runtime / pinned_memory.py: 94%

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 memory pool with bucket-based management. 

16 

17Provides a thread-safe pool of host page-locked CPU memory buffers for 

18efficient host-device data transfers during activation offloading. 

19""" 

20 

21import logging 

22import threading 

23 

24import torch 

25 

26logger = logging.getLogger(__name__) 

27 

28_BUCKET_SIZES = [2**i for i in range(10, 32)] 

29 

30 

31def _align_to_bucket(size: int) -> int: 

32 """Find the smallest bucket size >= size.""" 

33 for bucket in _BUCKET_SIZES: 

34 if bucket >= size: 

35 return bucket 

36 return size 

37 

38 

39def _bucket_for(size: int) -> int: 

40 """Return the bucket size a buffer belongs to.""" 

41 return _align_to_bucket(size) 

42 

43 

44class PinnedMemoryPool: 

45 """Thread-safe pool of pinned host CPU memory buffers with deferred recycling via CUDA events. 

46 

47 If an acquire request exceeds the pool's remaining capacity defined by 

48 ``max_host_bytes``, further acquires fall back to regular pageable CPU 

49 memory. 

50 """ 

51 

52 def __init__(self, max_host_bytes: int) -> None: 

53 self._pool: dict[int, list[torch.Tensor]] = {} 

54 self._pending: dict[int, list[tuple[torch.Tensor, torch.Event]]] = {} 

55 self._lock = threading.Lock() 

56 self._total_allocated = 0 

57 self._max_host_bytes = max_host_bytes 

58 

59 @property 

60 def total_allocated(self) -> int: 

61 """Total host bytes currently held by the allocator.""" 

62 return self._total_allocated 

63 

64 @property 

65 def max_host_bytes(self) -> int: 

66 """Hard limit on host memory in bytes.""" 

67 return self._max_host_bytes 

68 

69 def _reclaim_locked(self, bucket: int) -> None: 

70 """Move completed tensors from pending to available pool.""" 

71 if bucket not in self._pending: 

72 return 

73 

74 still_pending = [] 

75 for tensor, event in self._pending[bucket]: 

76 if event.query(): 

77 self._pool.setdefault(bucket, []).append(tensor) 

78 else: 

79 still_pending.append((tensor, event)) 

80 self._pending[bucket] = still_pending 

81 

82 def acquire(self, size: int) -> torch.Tensor: 

83 """Obtain a buffer of at least *size* bytes from the pool.""" 

84 bucket = _bucket_for(size) 

85 with self._lock: 

86 for bucket_size in (bucket, *(b for b in _BUCKET_SIZES if b > bucket)): 

87 self._reclaim_locked(bucket_size) 

88 entries = self._pool.get(bucket_size) 

89 if entries: 

90 return entries.pop()[:size] 

91 

92 aligned = _align_to_bucket(size) 

93 if self._total_allocated + aligned <= self._max_host_bytes: 

94 self._total_allocated += aligned 

95 logger.debug( 

96 "PinnedMemoryPool: allocate %d bytes (total=%d, limit=%d)", 

97 aligned, 

98 self._total_allocated, 

99 self._max_host_bytes, 

100 ) 

101 return torch.empty(aligned, dtype=torch.uint8, pin_memory=True)[:size] 

102 

103 self._reclaim_locked(bucket) 

104 if bucket in self._pending and self._pending[bucket]: 

105 tensor, event = self._pending[bucket].pop(0) 

106 event.synchronize() 

107 return tensor[:size] 

108 

109 raise RuntimeError( 

110 f"PinnedMemoryPool exhausted: total_allocated={self._total_allocated}, " 

111 f"max_host_bytes={self._max_host_bytes}, requested={size}" 

112 ) 

113 

114 def release(self, tensor: torch.Tensor, event: torch.Event | None = None) -> None: 

115 """Return a previously acquired buffer to the pool for reuse.""" 

116 if not tensor.is_pinned(): 

117 raise ValueError( 

118 "release() expects a pinned (page-locked) tensor, " 

119 f"got tensor on {tensor.device} with pin_memory={tensor.is_pinned()}" 

120 ) 

121 

122 storage = tensor.untyped_storage() 

123 full_tensor = torch.empty(0, dtype=torch.uint8, device="cpu") 

124 full_tensor.set_(storage, 0, (storage.size(),), (1,)) 

125 

126 bucket = _bucket_for(full_tensor.numel()) 

127 with self._lock: 

128 if event is not None: 

129 self._pending.setdefault(bucket, []).append((full_tensor, event)) 

130 else: 

131 self._pool.setdefault(bucket, []).append(full_tensor)