wandb

1
#!/usr/bin/env python
2
"""Parse tracelog output for analysis/diagrams.
3

4
NOTE: tracelog is still in development.
5

6
Usage:
7
    ./wandb/tools/tracelog-tool.py
8
    ./wandb/tools/tracelog-tool.py --logdir logdir/
9
    ./wandb/tools/tracelog-tool.py --format plantuml
10
"""
11

12
import argparse
13
import pathlib
14
import sys
15
from dataclasses import dataclass
16
from typing import List
17

18

19
@dataclass
20
class SequenceItem:
21
    ts: float
22
    src: str
23
    request: bool
24
    dst: str
25
    info: str
26

27

28
class TracelogParser:
29
    def __init__(self) -> None:
30
        self._items: List[SequenceItem] = []
31
        self._uuid_messages = dict()
32

33
    def _parse(self, line: str) -> None:
34
        line = line.strip()
35
        index = line.find("TRACELOG(")
36
        if index < 0:
37
            return
38
        line = line[index:]
39
        items = line.split()
40
        if len(items) != 10:
41
            return
42
        # ['TRACELOG(1)', '<-', '185542.522061', 'fd1e0e9f4d3f3520', 'dequeue', 'result_q', 'MsgRouterThr', 'poll_exit_response', '69aed18a893a49d182c7a13b498f805f', '-']
43
        magic, direct, ts, msg_id, op, resource, thr, msg, uuid, stream = items
44
        self._uuid_messages.setdefault(uuid, msg)
45
        if magic != "TRACELOG(1)":
46
            return
47
        thr = thr.replace("-", "_")
48
        if op == "queue":
49
            src = thr
50
            dst = resource
51
        elif op == "dequeue":
52
            dst = thr
53
            src = resource
54
        else:
55
            # TODO: handle this
56
            return
57
        request = True
58
        if direct == "<-":
59
            request = False
60
        ts = float(ts)
61
        if msg == "None":
62
            msg = "return_" + self._uuid_messages.get(uuid)
63
        item = SequenceItem(ts=ts, src=src, request=request, dst=dst, info=msg)
64
        self.add(item)
65

66
    def add(self, item: SequenceItem):
67
        self._items.append(item)
68

69
    def output_plantuml(self) -> None:
70
        lines = []
71
        for item in self._items:
72
            line = f"{item.src} --> {item.dst}: {item.info}"
73
            lines.append((item.ts, line))
74
        print("@startuml")
75
        header = """
76
!theme crt-amber
77
skinparam responseMessageBelowArrow true
78
box "User Process"
79
participant User as MainThread
80
control router as MsgRouterThr
81
control check_stop as ChkStopThr
82
control net_stat as NetStatThr
83
end box
84

85
queue record_q as record_q
86
queue result_q as result_q
87

88
box "Internal Process"
89
control handler as HandlerThread
90
control stats as StatsThr
91
queue send_q as send_q
92
queue write_q as write_q
93
control writer as WriterThread
94
control sender as SenderThread
95
end box
96
        """
97
        print(header)
98
        # TODO: move to common place (sorted sequence items)
99
        for _, line in sorted(lines):
100
            print(line)
101
        print("@enduml")
102

103
    def output_mermaid(self) -> None:
104
        lines = []
105
        for item in self._items:
106
            line = f"{item.src} ->> {item.dst}: {item.info}"
107
            lines.append((item.ts, line))
108

109
        header = """
110
sequenceDiagram
111
participant MainThread as User
112
participant MsgRouterThr as router
113
participant ChkStopThr as check_stop
114
participant NetStatThr as net_stat
115

116
participant record_q as record_q
117
participant result_q as result_q
118

119
participant HandlerThread as handler
120
participant StatsThr as stats
121
participant send_q as send_q
122
participant write_q as write_q
123
participant WriterThread as writer
124
participant SenderThread as sender
125
        """
126
        print(header)
127
        # TODO: move to common place (sorted sequence items)
128
        for _, line in sorted(lines):
129
            print(line)
130

131
    def load(self, fname: str) -> None:
132
        with open(fname) as f:
133
            for line in f.readlines():
134
                self._parse(line)
135

136
    def loaddir(self, dname: str) -> None:
137
        flist = []
138
        for p in pathlib.Path(dname).iterdir():
139
            if not p.is_file():
140
                continue
141
            flist.append(p)
142

143
        for f in flist:
144
            self.load(f)
145

146

147
def main():
148
    argparser = argparse.ArgumentParser()
149
    argparser.add_argument("--logdir", default="wandb/latest-run/logs/")
150
    argparser.add_argument("--format", default="mermaid")
151
    args = argparser.parse_args()
152

153
    parser = TracelogParser()
154
    parser.loaddir(args.logdir)
155
    if args.format == "plantuml":
156
        parser.output_plantuml()
157
    elif args.format == "mermaid":
158
        parser.output_mermaid()
159
    else:
160
        print(f"Unknown format: {args.format}")
161
        sys.exit(1)
162

163

164
if __name__ == "__main__":
165
    main()
166