aboutsummaryrefslogtreecommitdiff
path: root/afdo_tools/generate_afdo_from_tryjob.py
blob: b8a2d669aa64ec55b78e406e37fdef4e88f672c2 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
#!/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.

"""Given a tryjob and perf profile, generates an AFDO profile."""

from __future__ import print_function

import argparse
import distutils.spawn
import os
import os.path
import shutil
import subprocess
import sys
import tempfile

_CREATE_LLVM_PROF = 'create_llvm_prof'
_GS_PREFIX = 'gs://'


def _fetch_gs_artifact(remote_name, local_name):
  assert remote_name.startswith(_GS_PREFIX)
  subprocess.check_call(['gsutil', 'cp', remote_name, local_name])


def _fetch_and_maybe_unpack(remote_name, local_name):
  unpackers = [
      ('.tar.bz2', ['tar', 'xaf']),
      ('.bz2', ['bunzip2']),
      ('.tar.xz', ['tar', 'xaf']),
      ('.xz', ['xz', '-d']),
  ]

  unpack_ext = None
  unpack_cmd = None
  for ext, unpack in unpackers:
    if remote_name.endswith(ext):
      unpack_ext, unpack_cmd = ext, unpack
      break

  download_to = local_name + unpack_ext if unpack_ext else local_name
  _fetch_gs_artifact(remote_name, download_to)
  if unpack_cmd is not None:
    print('Unpacking', download_to)
    subprocess.check_output(unpack_cmd + [download_to])
    assert os.path.exists(local_name)


def _generate_afdo(perf_profile_loc, tryjob_loc, output_name):
  if perf_profile_loc.startswith(_GS_PREFIX):
    local_loc = 'perf.data'
    _fetch_and_maybe_unpack(perf_profile_loc, local_loc)
    perf_profile_loc = local_loc

  chrome_in_debug_loc = 'debug/opt/google/chrome/chrome.debug'
  debug_out = 'debug.tgz'
  _fetch_gs_artifact(os.path.join(tryjob_loc, 'debug.tgz'), debug_out)

  print('Extracting chrome.debug.')
  # This has tons of artifacts, and we only want Chrome; don't waste time
  # extracting the rest in _fetch_and_maybe_unpack.
  subprocess.check_call(['tar', 'xaf', 'debug.tgz', chrome_in_debug_loc])

  # Note that the AFDO tool *requires* a binary named `chrome` to be present if
  # we're generating a profile for chrome. It's OK for this to be split debug
  # information.
  os.rename(chrome_in_debug_loc, 'chrome')

  print('Generating AFDO profile.')
  subprocess.check_call([
      _CREATE_LLVM_PROF, '--out=' + output_name, '--binary=chrome',
      '--profile=' + perf_profile_loc
  ])


def _abspath_or_gs_link(path):
  if path.startswith(_GS_PREFIX):
    return path
  return os.path.abspath(path)


def _tryjob_arg(tryjob_arg):
  # Forward gs args through
  if tryjob_arg.startswith(_GS_PREFIX):
    return tryjob_arg

  # Clicking on the 'Artifacts' link gives us a pantheon link that's basically
  # a preamble and gs path.
  pantheon = 'https://pantheon.corp.google.com/storage/browser/'
  if tryjob_arg.startswith(pantheon):
    return _GS_PREFIX + tryjob_arg[len(pantheon):]

  # Otherwise, only do things with a tryjob ID (e.g. R75-11965.0.0-b3648595)
  if not tryjob_arg.startswith('R'):
    raise ValueError('Unparseable tryjob arg; give a tryjob ID, pantheon '
                     'link, or gs:// link. Please see source for more.')

  chell_path = 'chromeos-image-archive/chell-chrome-pfq-tryjob/'
  # ...And assume it's from chell, since that's the only thing we generate
  # profiles with today.
  return _GS_PREFIX + chell_path + tryjob_arg


def main():
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument(
      '--perf_profile',
      required=True,
      help='Path to our perf profile. Accepts either a gs:// path or local '
      'filepath.')
  parser.add_argument(
      '--tryjob',
      required=True,
      type=_tryjob_arg,
      help='Path to our tryjob\'s artifacts. Accepts a gs:// path, pantheon '
      'link, or tryjob ID, e.g. R75-11965.0.0-b3648595. In the last case, '
      'the assumption is that you ran a chell-chrome-pfq-tryjob.')
  parser.add_argument(
      '-o',
      '--output',
      default='afdo.prof',
      help='Where to put the AFDO profile. Default is afdo.prof.')
  parser.add_argument(
      '-k',
      '--keep_artifacts_on_failure',
      action='store_true',
      help='Don\'t remove the tempdir on failure')
  args = parser.parse_args()

  if not distutils.spawn.find_executable(_CREATE_LLVM_PROF):
    sys.exit(_CREATE_LLVM_PROF + ' not found; are you in the chroot?')

  profile = _abspath_or_gs_link(args.perf_profile)
  afdo_output = os.path.abspath(args.output)

  initial_dir = os.getcwd()
  temp_dir = tempfile.mkdtemp(prefix='generate_afdo')
  success = True
  try:
    os.chdir(temp_dir)
    _generate_afdo(profile, args.tryjob, afdo_output)

    # The AFDO tooling is happy to generate essentially empty profiles for us.
    # Chrome's profiles are often 8+ MB; if we only see a small fraction of
    # that, something's off. 512KB was arbitrarily selected.
    if os.path.getsize(afdo_output) < 512 * 1024:
      raise ValueError('The AFDO profile is suspiciously small for Chrome. '
                       'Something might have gone wrong.')
  except:
    success = False
    raise
  finally:
    os.chdir(initial_dir)

    if success or not args.keep_artifacts_on_failure:
      shutil.rmtree(temp_dir, ignore_errors=True)
    else:
      print('Artifacts are available at', temp_dir)


if __name__ == '__main__':
  sys.exit(main())