1#!/usr/bin/env python3
2# Copyright (C) 2023 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import sys
17if __name__ == "__main__":
18    sys.dont_write_bytecode = True
19
20import argparse
21import dataclasses
22import datetime
23import json
24import os
25import pathlib
26import statistics
27import zoneinfo
28import csv
29
30import pretty
31import utils
32
33# TODO:
34# - Flag if the last postroll build was more than 15 seconds or something. That's
35#   an indicator that something is amiss.
36# - Add a mode to print all of the values for multi-iteration runs
37# - Add a flag to reorder the tags
38# - Add a flag to reorder the headers in order to show grouping more clearly.
39
40
41def FindSummaries(args):
42    def find_summaries(directory):
43        return [str(p.resolve()) for p in pathlib.Path(directory).glob("**/summary.json")]
44    if not args:
45        # If they didn't give an argument, use the default dir
46        root = utils.get_root()
47        if not root:
48            return []
49        return find_summaries(root.joinpath("..", utils.DEFAULT_REPORT_DIR))
50    results = list()
51    for arg in args:
52        if os.path.isfile(arg):
53            # If it's a file add that
54            results.append(arg)
55        elif os.path.isdir(arg):
56            # If it's a directory, find all of the files there
57            results += find_summaries(arg)
58        else:
59            sys.stderr.write(f"Invalid summary argument: {arg}\n")
60            sys.exit(1)
61    return sorted(list(results))
62
63
64def LoadSummary(filename):
65    with open(filename) as f:
66        return json.load(f)
67
68# Columns:
69#   Date
70#   Branch
71#   Tag
72#   --
73#   Lunch
74# Rows:
75#   Benchmark
76
77def lunch_str(d):
78    "Convert a lunch dict to a string"
79    return f"{d['TARGET_PRODUCT']}-{d['TARGET_RELEASE']}-{d['TARGET_BUILD_VARIANT']}"
80
81def group_by(l, key):
82    "Return a list of tuples, grouped by key, sorted by key"
83    result = {}
84    for item in l:
85        result.setdefault(key(item), []).append(item)
86    return [(k, v) for k, v in result.items()]
87
88
89class Table:
90    def __init__(self, row_title, fixed_titles=[]):
91        self._data = {}
92        self._rows = []
93        self._cols = []
94        self._fixed_cols = {}
95        self._titles = [row_title] + fixed_titles
96
97    def Set(self, column_key, row_key, data):
98        self._data[(column_key, row_key)] = data
99        if not column_key in self._cols:
100            self._cols.append(column_key)
101        if not row_key in self._rows:
102            self._rows.append(row_key)
103
104    def SetFixedCol(self, row_key, columns):
105        self._fixed_cols[row_key] = columns
106
107    def Write(self, out, fmt):
108        table = []
109        # Expand the column items
110        for row in zip(*self._cols):
111            if row.count(row[0]) == len(row):
112                continue
113            table.append([""] * len(self._titles) + [col for col in row])
114        if table:
115            # Update the last row of the header with title and add separator
116            for i in range(len(self._titles)):
117                table[len(table)-1][i] = self._titles[i]
118            if fmt == "table":
119                table.append(pretty.SEPARATOR)
120        # Populate the data
121        for row in self._rows:
122            table.append([str(row)]
123                         + self._fixed_cols[row]
124                         + [str(self._data.get((col, row), "")) for col in self._cols])
125        if fmt == "csv":
126            csv.writer(sys.stdout, quoting=csv.QUOTE_MINIMAL).writerows(table)
127        else:
128            out.write(pretty.FormatTable(table, alignments="LL"))
129
130
131def format_duration_sec(ns, fmt_sec):
132    "Format a duration in ns to second precision"
133    sec = round(ns / 1000000000)
134    if fmt_sec:
135        return f"{sec}"
136    else:
137        h, sec = divmod(sec, 60*60)
138        m, sec = divmod(sec, 60)
139        result = ""
140        if h > 0:
141            result += f"{h:2d}h "
142        if h > 0 or m > 0:
143            result += f"{m:2d}m "
144        return result + f"{sec:2d}s"
145
146
147def main(argv):
148    parser = argparse.ArgumentParser(
149            prog="format_benchmarks",
150            allow_abbrev=False, # Don't let people write unsupportable scripts.
151            description="Print analysis tables for benchmarks")
152
153    parser.add_argument("--csv", action="store_true",
154                        help="Print in CSV instead of table.")
155
156    parser.add_argument("--sec", action="store_true",
157                        help="Print in seconds instead of minutes and seconds")
158
159    parser.add_argument("--tags", nargs="*",
160                        help="The tags to print, in order.")
161
162    parser.add_argument("summaries", nargs="*",
163                        help="A summary.json file or a directory in which to look for summaries.")
164
165    args = parser.parse_args()
166
167    # Load the summaries
168    summaries = [(s, LoadSummary(s)) for s in FindSummaries(args.summaries)]
169
170    # Convert to MTV time
171    for filename, s in summaries:
172        dt = datetime.datetime.fromisoformat(s["start_time"])
173        dt = dt.astimezone(zoneinfo.ZoneInfo("America/Los_Angeles"))
174        s["datetime"] = dt
175        s["date"] = datetime.date(dt.year, dt.month, dt.day)
176
177    # Filter out tags we don't want
178    if args.tags:
179        summaries = [(f, s) for f, s in summaries if s.get("tag", "") in args.tags]
180
181    # If they supplied tags, sort in that order, otherwise sort by tag
182    if args.tags:
183        tagsort = lambda tag: args.tags.index(tag)
184    else:
185        tagsort = lambda tag: tag
186
187    # Sort the summaries
188    summaries.sort(key=lambda s: (s[1]["date"], s[1]["branch"], tagsort(s[1]["tag"])))
189
190    # group the benchmarks by column and iteration
191    def bm_key(b):
192        return (
193            lunch_str(b["lunch"]),
194        )
195    for filename, summary in summaries:
196        summary["columns"] = [(key, group_by(bms, lambda b: b["id"])) for key, bms
197                              in group_by(summary["benchmarks"], bm_key)]
198
199    # Build the table
200    table = Table("Benchmark", ["Rebuild"])
201    for filename, summary in summaries:
202        for key, column in summary["columns"]:
203            for id, cell in column:
204                duration_ns = statistics.median([b["duration_ns"] for b in cell])
205                modules = cell[0]["modules"]
206                if not modules:
207                    modules = ["---"]
208                table.SetFixedCol(cell[0]["title"], [" ".join(modules)])
209                table.Set(tuple([summary["date"].strftime("%Y-%m-%d"),
210                                 summary["branch"],
211                                 summary["tag"]]
212                                + list(key)),
213                          cell[0]["title"], format_duration_sec(duration_ns, args.sec))
214
215    table.Write(sys.stdout, "csv" if args.csv else "table")
216
217if __name__ == "__main__":
218    main(sys.argv)
219
220