Pirama Arumuga Nainar | f767856 | 2021-04-21 11:36:00 -0700 | [diff] [blame] | 1 | #!/usr/bin/env python3 |
Stephen Hines | 6c6dede | 2020-03-20 19:19:44 -0700 | [diff] [blame] | 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 | |
| 18 | try: |
| 19 | import apiclient.discovery |
| 20 | import apiclient.http |
| 21 | from oauth2client import client as oauth2_client |
| 22 | except ImportError: |
Pirama Arumuga Nainar | f767856 | 2021-04-21 11:36:00 -0700 | [diff] [blame] | 23 | missingImportString = """ |
Stephen Hines | 6c6dede | 2020-03-20 19:19:44 -0700 | [diff] [blame] | 24 | Missing necessary libraries. Try doing the following: |
Pirama Arumuga Nainar | f767856 | 2021-04-21 11:36:00 -0700 | [diff] [blame] | 25 | $ sudo apt-get install python3-pip |
| 26 | $ pip install --user --upgrade google-api-python-client |
| 27 | $ pip install --user --upgrade oauth2client |
| 28 | """ |
Stephen Hines | 6c6dede | 2020-03-20 19:19:44 -0700 | [diff] [blame] | 29 | raise ImportError(missingImportString) |
| 30 | |
Pirama Arumuga Nainar | 19a588f | 2019-11-04 13:22:07 -0800 | [diff] [blame] | 31 | import io |
Pirama Arumuga Nainar | f767856 | 2021-04-21 11:36:00 -0700 | [diff] [blame] | 32 | import getpass |
Pirama Arumuga Nainar | 19a588f | 2019-11-04 13:22:07 -0800 | [diff] [blame] | 33 | import os |
| 34 | |
Pirama Arumuga Nainar | f767856 | 2021-04-21 11:36:00 -0700 | [diff] [blame] | 35 | import utils |
| 36 | |
Pirama Arumuga Nainar | 19a588f | 2019-11-04 13:22:07 -0800 | [diff] [blame] | 37 | ANDROID_BUILD_API_SCOPE = ( |
| 38 | 'https://www.googleapis.com/auth/androidbuild.internal') |
| 39 | ANDROID_BUILD_API_NAME = 'androidbuildinternal' |
| 40 | ANDROID_BUILD_API_VERSION = 'v2beta1' |
Pirama Arumuga Nainar | 19a588f | 2019-11-04 13:22:07 -0800 | [diff] [blame] | 41 | CHUNK_SIZE = 10 * 1024 * 1024 # 10M |
| 42 | |
Yi Kong | f8b68a2 | 2020-07-03 15:46:38 +0800 | [diff] [blame] | 43 | ANDROID_PGO_BUILD = 'pgo-coral-config1' |
Pirama Arumuga Nainar | 19a588f | 2019-11-04 13:22:07 -0800 | [diff] [blame] | 44 | |
Pirama Arumuga Nainar | f767856 | 2021-04-21 11:36:00 -0700 | [diff] [blame] | 45 | STUBBY_COMMAND_PATH = '/google/data/ro/teams/android-llvm/tests/sso_stubby_cmd.sh' |
| 46 | STUBBY_REQUEST = """ |
| 47 | target: {{ |
| 48 | scope: GAIA_USER |
| 49 | name: "{user}@google.com" |
| 50 | }} |
| 51 | target_credential: {{ |
| 52 | type: OAUTH2_TOKEN |
| 53 | oauth2_attributes: {{ |
| 54 | scope: '{scope}' |
| 55 | }} |
| 56 | }} |
| 57 | """ |
| 58 | |
| 59 | |
| 60 | def _get_oauth2_token(): |
| 61 | request = STUBBY_REQUEST.format( |
| 62 | user=getpass.getuser(), scope=ANDROID_BUILD_API_SCOPE) |
| 63 | with open(STUBBY_COMMAND_PATH) as stubby_command_file: |
| 64 | stubby_command = stubby_command_file.read().strip().split() |
| 65 | output = utils.check_output(stubby_command, input=request) |
| 66 | # output is of the format: |
| 67 | # oauth2_token: "<TOKEN>" |
| 68 | return output.split('"')[1] |
| 69 | |
Pirama Arumuga Nainar | 19a588f | 2019-11-04 13:22:07 -0800 | [diff] [blame] | 70 | |
| 71 | class AndroidBuildClient(object): |
| 72 | |
| 73 | def __init__(self): |
Pirama Arumuga Nainar | f767856 | 2021-04-21 11:36:00 -0700 | [diff] [blame] | 74 | creds = oauth2_client.AccessTokenCredentials( |
| 75 | access_token=_get_oauth2_token(), user_agent='unused/1.0') |
Pirama Arumuga Nainar | 19a588f | 2019-11-04 13:22:07 -0800 | [diff] [blame] | 76 | |
| 77 | self.client = apiclient.discovery.build( |
| 78 | ANDROID_BUILD_API_NAME, |
| 79 | ANDROID_BUILD_API_VERSION, |
| 80 | credentials=creds, |
| 81 | discoveryServiceUrl=apiclient.discovery.DISCOVERY_URI) |
| 82 | |
| 83 | # Get the latest test invocation for a given test_tag for a given build. |
| 84 | def get_invocation_id(self, build, test_tag): |
| 85 | request = self.client.testresult().list( |
| 86 | buildId=build, target=ANDROID_PGO_BUILD, attemptId='latest') |
| 87 | |
| 88 | response = request.execute() |
| 89 | testResultWithTag = [ |
| 90 | r for r in response['testResults'] if r['testTag'] == test_tag |
| 91 | ] |
| 92 | if len(testResultWithTag) != 1: |
| 93 | raise RuntimeError( |
| 94 | 'Expected one test with tag {} for build {}. Found {}. Full response is {}' |
| 95 | .format(test_tag, build, len(testResultWithTag), response)) |
| 96 | return testResultWithTag[0]['id'] |
| 97 | |
| 98 | # Get the full artifact name for the zipped PGO profiles |
| 99 | # (_data_local_tmp_pgo_<hash>.zip) for a given <build, test_tag, |
| 100 | # invocation_id>. |
| 101 | def get_test_artifact_name(self, build, test_tag, invocation_id): |
| 102 | request = self.client.testartifact().list( |
| 103 | buildType='submitted', |
| 104 | buildId=build, |
| 105 | target=ANDROID_PGO_BUILD, |
| 106 | attemptId='latest', |
| 107 | testResultId=invocation_id, |
| 108 | maxResults=100) |
| 109 | |
| 110 | response = request.execute() |
| 111 | profile_zip = [ |
| 112 | f for f in response['test_artifacts'] |
| 113 | if f['name'].endswith('zip') and '_data_local_tmp_pgo_' in f['name'] |
| 114 | ] |
| 115 | if len(profile_zip) != 1: |
| 116 | raise RuntimeError( |
| 117 | 'Expected one matching zipfile for invocation {} of {} for build {}. Found {} ({})' |
| 118 | .format(invocation_id, test_tag, build, len(profile_zip), |
| 119 | ', '.join(profile_zip))) |
| 120 | return profile_zip[0]['name'] |
| 121 | |
| 122 | # Download the zipped PGO profiles for a given <build, test_tag, |
| 123 | # invocation_id, artifact_name> into <output_zip>. |
| 124 | def download_test_artifact(self, build, invocation_id, artifact_name, |
| 125 | output_zip): |
| 126 | request = self.client.testartifact().get_media( |
| 127 | buildType='submitted', |
| 128 | buildId=build, |
| 129 | target=ANDROID_PGO_BUILD, |
| 130 | attemptId='latest', |
| 131 | testResultId=invocation_id, |
| 132 | resourceId=artifact_name) |
| 133 | |
| 134 | f = io.FileIO(output_zip, 'wb') |
| 135 | try: |
| 136 | downloader = apiclient.http.MediaIoBaseDownload( |
| 137 | f, request, chunksize=CHUNK_SIZE) |
| 138 | done = False |
| 139 | while not done: |
| 140 | status, done = downloader.next_chunk() |
| 141 | except apiclient.errors.HttpError as e: |
| 142 | if e.resp.status == 404: |
| 143 | raise RuntimeError( |
| 144 | 'Artifact {} does not exist for invocation {} for build {}.' |
| 145 | .format(artifact_name, invocation_id, build)) |
| 146 | |
| 147 | # For a <build, test_tag>, find the invocation_id, artifact_name and |
| 148 | # download the artifact into <output_dir>/pgo_profiles.zip. |
| 149 | def download_pgo_zip(self, build, test_tag, output_dir): |
| 150 | output_zip = os.path.join(output_dir, 'pgo_profiles.zip') |
| 151 | |
| 152 | invocation_id = self.get_invocation_id(build, test_tag) |
| 153 | artifact_name = self.get_test_artifact_name(build, test_tag, |
| 154 | invocation_id) |
| 155 | self.download_test_artifact(build, invocation_id, artifact_name, |
| 156 | output_zip) |
| 157 | return output_zip |