« 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"""Warmup-phase executor.
17During warmup, the executor records the execution trace while applying
18online greedy eviction to keep device memory within budget. Every
19input tensor that was evicted in an earlier op is faulted back
20synchronously (demand-paging) before dispatch.
22Physical state transitions are delegated to
23:class:`~offload.runtime.residency.ResidencyManager`.
24"""
26from __future__ import annotations
28import logging
29from typing import Any
31import torch
32from torch.utils._pytree import tree_flatten
34from hyper_parallel.auto_parallel.hyper_offload.execution.base import BaseExecutor
35from hyper_parallel.auto_parallel.hyper_offload.runtime.timer import DeviceTimer
36from hyper_parallel.auto_parallel.hyper_offload.execution.warmup.tracker import ActivationTracker
37from hyper_parallel.auto_parallel.hyper_offload.ir.replay import OpGuide
38from hyper_parallel.auto_parallel.hyper_offload.ir.trace import AccessKind, ActivationTrace, StorageAccess, TraceOp
39from hyper_parallel.auto_parallel.hyper_offload.runtime.bandwidth import profile_transfer_bandwidth
40from hyper_parallel.auto_parallel.hyper_offload.runtime.residency import ResidencyManager
42logger = logging.getLogger(__name__)
45def iter_tensors(value: Any) -> list[torch.Tensor]:
46 """Return tensor leaves from an arbitrary pytree."""
47 leaves, _ = tree_flatten(value)
48 return [leaf for leaf in leaves if isinstance(leaf, torch.Tensor)]
51class WarmupExecutor(BaseExecutor):
52 """Executor for the warmup phase.
54 Records the execution trace while applying online greedy eviction
55 to keep device memory within the configured budget. Evicts the
56 **oldest** activations first when the memory budget is exceeded
57 (oldest-first within the same op, largest-sized entries are
58 preferred as tie-breaker to minimise eviction count).
59 """
61 def __init__(
62 self,
63 residency_manager: ResidencyManager,
64 memory_limit_bytes: int,
65 ) -> None:
66 super().__init__(residency_manager)
67 self._memory_limit_bytes = memory_limit_bytes
68 self._tracker = ActivationTracker()
69 self._timer = DeviceTimer()
70 self._guide: list[OpGuide] = []
71 self._ops: list[TraceOp] = []
72 #: sid -> the op index that first produced this activation.
73 self._sid_produced_at_op: dict[int, int] = {}
75 # ------------------------------------------------------------------
76 # Eviction policy
77 # ------------------------------------------------------------------
79 def _enforce_budget(self, protected_sids: set[int]) -> None:
80 """Evict warmup activations until resident bytes fit the configured budget."""
81 while self.residency_manager.resident_bytes > self._memory_limit_bytes:
82 # Greedy: evict oldest first; within the same op, evict largest first.
83 victim_sid: int | None = None
84 victim_key: tuple[int, int] | None = None
86 for sid, produced_at_op in self._sid_produced_at_op.items():
87 if sid in protected_sids:
88 continue
89 size = self.residency_manager.device_resident_size(sid)
90 if size is None:
91 continue
92 key = (produced_at_op, -size)
93 if victim_key is None or key < victim_key:
94 victim_key = key
95 victim_sid = sid
97 if victim_sid is None:
98 raise RuntimeError(
99 "Warmup memory budget exceeded but no evictable activation found. "
100 f"resident_bytes={self.residency_manager.resident_bytes}, "
101 f"limit={self._memory_limit_bytes}, "
102 f"protected_sids={protected_sids}"
103 )
105 self.residency_manager.copy_d2h(victim_sid)
106 self.residency_manager.release_device(victim_sid)
108 # ------------------------------------------------------------------
109 # Lifecycle hooks
110 # ------------------------------------------------------------------
112 def on_op_begin(self, func, args, kwargs) -> None:
113 """Before op: enforce memory budget and fault inputs back to device."""
114 super().on_op_begin(func, args, kwargs)
116 protected_sids: set[int] = set()
117 for t in iter_tensors((args, kwargs)):
118 if (sid := self._tracker.get_activation_sid(t)) is not None:
119 protected_sids.add(sid)
120 self._enforce_budget(protected_sids)
122 self._timer.start()
124 def on_op_end(self, result) -> Any:
125 """After op: record trace, residency metadata, and return shadowed result."""
126 op_duration_ms = self._timer.stop()
128 func, args, kwargs = self._last_func, self._last_args, self._last_kwargs
130 self._tracker.register_op_activations(iter_tensors((args, kwargs)), iter_tensors(result))
132 op = TraceOp(name=func.__name__, duration_ms=op_duration_ms)
134 # --- Detect mutated (write) input tensors via func._schema ---
135 mutated_tensor_ids: set[int] = set()
136 if hasattr(func, "_schema") and func._schema.is_mutable: # pylint: disable=protected-access
137 flat_args = iter_tensors((args, {}))
138 for idx, arg_info in enumerate(func._schema.arguments): # pylint: disable=protected-access
139 if arg_info.alias_info is not None and arg_info.alias_info.is_write and idx < len(flat_args): # pylint: disable=protected-access
140 mutated_tensor_ids.add(id(flat_args[idx]))
142 seen: set[tuple[int, AccessKind]] = set()
144 # --- Input accesses ---
145 for t in iter_tensors((args, kwargs)):
146 sid = self._tracker.get_activation_sid(t)
147 if sid is None:
148 continue
149 is_mutated = id(t) in mutated_tensor_ids
150 kind = AccessKind.WRITE if is_mutated else AccessKind.READ
151 if (sid, kind) not in seen:
152 seen.add((sid, kind))
153 op.accesses.append(StorageAccess(self.op_idx, sid, kind))
155 # --- Output accesses + bindings ---
156 leaves, _ = tree_flatten(result)
157 output_bindings: dict[int, int] = {}
159 for leaf_index, t in enumerate(leaves):
160 if not isinstance(t, torch.Tensor):
161 continue
162 sid = self._tracker.get_activation_sid(t)
163 if sid is None:
164 continue
166 if (sid, AccessKind.WRITE) not in seen:
167 seen.add((sid, AccessKind.WRITE))
168 op.accesses.append(StorageAccess(self.op_idx, sid, AccessKind.WRITE))
170 # Track the op that first produced this activation (used by
171 # eviction policy: oldest-first).
172 if sid not in self._sid_produced_at_op:
173 self._sid_produced_at_op[sid] = self.op_idx
175 output_bindings[leaf_index] = sid
177 guide = OpGuide(
178 name=func.__name__,
179 output_leaf_count=len(leaves),
180 output_bindings=output_bindings,
181 )
183 self._ops.append(op)
184 self._guide.append(guide)
186 return self.apply_shadows(result, output_bindings)
188 def finish(self) -> tuple[ActivationTrace, list[OpGuide]]:
189 """Finish warmup and return recorded trace and replay guide."""
190 d2h, h2d = profile_transfer_bandwidth()
192 trace = ActivationTrace(
193 ops=list(self._ops),
194 storage_sizes=self._tracker.storage_sizes,
195 retained_sids=set(self.retained_sids),
196 memory_limit_bytes=self._memory_limit_bytes,
197 d2h_bandwidth_gbps=d2h,
198 h2d_bandwidth_gbps=h2d,
199 )
200 guide = self._guide
202 self.reset()
203 return trace, guide
205 def reset(self) -> None:
206 """Reset."""
207 self._sid_produced_at_op.clear()
208 self._ops = []
209 self._guide = []
210 self._tracker.clear_activations()
211 super().reset()