1# Copyright 2024 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 json 9import logging 10from collections import defaultdict 11from typing import (TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, 12 cast) 13 14from crossbench.benchmarks.base import BenchmarkProbeMixin, PressBenchmark 15from crossbench.probes.json import JsonResultProbe 16from crossbench.probes.metric import (CSVFormatter, Metric, MetricsMerger, 17 geomean) 18from crossbench.probes.results import ProbeResult, ProbeResultDict 19 20if TYPE_CHECKING: 21 import argparse 22 23 from crossbench.cli.parser import CrossBenchArgumentParser 24 from crossbench.path import LocalPath 25 from crossbench.runner.actions import Actions 26 from crossbench.runner.groups.browsers import BrowsersRunGroup 27 from crossbench.runner.groups.stories import StoriesRunGroup 28 from crossbench.runner.run import Run 29 from crossbench.stories.story import Story 30 from crossbench.types import Json 31 32 33class JetStreamProbe( 34 BenchmarkProbeMixin, JsonResultProbe, metaclass=abc.ABCMeta): 35 """ 36 JetStream-specific Probe. 37 Extracts all JetStream times and scores. 38 """ 39 FLATTEN: bool = False 40 JS: str = """ 41 let results = Object.create(null); 42 let benchmarks = [] 43 for (let benchmark of JetStream.benchmarks) { 44 const data = { score: benchmark.score }; 45 if ("worst4" in benchmark) { 46 data.firstIteration = benchmark.firstIteration; 47 data.average = benchmark.average; 48 data.worst4 = benchmark.worst4; 49 } else if ("runTime" in benchmark) { 50 data.runTime = benchmark.runTime; 51 data.startupTime = benchmark.startupTime; 52 } else if ("mainRun" in benchmark) { 53 data.mainRun = benchmark.mainRun; 54 data.stdlib = benchmark.stdlib; 55 } 56 results[benchmark.plan.name] = data; 57 benchmarks.push(benchmark); 58 }; 59 return results; 60""" 61 62 @property 63 def jetstream(self) -> JetStreamBenchmark: 64 return cast(JetStreamBenchmark, self.benchmark) 65 66 def to_json(self, actions: Actions) -> Dict[str, float]: 67 data = actions.js(self.JS) 68 assert len(data) > 0, "No benchmark data generated" 69 return data 70 71 def process_json_data(self, json_data: Dict[str, Any]) -> Dict[str, Any]: 72 assert "Total" not in json_data, ( 73 "JSON result data already contains a ['Total'] entry.") 74 json_data["Total"] = self._compute_total_metrics(json_data) 75 return json_data 76 77 def _compute_total_metrics(self, json_data: Dict[str, 78 Any]) -> Dict[str, float]: 79 # Manually add all total scores 80 accumulated_metrics = defaultdict(list) 81 for _, metrics in json_data.items(): 82 for metric, value in metrics.items(): 83 accumulated_metrics[metric].append(value) 84 total: Dict[str, float] = {} 85 for metric, values in accumulated_metrics.items(): 86 total[metric] = geomean(values) 87 return total 88 89 def log_run_result(self, run: Run) -> None: 90 self._log_result(run.results, single_result=True) 91 92 def log_browsers_result(self, group: BrowsersRunGroup) -> None: 93 self._log_result(group.results, single_result=False) 94 95 def _log_result(self, result_dict: ProbeResultDict, 96 single_result: bool) -> None: 97 if self not in result_dict: 98 return 99 results_json: LocalPath = result_dict[self].json 100 logging.info("-" * 80) 101 logging.critical("JetStream results:") 102 if not single_result: 103 logging.critical(" %s", result_dict[self].csv) 104 logging.info("- " * 40) 105 106 with results_json.open(encoding="utf-8") as f: 107 data = json.load(f) 108 if single_result: 109 logging.critical("Score %s", data["Total"]["score"]) 110 else: 111 self._log_result_metrics(data) 112 113 def _extract_result_metrics_table(self, metrics: Dict[str, Any], 114 table: Dict[str, List[str]]) -> None: 115 for metric_key, metric_value in metrics.items(): 116 if not self._is_valid_metric_key(metric_key): 117 continue 118 table[metric_key].append( 119 Metric.format(metric_value["average"], metric_value["stddev"])) 120 # Separate runs don't produce a score 121 if "Total/score" in metrics: 122 metric_value = metrics["Total/score"] 123 table["Score"].append( 124 Metric.format(metric_value["average"], metric_value["stddev"])) 125 126 def merge_stories(self, group: StoriesRunGroup) -> ProbeResult: 127 merged = MetricsMerger.merge_json_list( 128 story_group.results[self].json 129 for story_group in group.repetitions_groups) 130 return self.write_group_result(group, merged, JetStreamCSVFormatter) 131 132 def merge_browsers(self, group: BrowsersRunGroup) -> ProbeResult: 133 return self.merge_browsers_json_list(group).merge( 134 self.merge_browsers_csv_list(group)) 135 136 def _is_valid_metric_key(self, metric_key: str) -> bool: 137 parts = metric_key.split("/") 138 if len(parts) != 2: 139 return False 140 if self.jetstream.detailed_metrics: 141 return True 142 return parts[0] != "Total" and parts[1] == "score" 143 144class JetStreamCSVFormatter(CSVFormatter): 145 146 def format_items(self, data: Dict[str, Json], 147 sort: bool) -> Sequence[Tuple[str, Json]]: 148 items = list(data.items()) 149 if sort: 150 items.sort() 151 # Copy all /score items to the top: 152 total_key = "Total/score" 153 score_items = [] 154 for key, value in items: 155 if key != total_key and key.endswith("/score"): 156 score_items.append((key, value)) 157 total_item = [(total_key, data[total_key])] 158 return total_item + score_items + items 159 160 161class JetStreamBenchmark(PressBenchmark, metaclass=abc.ABCMeta): 162 163 @classmethod 164 def short_base_name(cls) -> str: 165 return "js" 166 167 @classmethod 168 def base_name(cls) -> str: 169 return "jetstream" 170 171 @classmethod 172 def add_cli_parser( 173 cls, subparsers: argparse.ArgumentParser, aliases: Sequence[str] = () 174 ) -> CrossBenchArgumentParser: 175 parser = super().add_cli_parser(subparsers, aliases) 176 parser.add_argument( 177 "--detailed-metrics", 178 "--details", 179 default=False, 180 action="store_true", 181 help="Report more detailed internal metrics.") 182 return parser 183 184 @classmethod 185 def kwargs_from_cli(cls, args: argparse.Namespace) -> Dict[str, Any]: 186 kwargs = super().kwargs_from_cli(args) 187 kwargs["detailed_metrics"] = args.detailed_metrics 188 return kwargs 189 190 def __init__(self, 191 stories: Sequence[Story], 192 custom_url: Optional[str] = None, 193 detailed_metrics: bool = False): 194 self._detailed_metrics = detailed_metrics 195 super().__init__(stories, custom_url) 196 197 @property 198 def detailed_metrics(self) -> bool: 199 return self._detailed_metrics 200