Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / hyper_offload / runtime / residency.py: 86%
153 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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"""Pure physical control plane for byte-level residency.
17This module provides the data-plane (:class:`PhysicalBuffer`) and
18the control-plane (:class:`ResidencyManager`) for raw byte buffers.
19"""
21from __future__ import annotations
23import logging
24from dataclasses import dataclass
26import torch
28from hyper_parallel.auto_parallel.hyper_offload.runtime.pinned_memory import PinnedMemoryPool
30logger = logging.getLogger(__name__)
33@dataclass
34class PhysicalBuffer:
35 """Pure physical memory block tracking host/device byte buffers.
37 This dataclass is intentionally minimal — it carries **no** knowledge
38 of logical tensors or ShadowTensor objects. Every field is
39 ``None`` when the corresponding resource is not held.
41 The **control plane** (:class:`ResidencyManager`) orchestrates
42 allocations and copies.
43 """
45 device: torch.device | None = None
46 """Target accelerator device (set on first registration)."""
48 host_buffer: torch.Tensor | None = None
49 """1-D uint8 pinned CPU buffer, or ``None`` when not resident on host."""
51 host_event: torch.Event | None = None
52 """:class:`torch.Event` recorded after the latest D2H copy completes."""
54 device_buffer: torch.Tensor | None = None
55 """1-D uint8 device buffer, or ``None`` when not resident on device."""
57 device_event: torch.Event | None = None
58 """:class:`torch.Event` recorded after the latest H2D copy completes."""
60 # ------------------------------------------------------------------
61 # Device-storage access (called by ShadowTensor)
62 # ------------------------------------------------------------------
64 def device_storage(self) -> torch.UntypedStorage:
65 """Return the device storage, waiting for any in-flight H2D event.
67 If the device buffer is not resident but host data is available,
68 this will synchronously demand-page the data back to the device.
70 Returns:
71 The underlying :class:`torch.UntypedStorage` of ``device_buffer``.
73 """
74 if self.device_buffer is None:
75 if self.host_buffer is not None and self.device is not None:
76 # Synchronous demand-paging fallback.
77 # Wait for any in-flight D2H copy to complete before reading
78 # from host_buffer (race-condition avoidance).
79 if self.host_event is not None:
80 self.host_event.synchronize()
81 self.host_event = None
83 size_bytes = self.host_buffer.numel()
84 dev_bytes = torch.empty(size_bytes, dtype=torch.uint8, device=self.device)
85 dev_bytes.copy_(self.host_buffer, non_blocking=False)
86 self.device_buffer = dev_bytes
87 else:
88 raise RuntimeError(
89 "No device buffer available and cannot demand-page. "
90 "Ensure host data is available if device data is evicted."
91 )
92 if self.device_event is not None:
93 current_stream = torch.accelerator.current_stream()
94 self.device_event.wait(current_stream)
95 self.device_buffer.record_stream(current_stream)
96 self.device_event = None
97 return self.device_buffer.untyped_storage()
100class ResidencyManager:
101 """Pure physical controller for byte-level tensor residency.
103 Owns
104 ----
105 * Physical residency table (``storage_id → PhysicalBuffer``).
106 * High-level state transitions (``copy_d2h``, ``copy_h2d``).
108 All public methods accept ``storage_id: int``.
109 """
111 def __init__(
112 self,
113 max_host_bytes: int,
114 ) -> None:
115 self._copy_stream = None
116 self._host_pool = PinnedMemoryPool(max_host_bytes)
117 self._residency: dict[int, PhysicalBuffer] = {}
119 def _get_copy_stream(self) -> torch.Stream:
120 """Return the internal copy stream, creating it lazily on first access.
122 Delaying stream creation avoids requiring an accelerator device context
123 at construction time (e.g. when running on CPU-only hosts or before
124 NPU/CUDA is initialised).
125 """
126 if self._copy_stream is None:
127 self._copy_stream = torch.Stream()
128 return self._copy_stream
130 # ------------------------------------------------------------------
131 # Device-side memory query
132 # ------------------------------------------------------------------
134 @property
135 def resident_bytes(self) -> int:
136 """Total bytes currently resident on the device side across all storage IDs."""
137 total = 0
138 for buf in self._residency.values():
139 if buf.device_buffer is not None:
140 total += buf.device_buffer.numel()
141 return total
143 def device_resident_size(self, sid: int) -> int | None:
144 """Return the device buffer size in bytes, or ``None`` if not device-resident."""
145 buffer = self._residency.get(sid)
146 if buffer is None or buffer.device_buffer is None:
147 return None
148 return buffer.device_buffer.numel()
150 # ------------------------------------------------------------------
151 # Stream synchronisation
152 # ------------------------------------------------------------------
154 def wait_for_transfers(self) -> None:
155 """Make the current accelerator stream wait for pending async transfers on the copy stream."""
156 torch.accelerator.current_stream().wait_stream(self._get_copy_stream())
158 def sync_all_transfers(self) -> None:
159 """Synchronise streams."""
160 self._get_copy_stream().synchronize()
162 # ------------------------------------------------------------------
163 # Registration: bind a storage ID to a tensor's device storage
164 # ------------------------------------------------------------------
166 def bind(self, sid: int, tensor: torch.Tensor) -> PhysicalBuffer:
167 """Point the physical buffer for *sid* at *tensor*'s device storage.
169 Returns the :class:`PhysicalBuffer` so that the caller can pass
170 it to a new :class:`~offload.execution.tensor.ShadowTensor`.
171 """
172 if sid not in self._residency:
173 self._residency[sid] = PhysicalBuffer()
174 buffer = self._residency[sid]
175 buffer.device = tensor.device
177 storage = tensor.untyped_storage()
178 dev_view = torch.empty(0, dtype=torch.uint8, device=tensor.device)
179 dev_view.set_(storage, 0, (storage.size(),), (1,))
180 buffer.device_buffer = dev_view
181 return buffer
183 # ------------------------------------------------------------------
184 # State transition: copy D2H
185 # ------------------------------------------------------------------
187 def copy_d2h(self, sid: int) -> None:
188 """Copy the physical storage for *sid* from device to host.
190 1. Look up the physical buffer for ``sid``.
191 2. If ``host_buffer`` is already present → no-op.
192 3. Launch an async D2H copy.
193 4. Keep the device buffer resident until ``release_device``.
194 """
195 buffer = self._residency.get(sid)
196 if buffer is None:
197 raise RuntimeError(
198 f"copy_d2h sid={sid}: no physical buffer registered"
199 )
200 if buffer.host_buffer is not None:
201 logger.debug("copy_d2h sid=%d: already on host, skip", sid)
202 return
203 if buffer.device_buffer is None:
204 raise RuntimeError(
205 f"copy_d2h sid={sid}: no device data to copy"
206 )
208 dev_src = buffer.device_buffer
209 size_bytes = dev_src.numel()
210 logger.debug(
211 "copy_d2h sid=%d: copying %d bytes (%.2f MiB) D2H",
212 sid,
213 size_bytes,
214 size_bytes / 1024**2,
215 )
217 # 1. Allocate pinned host buffer.
218 host_buf = self._host_pool.acquire(size_bytes)
220 # 2. Prevent the caching allocator from recycling the source
221 # memory while the copy stream reads it.
222 copy_stream = self._get_copy_stream()
223 if dev_src.device == copy_stream.device:
224 dev_src.record_stream(copy_stream)
226 # 3. Launch asynchronous D2H copy.
227 event = None
228 if dev_src.device != copy_stream.device:
229 host_buf.copy_(dev_src, non_blocking=False)
230 else:
231 producer_stream = torch.accelerator.current_stream()
232 event = torch.Event()
233 with copy_stream:
234 copy_stream.wait_stream(producer_stream)
235 host_buf.copy_(dev_src, non_blocking=True)
236 event.record(copy_stream)
238 # 4. Update physical buffer.
239 buffer.host_buffer = host_buf
240 buffer.host_event = event
242 logger.debug("copy_d2h sid=%d: done", sid)
244 # ------------------------------------------------------------------
245 # State transition: copy H2D
246 # ------------------------------------------------------------------
248 def copy_h2d(self, sid: int) -> None:
249 """Asynchronously copy (H2D) the physical storage for *sid* to device.
251 Allocates fresh device memory, launches an async H2D copy on the
252 copy stream, and updates the physical buffer with the
253 new ``device_buffer`` and ``device_event``. Returns immediately
254 without waiting for the copy to complete.
255 """
256 buffer = self._residency.get(sid)
257 if buffer is None:
258 raise RuntimeError(
259 f"copy_h2d sid={sid}: no physical buffer registered"
260 )
262 if buffer.device_buffer is not None:
263 logger.debug("copy_h2d sid=%d: already on device, skip", sid)
264 return
266 if buffer.host_buffer is None:
267 raise RuntimeError(
268 f"copy_h2d sid={sid}: no host data to copy"
269 )
271 if buffer.device is None:
272 raise RuntimeError(
273 f"copy_h2d sid={sid}: target device unknown"
274 )
276 size_bytes = buffer.host_buffer.numel()
277 logger.debug(
278 "copy_h2d sid=%d: copying %d bytes (%.2f MiB) H2D to %s",
279 sid,
280 size_bytes,
281 size_bytes / 1024**2,
282 buffer.device,
283 )
285 # 1. Allocate device memory.
286 dev_bytes = torch.empty(size_bytes, dtype=torch.uint8, device=buffer.device)
288 # 2. Launch async H2D copy (wait for prior D2H event if present).
289 copy_stream = self._get_copy_stream()
290 event = None
291 if dev_bytes.device != copy_stream.device:
292 dev_bytes.copy_(buffer.host_buffer, non_blocking=False)
293 else:
294 producer_stream = torch.accelerator.current_stream()
295 event = torch.Event()
296 with copy_stream:
297 copy_stream.wait_stream(producer_stream)
298 if buffer.host_event is not None:
299 buffer.host_event.wait(copy_stream)
300 dev_bytes.copy_(buffer.host_buffer, non_blocking=True)
301 event.record(copy_stream)
303 # 3. Update physical buffer.
304 if dev_bytes.device == copy_stream.device:
305 dev_bytes.record_stream(copy_stream)
306 buffer.device_buffer = dev_bytes
307 buffer.device_event = event
309 # ------------------------------------------------------------------
310 # Release helpers
311 # ------------------------------------------------------------------
313 def release_device(self, sid: int) -> None:
314 """Release device-resident bytes for a storage ID.
316 Frees the device buffer. If an H2D prefetch is still in flight,
317 waits for it before dropping the destination buffer reference.
318 Pending D2H copies are protected by ``record_stream`` during
319 ``copy_d2h``.
320 """
321 buffer = self._residency.get(sid)
322 if buffer is None or buffer.device_buffer is None:
323 return
325 if buffer.device_event is not None:
326 buffer.device_event.synchronize()
327 buffer.device_event = None
329 buffer.device_buffer = None
330 if buffer.host_buffer is None:
331 del self._residency[sid]
333 def release_host(self, sid: int) -> None:
334 """Release host-resident bytes for a storage ID.
336 Waits for any in-flight H2D copy that may be reading from the
337 host buffer before returning it to the pool.
338 """
339 buffer = self._residency.get(sid)
340 if buffer is None or buffer.host_buffer is None:
341 return
343 # Ensure any in-flight H2D copy that reads this host buffer
344 # has completed before we return it to the pool.
345 event_to_wait = buffer.device_event if buffer.device_event is not None else buffer.host_event
346 self._host_pool.release(buffer.host_buffer, event=event_to_wait)
347 buffer.host_buffer = None
348 buffer.host_event = None
349 if buffer.device_buffer is None:
350 del self._residency[sid]
352 # ------------------------------------------------------------------
353 # Runtime clear
354 # ------------------------------------------------------------------
356 def clear_runtime(self) -> None:
357 """Release all physical resources and reset tracking."""
358 for buffer in self._residency.values():
359 if buffer.host_buffer is not None:
360 self._host_pool.release(buffer.host_buffer, event=buffer.host_event)
361 self._residency.clear()