| #!/usr/bin/env python2 |
| # -*- coding: utf-8 -*- |
| # Copyright 2019 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Download profdata from different arches, merge them and upload to gs.""" |
| |
| from __future__ import print_function |
| |
| import argparse |
| import datetime |
| import distutils.spawn |
| import os |
| import os.path |
| import shutil |
| import subprocess |
| import sys |
| import tempfile |
| |
| _LLVM_PROFDATA = '/usr/bin/llvm-profdata' |
| _GS_PREFIX = 'gs://' |
| |
| |
| def _get_gs_latest(remote_lastest): |
| assert remote_lastest.startswith(_GS_PREFIX) |
| try: |
| return subprocess.check_output(['gsutil', 'cat', remote_lastest]) |
| except subprocess.CalledProcessError: |
| raise RuntimeError('Lastest artifacts not found: %s' % remote_lastest) |
| |
| |
| def _fetch_gs_artifact(remote_name, local_name): |
| assert remote_name.startswith(_GS_PREFIX) |
| subprocess.check_call(['gsutil', 'cp', remote_name, local_name]) |
| |
| |
| def _find_latest_profdata(arch): |
| remote_latest = ( |
| '%schromeos-image-archive/' |
| '%s-pgo-generate-llvm-next-toolchain/LATEST-master' % (_GS_PREFIX, arch)) |
| version = _get_gs_latest(remote_latest) |
| profdata = ('%s-pgo-generate-llvm-next-toolchain/%s/llvm_profdata.tar.xz' % |
| (arch, version)) |
| return profdata |
| |
| |
| def _get_gs_profdata(profdata): |
| remote_profdata = ('%schromeos-image-archive/%s' % (_GS_PREFIX, profdata)) |
| tar = 'llvm_profdata.tar.xz' |
| print('Downloading single profdata for: %s' % profdata) |
| _fetch_gs_artifact(remote_profdata, tar) |
| extract_cmd = ['tar', '-xf', tar] |
| |
| print('Extracting profdata tarball.\nCMD: %s\n' % extract_cmd) |
| subprocess.check_call(extract_cmd) |
| profdata = profdata.replace('llvm_profdata.tar.xz', 'llvm.profdata') |
| # Return directory to the llvm.profdata extracted. |
| return 'b/s/w/ir/cache/cbuild/repository/buildbot_archive/%s' % profdata |
| |
| |
| def _merge_profdata(profdata_list, output_name): |
| merge_cmd = [_LLVM_PROFDATA, 'merge', '-output', output_name] + profdata_list |
| print('Merging PGO profiles.\nCMD: %s\n' % merge_cmd) |
| subprocess.check_call(merge_cmd) |
| |
| |
| def _tar_and_upload_profdata(profdata): |
| timestamp = datetime.datetime.strftime(datetime.datetime.now(), '%Y%m%d') |
| tarball = 'llvm-profdata-%s.tar.xz' % timestamp |
| print('Making profdata tarball: %s' % tarball) |
| subprocess.check_call( |
| ['tar', '--sparse', '-I', 'xz', '-cf', tarball, profdata]) |
| |
| # TODO: it's better to create a subdir: distfiles/llvm_pgo_profile, but |
| # now llvm could only recognize distfiles. |
| upload_cmd = [ |
| 'gsutil', '-m', 'cp', '-n', '-a', 'public-read', tarball, |
| '%schromeos-localmirror/distfiles/%s' % (_GS_PREFIX, tarball) |
| ] |
| print('Uploading tarball to gs.\nCMD: %s\n' % upload_cmd) |
| subprocess.check_call(upload_cmd) |
| |
| |
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument( |
| '-l', |
| '--latest', |
| default=[], |
| action='append', |
| help='User can specify the profdata from which architectures to ' |
| 'download. By default, we merge profdata from arm, arm64, amd64.') |
| parser.add_argument( |
| '-t', |
| '--tryjob', |
| default=[], |
| action='append', |
| help='Extra pgo-generate-llvm-next-toolchain/tryjob results to be used. ' |
| 'Format should be ' |
| '{arch}-pgo-generate-llvm-next-toolchain(-tryjob)/{VERSION}.') |
| parser.add_argument( |
| '-o', |
| '--output', |
| default='llvm.profdata', |
| help='Where to put merged PGO profile. The default is to not save it ' |
| 'anywhere.') |
| args = parser.parse_args() |
| |
| # If no --latest specified, by default we collect from listed arches. |
| latest = ['arm', 'arm64', 'amd64'] if not args.latest else args.latest |
| |
| if not distutils.spawn.find_executable(_LLVM_PROFDATA): |
| sys.exit(_LLVM_PROFDATA + ' not found; are you in the chroot?') |
| |
| initial_dir = os.getcwd() |
| temp_dir = tempfile.mkdtemp(prefix='merge_pgo') |
| success = True |
| try: |
| os.chdir(temp_dir) |
| profdata_list = [] |
| |
| for arch in latest: |
| profdata = _find_latest_profdata(arch) |
| profdata_loc = _get_gs_profdata(profdata) |
| profdata_list.append(profdata_loc) |
| |
| if args.tryjob: |
| for tryjob in args.tryjob: |
| profdata = os.path.join(tryjob, 'llvm_profdata.tar.xz') |
| profdata_loc = _get_gs_profdata(profdata) |
| profdata_list.append(profdata_loc) |
| |
| for profdata in profdata_list: |
| if os.path.getsize(profdata_loc) < 512 * 1024: |
| raise RuntimeError('The PGO profile in %s is suspiciously small. ' |
| 'Something might have gone wrong.' % profdata) |
| |
| _merge_profdata(profdata_list, args.output) |
| print('Merged profdata locates at %s\n' % os.path.abspath(args.output)) |
| _tar_and_upload_profdata(args.output) |
| print('Merged profdata uploaded successfully.') |
| except: |
| success = False |
| raise |
| finally: |
| os.chdir(initial_dir) |
| if success: |
| print('Clearing temp directory.') |
| shutil.rmtree(temp_dir, ignore_errors=True) |
| else: |
| print('Script fails, temp directory is at: %s' % temp_dir) |
| |
| |
| if __name__ == '__main__': |
| sys.exit(main()) |