File size: 2,927 Bytes
9c6594c |
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 |
"""
Utilities for converting data types into structured JSON for dumping.
"""
import inspect
import os
import traceback
from collections.abc import Sequence
from typing import Any, Optional
import torch._logging._internal
INTERN_TABLE: dict[str, int] = {}
DUMPED_FILES: set[str] = set()
def intern_string(s: Optional[str]) -> int:
if s is None:
return -1
r = INTERN_TABLE.get(s, None)
if r is None:
r = len(INTERN_TABLE)
INTERN_TABLE[s] = r
torch._logging._internal.trace_structured(
"str", lambda: (s, r), suppress_context=True
)
return r
def dump_file(filename: str) -> None:
if "eval_with_key" not in filename:
return
if filename in DUMPED_FILES:
return
DUMPED_FILES.add(filename)
from torch.fx.graph_module import _loader
torch._logging._internal.trace_structured(
"dump_file",
metadata_fn=lambda: {
"name": filename,
},
payload_fn=lambda: _loader.get_source(filename),
)
def from_traceback(tb: Sequence[traceback.FrameSummary]) -> list[dict[str, Any]]:
# dict naming convention here coincides with
# python/combined_traceback.cpp
r = [
{
"line": frame.lineno,
"name": frame.name,
"filename": intern_string(frame.filename),
"loc": frame.line,
}
for frame in tb
]
return r
def get_user_stack(num_frames: int) -> list[dict[str, Any]]:
from torch._guards import TracingContext
from torch.utils._traceback import CapturedTraceback
user_tb = TracingContext.extract_stack()
if user_tb:
return from_traceback(user_tb[-1 * num_frames :])
tb = CapturedTraceback.extract().summary()
# Filter out frames that are within the torch/ codebase
torch_filepath = os.path.dirname(inspect.getfile(torch)) + os.path.sep
for i, frame in enumerate(reversed(tb)):
if torch_filepath not in frame.filename:
# Only display `num_frames` frames in the traceback
filtered_tb = tb[len(tb) - i - num_frames : len(tb) - i]
return from_traceback(filtered_tb)
return from_traceback(tb[-1 * num_frames :])
def get_framework_stack(
num_frames: int = 25, cpp: bool = False
) -> list[dict[str, Any]]:
"""
Returns the traceback for the user stack and the framework stack
"""
from torch.fx.experimental.symbolic_shapes import uninteresting_files
from torch.utils._traceback import CapturedTraceback
tb = CapturedTraceback.extract(cpp=cpp).summary()
tb = [
frame
for frame in tb
if (
(
frame.filename.endswith(".py")
and frame.filename not in uninteresting_files()
)
or ("at::" in frame.name or "torch::" in frame.name)
)
]
return from_traceback(tb[-1 * num_frames :])
|