wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
tensorflow--tensorflow/third_party/xla/build_tools/ci/build.py
T
wehub-resource-sync 8a852e4b4e
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details
chore: import upstream snapshot with attribution
2026-07-13 12:14:16 +08:00

1022 lines
31 KiB
Python
Executable File
Raw Permalink Blame History

#!/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()
TENSORFLOW_LINUX_X86_CPU_GITHUB_ACTIONS = enum.auto()
TENSORFLOW_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
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",
"-//xla/tools/benchmarks/utils:generate_benchmark_matrices_cc",
"-//xla/tools/benchmarks/utils:generate_benchmark_matrices_main",
"-//xla/tools/benchmarks/utils:generate_benchmark_matrices_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",
"HERMETIC_CUDA_VERSION": "12.8.0",
"HERMETIC_CUDNN_VERSION": "9.8.0",
},
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",
"HERMETIC_CUDA_VERSION": "12.8.0",
"HERMETIC_CUDNN_VERSION": "9.8.0",
},
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",
# TODO(alekstheod): remove when jax is migrated to
# latest the rules_ml_toolchain
"rules_ml_toolchain": f"{_GITHUB_WORKSPACE}/openxla/rules_ml_toolchain",
},
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_cuda",),
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"],),
)
tensorflow_tag_filters = (
"-no_oss",
"-tf_tosa",
"-oss_excluded",
"-oss_serial",
"-tpu",
"-benchmark-test",
"-v1only",
)
tensorflow_cpu_tag_filters = tensorflow_tag_filters + ("-gpu",)
tensorflow_gpu_tag_filters = tensorflow_tag_filters + (
"-no_gpu",
"-no_gpu_presubmit",
"-no_cuda11",
"+gpu",
)
Build(
type_=BuildType.TENSORFLOW_LINUX_X86_CPU_GITHUB_ACTIONS,
repo="tensorflow/tensorflow",
configs=(
"release_cpu_linux",
"rbe_linux_cpu",
),
target_patterns=(
"//tensorflow/compiler/...",
"-//tensorflow/compiler/tf2tensorrt/...",
"//tensorflow/python/...",
"-//tensorflow/python/distribute/...",
"-//tensorflow/python/kernel_tests/...",
"-//tensorflow/python/data/...",
"-//tensorflow/python/compiler/tensorrt/...",
),
build_tag_filters=tensorflow_cpu_tag_filters,
test_tag_filters=tensorflow_cpu_tag_filters,
options=dict(
verbose_failures=True,
test_output="errors",
profile="profile.json.gz",
test_lang_filters="cc,py",
color="yes",
),
override_repository=dict(
xla=f"{_GITHUB_WORKSPACE}/openxla/xla",
),
override_module=dict(
xla=f"{_GITHUB_WORKSPACE}/openxla/xla",
),
repo_env={"USE_PYWRAP_RULES": "True"},
)
Build(
type_=BuildType.TENSORFLOW_LINUX_X86_GPU_L4_GITHUB_ACTIONS,
repo="tensorflow/tensorflow",
configs=("release_gpu_linux", "rbe_linux_cuda", "hermetic_cuda_umd"),
target_patterns=(
"//tensorflow/compiler/...",
"-//tensorflow/compiler/tf2tensorrt/...",
"//tensorflow/python/...",
"-//tensorflow/python/distribute/...",
"-//tensorflow/python/kernel_tests/...",
"-//tensorflow/python/data/...",
"-//tensorflow/python/compiler/tensorrt/...",
),
build_tag_filters=tensorflow_gpu_tag_filters,
test_tag_filters=tensorflow_gpu_tag_filters,
override_repository=dict(
xla=f"{_GITHUB_WORKSPACE}/openxla/xla",
),
override_module=dict(
xla=f"{_GITHUB_WORKSPACE}/openxla/xla",
),
options=dict(
verbose_failures=True,
test_output="errors",
profile="profile.json.gz",
test_lang_filters="cc,py",
color="yes",
),
repo_env={"USE_PYWRAP_RULES": "True"},
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()
Reference in New Issue View Git Blame Copy Permalink