xref: /aosp_15_r20/tools/asuite/atest/profiler.py (revision c2e18aaa1096c836b086f94603d04f4eb9cf37f5)
1#!/usr/bin/env python3
2# Copyright 2022, The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""Script of Atest Profiler."""
17
18import argparse
19import os
20import shutil
21import subprocess
22import sys
23import tempfile
24import zipfile
25
26# This is mostly a copy of Soong's stub_template_host.txt, but with changes to
27# run soong python executables via a profiler tool. This is just a hack, soong
28# should really have the ability to build profiled binaries directly.
29
30
31def main():
32  """Main method that runs python profilers."""
33  parser = argparse.ArgumentParser(
34      description='Runs a soong-built python binary under a profiler.'
35  )
36  parser.add_argument(
37      'profiler',
38      choices=['pyinstrument', 'cProfile'],
39      help='The profiler to use',
40  )
41  parser.add_argument('profile_file', help='The output file of the profiler')
42  parser.add_argument(
43      'executable',
44      help='The soong-built python binary (with embedded_launcher: false)',
45  )
46  args, args_for_executable = parser.parse_known_args()
47
48  if not os.path.isfile(args.executable):
49    sys.exit(f'{args.executable}: File not found')
50  os.makedirs(os.path.dirname(args.profile_file), exist_ok=True)
51
52  runfiles_path = tempfile.mkdtemp(prefix='Soong.python_')
53  try:
54    _zf = zipfile.ZipFile(args.executable)
55    _zf.extractall(runfiles_path)
56    _zf.close()
57
58    sys.exit(
59        subprocess.call(
60            [
61                'python3',
62                '-m',
63                args.profiler,
64                '-o',
65                args.profile_file,
66                os.path.join(
67                    runfiles_path, '__soong_entrypoint_redirector__.py'
68                ),
69            ]
70            + args_for_executable,
71            close_fds=False,
72        )
73    )
74
75  finally:
76    shutil.rmtree(runfiles_path, ignore_errors=True)
77
78
79if __name__ == '__main__':
80  main()
81