« 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"""Activation identity and lifecycle management.
17:class:`ActivationTracker` is the single source of truth for:
18* Storage identity (``storage_id`` resolution).
19* Activation policy (which tensors should be shadowed/tracked).
20* Activation storage registration (activations produced inside the trace).
22It holds **no mutable runtime state** beyond the set of trace-created
23storage IDs. Callers that need physical state transitions should use
24:class:`~offload.runtime.residency.ResidencyManager` instead.
25"""
27from __future__ import annotations
29import logging
30import weakref
32import torch
34from hyper_parallel.auto_parallel.hyper_offload.execution.tensor import ShadowTensor
36logger = logging.getLogger(__name__)
39class ActivationTracker:
40 """Identity resolution + activation lifecycle management.
42 Owns
43 ----
44 * ``_storage_tracker`` (weak dictionary) for stable ``storage_id``.
45 * ``_storage_sizes`` — mapping from storage ID to size in bytes,
46 accumulated across all recorded ops.
47 * ``_activation_sids`` — storage IDs produced inside the current trace.
49 High-level API (:meth:`get_activation_sid`, :meth:`register_op_activations`)
50 should be preferred over the low-level private helpers.
51 """
53 def __init__(self) -> None:
54 self._storage_tracker: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
55 self._next_storage_id = 1
56 self._storage_sizes: dict[int, int] = {}
57 self._activation_sids: set[int] = set()
59 # ------------------------------------------------------------------
60 # Private low-level identity API
61 # ------------------------------------------------------------------
63 def _ensure_id(self, tensor: torch.Tensor) -> int | None:
64 """Get-or-create the unique storage ID for *tensor* (private).
66 If the storage has been seen before the existing ID is returned;
67 otherwise a fresh ID is assigned and recorded.
69 Returns ``None`` for tensors without a stable storage identity
70 (e.g. meta / quantized tensors that do not support
71 :meth:`torch.Tensor.untyped_storage`).
72 """
73 try:
74 s = tensor.untyped_storage()
75 except (AttributeError, RuntimeError):
76 return None
78 try:
79 sid = self._storage_tracker[s]
80 except KeyError:
81 sid = self._next_storage_id
82 self._storage_tracker[s] = sid
83 self._next_storage_id += 1
85 self._storage_sizes.setdefault(sid, s.size())
86 logger.debug("_ensure_id: sid=%d shape=%s", sid, tensor.shape)
87 return sid
89 # ------------------------------------------------------------------
90 # Unified SID resolution
91 # ------------------------------------------------------------------
93 def get_activation_sid(self, tensor: torch.Tensor) -> int | None:
94 """Return the tracked storage ID for *tensor*, or ``None``.
96 Handles both :class:`ShadowTensor` (which carries its SID inline)
97 and raw tensors via a read-only storage lookup.
99 This is the unified replacement for ``WarmupExecutor._sid_of``.
101 Notes on the raw-tensor eligibility heuristics:
102 * CPU tensors are never activations — the trace runs on CUDA.
103 * Tensors without ``untyped_storage`` (e.g. meta device) cannot
104 have a stable storage identity.
105 """
106 if isinstance(tensor, ShadowTensor):
107 return tensor.storage_id
108 if tensor.device.type == "cpu":
109 return None
110 if not hasattr(tensor, "untyped_storage"):
111 return None
112 try:
113 s = tensor.untyped_storage()
114 except (AttributeError, RuntimeError):
115 return None
116 sid = self._storage_tracker.get(s)
117 return sid if sid is not None and sid in self._activation_sids else None
119 # ------------------------------------------------------------------
120 # Step lifecycle
121 # ------------------------------------------------------------------
123 def register_op_activations(
124 self,
125 input_tensors: list[torch.Tensor],
126 output_tensors: list[torch.Tensor],
127 ) -> None:
128 """Register new activations produced by an op.
130 Compares the storage IDs of *output_tensors* against those of
131 *input_tensors* and marks any previously unseen output storage
132 as a trace-created activation.
134 """
135 # Collect storage IDs of all op inputs.
136 input_sids = {sid for t in input_tensors if (sid := self._ensure_id(t)) is not None}
138 # Register newly-created storages that appear for the first time
139 # as op outputs, collecting their sizes.
140 for t in output_tensors:
141 sid = self._ensure_id(t)
142 if sid is not None and sid not in input_sids:
143 self._activation_sids.add(sid)
145 # ------------------------------------------------------------------
146 # Lifecycle reset
147 # ------------------------------------------------------------------
149 @property
150 def storage_sizes(self) -> dict[int, int]:
151 """Return a copy of the accumulated storage size map."""
152 return dict(self._storage_sizes)
154 def clear_activations(self) -> None:
155 """Clear the activation set (called at the start of warmup)."""
156 self._activation_sids.clear()
157 self._storage_sizes.clear()
159 def __repr__(self) -> str:
160 """Return the string representation."""
161 return f"{type(self).__name__}(activations={len(self._activation_sids)})"