diff --git a/build_tools/ci/build.py b/build_tools/ci/build.py index c10e4023855f5..9507f0a5e38f1 100755 --- a/build_tools/ci/build.py +++ b/build_tools/ci/build.py @@ -93,6 +93,7 @@ class BuildType(enum.Enum): GPU_T4_SELF_HOSTED = enum.auto() MACOS_CPU_X86 = enum.auto() + MACOS_CPU_ARM64 = enum.auto() JAX_CPU_SELF_HOSTED = enum.auto() JAX_X86_GPU_T4_SELF_HOSTED = enum.auto() @@ -292,6 +293,39 @@ def nvidia_gpu_build_with_compute_capability( ), ) +_MACOS_ARM64_BUILD = Build( + type_=BuildType.MACOS_CPU_ARM64, + repo="openxla/xla", + configs=("nonccl",), + target_patterns=( + "//xla/...", + "-//xla/hlo/experimental/...", + "-//xla/python_api/...", + "-//xla/python/...", + "-//xla/service/gpu/...", + ), + options=dict( + **_DEFAULT_BAZEL_OPTIONS, + macos_minimum_os="10.15", + test_tmpdir="/Volumes/BuildData/bazel_output", + ), + 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-arm64", + ], + ["chmod", "+x", "/usr/local/bin/bazel"], + ["bazel", "--version"], # Sanity check due to strange failures + ["mkdir", "-p", "/Volumes/BuildData/bazel_output"], + ), +) + _JAX_CPU_SELF_HOSTED_BUILD = Build( type_=BuildType.JAX_CPU_SELF_HOSTED, repo="google/jax", @@ -397,6 +431,7 @@ def nvidia_gpu_build_with_compute_capability( _KOKORO_JOB_NAME_TO_BUILD_MAP = { "tensorflow/xla/macos/github_continuous/cpu_py39_full": _MACOS_X86_BUILD, + "tensorflow/xla/macos/cpu/cpu_py39_full": _MACOS_ARM64_BUILD, "xla-linux-x86-cpu": _CPU_X86_SELF_HOSTED_BUILD, "xla-linux-arm64-cpu": _CPU_ARM64_SELF_HOSTED_BUILD, "xla-linux-x86-gpu-t4": _GPU_T4_SELF_HOSTED_BUILD, diff --git a/build_tools/ci/golden_commands.txt b/build_tools/ci/golden_commands.txt index 4f876bb111a78..a832c888a0309 100644 --- a/build_tools/ci/golden_commands.txt +++ b/build_tools/ci/golden_commands.txt @@ -24,6 +24,15 @@ parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_fi bazel test --build_tag_filters=-multiaccelerator --test_tag_filters=-multiaccelerator --config=rbe_linux_x86_64_cuda --test_env=JAX_SKIP_SLOW_TESTS=1 --test_env=TF_CPP_MIN_LOG_LEVEL=0 --test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --repo_env=HERMETIC_PYTHON_VERSION=3.10 -- //tests:gpu_tests //tests:backend_independent_tests bazel analyze-profile profile.json.gz # END BuildType.JAX_X86_GPU_T4_SELF_HOSTED +# BEGIN BuildType.MACOS_CPU_ARM64 +sudo wget --no-verbose -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-darwin-arm64 +chmod +x /usr/local/bin/bazel +bazel --version +mkdir -p /Volumes/BuildData/bazel_output +parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd --test_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd --config=nonccl --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --macos_minimum_os=10.15 --test_tmpdir=/Volumes/BuildData/bazel_output --nobuild -- //xla/... -//xla/hlo/experimental/... -//xla/python_api/... -//xla/python/... -//xla/service/gpu/... +bazel test --build_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd --test_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd --config=nonccl --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --macos_minimum_os=10.15 --test_tmpdir=/Volumes/BuildData/bazel_output -- //xla/... -//xla/hlo/experimental/... -//xla/python_api/... -//xla/python/... -//xla/service/gpu/... +bazel analyze-profile profile.json.gz +# END BuildType.MACOS_CPU_ARM64 # BEGIN BuildType.MACOS_CPU_X86 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