infra/ci/worker/perf_metrics_uploader.py

#!/usr/bin/env python3
# Copyright (C) 2019 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import json
import hashlib
import sys

from config import DB, PROJECT
from common_utils import req, SCOPES
'''
Uploads the performance metrics of the Perfetto tests to StackDriver and
Firebase.

The expected format of the JSON is as follows:
{
  metrics: [
    {
      'metric': *metric name*,
      'value': *metric value*,
      'unit': *either s (seconds) or b (bytes)*,
      'tags': {
        *tag name*: *tag value*,
        ...
      },
      'labels': {
        *label name*: *label value*,
        ...
      }
    },
    ...
  ]
}
'''

STACKDRIVER_API = 'https://monitoring.googleapis.com/v3/projects/%s' % PROJECT
SCOPES.append('https://www.googleapis.com/auth/firebase.database')
SCOPES.append('https://www.googleapis.com/auth/userinfo.email')
SCOPES.append('https://www.googleapis.com/auth/monitoring.write')


def sha1(obj):
  hasher = hashlib.sha1()
  hasher.update(
      json.dumps(obj, sort_keys=True, separators=(',', ':')).encode('utf-8'))
  return hasher.hexdigest()


def metric_list_to_hash_dict(raw_metrics):
  metrics = {}
  for metric in raw_metrics:
    key = '%s-%s' % (metric['metric'], sha1(metric['tags']))
    metrics[key] = metric
  return metrics


def create_stackdriver_metrics(ts, metrics):
  # Chunk up metrics into 100 element chunks to comply with Stackdriver's
  # restrictions on the number of metrics in a request.
  metrics_list = list(metrics.values())
  metric_chunks = [metrics_list[x:x + 100] for x in range(0, len(metrics), 100)]
  desc_chunks = []

  for chunk in metric_chunks:
    desc = {'timeSeries': []}
    for metric in chunk:
      metric_name = metric['metric']
      desc['timeSeries'] += [{
          'metric': {
              'type':
                  'custom.googleapis.com/perfetto-ci/perf/%s' % metric_name,
              'labels':
                  dict(
                      list(metric.get('tags', {}).items()) +
                      list(metric.get('labels', {}).items())),
          },
          'resource': {
              'type': 'global'
          },
          'points': [{
              'interval': {
                  'endTime': ts
              },
              'value': {
                  'doubleValue': str(metric['value'])
              }
          }]
      }]
    desc_chunks.append(desc)
  return desc_chunks


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--job-id',
      type=str,
      required=True,
      help='The Perfetto CI job ID to tie this upload to')
  parser.add_argument(
      'metrics_file', type=str, help='File containing the metrics to upload')
  args = parser.parse_args()

  with open(args.metrics_file, 'r') as metrics_file:
    raw_metrics = json.loads(metrics_file.read())

  job = req('GET', '%s/jobs/%s.json' % (DB, args.job_id))
  ts = job['time_started']

  metrics = metric_list_to_hash_dict(raw_metrics['metrics'])
  req('PUT', '%s/perf/%s.json' % (DB, args.job_id), body=metrics)

  # Only upload Stackdriver metrics for post-submit runs.
  git_ref = job['env'].get('PERFETTO_TEST_GIT_REF')
  if git_ref == 'refs/heads/main':
    sd_metrics_chunks = create_stackdriver_metrics(ts, metrics)
    for sd_metrics in sd_metrics_chunks:
      req('POST', STACKDRIVER_API + '/timeSeries', body=sd_metrics)

  return 0


if __name__ == '__main__':
  sys.exit(main())