1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
|
# mypy: ignore-errors
import torch
from torch.utils._pytree import tree_map
from typing import Iterator, List, Optional
import logging
import contextlib
import itertools
from torch.utils._python_dispatch import TorchDispatchMode
from torch.utils.weak import WeakTensorKeyDictionary
import functools
from torch._C._profiler import gather_traceback, symbolize_tracebacks
logger = logging.getLogger("LoggingTensor")
_dtype_abbrs = {
torch.bfloat16: "bf16",
torch.float64: "f64",
torch.float32: "f32",
torch.float16: "f16",
torch.complex32: "c32",
torch.complex64: "c64",
torch.complex128: "c128",
torch.int8: "i8",
torch.int16: "i16",
torch.int32: "i32",
torch.int64: "i64",
torch.bool: "b8",
torch.uint8: "u8",
torch.float8_e4m3fn: "f8e4m3fn",
torch.float8_e5m2: "f8e5m2",
torch.float8_e4m3fnuz: "f8e4m3fnuz",
torch.float8_e5m2fnuz: "f8e5m2fnuz",
}
# How the chain of calls works for LoggingTensor:
# 1. Call torch.sin
# 2. Attempt __torch_function__. In LoggingTensor torch function is disabled so we bypass it entirely
# 3. Enter dispatcher, wind your way through Autograd
# 4. Hit Python dispatch key, call __torch_dispatch__
# This Tensor can work with autograd in two ways:
# - The wrapped Tensor does not require gradients. In that case, the LoggingTensor
# can require gradients if the user asks for it as a constructor kwarg.
# - The wrapped Tensor can require gradients. In that case autograd will be tracked
# for the wrapped Tensor and the LoggingTensor itself cannot require gradients.
# WARNING: We allow these two possibilities for testing purposes. You should NEVER use both in a single
# test or you might get surprising behavior.
# TODO: TensorBase should work
class LoggingTensor(torch.Tensor):
elem: torch.Tensor
__slots__ = ['elem']
context = contextlib.nullcontext
@staticmethod
def __new__(cls, elem, *args, **kwargs):
# The wrapping tensor (LoggingTensor) shouldn't hold any
# memory for the class in question, but it should still
# advertise the same device as before
r = torch.Tensor._make_wrapper_subclass( # type: ignore[attr-defined]
cls, elem.size(),
strides=elem.stride(), storage_offset=elem.storage_offset(),
# TODO: clone storage aliasing
dtype=elem.dtype, layout=elem.layout,
device=elem.device, requires_grad=kwargs.get("requires_grad", False)
)
# ...the real tensor is held as an element on the tensor.
r.elem = elem.detach() if r.requires_grad else elem
return r
def __repr__(self):
return super().__repr__(tensor_contents=f"{self.elem}")
@classmethod
def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
def unwrap(e):
return e.elem if isinstance(e, cls) else e
def wrap(e):
return cls(e) if isinstance(e, torch.Tensor) else e
with cls.context():
rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs)))
logging.getLogger("LoggingTensor").info(f"{func.__module__}.{func.__name__}", args, kwargs, rs) # noqa: G004
return rs
class LoggingTensorMode(TorchDispatchMode):
def __torch_dispatch__(self, func, types, args=(), kwargs=None):
if kwargs is None:
kwargs = {}
rs = func(*args, **kwargs)
logging.getLogger("LoggingTensor").info(f"{func.__module__}.{func.__name__}", args, kwargs, rs) # noqa: G004
return rs
class LoggingTensorReentrant(LoggingTensor):
context = torch.overrides.enable_reentrant_dispatch
# https://stackoverflow.com/questions/36408496/python-logging-handler-to-append-to-list
class LoggingTensorHandler(logging.Handler):
def __init__(
self, log_list: List[str], use_shortid_for_all_tensors: bool,
with_type: bool, tracebacks_list: Optional[List]) -> None:
logging.Handler.__init__(self)
self.log_list = log_list
self.use_shortid_for_all_tensors = use_shortid_for_all_tensors
self.tracebacks_list = tracebacks_list
self.memo = WeakTensorKeyDictionary()
self.next_id = 0
self.with_type = with_type
def _shortid(self, t: torch.Tensor) -> int:
if t not in self.memo:
self.memo[t] = self.next_id
self.next_id += 1
return self.memo[t]
def _fmt(self, a: object, with_type: bool = False) -> str:
cond_cls = torch.Tensor if self.use_shortid_for_all_tensors else LoggingTensor
if isinstance(a, cond_cls):
maybe_type = ""
if with_type and self.with_type:
maybe_type = f": {_dtype_abbrs[a.dtype]}[{', '.join(map(str, a.shape))}]"
x = f"${self._shortid(a)}{maybe_type}"
return x
else:
return repr(a)
def emit(self, record):
fmt_args = ", ".join(
itertools.chain(
(str(tree_map(self._fmt, a)) for a in record.args[0]),
(f"{k}={str(tree_map(self._fmt, v))}" for k, v in record.args[1].items()),
)
)
fmt_rets = tree_map(functools.partial(self._fmt, with_type=True), record.args[2])
self.log_list.append(f'{fmt_rets} = {record.msg}({fmt_args})')
if self.tracebacks_list is not None:
self.tracebacks_list.append(record.traceback)
def log_input(name: str, var: object) -> None:
logger.info("input", (name,), {}, var) # noqa: PLE1205
class GatherTraceback(logging.Filter):
def __init__(self, python=True, script=True, cpp=False):
self.python = python
self.script = script
self.cpp = cpp
def filter(self, record):
record.traceback = gather_traceback(python=self.python, script=self.script, cpp=self.cpp)
return True
@contextlib.contextmanager
def capture_logs(is_mode=False, python_tb=False, script_tb=False, cpp_tb=False) -> Iterator[List[str]]:
collect_traceback = python_tb or script_tb or cpp_tb
log_list: List[str] = []
tracebacks_list: List[str] = []
handler = LoggingTensorHandler(
log_list,
with_type=True,
use_shortid_for_all_tensors=is_mode,
tracebacks_list=tracebacks_list if collect_traceback else None
)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
logger.propagate = False
if collect_traceback:
logger.addFilter(GatherTraceback(python=python_tb, script=script_tb, cpp=cpp_tb))
try:
if collect_traceback:
yield log_list, tracebacks_list
else:
yield log_list
finally:
symbolized_tracebacks = symbolize_tracebacks(tracebacks_list)
tracebacks_list.clear()
tracebacks_list.extend(symbolized_tracebacks)
logger.removeHandler(handler)
@contextlib.contextmanager
def capture_logs_with_logging_tensor_mode(python_tb=False, script_tb=False, cpp_tb=False):
with LoggingTensorMode(), capture_logs(True, python_tb, script_tb, cpp_tb) as logs:
yield logs
|
