1#!/usr/bin/env python3 2# -*- coding: utf-8 -*- 3# Copyright 2019 The ChromiumOS Authors 4# Use of this source code is governed by a BSD-style license that can be 5# found in the LICENSE file. 6 7"""Given a tryjob and perf profile, generates an AFDO profile.""" 8 9 10import argparse 11import distutils.spawn 12import os 13import os.path 14import shutil 15import subprocess 16import sys 17import tempfile 18 19 20_CREATE_LLVM_PROF = "create_llvm_prof" 21_GS_PREFIX = "gs://" 22 23 24def _fetch_gs_artifact(remote_name, local_name): 25 assert remote_name.startswith(_GS_PREFIX) 26 subprocess.check_call(["gsutil", "cp", remote_name, local_name]) 27 28 29def _fetch_and_maybe_unpack(remote_name, local_name): 30 unpackers = [ 31 (".tar.bz2", ["tar", "xaf"]), 32 (".bz2", ["bunzip2"]), 33 (".tar.xz", ["tar", "xaf"]), 34 (".xz", ["xz", "-d"]), 35 ] 36 37 unpack_ext = None 38 unpack_cmd = None 39 for ext, unpack in unpackers: 40 if remote_name.endswith(ext): 41 unpack_ext, unpack_cmd = ext, unpack 42 break 43 44 download_to = local_name + unpack_ext if unpack_ext else local_name 45 _fetch_gs_artifact(remote_name, download_to) 46 if unpack_cmd is not None: 47 print("Unpacking", download_to) 48 subprocess.check_output(unpack_cmd + [download_to]) 49 assert os.path.exists(local_name) 50 51 52def _generate_afdo(perf_profile_loc, tryjob_loc, output_name): 53 if perf_profile_loc.startswith(_GS_PREFIX): 54 local_loc = "perf.data" 55 _fetch_and_maybe_unpack(perf_profile_loc, local_loc) 56 perf_profile_loc = local_loc 57 58 chrome_in_debug_loc = "debug/opt/google/chrome/chrome.debug" 59 debug_out = "debug.tgz" 60 _fetch_gs_artifact(os.path.join(tryjob_loc, "debug.tgz"), debug_out) 61 62 print("Extracting chrome.debug.") 63 # This has tons of artifacts, and we only want Chrome; don't waste time 64 # extracting the rest in _fetch_and_maybe_unpack. 65 subprocess.check_call(["tar", "xaf", "debug.tgz", chrome_in_debug_loc]) 66 67 # Note that the AFDO tool *requires* a binary named `chrome` to be present if 68 # we're generating a profile for chrome. It's OK for this to be split debug 69 # information. 70 os.rename(chrome_in_debug_loc, "chrome") 71 72 print("Generating AFDO profile.") 73 subprocess.check_call( 74 [ 75 _CREATE_LLVM_PROF, 76 "--out=" + output_name, 77 "--binary=chrome", 78 "--profile=" + perf_profile_loc, 79 ] 80 ) 81 82 83def _abspath_or_gs_link(path): 84 if path.startswith(_GS_PREFIX): 85 return path 86 return os.path.abspath(path) 87 88 89def _tryjob_arg(tryjob_arg): 90 # Forward gs args through 91 if tryjob_arg.startswith(_GS_PREFIX): 92 return tryjob_arg 93 94 # Clicking on the 'Artifacts' link gives us a pantheon link that's basically 95 # a preamble and gs path. 96 pantheon = "https://pantheon.corp.google.com/storage/browser/" 97 if tryjob_arg.startswith(pantheon): 98 return _GS_PREFIX + tryjob_arg[len(pantheon) :] 99 100 # Otherwise, only do things with a tryjob ID (e.g. R75-11965.0.0-b3648595) 101 if not tryjob_arg.startswith("R"): 102 raise ValueError( 103 "Unparseable tryjob arg; give a tryjob ID, pantheon " 104 "link, or gs:// link. Please see source for more." 105 ) 106 107 chell_path = "chromeos-image-archive/chell-chrome-pfq-tryjob/" 108 # ...And assume it's from chell, since that's the only thing we generate 109 # profiles with today. 110 return _GS_PREFIX + chell_path + tryjob_arg 111 112 113def main(): 114 parser = argparse.ArgumentParser(description=__doc__) 115 parser.add_argument( 116 "--perf_profile", 117 required=True, 118 help="Path to our perf profile. Accepts either a gs:// path or local " 119 "filepath.", 120 ) 121 parser.add_argument( 122 "--tryjob", 123 required=True, 124 type=_tryjob_arg, 125 help="Path to our tryjob's artifacts. Accepts a gs:// path, pantheon " 126 "link, or tryjob ID, e.g. R75-11965.0.0-b3648595. In the last case, " 127 "the assumption is that you ran a chell-chrome-pfq-tryjob.", 128 ) 129 parser.add_argument( 130 "-o", 131 "--output", 132 default="afdo.prof", 133 help="Where to put the AFDO profile. Default is afdo.prof.", 134 ) 135 parser.add_argument( 136 "-k", 137 "--keep_artifacts_on_failure", 138 action="store_true", 139 help="Don't remove the tempdir on failure", 140 ) 141 args = parser.parse_args() 142 143 if not distutils.spawn.find_executable(_CREATE_LLVM_PROF): 144 sys.exit(_CREATE_LLVM_PROF + " not found; are you in the chroot?") 145 146 profile = _abspath_or_gs_link(args.perf_profile) 147 afdo_output = os.path.abspath(args.output) 148 149 initial_dir = os.getcwd() 150 temp_dir = tempfile.mkdtemp(prefix="generate_afdo") 151 success = True 152 try: 153 os.chdir(temp_dir) 154 _generate_afdo(profile, args.tryjob, afdo_output) 155 156 # The AFDO tooling is happy to generate essentially empty profiles for us. 157 # Chrome's profiles are often 8+ MB; if we only see a small fraction of 158 # that, something's off. 512KB was arbitrarily selected. 159 if os.path.getsize(afdo_output) < 512 * 1024: 160 raise ValueError( 161 "The AFDO profile is suspiciously small for Chrome. " 162 "Something might have gone wrong." 163 ) 164 except: 165 success = False 166 raise 167 finally: 168 os.chdir(initial_dir) 169 170 if success or not args.keep_artifacts_on_failure: 171 shutil.rmtree(temp_dir, ignore_errors=True) 172 else: 173 print("Artifacts are available at", temp_dir) 174 175 176if __name__ == "__main__": 177 sys.exit(main()) 178