1#!/usr/bin/env python
2#
3# Copyright (C) 2019 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17# Sample Usage:
18# $ python update_profiles.py 500000 git_master ALL --profdata-suffix 2019-04-15
19#
20# Additional/frequently-used arguments:
21#   -b BUG adds a 'Bug: <BUG>' to the commit message when adding the profiles.
22#   --do-not-merge adds a 'DO NOT MERGE' tag to the commit message to restrict
23#                  automerge of profiles from release branches.
24#
25# Try '-h' for a full list of command line arguments.
26
27import argparse
28import os
29import shutil
30import subprocess
31import sys
32import tempfile
33import zipfile
34
35import utils
36
37X20_BASE_LOCATION = '/google/data/ro/teams/android-pgo-data'
38
39class Benchmark(object):
40    def __init__(self, name):
41        self.name = name
42
43    def x20_profile_location(self):
44        raise NotImplementedError()
45
46    def apct_job_name(self):
47        raise NotImplementedError()
48
49    def profdata_file(self, suffix=''):
50        profdata = os.path.join(self.name, '{}.profdata'.format(self.name))
51        if suffix:
52            profdata += '.' + suffix
53        return profdata
54
55    def profraw_files(self):
56        raise NotImplementedError()
57
58    def merge_profraws(self, profile_dir, output):
59        profraws = [os.path.join(profile_dir, p) for p in self.profraw_files(profile_dir)]
60        utils.run_llvm_profdata(profraws, output)
61
62
63class NativeExeBenchmark(Benchmark):
64    def apct_job_name(self):
65        return 'pgo-collector'
66
67    def x20_profile_location(self):
68        return os.path.join(X20_BASE_LOCATION, 'raw')
69
70    def profraw_files(self, profile_dir):
71        if self.name == 'hwui':
72            return ['hwuimacro.profraw', 'hwuimacro_64.profraw',
73                    'hwuimicro.profraw', 'hwuimicro_64.profraw',
74                    'skia_nanobench.profraw', 'skia_nanobench_64.profraw']
75        elif self.name == 'hwbinder':
76            return ['hwbinder.profraw', 'hwbinder_64.profraw']
77
78
79class APKBenchmark(Benchmark):
80    def apct_job_name(self):
81        return 'apk-pgo-collector'
82
83    def x20_profile_location(self):
84        return os.path.join(X20_BASE_LOCATION, 'apk-raw')
85
86    def profdata_file(self, suffix=''):
87        profdata = os.path.join('art', '{}_arm_arm64.profdata'.format(self.name))
88        if suffix:
89            profdata += '.' + suffix
90        return profdata
91
92    def profraw_files(self, profile_dir):
93        return os.listdir(profile_dir)
94
95
96def BenchmarkFactory(benchmark_name):
97    if benchmark_name == 'dex2oat':
98        return APKBenchmark(benchmark_name)
99    elif benchmark_name in ['hwui', 'hwbinder']:
100        return NativeExeBenchmark(benchmark_name)
101    else:
102        raise RuntimeError('Unknown benchmark ' + benchmark_name)
103
104
105def extract_profiles(benchmark, branch, build, output_dir):
106    # The APCT results are stored in
107    #   <x20_profile_base>/<branch>/<build>/<apct_job_name>/<arbitrary_invocation_dir>/
108    #
109    # The PGO files are in _data_local_tmp_<id>.zip in the above directory.
110
111    profile_base = os.path.join(benchmark.x20_profile_location(), branch, build,
112                                benchmark.apct_job_name())
113    invocation_dirs = os.listdir(profile_base)
114
115    if len(invocation_dirs) == 0:
116        raise RuntimeError('No invocations found in {}'.format(profile_base))
117    if len(invocation_dirs) > 1:
118        # TODO Add option to pick/select an invocation from the command line.
119        raise RuntimeError('More than one invocation found in {}'.format(profile_base))
120
121    profile_dir = os.path.join(profile_base, invocation_dirs[0])
122    zipfiles = [f for f in os.listdir(profile_dir) if f.startswith('_data_local_tmp')]
123
124    if len(zipfiles) != 1:
125        raise RuntimeError('Expected one zipfile in {}.  Found {}'.format(profile_dir,
126                                                                          len(zipfiles)))
127
128    zipfile_name = os.path.join(profile_dir, zipfiles[0])
129    zip_ref = zipfile.ZipFile(zipfile_name)
130    zip_ref.extractall(output_dir)
131    zip_ref.close()
132
133
134KNOWN_BENCHMARKS = ['ALL', 'dex2oat', 'hwui', 'hwbinder']
135
136def parse_args():
137    """Parses and returns command line arguments."""
138    parser = argparse.ArgumentParser()
139
140    parser.add_argument(
141        'build', metavar='BUILD',
142        help='Build number to pull from the build server.')
143
144    parser.add_argument(
145        '-b', '--bug', type=int,
146        help='Bug to reference in commit message.')
147
148    parser.add_argument(
149        '--use-current-branch', action='store_true',
150        help='Do not repo start a new branch for the update.')
151
152    parser.add_argument(
153        '--add-do-not-merge', action='store_true',
154        help='Add \'DO NOT MERGE\' to the commit message.')
155
156    parser.add_argument(
157        '--profdata-suffix', type=str, default='',
158        help='Suffix to append to merged profdata file')
159
160    parser.add_argument(
161        'branch', metavar='BRANCH',
162        help='Fetch profiles for BRANCH (e.g. git_qt-release)')
163
164    parser.add_argument(
165        'benchmark', metavar='BENCHMARK',
166        help='Update profiles for BENCHMARK.  Choices are {}'.format(KNOWN_BENCHMARKS))
167
168    parser.add_argument(
169        '--skip-cleanup', '-sc',
170        action='store_true',
171        default=False,
172        help='Skip the cleanup, and leave intermediate files (in /tmp/pgo-profiles-*)')
173
174    return parser.parse_args()
175
176
177def get_current_profile(benchmark):
178    profile = benchmark.profdata_file()
179    dirname, basename = os.path.split(profile)
180
181    old_profiles = [f for f in os.listdir(dirname) if f.startswith(basename)]
182    if len(old_profiles) == 0:
183        return ''
184    return os.path.join(dirname, old_profiles[0])
185
186
187def main():
188    args = parse_args()
189
190    if args.benchmark == 'ALL':
191        worklist = KNOWN_BENCHMARKS[1:]
192    else:
193        worklist  = [args.benchmark]
194
195    profiles_project = os.path.join(utils.android_build_top(), 'toolchain',
196                                    'pgo-profiles')
197    os.chdir(profiles_project)
198
199    if not args.use_current_branch:
200        branch_name = 'update-profiles-' + args.build
201        utils.check_call(['repo', 'start', branch_name, '.'])
202
203    for benchmark_name in worklist:
204        benchmark = BenchmarkFactory(benchmark_name)
205
206        # Existing profile file, which gets 'rm'-ed from 'git' down below.
207        current_profile = get_current_profile(benchmark)
208
209        # Extract profiles to a temporary directory.  After extraction, we
210        # expect to find one subdirectory with profraw files under the temporary
211        # directory.
212        extract_dir = tempfile.mkdtemp(prefix='pgo-profiles-'+benchmark_name)
213        extract_profiles(benchmark, args.branch, args.build, extract_dir)
214
215        if len(os.listdir(extract_dir)) != 1:
216            raise RuntimeError("Expected one subdir under {}".format(extract_dir))
217
218        extract_subdir = os.path.join(extract_dir, os.listdir(extract_dir)[0])
219
220        # Merge profiles.
221        profdata = benchmark.profdata_file(args.profdata_suffix)
222        benchmark.merge_profraws(extract_subdir, profdata)
223
224        # Construct 'git' commit message.
225        message_lines = [
226                'Update PGO profiles for {}'.format(benchmark_name), '',
227                'The profiles are from build {}.'.format(args.build), ''
228        ]
229
230        if args.add_do_not_merge:
231            message_lines[0] = '[DO NOT MERGE] ' + message_lines[0]
232
233        if args.bug:
234            message_lines.append('')
235            message_lines.append('Bug: http://b/{}'.format(args.bug))
236        message_lines.append('Test: Build (TH)')
237        message = '\n'.join(message_lines)
238
239        # Invoke git: Delete current profile, add new profile and commit these
240        # changes.
241        if current_profile:
242            utils.check_call(['git', 'rm', current_profile])
243        utils.check_call(['git', 'add', profdata])
244        utils.check_call(['git', 'commit', '-m', message])
245
246        if not args.skip_cleanup:
247            shutil.rmtree(extract_dir)
248
249
250if __name__ == '__main__':
251    main()
252