58 lines
1.3 KiB
Python
58 lines
1.3 KiB
Python
"""
|
|
Utilities for converting data types into structured JSON for dumping.
|
|
"""
|
|
|
|
import traceback
|
|
from typing import Any, Dict, List, Sequence, Set
|
|
|
|
import torch._logging._internal
|
|
|
|
|
|
INTERN_TABLE: Dict[str, int] = {}
|
|
|
|
|
|
DUMPED_FILES: Set[str] = set()
|
|
|
|
|
|
def intern_string(s: str) -> int:
|
|
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]]:
|
|
r = []
|
|
for frame in tb:
|
|
# dict naming convention here coincides with
|
|
# python/combined_traceback.cpp
|
|
r.append(
|
|
{
|
|
"line": frame.lineno,
|
|
"name": frame.name,
|
|
"filename": intern_string(frame.filename),
|
|
}
|
|
)
|
|
return r
|