255 lines
10 KiB
Python
255 lines
10 KiB
Python
|
from types import TracebackType
|
||
|
|
from typing import List, Optional
|
||
|
|
import tempfile
|
||
|
|
import traceback
|
||
|
|
import contextlib
|
||
|
|
import inspect
|
||
|
|
import os.path
|
||
|
|
|
||
|
|
# This file contains utilities for ensuring dynamically compile()'d
|
||
|
|
# code fragments display their line numbers in backtraces.
|
||
|
|
#
|
||
|
|
# The constraints:
|
||
|
|
#
|
||
|
|
# - We don't have control over the user exception printer (in particular,
|
||
|
|
# we cannot assume the linecache trick will work, c.f.
|
||
|
|
# https://stackoverflow.com/q/50515651/23845 )
|
||
|
|
#
|
||
|
|
# - We don't want to create temporary files every time we compile()
|
||
|
|
# some code; file creation should happen lazily only at exception
|
||
|
|
# time. Arguably, you *should* be willing to write out your
|
||
|
|
# generated Python code to file system, but in some situations
|
||
|
|
# (esp. library code) it would violate user expectation to write
|
||
|
|
# to the file system, so we try to avoid it. In particular, we'd
|
||
|
|
# like to keep the files around, so users can open up the files
|
||
|
|
# mentioned in the trace; if the file is invisible, we want to
|
||
|
|
# avoid clogging up the filesystem.
|
||
|
|
#
|
||
|
|
# If this is not a constraint for you, there is a substantially simpler
|
||
|
|
# way to implement the functionality in this PR: instead of using
|
||
|
|
# eval/exec directly, just always write a Python file to filesystem
|
||
|
|
# and compile that.
|
||
|
|
#
|
||
|
|
# - You have control over a context where the compiled code will get
|
||
|
|
# executed, so that we can interpose while the stack is unwinding
|
||
|
|
# (otherwise, we have no way to interpose on the exception printing
|
||
|
|
# process.)
|
||
|
|
#
|
||
|
|
# There are two things you have to do to make use of the utilities here:
|
||
|
|
#
|
||
|
|
# - When you compile your source code, you must save its string source
|
||
|
|
# in its f_globals under the magic name "__compile_source__"
|
||
|
|
#
|
||
|
|
# - Before running the compiled code, enter the
|
||
|
|
# report_compile_source_on_error() context manager.
|
||
|
|
|
||
|
|
@contextlib.contextmanager
|
||
|
|
def report_compile_source_on_error():
|
||
|
|
try:
|
||
|
|
yield
|
||
|
|
except Exception as exc:
|
||
|
|
tb = exc.__traceback__
|
||
|
|
|
||
|
|
# Walk the traceback, looking for frames that have
|
||
|
|
# source attached
|
||
|
|
stack = []
|
||
|
|
while tb is not None:
|
||
|
|
filename = tb.tb_frame.f_code.co_filename
|
||
|
|
source = tb.tb_frame.f_globals.get("__compile_source__")
|
||
|
|
|
||
|
|
if filename == "<string>" and source is not None:
|
||
|
|
# What black magic are we doing here? Intuitively, what
|
||
|
|
# we would like to do is overwrite the co_filename on any
|
||
|
|
# frames that were generated from exec/eval so that they
|
||
|
|
# point to a temporary file that has the actual line
|
||
|
|
# information, so Python's default error printer can print
|
||
|
|
# useful line information on it.
|
||
|
|
#
|
||
|
|
# Writing out the temporary file is easy. But overwriting
|
||
|
|
# co_filename is not! You can't modify the code object
|
||
|
|
# associated with a frame. You can, however, reconstruct
|
||
|
|
# a traceback with entirely new frames from scratch, so that's
|
||
|
|
# what we do. But there's another problem, which is how to
|
||
|
|
# make the frame?
|
||
|
|
#
|
||
|
|
# The black magic is we make a frankenstein frame and code
|
||
|
|
# object which resembles the original frame/code enough so
|
||
|
|
# that it will print properly under traceback and the default
|
||
|
|
# error printer, but IT IS NOT THE ORIGINAL FRAME (you
|
||
|
|
# couldn't, e.g., execute its code with different variables
|
||
|
|
# and expect it to work.)
|
||
|
|
|
||
|
|
# Don't delete the temporary file so the user can inspect it
|
||
|
|
# TODO: This creates a temporary file for every frame, but we
|
||
|
|
# technically only need one per distinct __compile_source__
|
||
|
|
with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix=".py") as f:
|
||
|
|
f.write(source)
|
||
|
|
# Create a frame. Python doesn't let you construct
|
||
|
|
# FrameType directly, so just make one with compile
|
||
|
|
frame = tb.tb_frame
|
||
|
|
code = compile('__inspect_currentframe()', f.name, 'eval')
|
||
|
|
code = code.replace(co_name=frame.f_code.co_name)
|
||
|
|
# Python 3.11 only
|
||
|
|
if hasattr(frame.f_code, 'co_linetable'):
|
||
|
|
# We can't copy ALL of the metadata over, because you
|
||
|
|
# can cause Python to segfault this way. What exactly
|
||
|
|
# do we need? We need enough information for
|
||
|
|
# traceback to be able to print the exception
|
||
|
|
# correctly. Code reading Lib/traceback.py reveals
|
||
|
|
# that traceback calls code.co_positions() in order to
|
||
|
|
# get the augmented line/col numbers. Objects/codeobject.c,
|
||
|
|
# specifically _PyCode_InitAddressRange, reveals that
|
||
|
|
# this iterator is initialized from co_linetable and
|
||
|
|
# co_firstfileno. So copy these we must!
|
||
|
|
code = code.replace( # type: ignore[call-arg]
|
||
|
|
co_linetable=frame.f_code.co_linetable, # type: ignore[attr-defined]
|
||
|
|
co_firstlineno=frame.f_code.co_firstlineno, # type: ignore[attr-defined]
|
||
|
|
)
|
||
|
|
fake_frame = eval(
|
||
|
|
code,
|
||
|
|
frame.f_globals,
|
||
|
|
{
|
||
|
|
**frame.f_locals,
|
||
|
|
'__inspect_currentframe': inspect.currentframe
|
||
|
|
}
|
||
|
|
)
|
||
|
|
fake_tb = TracebackType(
|
||
|
|
None, fake_frame, tb.tb_lasti, tb.tb_lineno
|
||
|
|
)
|
||
|
|
stack.append(fake_tb)
|
||
|
|
else:
|
||
|
|
stack.append(tb)
|
||
|
|
|
||
|
|
tb = tb.tb_next
|
||
|
|
|
||
|
|
# Reconstruct the linked list
|
||
|
|
tb_next = None
|
||
|
|
for tb in reversed(stack):
|
||
|
|
tb.tb_next = tb_next
|
||
|
|
tb_next = tb
|
||
|
|
|
||
|
|
raise exc.with_traceback(tb_next) # noqa: TRY200
|
||
|
|
|
||
|
|
def shorten_filename(fn, *, base=None):
|
||
|
|
"""Shorten a source filepath, with the assumption that torch/ subdirectories don't need to be shown to user."""
|
||
|
|
if base is None:
|
||
|
|
base = os.path.dirname(os.path.dirname(__file__))
|
||
|
|
# Truncate torch/foo.py to foo.py
|
||
|
|
try:
|
||
|
|
prefix = os.path.commonpath([fn, base])
|
||
|
|
except ValueError:
|
||
|
|
return fn
|
||
|
|
else:
|
||
|
|
return fn[len(prefix) + 1:]
|
||
|
|
|
||
|
|
def format_frame(frame, *, base=None, line=False):
|
||
|
|
"""
|
||
|
|
Format a FrameSummary in a short way, without printing full absolute path or code.
|
||
|
|
|
||
|
|
The idea is the result fits on a single line.
|
||
|
|
"""
|
||
|
|
extra_line = ""
|
||
|
|
if line:
|
||
|
|
extra_line = f"{frame.line} # "
|
||
|
|
return f"{extra_line}{shorten_filename(frame.filename, base=base)}:{frame.lineno} in {frame.name}"
|
||
|
|
|
||
|
|
def format_traceback_short(tb):
|
||
|
|
"""Format a TracebackType in a short way, printing only the inner-most frame."""
|
||
|
|
return format_frame(traceback.extract_tb(tb)[-1])
|
||
|
|
|
||
|
|
class CapturedTraceback:
|
||
|
|
__slots__ = ['tb', 'skip']
|
||
|
|
|
||
|
|
def __init__(self, tb, skip=0):
|
||
|
|
self.tb = tb
|
||
|
|
self.skip = skip
|
||
|
|
|
||
|
|
def cleanup(self):
|
||
|
|
self.tb = None
|
||
|
|
|
||
|
|
def summary(self):
|
||
|
|
import torch._C._profiler
|
||
|
|
|
||
|
|
if self.tb is None:
|
||
|
|
# TODO: Maybe indicate that the traceback was elided?
|
||
|
|
return traceback.StackSummary()
|
||
|
|
|
||
|
|
return _extract_symbolized_tb(
|
||
|
|
torch._C._profiler.symbolize_tracebacks([self.tb])[0],
|
||
|
|
self.skip
|
||
|
|
)
|
||
|
|
|
||
|
|
def __getstate__(self):
|
||
|
|
return (None, {
|
||
|
|
'tb': None, # TB is not pickleable
|
||
|
|
'skip': self.skip,
|
||
|
|
})
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def extract(*, script=False, cpp=False, skip=0):
|
||
|
|
"""
|
||
|
|
Like traceback.extract_stack(), but faster (approximately 20x faster); it
|
||
|
|
is fast enough that you can unconditionally log stacks this way as part of
|
||
|
|
normal execution. It returns a torch._C._profiler.CapturedTraceback
|
||
|
|
object that must be formatted specially with format_captured_tb.
|
||
|
|
|
||
|
|
By default, this only reports Python backtraces (like extract_stack). You
|
||
|
|
can set the script/cpp kwargs to also turn on TorchScript/C++ trace
|
||
|
|
reporting.
|
||
|
|
"""
|
||
|
|
import torch._C._profiler
|
||
|
|
|
||
|
|
if script or cpp:
|
||
|
|
assert skip == 0, "skip with script/cpp NYI"
|
||
|
|
|
||
|
|
return CapturedTraceback(
|
||
|
|
torch._C._profiler.gather_traceback(python=True, script=script, cpp=cpp),
|
||
|
|
# Elide extract() frame if we don't have script/cpp frames. If
|
||
|
|
# we do have those frames, it doesn't work so force zero.
|
||
|
|
0 if script or cpp else skip + 1
|
||
|
|
)
|
||
|
|
|
||
|
|
def format(self):
|
||
|
|
"""
|
||
|
|
Formats a single torch._C._profiler.CapturedTraceback into a list of
|
||
|
|
strings equivalent to the output of traceback.format_list. Note that if
|
||
|
|
pass it CapturedTraceback with C++ traces, it is better not to use this
|
||
|
|
function and use the batch formatting API format_captured_tbs to amortize
|
||
|
|
the cost of symbolization
|
||
|
|
"""
|
||
|
|
return traceback.format_list(self.summary())
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def format_all(tbs):
|
||
|
|
"""
|
||
|
|
Bulk version of CapturedTraceback.format. Returns a list of list of strings.
|
||
|
|
"""
|
||
|
|
import torch._C._profiler
|
||
|
|
|
||
|
|
# Directly populate tracebacks that already have cached summaries
|
||
|
|
rs: List[Optional[List[str]]] = []
|
||
|
|
delayed_idxs = []
|
||
|
|
for i, tb in enumerate(tbs):
|
||
|
|
if tb.tb is None:
|
||
|
|
rs.append([])
|
||
|
|
else:
|
||
|
|
rs.append(None)
|
||
|
|
delayed_idxs.append(i)
|
||
|
|
|
||
|
|
stbs = torch._C._profiler.symbolize_tracebacks([tbs[i].tb for i in delayed_idxs])
|
||
|
|
for i, stb in zip(delayed_idxs, stbs):
|
||
|
|
rs[i] = traceback.format_list(tbs[i].summary())
|
||
|
|
|
||
|
|
return rs
|
||
|
|
|
||
|
|
|
||
|
|
def _extract_symbolized_tb(tb, skip):
|
||
|
|
"""
|
||
|
|
Given a symbolized traceback from symbolize_tracebacks, return a StackSummary object of
|
||
|
|
pre-processed stack trace entries.
|
||
|
|
"""
|
||
|
|
stack = traceback.StackSummary()
|
||
|
|
for f in reversed(tb[skip:]):
|
||
|
|
stack.append(traceback.FrameSummary(f['filename'], f['line'], f['name']))
|
||
|
|
return stack
|