bazel/ci/ci.py

"""Module for running CI targets."""

import argparse
import logging
import os
import pathlib
import platform
import subprocess
import sys
from typing import Callable, List
import uuid

from tools.base.bazel.ci import bazel
from tools.base.bazel.ci import errors
from tools.base.bazel.ci import flake_reruns
from tools.base.bazel.ci import owners_checks
from tools.base.bazel.ci import query_checks
from tools.base.bazel.ci import studio_linux
from tools.base.bazel.ci import studio_mac
from tools.base.bazel.ci import studio_nightly
from tools.base.bazel.ci import studio_win
from tools.base.bazel.ci import studio_evals

_ARCH_ALIAS = {
    'amd64': 'x86_64',
    'aarch64': 'arm64',
}


class CI:
  """Continuous Integration wrapper.

  This is used to run multiple functions that may fail, but
  should be aggregated and displayed before exiting.
  """

  exceptions: List[BaseException] = []
  build_env: bazel.BuildEnv

  def __init__(self, build_env: bazel.BuildEnv):
    self.build_env = build_env
    self.exit_code = 0

  def run(self, func: Callable[[bazel.BuildEnv], None]):
    """Runs callables that might raise exceptions."""
    try:
      func(self.build_env)
    except errors.CIError as e:
      self.exceptions.append(e)
      self.exit_code = e.exit_code
    except subprocess.CalledProcessError as e:
      stderr = (e.stderr or b'').decode('utf-8')
      msg = f'#### internal subprocess error\n{e}\n{stderr}'
      self.exceptions.append(RuntimeError(msg))
      self.exit_code = 1

  def has_errors(self) -> bool:
    """Returns true if there were exceptions."""
    return bool(self.exceptions)

  def print_errors(self):
    """Write CI exceptions to stderr."""
    for err in self.exceptions:
      print(err, file=sys.stderr)


def find_workspace() -> str:
  """Returns the path of the WORKSPACE directory."""
  bazel_workspace = os.environ.get('BUILD_WORKSPACE_DIRECTORY')
  if bazel_workspace:
    return bazel_workspace
  raise RuntimeError('Missing environment variable: BUILD_WORKSPACE_DIRECTORY')


def studio_build_checks(ci: CI):
  """Runs checks against the build graph."""
  ci.run(query_checks.cquery_all)
  ci.run(query_checks.no_local_genrules)
  ci.run(query_checks.require_cpu_tags)
  ci.run(query_checks.gradle_requires_cpu4_or_more)
  ci.run(owners_checks.require_component_id)
  ci.run(query_checks.check_large_machine_allowlist)
  ci.run(query_checks.check_docker_network_allowlist)

  def validate_coverage_graph(env: bazel.BuildEnv):
    inv_id = uuid.uuid4()
    result = env.bazel_build(
        '--config=ci',
        '--nobuild',
        f'--invocation_id={inv_id}',
        '--',
        '@cov//:all.suite',
    )
    if result.returncode:
      raise RuntimeError((
          'Coverage build configuration is broken; you may need to update'
          ' tools/base/bazel/coverage/BUILD\n'
          f'\n See https://fusion2.corp.google.com/invocations/{inv_id}'
      ))

  # ci.run(validate_coverage_graph)
  if not ci.exceptions:
    return
  # Write the exceptions to a file, so Android Build shows a clear and
  # understandable failure message, instead of truncating bazel output.
  error_log = pathlib.Path(ci.build_env.dist_dir) / 'logs/build_error.log'
  with open(error_log, 'w') as f:
    for e in ci.exceptions:
      f.write(str(e))


def _get_bazel_path() -> str:
  host_os = platform.system().lower()
  host_arch = platform.machine().lower()
  host_arch = _ARCH_ALIAS.get(host_arch, host_arch)
  bazel_path = os.path.join(
      find_workspace(),
      'prebuilts',
      'tools',
      f'{host_os}-{host_arch}',
      'bazel',
      'bazelisk.exe' if host_os == 'windows' else 'bazelisk',
  )
  return bazel_path


def main():
  """Runs the CI target command.

  If any commands raise a CIError, this exits with the exit code of the last
  exception raised.
  """
  parser = argparse.ArgumentParser()
  parser.add_argument('target', help='The name of the CI target')
  args = parser.parse_args()

  build_env = bazel.make_build_env(bazel_path=_get_bazel_path())
  ci = CI(build_env=build_env)

  if build_env.dist_dir:
    logs_dir = os.path.join(build_env.dist_dir, 'logs')
    os.makedirs(logs_dir, exist_ok=True)
    logging.basicConfig(
        filename=os.path.join(logs_dir, 'ci.log'),
        format='%(asctime)s %(levelname)s:%(message)s',
        level=logging.INFO,
    )
  else:
    logging.basicConfig(level=logging.INFO)

  match args.target:
    case 'studio-build-checks':
      studio_build_checks(ci)
    case 'studio-linux':
      ci.run(studio_linux.studio_linux)
    case 'studio-linux-large':
      ci.run(studio_linux.studio_linux_large)
    case 'studio-linux_very_flaky':
      ci.run(studio_linux.studio_linux_very_flaky)
    case 'studio-linux-flake-reruns':
      ci.run(flake_reruns.studio_linux_flake_reruns)
    case 'studio-win':
      ci.run(studio_win.studio_win)
    case 'studio-win-canary':
      ci.run(studio_win.studio_win_canary)
    case 'studio-win-flake-reruns':
      ci.run(flake_reruns.studio_win_flake_reruns)
    case 'studio-mac':
      ci.run(studio_mac.studio_mac)
    case 'studio-mac-arm':
      ci.run(studio_mac.studio_mac_arm)
    case 'studio-nightly':
      ci.run(studio_nightly.studio_nightly)
    case 'studio-autobot':
      pass  # No-op. To be filled in later.
    case 'studio-evals':
        ci.run(studio_evals.studio_evals)
    case 'uitools_evals': # deprecated, to be removed
      ci.run(studio_evals.studio_evals)
    case _:
      raise NotImplementedError(f'target: "{args.target}" does not exist')

  if ci.has_errors():
    ci.print_errors()
    sys.exit(ci.exit_code)


if __name__ == '__main__':
  main()