pytorch

1
"""Scribe Uploader for Pytorch Benchmark Data
2

3
Currently supports data in pytest-benchmark format but can be extended.
4

5
New fields can be added just by modifying the schema in this file, schema
6
checking is only here to encourage reusing existing fields and avoiding typos.
7
"""
8

9
import argparse
10
import json
11
import os
12
import subprocess
13
import time
14
from collections import defaultdict
15

16
import requests
17

18

19
class ScribeUploader:
20
    def __init__(self, category):
21
        self.category = category
22

23
    def format_message(self, field_dict):
24
        assert "time" in field_dict, "Missing required Scribe field 'time'"
25
        message = defaultdict(dict)
26
        for field, value in field_dict.items():
27
            if field in self.schema["normal"]:
28
                message["normal"][field] = str(value)
29
            elif field in self.schema["int"]:
30
                message["int"][field] = int(value)
31
            elif field in self.schema["float"]:
32
                message["float"][field] = float(value)
33
            else:
34
                raise ValueError(
35
                    f"Field {field} is not currently used, be intentional about adding new fields"
36
                )
37
        return message
38

39
    def _upload_intern(self, messages):
40
        for m in messages:
41
            json_str = json.dumps(m)
42
            cmd = ["scribe_cat", self.category, json_str]
43
            subprocess.run(cmd)
44

45
    def upload(self, messages):
46
        if os.environ.get("SCRIBE_INTERN"):
47
            return self._upload_intern(messages)
48
        access_token = os.environ.get("SCRIBE_GRAPHQL_ACCESS_TOKEN")
49
        if not access_token:
50
            raise ValueError("Can't find access token from environment variable")
51
        url = "https://graph.facebook.com/scribe_logs"
52
        r = requests.post(
53
            url,
54
            data={
55
                "access_token": access_token,
56
                "logs": json.dumps(
57
                    [
58
                        {
59
                            "category": self.category,
60
                            "message": json.dumps(message),
61
                            "line_escape": False,
62
                        }
63
                        for message in messages
64
                    ]
65
                ),
66
            },
67
        )
68
        print(r.text)
69
        r.raise_for_status()
70

71

72
class PytorchBenchmarkUploader(ScribeUploader):
73
    def __init__(self):
74
        super().__init__("perfpipe_pytorch_benchmarks")
75
        self.schema = {
76
            "int": [
77
                "time",
78
                "rounds",
79
            ],
80
            "normal": [
81
                "benchmark_group",
82
                "benchmark_name",
83
                "benchmark_executor",
84
                "benchmark_fuser",
85
                "benchmark_class",
86
                "benchmark_time",
87
                "pytorch_commit_id",
88
                "pytorch_branch",
89
                "pytorch_commit_time",
90
                "pytorch_version",
91
                "pytorch_git_dirty",
92
                "machine_kernel",
93
                "machine_processor",
94
                "machine_hostname",
95
                "circle_build_num",
96
                "circle_project_reponame",
97
            ],
98
            "float": [
99
                "stddev",
100
                "min",
101
                "median",
102
                "max",
103
                "mean",
104
            ],
105
        }
106

107
    def post_pytest_benchmarks(self, pytest_json):
108
        machine_info = pytest_json["machine_info"]
109
        commit_info = pytest_json["commit_info"]
110
        upload_time = int(time.time())
111
        messages = []
112
        for b in pytest_json["benchmarks"]:
113
            test = b["name"].split("[")[0]
114
            net_name = b["params"]["net_name"]
115
            benchmark_name = f"{test}[{net_name}]"
116
            executor = b["params"]["executor"]
117
            fuser = b["params"]["fuser"]
118
            m = self.format_message(
119
                {
120
                    "time": upload_time,
121
                    "benchmark_group": b["group"],
122
                    "benchmark_name": benchmark_name,
123
                    "benchmark_executor": executor,
124
                    "benchmark_fuser": fuser,
125
                    "benchmark_class": b["fullname"],
126
                    "benchmark_time": pytest_json["datetime"],
127
                    "pytorch_commit_id": commit_info["id"],
128
                    "pytorch_branch": commit_info["branch"],
129
                    "pytorch_commit_time": commit_info["time"],
130
                    "pytorch_version": None,
131
                    "pytorch_git_dirty": commit_info["dirty"],
132
                    "machine_kernel": machine_info["release"],
133
                    "machine_processor": machine_info["processor"],
134
                    "machine_hostname": machine_info["node"],
135
                    "circle_build_num": os.environ.get("CIRCLE_BUILD_NUM"),
136
                    "circle_project_reponame": os.environ.get(
137
                        "CIRCLE_PROJECT_REPONAME"
138
                    ),
139
                    "stddev": b["stats"]["stddev"],
140
                    "rounds": b["stats"]["rounds"],
141
                    "min": b["stats"]["min"],
142
                    "median": b["stats"]["median"],
143
                    "max": b["stats"]["max"],
144
                    "mean": b["stats"]["mean"],
145
                }
146
            )
147
            messages.append(m)
148
        self.upload(messages)
149

150

151
if __name__ == "__main__":
152
    parser = argparse.ArgumentParser(description=__doc__)
153
    parser.add_argument(
154
        "--pytest-bench-json",
155
        "--pytest_bench_json",
156
        type=argparse.FileType("r"),
157
        help="Upload json data formatted by pytest-benchmark module",
158
    )
159
    args = parser.parse_args()
160
    if args.pytest_bench_json:
161
        benchmark_uploader = PytorchBenchmarkUploader()
162
        json_data = json.load(args.pytest_bench_json)
163
        benchmark_uploader.post_pytest_benchmarks(json_data)
164