Merge commit 'db579439eef970657f5ddbf05dc9b798cb748c51' into sync-231206

Conflicts:
	third_party/xla/xla/service/gpu/BUILD
	third_party/xla/xla/service/gpu/buffer_comparator_test.cc
	third_party/xla/xla/stream_executor/device_description.h
	third_party/xla/xla/stream_executor/rocm/hip_blas_lt.cc
	third_party/xla/xla/stream_executor/rocm/hip_blas_lt.h
	third_party/xla/xla/tests/BUILD

.bazelrc

@@ -543,34 +543,9 @@ build:rbe_linux_rocm_base --platforms="@ubuntu20.04-gcc9_manylinux2014-rocm_conf
 build:rbe_linux_rocm_base --action_env=TF_ROCM_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-rocm_config_rocm"
 build:rbe_linux_rocm_py3.9 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-rocm_config_python3.9"
 
-build:rbe_linux_cuda_nvcc --config=cuda
-build:rbe_linux_cuda_nvcc --action_env=TF_NVCC_CLANG="1"
+build:rbe_linux_cuda_nvcc --config=rbe_linux_cuda
 build:rbe_linux_cuda_nvcc --repo_env TF_NCCL_USE_STUB=1
-build:rbe_linux_cuda_nvcc --@local_xla//xla/python:enable_gpu=true
-build:rbe_linux_cuda_nvcc --@local_xla//xla/python:jax_cuda_pip_rpaths=true
-build:rbe_linux_cuda_nvcc --define=xla_python_enable_gpu=true
-build:rbe_linux_cuda_nvcc --config=tensorrt
-build:rbe_linux_cuda_nvcc --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_75,compute_80"
-build:rbe_linux_cuda_nvcc --action_env=TF_CUDA_VERSION="12"
-build:rbe_linux_cuda_nvcc --action_env=TF_CUDNN_VERSION="8"
-build:rbe_linux_cuda_nvcc --action_env=CUDA_TOOLKIT_PATH="/usr/local/cuda-12.2"
-build:rbe_linux_cuda_nvcc --action_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"
-build:rbe_linux_cuda_nvcc --crosstool_top="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_cuda//crosstool:toolchain"
-build:rbe_linux_cuda_nvcc --config=rbe_linux
-build:rbe_linux_cuda_nvcc --host_crosstool_top="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_cuda//crosstool:toolchain"
-build:rbe_linux_cuda_nvcc --extra_toolchains="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_cuda//crosstool:toolchain-linux-x86_64"
-build:rbe_linux_cuda_nvcc --extra_execution_platforms="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_platform//:platform"
-build:rbe_linux_cuda_nvcc --host_platform="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_platform//:platform"
-build:rbe_linux_cuda_nvcc --platforms="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_platform//:platform"
-build:rbe_linux_cuda_nvcc --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_python3.9"
-build:rbe_linux_cuda_nvcc --python_path="/usr/bin/python3"
-# These you may need to change for your own GCP project.
-common:rbe_linux_cuda_nvcc --remote_instance_name=projects/tensorflow-testing/instances/default_instance
-build:rbe_linux_cuda_nvcc --repo_env=REMOTE_GPU_TESTING=1
-build:rbe_linux_cuda_nvcc --repo_env=TF_CUDA_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-cuda12.2-cudnn8.9_config_cuda"
-build:rbe_linux_cuda_nvcc --repo_env=TF_TENSORRT_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-cuda12.2-cudnn8.9_config_tensorrt"
-build:rbe_linux_cuda_nvcc --repo_env=TF_NCCL_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-cuda12.2-cudnn8.9_config_nccl"
-test:rbe_linux_cuda_nvcc --test_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"
+build:rbe_linux_cuda_nvcc --action_env=TF_NVCC_CLANG="1"
 
 # TODO(kanglan): Remove rbe_win and rbe_win_py3* after b/289091160 is fixed
 build:rbe_win --config=rbe_base
@@ -713,6 +688,15 @@ test:release_cpu_macos --config=release_base
 build:release_macos_base --action_env  DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer
 build:release_macos_base --define=no_nccl_support=true --output_filter=^$
 
+# Build configs for macOS x86
+build:release_macos_x86 --config=release_macos_base
+# Build with the AVX instruction set when on macOS x86
+build:release_macos_x86 --config=avx_linux
+build:release_macos_x86 --cpu=darwin
+# Target Catalina as the minimum compatible OS version
+build:release_macos_x86 --macos_minimum_os=10.15
+build:release_macos_x86 --action_env MACOSX_DEPLOYMENT_TARGET=10.15
+
 # Build configs for macOS Arm64
 build:release_macos_arm64 --config=release_macos_base
 build:release_macos_arm64 --cpu=darwin_arm64
@@ -727,6 +711,9 @@ test:release_macos_base --test_timeout=300,450,1200,3600 --test_output=errors
 test:release_macos_base --build_tests_only --keep_going
 test:release_macos_base --flaky_test_attempts=3
 
+# Test configs for macOS x86
+test:release_macos_x86 --config=release_macos_base
+
 # Test configs for macOS Arm64
 test:release_macos_arm64 --config=release_macos_base
 
@@ -788,6 +775,11 @@ test:macos_arm64_wheel_test_filters --test_tag_filters=-no_oss,-oss_excluded,-os
 test:macos_arm64_wheel_test_filters --build_tag_filters=-no_oss,-oss_excluded,-oss_serial,-no_oss_py39,-no_oss_py310,-nomac,-no_mac,-mac_excluded,-v1only,-gpu,-tpu,-benchmark-test,-no_mac_arm64,-no_aarch64
 test:macos_arm64_wheel_test_filters --test_lang_filters=py --test_size_filters=small,medium
 test:macos_arm64_wheel_test --config=macos_arm64_wheel_test_filters -- //tensorflow/... -//tensorflow/python/integration_testing/... -//tensorflow/compiler/tf2tensorrt/... -//tensorflow/core/tpu/... -//tensorflow/lite/... -//tensorflow/tools/toolchains/... -//tensorflow/go/... -//tensorflow/java/... -//tensorflow/compiler/aot/...
+# MACOS X86 WHEEL
+test:macos_x86_wheel_test_filters --test_tag_filters=-no_oss,-oss_excluded,-oss_serial,-no_oss_py38,-no_oss_py39,-no_oss_py310,-nomac,-no_mac,-mac_excluded,-v1only,-gpu,-tpu,-benchmark-test
+test:macos_x86_wheel_test_filters --build_tag_filters=-no_oss,-oss_excluded,-oss_serial,-no_oss_py38,-no_oss_py39,-no_oss_py310,-nomac,-no_mac,-mac_excluded,-v1only,-gpu,-tpu,-benchmark-test
+test:macos_x86_wheel_test_filters --test_lang_filters=py --test_size_filters=small,medium
+test:macos_x86_wheel_test --config=macos_x86_wheel_test_filters -- //tensorflow/... -//tensorflow/python/integration_testing/... -//tensorflow/compiler/tf2tensorrt/... -//tensorflow/core/tpu/... -//tensorflow/lite/... -//tensorflow/tools/toolchains/... -//tensorflow/go/... -//tensorflow/java/... -//tensorflow/compiler/aot/...
 
 # PYCPP TESTS run a suite of Python and C++ tests to verify general correctness over
 # the whole TF code base. These are usually run continuously or upon presubmit.

.github/workflows/osv-scanner-scheduled.yml

@@ -27,6 +27,7 @@ permissions:
 
 jobs:
   scan-scheduled:
+    if: github.repository == 'tensorflow/tensorflow'
     uses: "google/osv-scanner/.github/workflows/osv-scanner-reusable.yml@main"
     with:
       scan-args: |-
@@ -36,4 +37,4 @@ jobs:
         --lockfile=requirements.txt:./requirements_lock_3_12.txt
         --lockfile=requirements.txt:./ci/official/containers/linux_arm64/devel.requirements.txt
         --lockfile=requirements.txt:./ci/official/containers/linux_arm64/jax.requirements.txt
-        --lockfile=requirements.txt:./ci/official/containers/linux_arm64/devel.usertools/test.requirements.txt
+        --lockfile=requirements.txt:./ci/official/containers/linux_arm64/devel.usertools/test.requirements.txt

.github/workflows/stale-issues.yml

@@ -31,7 +31,7 @@ jobs:
       pull-requests: write
     steps:
       - name: Awaiting response issues
-        uses: actions/stale@v7
+        uses: actions/stale@6f05e4244c9a0b2ed3401882b05d701dd0a7289b # v7.0.0
         with:
           #Comma separated list of labels that can be assigned to issues to exclude them from being marked as stale
           exempt-issue-labels: 'override-stale'
@@ -59,7 +59,7 @@ jobs:
           close-pr-message: "This PR was closed because it has been inactive for 14 days since being marked as stale. Please reopen if you'd like to work on this further."
           repo-token: ${{ secrets.GITHUB_TOKEN }}
       - name: Contribution issues
-        uses: actions/stale@v7
+        uses: actions/stale@6f05e4244c9a0b2ed3401882b05d701dd0a7289b # v7.0.0
         with:
           #Comma separated list of labels that can be assigned to issues to exclude them from being marked as stale
           exempt-issue-labels: 'override-stale'

RELEASE.md

@@ -9,6 +9,12 @@
 * <DOCUMENT BREAKING CHANGES HERE>
 * <THIS SECTION SHOULD CONTAIN API, ABI AND BEHAVIORAL BREAKING CHANGES>
 
+* `tf.summary.trace_on` now takes a `profiler_outdir` argument. This must be set
+  if `profiler` arg is set to `True`.
+    * `tf.summary.trace_export`'s `profiler_outdir` arg is now a no-op. Enabling
+      the profiler now requires setting `profiler_outdir` in `trace_on`.
+
+
 ### Known Caveats
 
 * <CAVEATS REGARDING THE RELEASE (BUT NOT BREAKING CHANGES).>

ci/official/envs/ci_default

@@ -30,4 +30,9 @@ TFCI_WHL_AUDIT_PLAT=
 TFCI_WHL_BAZEL_TEST_ENABLE=1
 TFCI_WHL_SIZE_LIMIT=
 TFCI_WHL_SIZE_LIMIT_ENABLE=1
-TFCI_PYENV_INSTALL_LOCAL_ENABLE=
+TFCI_MACOS_UPGRADE_PYENV_ENABLE=
+TFCI_MACOS_INSTALL_BAZELISK_ENABLE=
+TFCI_MACOS_INSTALL_BAZELISK_URL=
+TFCI_MACOS_PYENV_INSTALL_ENABLE=
+TFCI_MACOS_BAZEL_TEST_DIR_ENABLE=
+TFCI_MACOS_BAZEL_TEST_DIR_PATH=

ci/official/envs/continuous_linux_arm64_cpu_py39_cross_compile

@@ -0,0 +1,6 @@
+# This envrionment is experimental and should not yet be used for production jobs
+TFCI_BAZEL_COMMON_ARGS="--config rbe_cross_compile_linux_arm64 --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION"
+TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=cross_compile_linux_arm64
+TFCI_DOCKER_IMAGE=gcr.io/tensorflow-sigs/build-arm64:tf-latest-multi-python
+TFCI_DOCKER_REBUILD_ARGS="--target=tf ci/official/containers/linux_arm64"
+TFCI_PYTHON_VERSION=3.9

ci/official/envs/continuous_macos_arm64_py310

@@ -2,3 +2,4 @@ TFCI_BAZEL_COMMON_ARGS="--config release_macos_arm64 --config tf_public_macos_ca
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_arm64
 TFCI_DOCKER_ENABLE=0
 TFCI_PYTHON_VERSION=3.10
+TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"

ci/official/envs/continuous_macos_arm64_py311

@@ -2,3 +2,4 @@ TFCI_BAZEL_COMMON_ARGS="--config release_macos_arm64 --config tf_public_macos_ca
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_arm64
 TFCI_DOCKER_ENABLE=0
 TFCI_PYTHON_VERSION=3.11
+TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"

ci/official/envs/continuous_macos_arm64_py39

@@ -2,3 +2,4 @@ TFCI_BAZEL_COMMON_ARGS="--config release_macos_arm64 --config tf_public_macos_ca
 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_arm64
 TFCI_DOCKER_ENABLE=0
 TFCI_PYTHON_VERSION=3.9
+TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"

ci/official/envs/nightly_libtensorflow_macos_arm64

@@ -1,7 +1,7 @@
-source ci/official/envs/ci_nightly_uploads
-TFCI_BAZEL_COMMON_ARGS="--config release_macos_arm64 --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION"
+# Disable arm64 uploads while being worked on
+source ci/official/envs/disable_all_uploads
+TFCI_BAZEL_COMMON_ARGS="--config release_macos_arm64 --config tf_public_cache_push --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION"
 TFCI_DOCKER_ENABLE=0
-TFCI_LIB_SUFFIX="-cpu-macos-arm64"
+TFCI_LIB_SUFFIX="-cpu-darwin-arm64"
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
 TFCI_PYTHON_VERSION=3.10
-TFCI_UPLOAD_WHL_GCS_URI=1

ci/official/envs/nightly_libtensorflow_macos_x86

@@ -0,0 +1,7 @@
+# Disable macOS x86 uploads while being worked on
+source ci/official/envs/disable_all_uploads
+TFCI_BAZEL_COMMON_ARGS="--config release_macos_x86 --config tf_public_cache_push --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION"
+TFCI_DOCKER_ENABLE=0
+TFCI_LIB_SUFFIX="-cpu-darwin-x86_64"
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
+TFCI_PYTHON_VERSION=3.10

ci/official/envs/nightly_macos_arm64_py310

@@ -5,8 +5,7 @@ TFCI_BUILD_PIP_PACKAGE_ARGS="--cpu --nightly_flag"
 TFCI_DOCKER_ENABLE=0
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
 TFCI_PYTHON_VERSION=3.10
-TFCI_UPLOAD_WHL_GCS_ENABLE=1
-TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-nightly/macos-arm64/$(date -I)"
 TFCI_WHL_AUDIT_ENABLE=
 TFCI_WHL_SIZE_LIMIT=240M
-TFCI_PYENV_INSTALL_LOCAL_ENABLE=$TFCI_PYTHON_VERSION
+TFCI_MACOS_PYENV_INSTALL_ENABLE=1
+TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"

ci/official/envs/nightly_macos_arm64_py311

@@ -5,8 +5,6 @@ TFCI_BUILD_PIP_PACKAGE_ARGS="--cpu --nightly_flag"
 TFCI_DOCKER_ENABLE=0
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
 TFCI_PYTHON_VERSION=3.11
-TFCI_UPLOAD_WHL_GCS_ENABLE=1
-TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-nightly/macos-arm64/$(date -I)"
 TFCI_WHL_AUDIT_ENABLE=
 TFCI_WHL_SIZE_LIMIT=240M
-TFCI_PYENV_INSTALL_LOCAL_ENABLE=$TFCI_PYTHON_VERSION
+TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"

ci/official/envs/nightly_macos_arm64_py312

@@ -5,8 +5,7 @@ TFCI_BUILD_PIP_PACKAGE_ARGS="--cpu --nightly_flag"
 TFCI_DOCKER_ENABLE=0
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
 TFCI_PYTHON_VERSION=3.12
-TFCI_UPLOAD_WHL_GCS_ENABLE=
-TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-nightly/macos-arm64/$(date -I)"
 TFCI_WHL_AUDIT_ENABLE=
 TFCI_WHL_SIZE_LIMIT=240M
-TFCI_PYENV_INSTALL_LOCAL_ENABLE=$TFCI_PYTHON_VERSION
+TFCI_MACOS_PYENV_INSTALL_ENABLE=1
+TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"

ci/official/envs/nightly_macos_arm64_py39

@@ -7,8 +7,7 @@ TFCI_BUILD_PIP_PACKAGE_ARGS="--cpu --nightly_flag"
 TFCI_DOCKER_ENABLE=0
 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
 TFCI_PYTHON_VERSION=3.9
-TFCI_UPLOAD_WHL_GCS_ENABLE=1
-TFCI_UPLOAD_WHL_GCS_URI="gs://tensorflow-nightly/macos/arm64/$(date -I)"
 TFCI_WHL_AUDIT_ENABLE=
 TFCI_WHL_SIZE_LIMIT=240M
-TFCI_PYENV_INSTALL_LOCAL_ENABLE=$TFCI_PYTHON_VERSION
+TFCI_MACOS_PYENV_INSTALL_ENABLE=1
+TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"

ci/official/envs/nightly_macos_x86_py310

@@ -0,0 +1,16 @@
+# Disable macOS x86 uploads while being worked on
+source ci/official/envs/disable_all_uploads
+TFCI_BAZEL_COMMON_ARGS="--config release_macos_x86 --config tf_public_macos_cache_push --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION"
+TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_x86
+TFCI_BUILD_PIP_PACKAGE_ARGS="--cpu --nightly_flag"
+TFCI_DOCKER_ENABLE=0
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
+TFCI_PYTHON_VERSION=3.10
+TFCI_WHL_AUDIT_ENABLE=
+TFCI_WHL_SIZE_LIMIT=255M
+TFCI_MACOS_INSTALL_BAZELISK_ENABLE=1
+TFCI_MACOS_INSTALL_BAZELISK_URL="https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-darwin-amd64"
+TFCI_MACOS_UPGRADE_PYENV_ENABLE=1
+TFCI_MACOS_PYENV_INSTALL_ENABLE=1
+TFCI_MACOS_BAZEL_TEST_DIR_ENABLE=1
+TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"

ci/official/envs/nightly_macos_x86_py311

@@ -0,0 +1,16 @@
+# Disable macOS x86 uploads while being worked on
+source ci/official/envs/disable_all_uploads
+TFCI_BAZEL_COMMON_ARGS="--config release_macos_x86 --config tf_public_macos_cache_push --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION"
+TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_x86
+TFCI_BUILD_PIP_PACKAGE_ARGS="--cpu --nightly_flag"
+TFCI_DOCKER_ENABLE=0
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
+TFCI_PYTHON_VERSION=3.11
+TFCI_WHL_AUDIT_ENABLE=
+TFCI_WHL_SIZE_LIMIT=255M
+TFCI_MACOS_INSTALL_BAZELISK_ENABLE=1
+TFCI_MACOS_INSTALL_BAZELISK_URL="https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-darwin-amd64"
+TFCI_MACOS_UPGRADE_PYENV_ENABLE=1
+TFCI_MACOS_PYENV_INSTALL_ENABLE=1
+TFCI_MACOS_BAZEL_TEST_DIR_ENABLE=1
+TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"

ci/official/envs/nightly_macos_x86_py312

@@ -0,0 +1,16 @@
+# Disable macOS x86 uploads while being worked on
+source ci/official/envs/disable_all_uploads
+TFCI_BAZEL_COMMON_ARGS="--config release_macos_x86 --config tf_public_macos_cache_push --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION"
+TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_x86
+TFCI_BUILD_PIP_PACKAGE_ARGS="--cpu --nightly_flag"
+TFCI_DOCKER_ENABLE=0
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
+TFCI_PYTHON_VERSION=3.12
+TFCI_WHL_AUDIT_ENABLE=
+TFCI_WHL_SIZE_LIMIT=255M
+TFCI_MACOS_INSTALL_BAZELISK_ENABLE=1
+TFCI_MACOS_INSTALL_BAZELISK_URL="https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-darwin-amd64"
+TFCI_MACOS_UPGRADE_PYENV_ENABLE=1
+TFCI_MACOS_PYENV_INSTALL_ENABLE=1
+TFCI_MACOS_BAZEL_TEST_DIR_ENABLE=1
+TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"

ci/official/envs/nightly_macos_x86_py39

@@ -0,0 +1,14 @@
+# Disable macOS x86 uploads while being worked on
+source ci/official/envs/disable_all_uploads
+TFCI_BAZEL_COMMON_ARGS="--config release_macos_x86 --config tf_public_macos_cache_push --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION"
+TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_x86
+TFCI_BUILD_PIP_PACKAGE_ARGS="--cpu --nightly_flag"
+TFCI_DOCKER_ENABLE=0
+TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1
+TFCI_PYTHON_VERSION=3.9
+TFCI_WHL_AUDIT_ENABLE=
+TFCI_WHL_SIZE_LIMIT=255M
+TFCI_MACOS_INSTALL_BAZELISK_ENABLE=1
+TFCI_MACOS_INSTALL_BAZELISK_URL="https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-darwin-amd64"
+TFCI_MACOS_BAZEL_TEST_DIR_ENABLE=1
+TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"

ci/official/utilities/setup_macos.sh

@@ -34,31 +34,61 @@ else
   exit 1
 fi
 
-if [[ -n "${KOKORO_JOB_NAME}" ]]; then
-  # Mac builds need ~150 GB of disk space to be able to run all the tests. By
-  # default, Kokoro runs the Bazel commands in a partition that does not have
-  # enough free space so we need to set TEST_TMPDIR explicitly.
-  mkdir -p /Volumes/BuildData/bazel_output
-  export TEST_TMPDIR=/Volumes/BuildData/bazel_output
+# "TFCI_MACOS_BAZEL_TEST_DIR_PATH" specifies the directory that Bazel should use
+# when running tests. Each test will be executed in a separate subdirectory
+# inside this directory. TF Mac builds need ~150 GB of disk space to be able to
+# run all the tests. Since TFCI Mac VMs execute Bazel test commands in a
+# partition with insufficient storage, we specify the
+# 'TFCI_MACOS_BAZEL_TEST_DIR_PATH' environment variable to point to a partition
+# with ample storage. When this variable is empty (i.e by default), Bazel will
+# use the output base directory to run tests.
+if [[ "${TFCI_MACOS_BAZEL_TEST_DIR_ENABLE}" == 1 ]]; then
+  mkdir -p "${TFCI_MACOS_BAZEL_TEST_DIR_PATH}"
+  export TEST_TMPDIR="${TFCI_MACOS_BAZEL_TEST_DIR_PATH}"
+fi
+
+# "TFCI_MACOS_INSTALL_BAZELISK_ENABLE" is used to decide if we need to install
+# Bazelisk manually. We enable this for macOS x86 builds as those VMs do not
+# have Bazelisk pre-installed. "TFCI_MACOS_INSTALL_BAZELISK_URL" contains the
+# link to the Bazelisk binary which needs to be downloaded.
+if [[ "${TFCI_MACOS_INSTALL_BAZELISK_ENABLE}" == 1 ]]; then
+  sudo wget --no-verbose -O "/usr/local/bin/bazel" "${TFCI_MACOS_INSTALL_BAZELISK_URL}"
+  chmod +x "/usr/local/bin/bazel"
+fi
 
-  # Before uploading the nightly and release wheels, we install them in a
-  # virtual environment and run some smoke tests on it. The Kokoro Mac VMs
-  # only have Python 3.11 installed so we need to install the other Python
-  # versions manually.
-  if [[ -n "${TFCI_BUILD_PIP_PACKAGE_ARGS}" ]] && [[ "${TFCI_PYENV_INSTALL_LOCAL_ENABLE}" != 3.11 ]]; then
-    pyenv install "${TFCI_PYENV_INSTALL_LOCAL_ENABLE}"
-    pyenv local "${TFCI_PYENV_INSTALL_LOCAL_ENABLE}"
-  fi
-elif [[ "${TFCI_WHL_BAZEL_TEST_ENABLE}" == 1 ]]; then
-  echo '==TFCI==: Note: Mac builds need ~150 GB of disk space to be able to'
-  echo 'run all the tests. Please make sure your system has enough disk space'
-  echo 'You can control where Bazel stores test artifacts by setting the'
-  echo '`TEST_TMPDIR` environment variable.'
+# "TFCI_MACOS_UPGRADE_PYENV_ENABLE" is used to decide if we need to upgrade the
+# Pyenv version. We enable this for macOS x86 builds as the default Pyenv on
+# those VMs does not support installing Python 3.10 and above which we need
+# for running smoke tests in nightly/release wheel builds.
+if [[ "${TFCI_MACOS_UPGRADE_PYENV_ENABLE}" == 1 ]]; then
+  brew upgrade pyenv
 fi
 
-if [[ "${TFCI_PYTHON_VERSION}" == "3.12" ]]; then
+# "TFCI_MACOS_PYENV_INSTALL_ENABLE" controls whether to use Pyenv to install
+# the Python version set in "TFCI_PYTHON_VERSION" and use it as default.
+# We enable this in the nightly and release builds because before uploading the
+# wheels, we install them in a virtual environment and run some smoke tests on
+# it. TFCI Mac VMs only have one Python version installed so we need to install
+# the other versions manually.
+if [[ "${TFCI_MACOS_PYENV_INSTALL_ENABLE}" == 1 ]]; then
+  pyenv install "$TFCI_PYTHON_VERSION"
+  pyenv local "$TFCI_PYTHON_VERSION"
+  # Do a sanity check to make sure that we using the correct Python version
+  python --version
+fi
+
+if [[ "$TFCI_PYTHON_VERSION" == "3.12" ]]; then
   # dm-tree (Keras v3 dependency) doesn't have pre-built wheels for 3.12 yet.
   # Having CMake allows building them.
   # Once the wheels are added, this should be removed - b/308399490.
   brew install cmake
 fi
+
+# Scheduled nightly and release builds upload build artifacts (Pip packages,
+# Libtensorflow archives) to GCS buckets. TFCI Mac VMs need to authenticate as
+# a service account that has the right permissions to be able to do so.
+set +x
+if [[ -n "${GOOGLE_APPLICATION_CREDENTIALS:-}" ]]; then
+  gcloud auth activate-service-account
+fi
+set -x