• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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