• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#! /usr/bin/env python3
2
3"""
4Download test artifacts from ATI (Android Test Investigate).
5
6This script downloads imgdiag artifacts for specified ATI test runs.
7
8Usage:
9    # Download artifacts from specific runs:
10    ./ati_download_artifacts.py --invocation-id I79100010355578895 --invocation-id I84900010357346481
11
12    # Download latest 10 imgdiag runs:
13    ./ati_download_artifacts.py --imgdiag-atp 10
14
15    # Check all command line flags:
16    ./ati_download_artifacts.py --help
17"""
18
19import tempfile
20import argparse
21import subprocess
22import time
23import os
24import concurrent.futures
25import json
26
27try:
28  from tqdm import tqdm
29except:
30  def tqdm(x, *args, **kwargs):
31    return x
32
33LIST_ARTIFACTS_QUERY = """
34SELECT
35  CONCAT('https://android-build.corp.google.com/builds/', invocation_id, '/', name) AS url
36FROM android_build.testartifacts.latest
37WHERE invocation_id = '{invocation_id}';
38"""
39
40LIST_IMGDIAG_RUNS_QUERY = """
41SELECT t.invocation_id,
42 t.test.name,
43 t.timing.creation_timestamp
44FROM android_build.invocations.latest AS t
45WHERE t.test.name like '%imgdiag_top_100%'
46ORDER BY t.timing.creation_timestamp DESC
47LIMIT {};
48"""
49
50REQUIRED_IMGDIAG_FILES = [
51  "combined_imgdiag_data",
52  "all-dirty-objects",
53  "dirty-image-objects-art",
54  "dirty-image-objects-framework",
55  "dirty-page-counts",
56]
57
58def filter_artifacts(artifacts):
59  return [a for a in artifacts if any(x in a for x in REQUIRED_IMGDIAG_FILES)]
60
61def list_last_imgdiag_runs(run_count):
62  query_file = tempfile.NamedTemporaryFile()
63  out_file = tempfile.NamedTemporaryFile()
64  with open(query_file.name, 'w') as f:
65    f.write(LIST_IMGDIAG_RUNS_QUERY.format(run_count))
66  cmd = f'f1-sql --input_file={query_file.name} --output_file={out_file.name} --csv_output --print_queries=false'
67  res = subprocess.run(cmd, shell=True, check=True, capture_output=True)
68  with open(out_file.name) as f:
69    content = f.read()
70  content = content.split()[1:]
71  for i in range(len(content)):
72    content[i] = content[i].replace('"', '').split(',')
73  return content
74
75def list_artifacts(invocation_id):
76  if not invocation_id:
77    raise ValueError(f'Invalid invocation: {invocation_id}')
78
79  query_file = tempfile.NamedTemporaryFile()
80  out_file = tempfile.NamedTemporaryFile()
81  with open(query_file.name, 'w') as f:
82    f.write(LIST_ARTIFACTS_QUERY.format(invocation_id=invocation_id))
83  cmd = f'f1-sql --input_file={query_file.name} --output_file={out_file.name} --csv_output --print_queries=false --quiet'
84  execute_command(cmd)
85  with open(out_file.name) as f:
86    content = f.read()
87  content = content.split()
88  content = content[1:]
89  for i in range(len(content)):
90    content[i] = content[i].replace('"', '')
91  return content
92
93def execute_command(cmd):
94  for i in range(5):
95    try:
96      subprocess.run(cmd, shell=True, check=True)
97      return
98    except Exception as e:
99      print(f'Failed to run: {cmd}\nException: {e}')
100      time.sleep(2 ** i)
101
102  raise RuntimeError(f'Failed to run: {cmd}')
103
104def download_artifacts(res_dir, artifacts):
105  os.makedirs(res_dir, exist_ok=True)
106
107  commands = []
108  for url in artifacts:
109    filename = url.split('/')[-1]
110    out_path = os.path.join(res_dir, filename)
111    cmd = f'sso_client {url} --connect_timeout=120 --dns_timeout=120 --request_timeout=600 --location > {out_path}'
112    commands.append(cmd)
113
114  if not commands:
115    return
116
117  with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
118    res = list(tqdm(executor.map(execute_command, commands), total=len(commands), leave=False))
119
120
121def download_invocations(args, invocations):
122  for invoc_id in tqdm(invocations):
123    artifacts = list_artifacts(invoc_id)
124    if not args.download_all:
125      artifacts = filter_artifacts(artifacts)
126
127    res_dir = os.path.join(args.out_dir, invoc_id)
128    download_artifacts(res_dir, artifacts)
129
130
131def main():
132  parser = argparse.ArgumentParser(
133    description='Download artifacts from ATI',
134    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
135  )
136  parser.add_argument(
137    '--out-dir',
138    default='./',
139    help='Output dir for downloads',
140  )
141  parser.add_argument(
142    '--invocation-id',
143    action='append',
144    default=None,
145    help='Download artifacts from the specified invocations',
146  )
147  parser.add_argument(
148    '--imgdiag-atp',
149    metavar='N',
150    dest='atp_run_count',
151    type=int,
152    default=None,
153    help='Download latest N imgdiag runs',
154  )
155  parser.add_argument(
156    '--download-all',
157    action=argparse.BooleanOptionalAction,
158    default=False,
159    help='Whether to download all artifacts or combined imgdiag data only',
160  )
161  parser.add_argument(
162    '--overwrite',
163    action=argparse.BooleanOptionalAction,
164    default=False,
165    help='Download artifacts again even if the invocation_id dir already exists',
166  )
167  args = parser.parse_args()
168  if not args.invocation_id and not args.atp_run_count:
169    print('Must specify at least one of: --invocation-id or --imgdiag-atp')
170    return
171
172  invocations = set()
173  if args.invocation_id:
174    invocations.update(args.invocation_id)
175  if args.atp_run_count:
176    recent_runs = list_last_imgdiag_runs(args.atp_run_count)
177    invocations.update({invoc_id for invoc_id, name, timestamp in recent_runs})
178
179  if not args.overwrite:
180    existing_downloads = set()
181    for invoc_id in invocations:
182      res_dir = os.path.join(args.out_dir, invoc_id)
183      if os.path.isdir(res_dir):
184        existing_downloads.add(invoc_id)
185
186    if existing_downloads:
187      print(f'Skipping existing downloads: {existing_downloads}')
188      invocations = invocations - existing_downloads
189
190  print(f'Downloading: {invocations}')
191  download_invocations(args, invocations)
192
193
194if __name__ == '__main__':
195  main()
196