1#!/usr/bin/env python3 2# -*- coding: utf-8 -*- 3# Copyright 2019 The Chromium OS Authors. All rights reserved. 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 9from __future__ import print_function 10 11import argparse 12import distutils.spawn 13import os 14import os.path 15import shutil 16import subprocess 17import sys 18import tempfile 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 _CREATE_LLVM_PROF, '--out=' + output_name, '--binary=chrome', 75 '--profile=' + perf_profile_loc 76 ]) 77 78 79def _abspath_or_gs_link(path): 80 if path.startswith(_GS_PREFIX): 81 return path 82 return os.path.abspath(path) 83 84 85def _tryjob_arg(tryjob_arg): 86 # Forward gs args through 87 if tryjob_arg.startswith(_GS_PREFIX): 88 return tryjob_arg 89 90 # Clicking on the 'Artifacts' link gives us a pantheon link that's basically 91 # a preamble and gs path. 92 pantheon = 'https://pantheon.corp.google.com/storage/browser/' 93 if tryjob_arg.startswith(pantheon): 94 return _GS_PREFIX + tryjob_arg[len(pantheon):] 95 96 # Otherwise, only do things with a tryjob ID (e.g. R75-11965.0.0-b3648595) 97 if not tryjob_arg.startswith('R'): 98 raise ValueError('Unparseable tryjob arg; give a tryjob ID, pantheon ' 99 'link, or gs:// link. Please see source for more.') 100 101 chell_path = 'chromeos-image-archive/chell-chrome-pfq-tryjob/' 102 # ...And assume it's from chell, since that's the only thing we generate 103 # profiles with today. 104 return _GS_PREFIX + chell_path + tryjob_arg 105 106 107def main(): 108 parser = argparse.ArgumentParser(description=__doc__) 109 parser.add_argument( 110 '--perf_profile', 111 required=True, 112 help='Path to our perf profile. Accepts either a gs:// path or local ' 113 'filepath.') 114 parser.add_argument( 115 '--tryjob', 116 required=True, 117 type=_tryjob_arg, 118 help="Path to our tryjob's artifacts. Accepts a gs:// path, pantheon " 119 'link, or tryjob ID, e.g. R75-11965.0.0-b3648595. In the last case, ' 120 'the assumption is that you ran a chell-chrome-pfq-tryjob.') 121 parser.add_argument( 122 '-o', 123 '--output', 124 default='afdo.prof', 125 help='Where to put the AFDO profile. Default is afdo.prof.') 126 parser.add_argument( 127 '-k', 128 '--keep_artifacts_on_failure', 129 action='store_true', 130 help="Don't remove the tempdir on failure") 131 args = parser.parse_args() 132 133 if not distutils.spawn.find_executable(_CREATE_LLVM_PROF): 134 sys.exit(_CREATE_LLVM_PROF + ' not found; are you in the chroot?') 135 136 profile = _abspath_or_gs_link(args.perf_profile) 137 afdo_output = os.path.abspath(args.output) 138 139 initial_dir = os.getcwd() 140 temp_dir = tempfile.mkdtemp(prefix='generate_afdo') 141 success = True 142 try: 143 os.chdir(temp_dir) 144 _generate_afdo(profile, args.tryjob, afdo_output) 145 146 # The AFDO tooling is happy to generate essentially empty profiles for us. 147 # Chrome's profiles are often 8+ MB; if we only see a small fraction of 148 # that, something's off. 512KB was arbitrarily selected. 149 if os.path.getsize(afdo_output) < 512 * 1024: 150 raise ValueError('The AFDO profile is suspiciously small for Chrome. ' 151 'Something might have gone wrong.') 152 except: 153 success = False 154 raise 155 finally: 156 os.chdir(initial_dir) 157 158 if success or not args.keep_artifacts_on_failure: 159 shutil.rmtree(temp_dir, ignore_errors=True) 160 else: 161 print('Artifacts are available at', temp_dir) 162 163 164if __name__ == '__main__': 165 sys.exit(main()) 166