#!/usr/bin/python3 # Copyright 2024 The OpenXLA Authors. # # 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. # ============================================================================== r"""XLA build script for use in CI. This build script aims to be completely agnostic to the specifics of the VM, the exceptions are uses of `KOKORO_ARTIFACTS_DIR` and `GITHUB_WORKSPACE` to know where JAX or TensorFlow lives depending on which build is being executed. To update the goldens associated with this file, run: ```PYTHONDONTWRITEBYTECODE=1 python3 build.py \ --dump_commands > golden_commands.txt``` """ import argparse import dataclasses import enum import logging import os import subprocess import sys from typing import Any, ClassVar, Dict, List, Tuple # TODO(ddunleavy): move this to the bazelrc _DEFAULT_BAZEL_OPTIONS = dict( color="yes", test_output="errors", verbose_failures=True, keep_going=True, nobuild_tests_only=True, profile="profile.json.gz", flaky_test_attempts=3, jobs=150, bes_upload_mode="fully_async", ) _KW_ONLY_IF_PYTHON310 = {"kw_only": True} if sys.version_info >= (3, 10) else {} _XLA_DEFAULT_TARGET_PATTERNS = ( "//xla/...", "//build_tools/...", "@tsl//tsl/...", ) _XLA_CPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS = ( "//xla/tools/multihost_hlo_runner:hlo_runner_main", "//xla/tools:compute_xspace_stats_main", ) _XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS = ( "//xla/tools/multihost_hlo_runner:hlo_runner_main_gpu", "//xla/tools:compute_xspace_stats_main_gpu", ) _GITHUB_WORKSPACE = os.environ.get("GITHUB_WORKSPACE", "$GITHUB_WORKSPACE") def retry( args: List[str], delay_seconds: int = 15, retries: int = 3 ) -> List[str]: # Possibly a slight abuse of `parallel` as nothing happens in parallel, just # retries with delay if the command fails. # pyformat:disable return [ "parallel", "--ungroup", "--retries", str(retries), "--delay", str(delay_seconds), "--nonall", "--", *args, ] def sh(args, check=True, **kwargs): logging.info("Starting process: %s", " ".join(args)) return subprocess.run(args, check=check, **kwargs) def _dict_to_cli_options(d: Dict[str, Any]) -> List[str]: # pylint: disable=g-bool-id-comparison return [f"--{k}" if v is True else f"--{k}={v}" for k, v in d.items()] class BuildType(enum.Enum): """Enum representing all types of builds. Should be named as `REPO,OS,HOST_TYPE,BACKEND,GPU_TYPE,CI_TYPE`. """ XLA_LINUX_X86_CPU_GITHUB_ACTIONS = enum.auto() XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS = enum.auto() XLA_LINUX_ARM64_CPU_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_GPU_L4_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_GPU_8X_H100_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_GPU_ONEAPI_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_GPU_HERMETIC_ROCM_GITHUB_ACTIONS = enum.auto() # Presubmit builds for regression testing. XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS = enum.auto() XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS = enum.auto() XLA_MACOS_X86_CPU_KOKORO = enum.auto() XLA_MACOS_ARM64_CPU_KOKORO = enum.auto() JAX_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS = enum.auto() JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS = enum.auto() JAX_LINUX_X86_GPU_L4_GITHUB_ACTIONS = enum.auto() @classmethod def from_str(cls, s): try: return cls[s.replace(" ", "_").upper()] except KeyError: # Sloppy looking exception handling, but argparse will catch ValueError # and give a pleasant error message. KeyError would not work here. raise ValueError # pylint: disable=raise-missing-from @dataclasses.dataclass(frozen=True, **_KW_ONLY_IF_PYTHON310) class Build: """Class representing a build of XLA.""" _builds: ClassVar[Dict[BuildType, "Build"]] = {} type_: BuildType repo: str target_patterns: Tuple[str, ...] subcommand: str = "test" configs: Tuple[str, ...] = () build_tag_filters: Tuple[str, ...] = () test_tag_filters: Tuple[str, ...] = () action_env: Dict[str, Any] = dataclasses.field(default_factory=dict) test_env: Dict[str, Any] = dataclasses.field(default_factory=dict) repo_env: Dict[str, Any] = dataclasses.field(default_factory=dict) override_repository: Dict[str, str] = dataclasses.field(default_factory=dict) override_module: Dict[str, str] = dataclasses.field(default_factory=dict) options: Dict[str, Any] = dataclasses.field(default_factory=dict) startup_options: Dict[str, Any] = dataclasses.field(default_factory=dict) extra_setup_commands: Tuple[List[str], ...] = () def __post_init__(self): # pylint: disable=protected-access assert ( self.type_ not in self.__class__._builds ), "Can't have multiple builds of same BuildType!" assert ( self.repo == "openxla/xla" or self.override_repository or self.override_module ), "Must override repo if repo under test isn't XLA!" self.__class__._builds[self.type_] = self @classmethod def all_builds(cls): return cls._builds def bazel_command( self, subcommand: str = "test", extra_options: Tuple[str, ...] = () ) -> List[str]: """Returns a bazel test command for this build. Args: subcommand: The subcommand to give to bazel. `test` by default. extra_options: Extra options. For now just used to pass in `--nobuild`. Returns: List of command line arguments """ options = _dict_to_cli_options(self.options) startup_options = _dict_to_cli_options(self.startup_options) configs = [f"--config={config}" for config in self.configs] build_tag_filters = ( f"--build_tag_filters={','.join(self.build_tag_filters)}" ) test_tag_filters = f"--test_tag_filters={','.join(self.test_tag_filters)}" action_env = [f"--action_env={k}={v}" for k, v in self.action_env.items()] test_env = [f"--test_env={k}={v}" for k, v in self.test_env.items()] repo_env = [f"--repo_env={k}={v}" for k, v in self.repo_env.items()] override_repository = [ f"--override_repository={k}={v}" for k, v in self.override_repository.items() ] override_module = [ f"--override_module={k}={v}" for k, v in self.override_module.items() ] tag_filters = [build_tag_filters, test_tag_filters] all_options = ( tag_filters + configs + action_env + test_env + repo_env + override_repository + override_module + options + list(extra_options) ) return [ "bazel", *startup_options, subcommand, *all_options, "--", *self.target_patterns, ] def commands(self) -> List[List[str]]: """Returns list of commands for a build.""" cmds = [] cmds.extend(self.extra_setup_commands) # We really want `bazel fetch` here, but it uses `bazel query` and not # `cquery`, which means that it fails due to config issues that aren't # problems in practice. # TODO(ddunleavy): Remove the condition here. Need to get parallel on the # MacOS VM. macos_build = ( self.type_ == BuildType.XLA_MACOS_X86_CPU_KOKORO or self.type_ == BuildType.XLA_MACOS_ARM64_CPU_KOKORO ) windows_build = ( self.type_ == BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS or self.type_ == BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS ) if not (macos_build or windows_build): cmds.append( retry( self.bazel_command( subcommand="build", extra_options=("--nobuild",) ) ) ) cmds.append(self.bazel_command(subcommand=self.subcommand)) cmds.append(["bazel", "analyze-profile", "profile.json.gz"]) return cmds _CUDA_COMPUTE_CAPABILITIES = (60, 70, 80, 90, 100, 103, 120) def _tag_filters_only_for_compute_capability( compute_capability: int, ) -> Tuple[str, ...]: """Returns the tag exclude filters for the given compute capability.""" tag_filters = () for cc in _CUDA_COMPUTE_CAPABILITIES: if compute_capability != cc: tag_filters += (f"-requires-gpu-sm{cc}",) tag_filters += (f"-requires-gpu-sm{cc}-only",) tag_filters += ("-requires-gpu-amd",) tag_filters += ("-requires-gpu-intel",) return tag_filters def _tag_filters_for_compute_capability( compute_capability: int, ) -> Tuple[str, ...]: """Returns the tag filters for the given compute capability.""" tag_filters = (f"requires-gpu-sm{compute_capability}-only",) for cc in _CUDA_COMPUTE_CAPABILITIES: if compute_capability >= cc: tag_filters += (f"requires-gpu-sm{cc}",) else: tag_filters += (f"-requires-gpu-sm{cc}",) tag_filters += (f"-requires-gpu-sm{cc}-only",) tag_filters += ("-requires-gpu-amd",) tag_filters += ("-requires-gpu-intel",) return tag_filters nvidia_single_gpu_build_filters = ( "-no_oss", "requires-gpu-nvidia", "-rocm-only", "-oneapi-only", "gpu", ) nvidia_single_gpu_test_filters = ( *nvidia_single_gpu_build_filters, "-multi_gpu", ) def nvidia_gpu_build_with_compute_capability( *, type_: BuildType, configs: Tuple[str, ...], compute_capability: int, multi_gpu: bool = False, ) -> Build: """Returns a build for a Nvidia GPU build with the given compute capability.""" repo_env = {"TF_CUDA_COMPUTE_CAPABILITIES": f"{compute_capability/10}"} if multi_gpu: options = { "//xla/tsl:ci_build": True, **_DEFAULT_BAZEL_OPTIONS, # TODO(b/464116734): Remove once we have more machines "test_sharding_strategy": "disabled", "test_strategy": "exclusive", } nvidia_only_multi_gpu_filters = ( "-no_oss", "-rocm-only", "-oneapi-only", "multi_gpu", "-xla_nvgpu_any", ) build_tag_filters = nvidia_only_multi_gpu_filters test_tag_filters = ( nvidia_only_multi_gpu_filters + _tag_filters_only_for_compute_capability(compute_capability) ) repo_env["REMOTE_GPU_TESTING"] = 0 # pyrefly: ignore[bad-assignment] else: options = { "run_under": "//build_tools/ci:parallel_gpu_execute", "//xla/tsl:ci_build": True, **_DEFAULT_BAZEL_OPTIONS, } build_tag_filters = nvidia_single_gpu_build_filters test_tag_filters = ( nvidia_single_gpu_test_filters + _tag_filters_for_compute_capability(compute_capability) ) return Build( type_=type_, repo="openxla/xla", target_patterns=_XLA_DEFAULT_TARGET_PATTERNS, configs=configs, test_tag_filters=test_tag_filters, build_tag_filters=build_tag_filters, options=options, repo_env=repo_env, extra_setup_commands=(["nvidia-smi"],), ) cpu_x86_tag_filter = ( "-no_oss", "-gpu", "-requires-gpu-nvidia", "-requires-gpu-amd", "-requires-gpu-intel", ) Build( type_=BuildType.XLA_LINUX_X86_CPU_GITHUB_ACTIONS, repo="openxla/xla", configs=("warnings", "nonccl", "rbe_linux_cpu"), target_patterns=_XLA_DEFAULT_TARGET_PATTERNS, build_tag_filters=cpu_x86_tag_filter, test_tag_filters=cpu_x86_tag_filter, options={**_DEFAULT_BAZEL_OPTIONS, "//xla/tsl:ci_build": True}, ) windows_x86_tag_filter = ( "-no_windows", "-no_oss", "-gpu", "-tpu", "-requires-gpu-nvidia", "-requires-gpu-amd", "-requires-gpu-intel", ) Build( type_=BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS, repo="openxla/xla", configs=( "warnings", "nonccl", "rbe_xla_windows_x86_cpu_2022", ), target_patterns=( "//xla/...", "//build_tools/...", "@tsl//tsl/...", "-//xla/tpu/...", # mpitrampoline and gloo are not windows compatible "-//xla/backends/cpu/collectives:gloo_collectives_test", "-//xla/backends/cpu/collectives:mpi_collectives", "-//xla/backends/cpu/collectives:mpi_communicator", # implementation is not windows compatible "-//xla/python/transfer/...", "-//xla/backends/profiler/subprocess:subprocess_profiling_session", "-//xla/backends/profiler/subprocess:subprocess_profiling_session_test", "-//xla/backends/profiler/subprocess:subprocess_registry", "-//xla/backends/profiler/subprocess:subprocess_registry_test", # xnnpack is not windows compatible "-//xla/backends/cpu/runtime/ynnpack:ynn_fusion_thunk", "-//xla/backends/cpu/runtime/ynnpack:ynn_interop", "-//xla/backends/cpu/runtime/ynnpack:ynn_threadpool", # triton is not windows compatible "-//xla/backends/gpu/...", "-//xla/codegen/emitters/tests/...", "-//xla/service/gpu/...", # eigen is not windows compatible "-//xla/codegen/intrinsic/cpp:eigen_unary_test", ), build_tag_filters=windows_x86_tag_filter, test_tag_filters=windows_x86_tag_filter, options={ **_DEFAULT_BAZEL_OPTIONS, "//xla/tsl:ci_build": True, }, subcommand="build", startup_options={ "output_user_root": "C:/x", }, ) Build( type_=BuildType.XLA_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS, repo="openxla/xla", configs=("warnings", "nonccl", "rbe_linux_cpu", "bzlmod"), target_patterns=_XLA_DEFAULT_TARGET_PATTERNS, build_tag_filters=cpu_x86_tag_filter, test_tag_filters=cpu_x86_tag_filter, options={**_DEFAULT_BAZEL_OPTIONS, "//xla/tsl:ci_build": True}, ) cpu_arm_tag_filter = ( "-no_oss", "-gpu", "-requires-gpu-nvidia", "-requires-gpu-amd", "-requires-gpu-intel", "-not_run:arm", ) Build( type_=BuildType.XLA_LINUX_ARM64_CPU_GITHUB_ACTIONS, repo="openxla/xla", configs=("warnings", "rbe_cross_compile_linux_arm64", "nonccl"), target_patterns=_XLA_DEFAULT_TARGET_PATTERNS, options={ **_DEFAULT_BAZEL_OPTIONS, "build_tests_only": True, "//xla/tsl:ci_build": True, }, build_tag_filters=cpu_arm_tag_filter, test_tag_filters=cpu_arm_tag_filter, ) nvidia_gpu_build_with_compute_capability( type_=BuildType.XLA_LINUX_X86_GPU_L4_GITHUB_ACTIONS, configs=("warnings", "rbe_linux_cuda_nvcc", "hermetic_cuda_umd"), compute_capability=75, multi_gpu=False, ) nvidia_gpu_build_with_compute_capability( type_=BuildType.XLA_LINUX_X86_GPU_8X_H100_GITHUB_ACTIONS, configs=("warnings", "rbe_linux_cuda_nvcc", "hermetic_cuda_umd"), compute_capability=90, multi_gpu=True, ) oneapi_build_tag_filter = ( "oneapi-only", "requires-gpu-intel", "-requires-gpu-amd", "-requires-gpu-nvidia", "-no_oss", "-cuda-only", "-rocm-only", "-no-oneapi", "gpu", ) oneapi_test_tag_filter = ( "oneapi-only", # This build of oneAPI backend runs on X86 host without an Intel GPU,so # we are excluding the tests requiring Intel GPU "-requires-gpu-intel", "-requires-gpu-amd", "-requires-gpu-nvidia", "-no_oss", "-cuda-only", "-rocm-only", "-no-oneapi", "gpu", ) Build( type_=BuildType.XLA_LINUX_X86_GPU_ONEAPI_GITHUB_ACTIONS, repo="openxla/xla", configs=( "nonccl", "rbe_linux_cpu", "sycl", "sycl_hermetic", "icpx_clang", ), target_patterns=_XLA_DEFAULT_TARGET_PATTERNS, build_tag_filters=oneapi_build_tag_filter, test_tag_filters=oneapi_test_tag_filter, options={**_DEFAULT_BAZEL_OPTIONS, "//xla/tsl:ci_build": True}, subcommand="build", ) # ROCm builds - hermetic LLVM # Tag filters from build_tools/rocm/rocm_tag_filters.sh + gpu rocm_tag_filter = ( "-no_gpu", "-requires-gpu-intel", "-requires-gpu-nvidia", "-cuda-only", "-oneapi-only", "-requires-gpu-sm60", "-requires-gpu-sm60-only", "-requires-gpu-sm70", "-requires-gpu-sm70-only", "-requires-gpu-sm80", "-requires-gpu-sm80-only", "-requires-gpu-sm86", "-requires-gpu-sm86-only", "-requires-gpu-sm89", "-requires-gpu-sm89-only", "-requires-gpu-sm90", "-requires-gpu-sm90-only", "-skip_rocprofiler_sdk", "-no_oss", "-oss_excluded", "-oss_serial", "gpu", ) Build( type_=BuildType.XLA_LINUX_X86_GPU_HERMETIC_ROCM_GITHUB_ACTIONS, repo="openxla/xla", configs=( "warnings", "rbe_linux_cpu", "rocm_clang_hermetic", "rocm_ci_hermetic", ), target_patterns=_XLA_DEFAULT_TARGET_PATTERNS, build_tag_filters=rocm_tag_filter, test_tag_filters=rocm_tag_filter, options={**_DEFAULT_BAZEL_OPTIONS, "//xla/tsl:ci_build": True}, subcommand="build", ) Build( type_=BuildType.XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS, repo="openxla/xla", configs=("warnings", "nonccl", "rbe_linux_cpu"), target_patterns=_XLA_CPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS, build_tag_filters=cpu_x86_tag_filter, test_tag_filters=cpu_x86_tag_filter, options={**_DEFAULT_BAZEL_OPTIONS, "//xla/tsl:ci_build": True}, subcommand="build", ) Build( type_=BuildType.XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS, repo="openxla/xla", configs=("warnings", "rbe_cross_compile_linux_arm64", "nonccl"), target_patterns=_XLA_CPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS, options={ **_DEFAULT_BAZEL_OPTIONS, "build_tests_only": False, "//xla/tsl:ci_build": True, }, build_tag_filters=cpu_arm_tag_filter, test_tag_filters=cpu_arm_tag_filter, subcommand="build", ) Build( type_=BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS, repo="openxla/xla", target_patterns=_XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS, configs=("warnings", "rbe_linux_cuda_nvcc", "hermetic_cuda_umd"), test_tag_filters=nvidia_single_gpu_test_filters + _tag_filters_for_compute_capability(compute_capability=75), build_tag_filters=nvidia_single_gpu_build_filters, options={ "run_under": "//build_tools/ci:parallel_gpu_execute", "//xla/tsl:ci_build": True, **_DEFAULT_BAZEL_OPTIONS, }, repo_env={ "TF_CUDA_COMPUTE_CAPABILITIES": "7.5", }, extra_setup_commands=(["nvidia-smi"],), subcommand="build", ) Build( type_=BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS, repo="openxla/xla", target_patterns=_XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS, configs=( "warnings", "rbe_linux_cuda_nvcc", "hermetic_cuda_umd", "cuda_libraries_from_stubs", ), test_tag_filters=nvidia_single_gpu_test_filters + _tag_filters_for_compute_capability(compute_capability=75), build_tag_filters=nvidia_single_gpu_build_filters, options={ "run_under": "//build_tools/ci:parallel_gpu_execute", "//xla/tsl:ci_build": True, **_DEFAULT_BAZEL_OPTIONS, }, repo_env={ "TF_CUDA_COMPUTE_CAPABILITIES": "7.5", }, extra_setup_commands=(["nvidia-smi"],), subcommand="build", ) Build( type_=BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS, repo="openxla/xla", configs=("warnings", "rbe_linux_cuda_nvcc", "hermetic_cuda_umd"), target_patterns=_XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS, test_tag_filters=nvidia_single_gpu_test_filters + _tag_filters_for_compute_capability(compute_capability=75), build_tag_filters=nvidia_single_gpu_build_filters, options={ "run_under": "//build_tools/ci:parallel_gpu_execute", "//xla/tsl:ci_build": True, **_DEFAULT_BAZEL_OPTIONS, }, repo_env={ "TF_CUDA_COMPUTE_CAPABILITIES": "7.5", }, extra_setup_commands=(["nvidia-smi"],), subcommand="build", ) Build( type_=BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS, repo="openxla/xla", configs=( "warnings", "rbe_linux_cuda_nvcc", "hermetic_cuda_umd", "cuda_libraries_from_stubs", ), target_patterns=_XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS, test_tag_filters=nvidia_single_gpu_test_filters + _tag_filters_for_compute_capability(compute_capability=75), build_tag_filters=nvidia_single_gpu_build_filters, options={ "run_under": "//build_tools/ci:parallel_gpu_execute", "//xla/tsl:ci_build": True, **_DEFAULT_BAZEL_OPTIONS, }, repo_env={ "TF_CUDA_COMPUTE_CAPABILITIES": "7.5", }, extra_setup_commands=(["nvidia-smi"],), subcommand="build", ) Build( type_=BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS, repo="openxla/xla", configs=(), target_patterns=_XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS, test_tag_filters=nvidia_single_gpu_test_filters + _tag_filters_for_compute_capability(compute_capability=100), build_tag_filters=nvidia_single_gpu_build_filters, options={ "run_under": "//build_tools/ci:parallel_gpu_execute", # Use User Mode and Kernel Mode Drivers pre-installed on the system. "//xla/tsl:ci_build": True, **_DEFAULT_BAZEL_OPTIONS, }, repo_env={ "TF_CUDA_COMPUTE_CAPABILITIES": "10", }, extra_setup_commands=(["nvidia-smi"],), subcommand="build", ) Build( type_=BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS, repo="openxla/xla", configs=("cuda_libraries_from_stubs",), target_patterns=_XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS, test_tag_filters=nvidia_single_gpu_test_filters + _tag_filters_for_compute_capability(compute_capability=100), build_tag_filters=nvidia_single_gpu_build_filters, options={ "run_under": "//build_tools/ci:parallel_gpu_execute", # Use User Mode and Kernel Mode Drivers pre-installed on the system. "//xla/tsl:ci_build": True, **_DEFAULT_BAZEL_OPTIONS, }, repo_env={ "TF_CUDA_COMPUTE_CAPABILITIES": "10", }, extra_setup_commands=(["nvidia-smi"],), subcommand="build", ) macos_tag_filter = ( "-no_oss", "-gpu", "-no_mac", "-mac_excluded", "-requires-gpu-nvidia", "-requires-gpu-amd", "-requires-gpu-intel", ) Build( type_=BuildType.XLA_MACOS_X86_CPU_KOKORO, repo="openxla/xla", configs=("nonccl",), target_patterns=( "//xla/...", "-//xla/python_api/...", "-//xla/python/...", "-//xla/service/gpu/...", ), options={ **_DEFAULT_BAZEL_OPTIONS, "macos_minimum_os": "10.15", "test_tmpdir": "/Volumes/BuildData/bazel_output", "//xla/tsl:ci_build": True, }, build_tag_filters=macos_tag_filter, test_tag_filters=macos_tag_filter, extra_setup_commands=( [ "sudo", "wget", "--no-verbose", "-O", "/usr/local/bin/bazel", "https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-darwin-amd64", ], ["chmod", "+x", "/usr/local/bin/bazel"], ["bazel", "--version"], # Sanity check due to strange failures ["mkdir", "-p", "/Volumes/BuildData/bazel_output"], ), ) Build( type_=BuildType.XLA_MACOS_ARM64_CPU_KOKORO, repo="openxla/xla", configs=("nonccl",), target_patterns=( "//xla/...", "-//xla/python_api/...", "-//xla/python/...", "-//xla/service/gpu/...", ), options={ **_DEFAULT_BAZEL_OPTIONS, "macos_minimum_os": "10.15", "test_tmpdir": "/tmpfs/bazel_output", "test_size_filters": "small,medium", "//xla/tsl:ci_build": True, }, build_tag_filters=macos_tag_filter, test_tag_filters=macos_tag_filter, extra_setup_commands=( ["df", "-h"], # Debug "No space left on device" error: b/396611909. ["bazel", "--version"], # Sanity check due to strange failures ["mkdir", "-p", "/tmpfs/bazel_output"], ), ) Build( type_=BuildType.JAX_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS, repo="google/jax", configs=("rbe_linux_x86_64", "bzlmod"), target_patterns=( "//tests:cpu_tests", "//tests:backend_independent_tests", "//jax/experimental/jax2tf/tests:jax2tf_test_cpu", "//tests/multiprocess:cpu_tests", "//jax/experimental/jax2tf/tests/multiprocess:cpu_tests", "//jaxlib/tools:check_cpu_wheel_sources_test", "//jaxlib/tools:jaxlib_wheel_size_test", "//:jax_wheel_size_test", ), test_env=dict( JAX_NUM_GENERATED_CASES=25, JAX_SKIP_SLOW_TESTS=1, ), override_module={ "xla": f"{_GITHUB_WORKSPACE}/openxla/xla", }, options=_DEFAULT_BAZEL_OPTIONS, repo_env={"HERMETIC_PYTHON_VERSION": "3.12"}, ) Build( type_=BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS, repo="google/jax", configs=("rbe_windows_amd64",), target_patterns=( "//tests:cpu_tests", "//tests:backend_independent_tests", "//jax/experimental/jax2tf/tests:jax2tf_test_cpu", "//tests/multiprocess:cpu_tests", "//jax/experimental/jax2tf/tests/multiprocess:cpu_tests", "//jaxlib/tools:check_cpu_wheel_sources_test", "//jaxlib/tools:jaxlib_wheel_size_test", "//:jax_wheel_size_test", ), test_env=dict( JAX_NUM_GENERATED_CASES=25, JAX_SKIP_SLOW_TESTS=1, ), override_repository=dict( xla=f"{_GITHUB_WORKSPACE}\\openxla\\xla", ), override_module=dict( xla=f"{_GITHUB_WORKSPACE}\\openxla\\xla", ), options={**_DEFAULT_BAZEL_OPTIONS, "build_runfile_links": False}, repo_env={"HERMETIC_PYTHON_VERSION": "3.12"}, subcommand="build", startup_options={ "output_base": f"{_GITHUB_WORKSPACE}\\bazel_output_base", }, ) Build( type_=BuildType.JAX_LINUX_X86_GPU_L4_GITHUB_ACTIONS, repo="google/jax", configs=("rbe_linux_x86_64_cuda13",), target_patterns=( "//tests:gpu_tests", "//tests:backend_independent_tests", "//tests/pallas:gpu_tests", "//tests/pallas:backend_independent_tests", "//jaxlib/tools:check_gpu_wheel_sources_test", "//jaxlib/tools:jax_cuda_plugin_wheel_size_test", "//jaxlib/tools:jax_cuda_pjrt_wheel_size_test", "//jaxlib/tools:jaxlib_wheel_size_test", "//:jax_wheel_size_test", ), build_tag_filters=("-multiaccelerator",), test_tag_filters=("-multiaccelerator",), test_env=dict( JAX_SKIP_SLOW_TESTS=1, TF_CPP_MIN_LOG_LEVEL=0, JAX_EXCLUDE_TEST_TARGETS="PmapTest.testSizeOverflow", ), override_repository=dict( xla=f"{_GITHUB_WORKSPACE}/openxla/xla", ), override_module=dict( xla=f"{_GITHUB_WORKSPACE}/openxla/xla", ), options={ **_DEFAULT_BAZEL_OPTIONS, "@local_config_cuda//cuda:override_include_cuda_libs": True, }, repo_env={"HERMETIC_PYTHON_VERSION": "3.12"}, extra_setup_commands=(["nvidia-smi"],), ) def dump_all_build_commands(): """Used to generate what commands are run for each build.""" # Awkward workaround b/c Build instances are not hashable for build in sorted(Build.all_builds().values(), key=lambda b: str(b.type_)): sys.stdout.write(f"# BEGIN {build.type_}\n") for cmd in build.commands(): sys.stdout.write(" ".join(cmd) + "\n") sys.stdout.write(f"# END {build.type_}\n") def _parse_args(): """Defines flags and parses args.""" parser = argparse.ArgumentParser(allow_abbrev=False) group = parser.add_mutually_exclusive_group(required=True) group.add_argument( "--build", type=BuildType.from_str, choices=list(BuildType), ) group.add_argument( "--dump_commands", action="store_true", ) return parser.parse_args() def main(): logging.basicConfig() logging.getLogger().setLevel(logging.INFO) args = _parse_args() if args.dump_commands: dump_all_build_commands() return else: for cmd in Build.all_builds()[args.build].commands(): sh(cmd) if __name__ == "__main__": main()