build_util.py

import os
from typing import Iterable
import subprocess
import sys
import shutil


def check_env_flag(name: str, default: str = '') -> bool:
  return os.getenv(name, default).upper() in ['ON', '1', 'YES', 'TRUE', 'Y']


def bazel_options_from_env() -> Iterable[str]:
  bazel_flags = []

  if check_env_flag('DEBUG'):
    bazel_flags.append('--config=dbg')

  if check_env_flag('TPUVM_MODE'):
    bazel_flags.append('--config=tpu')

  gcloud_key_file = os.getenv('GCLOUD_SERVICE_KEY_FILE', default='')
  # Remote cache authentication.
  if gcloud_key_file:
    # Temporary workaround to allow PRs from forked repo to run CI. See details at (#5259).
    # TODO: Remove the check once self-hosted GHA workers are available to CPU/GPU CI.
    gcloud_key_file_size = os.path.getsize(gcloud_key_file)
    if gcloud_key_file_size > 1:
      bazel_flags.append('--google_credentials=%s' % gcloud_key_file)
      bazel_flags.append('--config=remote_cache')
  else:
    if check_env_flag('BAZEL_REMOTE_CACHE'):
      bazel_flags.append('--config=remote_cache')

  cache_silo_name = os.getenv('SILO_NAME', default='dev')
  if cache_silo_name:
    bazel_flags.append('--remote_default_exec_properties=cache-silo-key=%s' %
                       cache_silo_name)

  if check_env_flag('BUILD_CPP_TESTS', default='0'):
    bazel_flags.append('//test/cpp:all')
    bazel_flags.append('//torch_xla/csrc/runtime:all')

  bazel_jobs = os.getenv('BAZEL_JOBS', default='')
  if bazel_jobs:
    bazel_flags.append('--jobs=%s' % bazel_jobs)

  # Build configuration.
  if check_env_flag('BAZEL_VERBOSE'):
    bazel_flags.append('-s')
  if check_env_flag('XLA_CUDA'):
    bazel_flags.append('--config=cuda')
  if check_env_flag('XLA_CPU_USE_ACL'):
    bazel_flags.append('--config=acl')

  return bazel_flags


def bazel_build(bazel_target: str,
                destination_dir: str,
                options: Iterable[str] = []):
  bazel_argv = [
      'bazel', 'build', bazel_target,
      f"--symlink_prefix={os.path.join(os.getcwd(), 'bazel-')}"
  ]

  # Remove duplicated flags because they confuse bazel
  flags = set(bazel_options_from_env() + options)
  bazel_argv.extend(flags)

  print(' '.join(bazel_argv), flush=True)
  subprocess.check_call(bazel_argv, stdout=sys.stdout, stderr=sys.stderr)

  target_path = bazel_target.replace('@xla//', 'external/xla/').replace(
      '//', '').replace(':', '/')
  output_path = os.path.join('bazel-bin', target_path)
  output_filename = os.path.basename(output_path)

  if not os.path.exists(destination_dir):
    os.makedirs(destination_dir)

  shutil.copyfile(output_path, os.path.join(destination_dir, output_filename))