xref: /aosp_15_r20/external/pytorch/torch/_logging/structured.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1"""
2Utilities for converting data types into structured JSON for dumping.
3"""
4
5import traceback
6from typing import Any, Dict, List, Sequence, Set
7
8import torch._logging._internal
9
10
11INTERN_TABLE: Dict[str, int] = {}
12
13
14DUMPED_FILES: Set[str] = set()
15
16
17def intern_string(s: str) -> int:
18    r = INTERN_TABLE.get(s, None)
19    if r is None:
20        r = len(INTERN_TABLE)
21        INTERN_TABLE[s] = r
22        torch._logging._internal.trace_structured(
23            "str", lambda: (s, r), suppress_context=True
24        )
25    return r
26
27
28def dump_file(filename: str) -> None:
29    if "eval_with_key" not in filename:
30        return
31    if filename in DUMPED_FILES:
32        return
33    DUMPED_FILES.add(filename)
34    from torch.fx.graph_module import _loader
35
36    torch._logging._internal.trace_structured(
37        "dump_file",
38        metadata_fn=lambda: {
39            "name": filename,
40        },
41        payload_fn=lambda: _loader.get_source(filename),
42    )
43
44
45def from_traceback(tb: Sequence[traceback.FrameSummary]) -> List[Dict[str, Any]]:
46    r = []
47    for frame in tb:
48        # dict naming convention here coincides with
49        # python/combined_traceback.cpp
50        r.append(
51            {
52                "line": frame.lineno,
53                "name": frame.name,
54                "filename": intern_string(frame.filename),
55            }
56        )
57    return r
58