• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2022 The Chromium Authors
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5from __future__ import annotations
6
7import abc
8import csv
9import json
10import logging
11from collections import defaultdict
12from typing import (TYPE_CHECKING, Any, Callable, Dict, Generic, List, Optional,
13                    Type, TypeVar, Union)
14
15from tabulate import tabulate
16
17from crossbench.probes import helper
18from crossbench.probes.metric import (CSVFormatter, MetricsMerger,
19                                      metric_geomean)
20from crossbench.probes.probe import Probe, ProbeContext, ProbeMissingDataError
21from crossbench.probes.results import (EmptyProbeResult, LocalProbeResult,
22                                       ProbeResult)
23
24if TYPE_CHECKING:
25  from crossbench.path import LocalPath
26  from crossbench.runner.actions import Actions
27  from crossbench.runner.groups.base import RunGroup
28  from crossbench.runner.groups.browsers import BrowsersRunGroup
29  from crossbench.runner.groups.repetitions import RepetitionsRunGroup
30  from crossbench.runner.run import Run
31  from crossbench.types import Json
32
33class JsonResultProbe(Probe, metaclass=abc.ABCMeta):
34  """
35  Abstract Probe that stores a Json result extracted by the `to_json` method
36
37  Tje `to_json` is provided by subclasses. A typical examples includes just
38  running a JS script on the page.
39  Multiple Json result files for RepetitionsRunGroups are merged with the
40  MetricsMerger. Custom merging for other RunGroups can be defined in the
41  subclass.
42  """
43
44  FLATTEN = True
45  SORT_KEYS = True
46
47  @property
48  def result_path_name(self) -> str:
49    return f"{self.name}.json"
50
51  @abc.abstractmethod
52  def to_json(self, actions: Actions) -> Json:
53    """
54    Override in subclasses.
55    Returns json-serializable data.
56    """
57    return None
58
59  def flatten_json_data(self, json_data: Any) -> Json:
60    return helper.Flatten(json_data).data
61
62  def process_json_data(self, json_data) -> Any:
63    return json_data
64
65  def get_context(self, run: Run) -> JsonResultProbeContext:
66    return JsonResultProbeContext(self, run)
67
68  def merge_repetitions(
69      self,
70      group: RepetitionsRunGroup,
71  ) -> ProbeResult:
72    merger = MetricsMerger()
73    for run in group.runs:
74      if self not in run.results:
75        raise ProbeMissingDataError(
76            f"Probe {self.NAME} produced no data to merge.")
77      source_file = run.results[self].json
78      assert source_file.is_file(), (
79          f"{source_file} from {run} is not a file or doesn't exist.")
80      with source_file.open(encoding="utf-8") as f:
81        merger.add(json.load(f))
82    return self.write_group_result(group, merger, csv_formatter=CSVFormatter)
83
84  def merge_browsers_json_list(self, group: BrowsersRunGroup) -> ProbeResult:
85    merged_json: Dict[str, Dict[str, Any]] = {}
86    for story_group in group.story_groups:
87      browser_result: Dict[str, Any] = {}
88      merged_json[story_group.browser.unique_name] = browser_result
89      browser_result["info"] = story_group.info
90      browser_json_path = story_group.results[self].json
91      assert browser_json_path.is_file(), (
92          f"{browser_json_path} from {story_group} "
93          "is not a file or doesn't exist.")
94      with browser_json_path.open(encoding="utf-8") as f:
95        browser_result["data"] = json.load(f)
96    merged_json_path = group.get_local_probe_result_path(self)
97    assert not merged_json_path.exists(), (
98        f"Cannot override existing Json result: {merged_json_path}")
99    with merged_json_path.open("w", encoding="utf-8") as f:
100      json.dump(merged_json, f, indent=2)
101      # TODO(375390958): figure out why files aren't fully written to
102      # pyfakefs here.
103      f.write("\n")
104    return LocalProbeResult(json=(merged_json_path,))
105
106  def merge_browsers_csv_list(self, group: BrowsersRunGroup) -> ProbeResult:
107    csv_file_list: List[LocalPath] = []
108    for story_group in group.story_groups:
109      csv_file_list.append(story_group.results[self].csv)
110    merged_table = helper.merge_csv(csv_file_list, row_header_len=-1)
111    merged_json_path = group.get_local_probe_result_path(self, exists_ok=True)
112    merged_csv_path = merged_json_path.with_suffix(".csv")
113    assert not merged_csv_path.exists(), (
114        f"Cannot override existing CSV result: {merged_csv_path}")
115    with merged_csv_path.open("w", newline="", encoding="utf-8") as f:
116      csv.writer(f, delimiter="\t").writerows(merged_table)
117    return LocalProbeResult(csv=(merged_csv_path,))
118
119  def write_group_result(
120      self,
121      group: RunGroup,
122      merged_data: Union[Dict, List, MetricsMerger],
123      csv_formatter: Optional[Type[CSVFormatter]] = CSVFormatter,
124      value_fn: Callable[[Any], Any] = metric_geomean) -> ProbeResult:
125    merged_json_path = group.get_local_probe_result_path(self)
126    with merged_json_path.open("w", encoding="utf-8") as f:
127      if isinstance(merged_data, (dict, list)):
128        json.dump(merged_data, f, indent=2)
129      else:
130        json.dump(merged_data.to_json(sort=self.SORT_KEYS), f, indent=2)
131      # TODO(375390958): figure out why files aren't fully written to
132      # pyfakefs here.
133      f.write("\n")
134    if not csv_formatter:
135      return LocalProbeResult(json=(merged_json_path,))
136    if not isinstance(merged_data, MetricsMerger):
137      raise ValueError("write_csv is only supported for MetricsMerger, "
138                       f"but found {type(merged_data)}'.")
139    return self.write_group_csv_result(group, merged_data, merged_json_path,
140                                       csv_formatter, value_fn)
141
142  def write_group_csv_result(self, group: RunGroup, merged_data: MetricsMerger,
143                             merged_json_path: LocalPath,
144                             csv_formatter: Type[CSVFormatter],
145                             value_fn: Callable[[Any], Any]) -> ProbeResult:
146    merged_csv_path = merged_json_path.with_suffix(".csv")
147    assert not merged_csv_path.exists(), (
148        f"Cannot override existing CSV result: {merged_csv_path}")
149    # Create a CSV table:
150    # 0 | info label 0,                                          info_value 0
151    #     ...                                                    ...
152    # N | info label N,                                          info_value N
153    # 0 | metric 0 full path, metric path[0] ... metric path[N], metric 0 value
154    #     ...                                                    ...
155    # M | metric M full path, ...                                metric M value
156    headers = []
157    for label, info_value in group.info.items():
158      headers.append((label, info_value))
159    csv_data = csv_formatter(
160        merged_data, value_fn, headers=headers, sort=self.SORT_KEYS).table
161    with merged_csv_path.open("w", newline="", encoding="utf-8") as f:
162      writer = csv.writer(f, delimiter="\t")
163      writer.writerows(csv_data)
164    return LocalProbeResult(json=(merged_json_path,), csv=(merged_csv_path,))
165
166  LOG_SUMMARY_KEYS = ("label", "browser", "version", "os", "device", "cpu",
167                      "runs", "failed runs")
168
169  def _log_result_metrics(self, data: Dict) -> None:
170    table: Dict[str, List[str]] = defaultdict(list)
171    for browser_result in data.values():
172      for info_key in self.LOG_SUMMARY_KEYS:
173        table[info_key].append(browser_result["info"][info_key])
174      data = browser_result["data"]
175      self._extract_result_metrics_table(data, table)
176    flattened: List[List[str]] = list(
177        [label] + values for label, values in table.items())
178    logging.critical(tabulate(flattened, tablefmt="plain"))
179
180  def _extract_result_metrics_table(self, metrics: Dict[str, Any],
181                                    table: Dict[str, List[str]]) -> None:
182    """Add individual metrics to the table in here.
183    Typically you only add score and total values for each benchmark or
184    benchmark item."""
185    del metrics
186    del table
187
188
189JsonResultProbeT = TypeVar("JsonResultProbeT", bound="JsonResultProbe")
190
191
192class JsonResultProbeContext(ProbeContext[JsonResultProbeT],
193                             Generic[JsonResultProbeT]):
194
195  def __init__(self, probe: JsonResultProbeT, run: Run) -> None:
196    super().__init__(probe, run)
197    self._json_data: Json = None
198
199  @property
200  def probe(self) -> JsonResultProbeT:
201    return super().probe
202
203  def to_json(self, actions: Actions) -> Json:
204    return self.probe.to_json(actions)
205
206  def start(self) -> None:
207    pass
208
209  def stop(self) -> None:
210    self._json_data = self.extract_json(self.run)
211
212  def teardown(self) -> ProbeResult:
213    if self._json_data is None:
214      return EmptyProbeResult()
215    self._json_data = self.process_json_data(self._json_data)
216    return self.write_json(self.run, self._json_data)
217
218  def extract_json(self, run: Run) -> Json:
219    with run.actions(f"Extracting Probe({self.probe.name})") as actions:
220      json_data = self.to_json(actions)
221      assert json_data is not None, (
222          f"Probe({self.probe.name}) produced no data")
223      return json_data
224
225  def write_json(self, run: Run, json_data: Json) -> ProbeResult:
226    flattened_file = None
227    with run.actions(f"Writing Probe({self.probe.name})"):
228      assert json_data is not None, (
229          f"Probe({self.probe.name}) produced no Json data.")
230      raw_file = self.local_result_path
231      if self.probe.FLATTEN:
232        raw_file = raw_file.with_suffix(".json.nested")
233        flattened_file = self.local_result_path
234        flat_json_data = self.flatten_json_data(json_data)
235        with flattened_file.open("w", encoding="utf-8") as f:
236          json.dump(flat_json_data, f, indent=2)
237          # TODO(375390958): figure out why files aren't fully written to
238          # pyfakefs here.
239          f.write("\n")
240      with raw_file.open("w", encoding="utf-8") as f:
241        json.dump(json_data, f, indent=2)
242        # TODO(375390958): figure out why files aren't fully written to
243        # pyfakefs here.
244        f.write("\n")
245    if flattened_file:
246      return LocalProbeResult(json=(flattened_file,), file=(raw_file,))
247    return LocalProbeResult(json=(raw_file,))
248
249  def process_json_data(self, json_data: Json) -> Json:
250    return self.probe.process_json_data(json_data)
251
252  def flatten_json_data(self, json_data: Any) -> Json:
253    return self.probe.flatten_json_data(json_data)
254