xref: /aosp_15_r20/tools/treble/build/sandbox/rbe_action.py (revision 105f628577ac4ba0e277a494fbb614ed8c12a994)
1# Copyright 2020 Google LLC
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""Runs a single action remotely with RBE."""
16
17import argparse
18import os
19import rbe
20import subprocess
21import sys
22
23
24def main():
25  parser = argparse.ArgumentParser(
26      description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
27  parser.add_argument(
28      '--command',
29      default='echo RBE check successful.',
30      help='Command to run remotely with RBE.')
31  parser.add_argument(
32      '--print', '-p',
33      action='store_true',
34      help='Prints the executed commands')
35  args = parser.parse_args()
36
37  env = []
38  cleanup = rbe.setup(env, sys.stdout if args.print else subprocess.DEVNULL)
39  src_root = os.path.normpath(
40      os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../../..'))
41  env = rbe.env_array_to_dict(rbe.prepare_env(env))
42  env['PATH'] = os.getenv('PATH')
43  env['HOME'] = os.getenv('HOME')
44  for d in ['FLAG_log_dir', 'RBE_output_dir', 'RBE_proxy_log_dir']:
45    env[d] = '/tmp'  # We want the logs in /tmp instead of out.
46  try:
47    # Bootstrap the RBE proxy.
48    bootstrap_cmd = rbe.get_nsjail_bin_wrapper() + \
49      [os.path.join(rbe.TOOLS_DIR, 'bootstrap')]
50    shell_env = ' '.join(['%s=%s' % (k,v) for k, v in env.items()])
51    if args.print:
52      print('Bootstrap RBE reproxy:')
53      print('cd ' + src_root)
54      print('%s %s' % (shell_env, ' '.join(bootstrap_cmd)))
55    subprocess.check_call(
56        bootstrap_cmd, env=env, cwd=src_root, stdout=subprocess.DEVNULL)
57    # Execute the remote command.
58    rewrapper_cmd = rbe.get_nsjail_bin_wrapper() + [
59        os.path.join(rbe.TOOLS_DIR, 'rewrapper'),
60        '--platform=container-image=docker://gcr.io/androidbuild-re-dockerimage/android-build-remoteexec-image@sha256:582efb38f0c229ea39952fff9e132ccbe183e14869b39888010dacf56b360d62', \
61        '--labels=type=tool',
62        '--exec_strategy=remote',
63        '--dial_timeout=5s',
64        '--exec_root=' + src_root,
65        '--',
66    ] + args.command.split()
67    if args.print:
68      print('Run remote command with RBE:')
69      print('%s %s' % (shell_env, ' '.join(rewrapper_cmd)))
70    subprocess.check_call(rewrapper_cmd, env=env, cwd=src_root)
71  finally:
72    # Shut down the RBE proxy.
73    if args.print:
74      print('RBE proxy shutdown:')
75      print('killall reproxy')
76    subprocess.call(
77        ['killall', 'reproxy'],
78        stdout=subprocess.DEVNULL,
79        stderr=subprocess.DEVNULL)
80    cleanup()
81
82
83if __name__ == '__main__':
84  main()
85