chore: import upstream snapshot with attribution

ci/official/utilities/cleanup_docker.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+cat <<EOF
+IMPORTANT: These tests ran under docker. This script does not clean up the
+container for you! You can delete the container with:
+
+$ docker rm -f tf
+
+You can also execute more commands within the container with e.g.:
+
+$ docker exec tf bazel clean
+$ docker exec -it tf bash
+EOF
+
+docker ps

ci/official/utilities/cleanup_summary.sh

@@ -0,0 +1,58 @@
+#!/bin/bash
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+set -exo pipefail
+
+function resultstore_extract_fallback {
+  # In case the main script fails somehow.
+  cat <<EOF
+IMPORTANT: For bazel invocations that uploaded to ResultStore (e.g. RBE), you
+can view more detailed results that are probably easier to read than this log.
+Try the links below:
+EOF
+  # Find any "Streaming build results to" lines,
+  # de-duplicate,
+  # and print the last word from each
+  awk '/Streaming build results to/ {print $NF}' "$TFCI_OUTPUT_DIR/script.log" | uniq
+}
+
+# Print out any ResultStore URLs for Bazel invocations' results.
+# Each failed target there will have its own representation, making failures
+# easier to find and read.
+function resultstore_extract {
+  local PYTHON_BIN XML_PATH
+  PYTHON_BIN=$(which python3 2>/dev/null || which python)
+  XML_PATH="$TFCI_OUTPUT_DIR/Bazel_Test_and_Build_Results/sponge_log.xml"
+
+  local script_path="$TFCI_GIT_DIR/ci/official/utilities/extract_resultstore_links.py"
+  local log_path="$TFCI_OUTPUT_DIR/script.log"
+
+  if [[ "$TFCI_GITHUB_ACTIONS" == "true" ]] && { [[ $(uname -s) == "MSYS_NT"* ]] || [[ $(uname -s) == "MINGW64_NT"* ]]; }; then
+    script_path=$(cygpath -w "$script_path")
+    log_path=$(cygpath -w "$log_path")
+    XML_PATH=$(cygpath -w "$XML_PATH")
+  fi
+
+  "$PYTHON_BIN" \
+    "$script_path" \
+    "$log_path" \
+    --print \
+    --xml-out-path "$XML_PATH"
+}
+
+if grep -q "Streaming build results to" "$TFCI_OUTPUT_DIR/script.log"; then
+  resultstore_extract || resultstore_extract_fallback
+fi

ci/official/utilities/code_check_changed_files.bats

@@ -0,0 +1,84 @@
+# vim: filetype=bash
+#
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+setup_file() {
+    bazel version  # Start the bazel server
+
+    # Fixes "fatal: detected dubious ownership in repository" for Docker.
+    git config --system --add safe.directory '*'
+    git config --system protocol.file.allow always
+
+    # Note that you could generate a list of all the affected targets with e.g.:
+    # bazel query $(paste -sd "+" $BATS_FILE_TMPDIR/changed_files) --keep_going
+    # Only shows Added, Changed, Modified, Renamed, and Type-changed files
+    if [[ "$(git rev-parse --abbrev-ref HEAD)" == "pull_branch" ]]; then
+        # TF's CI runs 'git fetch origin "pull/PR#/merge:pull_branch"'
+        # To get the as-merged branch during the CI tests
+        git diff --diff-filter ACMRT --name-only pull_branch^ pull_branch > $BATS_FILE_TMPDIR/changed_files
+    else
+        # If the branch is not present, then diff against origin/master
+        git diff --diff-filter ACMRT --name-only origin/master > $BATS_FILE_TMPDIR/changed_files
+    fi
+}
+
+# Note: this is excluded on the full code base, since any submitted code must
+# have passed Google's internal style guidelines.
+@test "Check buildifier formatting on BUILD files" {
+    echo "buildifier formatting is recommended. Here are the suggested fixes:"
+    echo "============================="
+    grep -e 'BUILD' $BATS_FILE_TMPDIR/changed_files \
+        | xargs buildifier -v -mode=diff -diff_command="git diff --no-index"
+}
+
+# Note: this is excluded on the full code base, since any submitted code must
+# have passed Google's internal style guidelines.
+@test "Check formatting for C++ files" {
+    skip "clang-format doesn't match internal clang-format checker"
+    echo "clang-format is recommended. Here are the suggested changes:"
+    echo "============================="
+    grep -e '\.h$' -e '\.cc$' $BATS_FILE_TMPDIR/changed_files > $BATS_TEST_TMPDIR/files || true
+    if [[ ! -s $BATS_TEST_TMPDIR/files ]]; then return 0; fi
+    xargs -a $BATS_TEST_TMPDIR/files -i -n1 -P $(nproc --all) \
+        bash -c 'clang-format-12 --style=Google {} | git diff --no-index {} -' \
+        | tee $BATS_TEST_TMPDIR/needs_help.txt
+    echo "You can use clang-format --style=Google -i <file> to apply changes to a file."
+    [[ ! -s $BATS_TEST_TMPDIR/needs_help.txt ]]
+}
+
+# Note: this is excluded on the full code base, since any submitted code must
+# have passed Google's internal style guidelines.
+@test "Check pylint for Python files" {
+    echo "Python formatting is recommended. Here are the pylint errors:"
+    echo "============================="
+    grep -e "\.py$" $BATS_FILE_TMPDIR/changed_files > $BATS_TEST_TMPDIR/files || true
+    if [[ ! -s $BATS_TEST_TMPDIR/files ]]; then return 0; fi
+    xargs -a $BATS_TEST_TMPDIR/files -n1 -P $(nproc --all) \
+        python -m pylint --rcfile=tensorflow/tools/ci_build/pylintrc --score false \
+        | grep -v "**** Module" \
+        | tee $BATS_TEST_TMPDIR/needs_help.txt
+    [[ ! -s $BATS_TEST_TMPDIR/needs_help.txt ]]
+}
+
+@test "API compatibility test passes, ensuring no unexpected changes to the TF API" {
+    bazel test $TFCI_BAZEL_COMMON_ARGS //tensorflow/tools/api/tests:api_compatibility_test
+    echo "You have to re-generate the TF API goldens and have the API changes reviewed."
+    echo "Look at the instructions for ':api_compatibility_test -- --update_goldens=True'"
+}
+
+teardown_file() {
+    bazel shutdown
+}

ci/official/utilities/code_check_full.bats

@@ -0,0 +1,330 @@
+# vim: filetype=bash
+#
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+setup_file() {
+    bazel version  # Start the bazel server
+}
+
+# Do a bazel query specifically for the licenses checker. It searches for
+# targets matching the provided query, which start with // or @ but not
+# //tensorflow (so it looks for //third_party, //external, etc.), and then
+# gathers the list of all packages (i.e. directories) which contain those
+# targets.
+license_query() {
+ bazel cquery --experimental_cc_shared_library "$1" --keep_going \
+  | grep -e "^//" -e "^@" \
+  | grep -E -v "^//tensorflow" \
+  | sed -e 's|:.*||' \
+  | sort -u
+}
+
+# Verify that, given a build target and a license-list generator target, all of
+# the dependencies of that target which include a license notice file are then
+# included when generating that license. Necessary because the license targets
+# in TensorFlow are manually enumerated rather than generated automatically.
+do_external_licenses_check(){
+  BUILD_TARGET="$1"
+  LICENSES_TARGET="$2"
+
+  # grep patterns for targets which are allowed to be missing from the licenses
+  cat > $BATS_TEST_TMPDIR/allowed_to_be_missing <<EOF
+@absl_py//absl
+@bazel_tools//platforms
+@bazel_tools//third_party/
+@bazel_tools//tools
+@jpegxl//lib
+@local
+@xla//xla
+@tsl//tsl
+@com_google_absl//absl
+@pybind11_abseil//pybind11_abseil
+@org_tensorflow//
+@com_github_googlecloudplatform_google_cloud_cpp//google
+@com_github_grpc_grpc//src/compiler
+@com_google_protobuf//upb_generator
+@llvm-project//third-party/siphash
+@platforms//os
+@ml_dtypes_py//ml_dtypes
+@ruy//
+@rules_java_builtin//toolchains
+@rules_ml_toolchain//
+@rules_python//
+@stablehlo//stablehlo/experimental
+EOF
+
+  # grep patterns for targets which are allowed to be extra licenses
+  cat > $BATS_TEST_TMPDIR/allowed_to_be_extra <<EOF
+@xla//third_party/mkl
+@xla//third_party/mkl_dnn
+@absl_py//
+@bazel_tools//src
+@bazel_tools//platforms
+@bazel_tools//tools
+@org_tensorflow//tensorflow
+@com_google_absl//
+@com_google_protobuf//
+@internal_platforms_do_not_use//host
+@jpegxl//
+@pybind11_abseil//pybind11_abseil
+//external
+@local
+@xla//xla
+@tsl//tsl
+@com_github_googlecloudplatform_google_cloud_cpp//
+@embedded_jdk//
+^//$
+@ml_dtypes_py//
+@ruy//
+@rules_ml_toolchain//
+EOF
+
+  license_query "attr('licenses', 'notice', deps($BUILD_TARGET))" > $BATS_TEST_TMPDIR/expected_licenses
+  license_query "deps($LICENSES_TARGET)" > $BATS_TEST_TMPDIR/actual_licenses
+
+  # Column 1 is left only, Column 2 is right only, Column 3 is shared lines
+  # Select lines unique to actual_licenses, i.e. extra licenses.
+  comm -1 -3 $BATS_TEST_TMPDIR/expected_licenses $BATS_TEST_TMPDIR/actual_licenses | grep -v -f $BATS_TEST_TMPDIR/allowed_to_be_extra > $BATS_TEST_TMPDIR/actual_extra_licenses || true
+  # Select lines unique to expected_licenses, i.e. missing licenses
+  comm -2 -3 $BATS_TEST_TMPDIR/expected_licenses $BATS_TEST_TMPDIR/actual_licenses | grep -v -f $BATS_TEST_TMPDIR/allowed_to_be_missing > $BATS_TEST_TMPDIR/actual_missing_licenses || true
+
+  if [[ -s $BATS_TEST_TMPDIR/actual_extra_licenses ]]; then
+    echo "Please remove the following extra licenses from $LICENSES_TARGET:"
+    cat $BATS_TEST_TMPDIR/actual_extra_licenses
+  fi
+
+  if [[ -s $BATS_TEST_TMPDIR/actual_missing_licenses ]]; then
+    echo "Please include the missing licenses for the following packages in $LICENSES_TARGET:"
+    cat $BATS_TEST_TMPDIR/actual_missing_licenses
+  fi
+
+  # Fail if either of the two "extras" or "missing" lists are present. If so,
+  # then the user will see the above error messages.
+  [[ ! -s $BATS_TEST_TMPDIR/actual_extra_licenses ]] && [[ ! -s $BATS_TEST_TMPDIR/actual_missing_licenses ]]
+}
+
+@test "Pip package generated license includes all dependencies' licenses" {
+  do_external_licenses_check \
+    "//tensorflow/tools/pip_package:wheel" \
+    "//tensorflow/tools/pip_package:licenses"
+}
+
+# This test ensures that all the targets built into the Python package include
+# their dependencies. It's a rewritten version of the "smoke test", an older
+# Python script that was very difficult to understand. See
+# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/pip_package/pip_smoke_test.py
+@test "Pip package includes all required //tensorflow dependencies" {
+  # grep patterns for packages whose dependencies can be ignored
+  cat > $BATS_TEST_TMPDIR/ignore_deps_for_these_packages <<EOF
+//tensorflow/lite
+//tensorflow/compiler/mlir/lite
+//tensorflow/compiler/mlir/tfrt
+//tensorflow/core/runtime_fallback
+//tensorflow/core/tfrt
+//tensorflow/python/kernel_tests/signal
+//tensorflow/examples
+//tensorflow/tools/android
+//tensorflow/python/eager/benchmarks
+EOF
+
+  # grep patterns for files and targets which don't need to be in the pip
+  # package, ever.
+  cat > $BATS_TEST_TMPDIR/ignore_these_deps <<EOF
+benchmark
+_test$
+_test.py$
+_test_cpu$
+_test_cpu.py$
+_test_gpu$
+_test_gpu.py$
+_test_lib$
+//tensorflow/cc/saved_model:saved_model_test_files
+//tensorflow/cc/saved_model:saved_model_half_plus_two
+//tensorflow:no_tensorflow_py_deps
+//tensorflow/tools/pip_package:win_pip_package_marker
+//tensorflow/core:image_testdata
+//tensorflow/core/kernels/cloud:bigquery_reader_ops
+//tensorflow/python:extra_py_tests_deps
+//tensorflow/python:mixed_precision
+//tensorflow/python:tf_optimizer
+//tensorflow/python:compare_test_proto_py
+//tensorflow/python/framework:test_file_system.so
+//tensorflow/python/debug:grpc_tensorflow_server.par
+//tensorflow/python/feature_column:vocabulary_testdata
+//tensorflow/python/util:nest_test_main_lib
+//tensorflow/lite/experimental/examples/lstm:rnn_cell
+//tensorflow/lite/experimental/examples/lstm:rnn_cell.py
+//tensorflow/lite/experimental/examples/lstm:unidirectional_sequence_lstm_test
+//tensorflow/lite/experimental/examples/lstm:unidirectional_sequence_lstm_test.py
+//tensorflow/lite/python:interpreter
+//tensorflow/lite/python:interpreter_test
+//tensorflow/lite/python:interpreter.py
+//tensorflow/lite/python:interpreter_test.py
+EOF
+
+  # Get the full list of files and targets which get included into the pip
+  # package
+  bazel cquery --keep_going 'deps(//tensorflow/tools/pip_package:wheel)' | sort -u > $BATS_TEST_TMPDIR/pip_deps
+  # Find all Python py_test targets not tagged "no_pip" or "manual", excluding
+  # any targets in ignored packages. Combine this list of targets into a bazel
+  # query list (e.g. the list becomes "target+target2+target3")
+  bazel cquery --keep_going 'kind(py_test, //tensorflow/python/...) - attr("tags", "no_pip|manual", //tensorflow/python/...)' | grep -v -f $BATS_TEST_TMPDIR/ignore_deps_for_these_packages | paste -sd "+" - > $BATS_TEST_TMPDIR/deps
+  # Find all one-step dependencies of those tests which are from //tensorflow
+  # (since external deps will come from Python-level pip dependencies),
+  # excluding dependencies and files that are known to be unneccessary.
+  # This creates a list of targets under //tensorflow that are required for
+  # TensorFlow python tests.
+  bazel cquery --keep_going "deps($(cat $BATS_TEST_TMPDIR/deps), 1)" | grep "^//tensorflow" | grep -v -f $BATS_TEST_TMPDIR/ignore_these_deps | sort -u > $BATS_TEST_TMPDIR/required_deps
+
+
+  # Find if any required dependencies are missing from the list of dependencies
+  # included in the pip package.
+  # (comm: Column 1 is left, Column 2 is right, Column 3 is shared lines)
+  comm -2 -3 $BATS_TEST_TMPDIR/required_deps $BATS_TEST_TMPDIR/pip_deps > $BATS_TEST_TMPDIR/missing_deps || true
+
+  if [[ -s $BATS_TEST_TMPDIR/missing_deps ]]; then
+    cat <<EOF
+One or more test dependencies are not in the pip package.
+If these test dependencies need to be in the TensorFlow pip package, please
+add them to //tensorflow/tools/pip_package/BUILD. Otherwise, add the no_pip tag
+to the test, or change code_check_full.bats in the SIG Build repository. That's
+https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/tf_sig_build_dockerfiles/devel.usertools/code_check_full.bats
+Here are the affected tests:
+EOF
+    while read dep; do
+      echo "For dependency $dep:"
+      # For every missing dependency, find the tests which directly depend on
+      # it, and print that list for debugging. Not really clear if this is
+      # helpful since the only examples I've seen are enormous.
+      bazel cquery "rdeps(kind(py_test, $(cat $BATS_TEST_TMPDIR/deps)), $dep, 1)"
+    done < $BATS_TEST_TMPDIR/missing_deps
+    exit 1
+  fi
+}
+
+# The Python package is not allowed to depend on any CUDA packages.
+@test "Pip package doesn't depend on CUDA" {
+  bazel cquery \
+    --experimental_cc_shared_library \
+    --@local_config_cuda//:enable_cuda \
+    --@local_config_cuda//cuda:include_cuda_libs=false \
+    --repo_env=HERMETIC_CUDA_VERSION="12.5.1" \
+    --repo_env=HERMETIC_CUDNN_VERSION="9.3.0" \
+    --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5" \
+    --repo_env=HERMETIC_NCCL_VERSION="2.29.7" \
+    "somepath(//tensorflow/tools/pip_package:wheel, " \
+    "@local_config_cuda//cuda:cudart + "\
+    "@local_config_cuda//cuda:cudart + "\
+    "@local_config_cuda//cuda:cuda_driver + "\
+    "@local_config_cuda//cuda:cudnn + "\
+    "@local_config_cuda//cuda:curand + "\
+    "@local_config_cuda//cuda:cusolver + "\
+    "@local_config_tensorrt//:tensorrt)" --keep_going > $BATS_TEST_TMPDIR/out
+
+  cat <<EOF
+There was a path found connecting //tensorflow/tools/pip_package:wheel
+to a banned CUDA dependency. Here's the output from bazel query:
+EOF
+  cat $BATS_TEST_TMPDIR/out
+  [[ ! -s $BATS_TEST_TMPDIR/out ]]
+}
+
+@test "Pip package doesn't depend on CUDA for static builds (i.e. Windows)" {
+  bazel cquery \
+    --experimental_cc_shared_library \
+    --@local_config_cuda//:enable_cuda \
+    --@local_config_cuda//cuda:include_cuda_libs=false \
+    --repo_env=HERMETIC_CUDA_VERSION="12.5.1" \
+    --repo_env=HERMETIC_CUDNN_VERSION="9.3.0" \
+    --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5" \
+    --repo_env=HERMETIC_NCCL_VERSION="2.29.7" \
+    --define framework_shared_object=false \
+    "somepath(//tensorflow/tools/pip_package:wheel, " \
+    "@local_config_cuda//cuda:cudart + "\
+    "@local_config_cuda//cuda:cudart + "\
+    "@local_config_cuda//cuda:cuda_driver + "\
+    "@local_config_cuda//cuda:cudnn + "\
+    "@local_config_cuda//cuda:curand + "\
+    "@local_config_cuda//cuda:cusolver + "\
+    "@local_config_tensorrt//:tensorrt)" --keep_going > $BATS_TEST_TMPDIR/out
+
+  cat <<EOF
+There was a path found connecting //tensorflow/tools/pip_package:wheel
+to a banned CUDA dependency when '--define framework_shared_object=false' is set.
+This means that a CUDA target was probably included via an is_static condition,
+used when targeting platforms like Windows where we build statically instead
+of dynamically. Here's the output from bazel query:
+EOF
+  cat $BATS_TEST_TMPDIR/out
+  [[ ! -s $BATS_TEST_TMPDIR/out ]]
+}
+
+@test "All tensorflow.org/code links point to real files" {
+    for i in $(grep -onI 'https://www.tensorflow.org/code/[a-zA-Z0-9/._-]\+' -r tensorflow); do
+        target=$(echo $i | sed 's!.*https://www.tensorflow.org/code/!!g')
+
+        if [[ ! -f $target ]] && [[ ! -d $target ]]; then
+            echo "$i" >> errors.txt
+        fi
+        if [[ -e errors.txt ]]; then
+            echo "Broken links found:"
+            cat errors.txt
+            rm errors.txt
+            false
+        fi
+    done
+}
+
+@test "No duplicate files on Windows" {
+    cat <<EOF
+Please rename files so there are no repeats. For example, README.md and
+Readme.md would be the same file on Windows. In this test, you would get a
+warning for "readme.md" because it makes everything lowercase. There are
+repeats of these filename(s) with different casing:
+EOF
+    find . | tr '[A-Z]' '[a-z]' | sort | uniq -d | tee $BATS_FILE_TMPDIR/repeats
+    [[ ! -s $BATS_FILE_TMPDIR/repeats ]]
+}
+
+# It's unclear why, but running this on //tensorflow/... is faster than running
+# only on affected targets, usually. There are targets in //tensorflow/lite that
+# don't pass --nobuild, so they're on their own.
+#
+# Although buildifier checks for formatting as well, "bazel build nobuild"
+# checks for cross-file issues like bad includes or missing BUILD definitions.
+#
+# We can't test on the windows toolchains because they're using a legacy
+# toolchain format (or something) that specifies the toolchain directly instead
+# of as a "repository". They can't be valid on Linux because Linux can't do
+# anything with a Windows-only toolchain, and bazel errors if trying to build
+# that directory.
+@test "bazel nobuild passes on all of TF except TF Lite and win toolchains" {
+    bazel build --experimental_cc_shared_library --incompatible_enable_android_toolchain_resolution=false --skip_incompatible_explicit_targets --nobuild --keep_going -- //tensorflow/core/... //tensorflow/python/... //tensorflow/compiler/... //tensorflow/tools/pip_package:build_pip_package_py -//tensorflow/lite/... -//tensorflow/tools/lib_package/...
+}
+
+@test "API compatibility test passes, ensuring no unexpected changes to the TF API" {
+    bazel test $TFCI_BAZEL_COMMON_ARGS //tensorflow/tools/api/tests:api_compatibility_test
+    echo "You have to re-generate the TF API goldens and have the API changes reviewed."
+    echo "Look at the instructions for ':api_compatibility_test -- --update_goldens=True'"
+}
+
+# See b/279852433 (internal).
+@test "Verify that it's possible to query every TensorFlow target without BUILD errors" {
+    bazel query --keep_going --deleted_packages=tensorflow/tools/toolchains/win/20240424,tensorflow/tools/toolchains/win/tf_win_05022023,tensorflow/tools/toolchains/win/bazel_211,tensorflow/tools/toolchains/win2022/20260322,tensorflow/tools/toolchains/win2022/20241118,tensorflow/tools/lib_package "deps(//tensorflow/... - attr(tags, 'manual|no_oss|gpu|metal|tflite|android|ios', //tensorflow/...) - //tensorflow/tools/pip_package:prebuilt_tf_py_import - //tensorflow/tools/pip_package:prebuilt_tf_py_import_unpacked_wheel)" > /dev/null
+}
+
+teardown_file() {
+    bazel shutdown
+}

ci/official/utilities/convert_msys_paths_to_win_paths.py

@@ -0,0 +1,76 @@
+#!/usr/bin/python3
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+
+"""Converts MSYS Linux-like paths stored in env variables to Windows paths.
+
+This is necessary on Windows, because some applications do not understand/handle
+Linux-like paths MSYS uses, for example, Docker.
+"""
+
+import argparse
+import os
+
+
+def should_convert(var_name: str,
+                   blacklist: list[str] | None,
+                   whitelist_prefix: list[str] | None):
+  """Check the variable name against white/black lists."""
+  if blacklist and var_name in blacklist:
+    return False
+  if not whitelist_prefix:
+    return True
+
+  for prefix in whitelist_prefix:
+    if var_name.startswith(prefix):
+      return True
+  return False
+
+
+def main(parsed_args: argparse.Namespace):
+  converted_vars = {}
+
+  for var, value in os.environ.items():
+    if not value or not should_convert(var,
+                                       parsed_args.blacklist,
+                                       parsed_args.whitelist_prefix):
+      continue
+
+    # In Python, MSYS, Linux-like paths are automatically read as Windows paths
+    # with forward slashes, e.g. 'C:/Program Files', instead of
+    # '/c/Program Files', thus becoming converted simply by virtue of having
+    # been read.
+    converted_vars[var] = value
+
+  var_str = '\n'.join(f'{k}="{v}"'
+                      for k, v in converted_vars.items())
+  # The string can then be piped into `source`, to re-set the
+  # 'converted' variables.
+  print(var_str)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser(description=(
+      'Convert MSYS paths in environment variables to Windows paths.'))
+  parser.add_argument('--blacklist',
+                      nargs='*',
+                      help='List of variables to ignore')
+  parser.add_argument('--whitelist-prefix',
+                      nargs='*',
+                      help='Prefix for variables to include')
+  args = parser.parse_args()
+
+  main(args)

ci/official/utilities/extract_resultstore_links.py

@@ -0,0 +1,294 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+"""Extracts ResultStore links from a log containing Bazel invocations.
+
+The links and the invocations' status can then be printed out, or output in the
+form of JUnit-based XML.
+"""
+import argparse
+import datetime
+import os
+import re
+from typing import Dict, Union
+import xml.etree.ElementTree as ElemTree
+
+
+ResultDictType = Dict[str, Dict[str, Union[str, int]]]
+
+RESULT_STORE_LINK_RE = re.compile(
+    r'^INFO: Streaming build results to: (https://[\w./\-]+)')
+FAILED_BUILD_LINE = 'FAILED: Build did NOT complete successfully'
+BUILD_STATUS_LINE = 'INFO: Build'
+TESTS_FAILED_RE = re.compile(r'^INFO: Build completed, \d+ tests? FAILED')
+BAZEL_COMMAND_RE = re.compile(
+    r'(^| )(?P<command>bazel (.*? )?(?P<type>test|build) .+)')
+
+
+class InvokeStatus:
+  tests_failed = 'tests_failed'
+  build_failed = 'build_failed'
+  passed = 'passed'
+
+
+def parse_args() -> argparse.Namespace:
+  """Parses the commandline args."""
+  parser = argparse.ArgumentParser(
+      description='Extracts ResultStore links from a build log.\n'
+                  'These can be then printed out, and/or output into a '
+                  'JUnit-based XML file inside a specified directory.')
+
+  parser.add_argument('build_log',
+                      help='Path to a build log.')
+  parser.add_argument('--xml-out-path',
+                      required=False,
+                      help='Path to which to output '
+                           'the JUnit-based XML with ResultStore links.')
+  parser.add_argument('--print',
+                      action='store_true', dest='print', default=False,
+                      help='Whether to print out a short summary with the '
+                           'found ResultStore links (if any).')
+  parser.add_argument('-v', '--verbose',
+                      action='store_true', dest='verbose', default=False,
+                      help='Prints out lines helpful for debugging.')
+  parsed_args = parser.parse_args()
+  if not parsed_args.print and not parsed_args.xml_out_path:
+    raise TypeError('`--print` or `--xml-out-path` must be specified')
+
+  return parsed_args
+
+
+def parse_log(file_path: str,
+              verbose: bool = False) -> ResultDictType:
+  """Finds ResultStore links, and tries to determine their status."""
+  with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
+    log_lines = f.read().splitlines()
+
+  result_store_links: ResultDictType = {}
+  current_url = None
+  for i in range(len(log_lines)):
+    line = log_lines[i]
+    result_store_line_match = re.search(RESULT_STORE_LINK_RE, line)
+    if not result_store_line_match:
+      continue
+
+    url = result_store_line_match.group(1)
+    url_lines = result_store_links.setdefault(url, {})
+    # Each bazel RBE invocation should produce two
+    # 'Streaming build results to: ...' lines, one at the start, and one at the
+    # end of the invocation.
+    # If there's a failure message, it will be found in-between these two.
+
+    if not current_url:
+      url_lines['start'] = i
+    elif current_url == url:
+      url_lines['end'] = i
+    else:
+      result_store_links[current_url]['next_url'] = i
+      url_lines['start'] = i
+    current_url = url
+
+  previous_end_line = None
+  for url, lines in result_store_links.items():
+    lines['status'] = InvokeStatus.passed  # default to passed
+    start_line = lines['start']
+    end_line = lines.get('end', lines.get('next_url', len(log_lines))) - 1
+    k = end_line
+    while k > start_line:
+      backtrack_line = log_lines[k]
+      build_failed = backtrack_line.startswith(FAILED_BUILD_LINE)
+      if build_failed or not backtrack_line.startswith(BUILD_STATUS_LINE):
+        tests_failed = False
+      else:
+        tests_failed = re.search(TESTS_FAILED_RE, backtrack_line)
+      if build_failed or tests_failed:
+        log_fragment = '\n'.join(
+            log_lines[max(k - 20, 0):min(end_line + 1, len(log_lines) - 1)])
+        lines['log_fragment'] = log_fragment
+        lines['status'] = (InvokeStatus.build_failed if build_failed
+                           else InvokeStatus.tests_failed)
+        if verbose:
+          print(f'Found failed invocation: {url.rsplit("/")[-1]}\n'
+                f'Log fragment:\n'
+                f'```\n{log_fragment}\n```\n'
+                f'{"=" * 140}')
+        break
+      k -= 1
+
+    # A low-effort attempt to find the bazel command that triggered the
+    # invocation.
+    bazel_comm_min_line_i = (previous_end_line if previous_end_line is not None
+                             else 0)
+    while k > bazel_comm_min_line_i:
+      backtrack_line = log_lines[k]
+      # Don't attempt to parse multi-line commands broken up by backslashes
+      if 'bazel ' in backtrack_line and not backtrack_line.endswith('\\'):
+        bazel_line = BAZEL_COMMAND_RE.search(backtrack_line)
+        if bazel_line:
+          lines['command'] = bazel_line.group('command')
+          lines['command_type'] = bazel_line.group('type')
+          break
+      k -= 1
+      continue
+    previous_end_line = lines.get('end') or start_line
+
+  return result_store_links
+
+
+def indent_xml(elem, level=0) -> None:
+  """Indents and newlines the XML for better output."""
+  indent_str = '\n' + level * '  '
+  if len(elem):  # pylint: disable=g-explicit-length-test  # `if elem` not valid
+    if not elem.text or not elem.text.strip():
+      elem.text = indent_str + '  '
+    if not elem.tail or not elem.tail.strip():
+      elem.tail = indent_str
+    for elem in elem:
+      indent_xml(elem, level + 1)
+    if not elem.tail or not elem.tail.strip():
+      elem.tail = indent_str
+  else:
+    if level and (not elem.tail or not elem.tail.strip()):
+      elem.tail = indent_str
+
+
+def create_xml_file(result_store_dict: ResultDictType,
+                    output_path: str,
+                    verbose: bool = False):
+  """Creates a JUnit-based XML file, with each invocation as a testcase."""
+  os.makedirs(os.path.dirname(output_path), exist_ok=True)
+  failure_count = 0
+  error_count = 0
+
+  date_time = datetime.datetime
+  attrib = {'name': 'Bazel Invocations', 'time': '0.0',
+            'timestamp': date_time.isoformat(date_time.utcnow())}
+  testsuites = ElemTree.Element('testsuites')
+  testsuite = ElemTree.SubElement(testsuites, 'testsuite')
+  for url, invocation_results in result_store_dict.items():
+    invocation_id = url.rsplit('/')[-1]
+    if verbose:
+      print(f'Creating testcase for invocation {invocation_id}')
+    status = invocation_results['status']
+    command = invocation_results.get('command')
+    command_type = invocation_results.get('command_type')
+
+    case_attrib = attrib.copy()
+    if command_type:
+      command_type = command_type.title()
+      case_name = f'{command_type} invocation {invocation_id}'
+    else:
+      case_name = f' Invocation {invocation_id}'
+    case_attrib.update({'name': case_name,
+                        'status': 'run', 'result': 'completed'})
+
+    testcase = ElemTree.SubElement(testsuite, 'testcase', attrib=case_attrib)
+    if status in (InvokeStatus.tests_failed, InvokeStatus.build_failed):
+      if status == InvokeStatus.tests_failed:
+        failure_count += 1
+        elem_name = 'failure'
+      else:
+        error_count += 1
+        elem_name = 'error'
+      if command:
+        failure_msg = (f'\nThe command was:\n\n'
+                       f'{command}\n\n')
+      else:
+        failure_msg = ('\nCouldn\'t parse a bazel command '
+                       'matching the invocation, inside the log. '
+                       'Please look for it in the build log.\n\n')
+      failure_msg += (
+          f'See the ResultStore link for a detailed view of failed targets:\n'
+          f'{url}\n\n')
+      failure_msg += (
+          f'Here\'s a fragment of the log containing the failure:\n\n'
+          f'[ ... TRUNCATED ... ]\n\n'
+          f'{invocation_results["log_fragment"]}\n'
+          f'\n[ ... TRUNCATED ... ]\n'
+      )
+      failure = ElemTree.SubElement(
+          testcase, elem_name,
+          message=f'Bazel invocation {invocation_id} failed.')
+      failure.text = failure_msg
+    else:
+      properties = ElemTree.SubElement(testcase, 'properties')
+      success_msg = 'Build completed successfully.\n' f'See {url} for details.'
+      ElemTree.SubElement(properties, 'property',
+                          name='description',
+                          value=success_msg)
+      if command:
+        ElemTree.SubElement(properties, 'property',
+                            name='bazel_command',
+                            value=command)
+
+    suite_specific = {'tests': str(len(result_store_dict)),
+                      'errors': str(error_count),
+                      'failures': str(failure_count)}
+    suite_attrib = attrib.copy()
+    suite_attrib.update(suite_specific)
+    testsuites.attrib = suite_attrib
+    testsuite.attrib = suite_attrib
+    indent_xml(testsuites)
+
+  tree = ElemTree.ElementTree(testsuites)
+  file_path = os.path.join(output_path)
+  with open(file_path, 'wb') as f:
+    f.write(b'<?xml version="1.0"?>\n')
+    tree.write(f)
+    if verbose:
+      print(f'\nWrote XML with Bazel invocation results to {file_path}')
+
+
+def print_invocation_results(result_store_dict: ResultDictType):
+  """Prints out a short summary of the found ResultStore links (if any)."""
+  print()
+  if not result_store_dict:
+    print('Found no ResultStore links for Bazel build/test invocations.')
+  else:
+    print(f'Found {len(result_store_dict)} ResultStore link(s) for '
+          f'Bazel invocations.\n'
+          f'ResultStore contains individual representations of each target '
+          f'that were run/built during the invocation.\n'
+          f'These results are generally easier to read than looking through '
+          f'the entire build log:\n')
+  i = 1
+  for url, invocation_results in result_store_dict.items():
+    line_str = f'Invocation #{i} ({invocation_results["status"]}):\n'
+    command = invocation_results.get('command')
+    if command:
+      line_str += command
+    else:
+      line_str += ('Couldn\'t parse the bazel command, '
+                   'check inside the build log instead')
+    line_str += f'\n{url}\n'
+    print(line_str)
+    i += 1
+
+
+def main():
+  args = parse_args()
+  verbose = args.verbose
+  build_log_path = os.path.expandvars(args.build_log)
+  links = parse_log(build_log_path, verbose=verbose)
+
+  if args.xml_out_path:
+    output_path = os.path.expandvars(args.xml_out_path)
+    create_xml_file(links, output_path, verbose=verbose)
+  if args.print:
+    print_invocation_results(links)
+
+
+if __name__ == '__main__':
+  main()

ci/official/utilities/generate_index_html.sh

@@ -0,0 +1,52 @@
+#!/bin/bash
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+# Generates a handy index.html with a bunch of Kokoro links for GitHub
+# presubmits.
+# Usage: generate_index_html.sh /path/to/output/index.html
+
+cat > "$1" <<EOF
+<html>
+<head>
+<title>$(basename "$KOKORO_JOB_NAME")</title>
+</head>
+<body>
+<h1>TensorFlow Job Logs and Links</h1>
+<h2>Job Details</h2>
+<ul>
+<li>Job name: $KOKORO_JOB_NAME</li>
+<li>Job pool: $KOKORO_JOB_POOL</li>
+<li>Job ID: $KOKORO_BUILD_ID</li>
+<li>Current HEAD Piper Changelist, if any: cl/${KOKORO_PIPER_CHANGELIST:-not available}</li>
+<li>Pull Request Number, if any: ${KOKORO_GITHUB_PULL_REQUEST_NUMBER_tensorflow:- none}</li>
+<li>Pull Request Link, if any: <a href="$KOKORO_GITHUB_PULL_REQUEST_URL_tensorflow">${KOKORO_GITHUB_PULL_REQUEST_URL_tensorflow:-none}</a></li>
+<li>Commit: $KOKORO_GIT_COMMIT_tensorflow</li>
+</ul>
+<h2>Googlers-Only Links</h2>
+<ul>
+<li><a href="http://sponge2/$KOKORO_BUILD_ID">Sponge2</a></li>
+<li><a href="http://sponge/target:$KOKORO_JOB_NAME">Sponge - recent jobs</a></li>
+<li><a href="http://fusion2/ci/kokoro/prod:$(echo "$KOKORO_JOB_NAME" | sed 's!/!%2F!g')">Test Fusion - recent jobs</a></li>
+<li><a href="http://cs/f:devtools/kokoro/config/prod/$KOKORO_JOB_NAME">Codesearch - job definition</a></li>
+<li><a href="http://cs/f:learning/brain/testing/kokoro/$(echo "$KOKORO_JOB_NAME" | sed 's!tensorflow/!!g')">Codesearch - build definition & scripts</a></li>
+<li><a href="http://cs/$KOKORO_JOB_NAME">Codesearch - All references to this job</a></li>
+</ul>
+<h2>Non-Googler Links</h2>
+<ul>
+<li><a href="https://source.cloud.google.com/results/invocations/$KOKORO_BUILD_ID">ResultStore</a></li>
+</ul>
+</body></html>
+EOF

ci/official/utilities/rename_and_verify_wheels.sh

@@ -0,0 +1,196 @@
+#!/usr/bin/env bash
+#
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+#
+# Usage: rename_and_verify_wheels.sh
+# This script is aware of TFCI_ variables, so it doesn't need any arguments.
+# Puts new wheel through auditwheel to rename and verify it, deletes the old
+# one, checks the filesize, and then ensures the new wheel is installable.
+set -exo pipefail
+
+cd "$TFCI_OUTPUT_DIR"
+
+# Install UMD compat library if enabled
+if [[ "$TFCI_BAZEL_HERMETIC_CUDA_UMD_ENABLE" == 1 ]]; then
+  # Extract the UMD version resolved for this build directly from .bazelrc
+  export HERMETIC_CUDA_UMD_VERSION=""
+  TEST_CONFIG="${TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX}_wheel_test"
+  CONFIG_LINE=$(grep "^test:${TEST_CONFIG} " "$TFCI_GIT_DIR/.bazelrc" || true)
+  if [[ "$CONFIG_LINE" =~ HERMETIC_CUDA_UMD_VERSION=\"?([0-9]+\.[0-9]+\.[0-9]+) ]]; then
+    export HERMETIC_CUDA_UMD_VERSION="${BASH_REMATCH[1]}"
+  else
+    for conf in $(echo "$CONFIG_LINE" | grep -o -e '--config=[a-zA-Z0-9_-]*' | sed 's/--config=//'); do
+      SUBCONFIG_LINE=$(grep "^test:${conf} " "$TFCI_GIT_DIR/.bazelrc" || true)
+      if [[ "$SUBCONFIG_LINE" =~ HERMETIC_CUDA_UMD_VERSION=\"?([0-9]+\.[0-9]+\.[0-9]+) ]]; then
+        export HERMETIC_CUDA_UMD_VERSION="${BASH_REMATCH[1]}"
+        break
+      fi
+    done
+  fi
+
+  if [[ -n "$HERMETIC_CUDA_UMD_VERSION" ]]; then
+    echo "Installing UMD compat library inside the container..."
+    if [[ "$HERMETIC_CUDA_UMD_VERSION" =~ ^([0-9]+)\.([0-9]+) ]]; then
+      COMPAT_VERSION="${BASH_REMATCH[1]}-${BASH_REMATCH[2]}"
+    else
+      echo "Error: Invalid HERMETIC_CUDA_UMD_VERSION format ($HERMETIC_CUDA_UMD_VERSION)."
+      exit 1
+    fi
+
+    echo "Setting up NVIDIA apt repository..."
+    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub || true
+    echo "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /" > /etc/apt/sources.list.d/nvidia.list
+
+    echo "Running apt-get to install cuda-compat-${COMPAT_VERSION}..."
+    apt-get update -y
+
+    # Fetch the exact UMD version that Bazel will download from NVIDIA redistrib JSON.
+    BAZEL_UMD_VERSION=$(curl -s "https://developer.download.nvidia.com/compute/cuda/redist/redistrib_${HERMETIC_CUDA_UMD_VERSION}.json" | python3 -c 'import json, sys; print(json.load(sys.stdin).get("nvidia_driver", {}).get("version", ""))')
+
+    if [[ -z "$BAZEL_UMD_VERSION" ]]; then
+      echo "Warning: Could not extract Bazel UMD version from redistrib JSON. Falling back to default cuda-compat-${COMPAT_VERSION}."
+      apt-get install -y --no-install-recommends "cuda-compat-${COMPAT_VERSION}"
+    else
+      echo "Hermetic UMD version from Bazel redistrib json: $BAZEL_UMD_VERSION"
+      # Find the exact package version in apt-cache that matches the Bazel UMD version.
+      EXACT_PKG=$(apt-cache madison cuda-compat-${COMPAT_VERSION} | grep "$BAZEL_UMD_VERSION" | head -n1 | cut -d"|" -f2 | tr -d " " || true)
+
+      if [[ -n "$EXACT_PKG" ]]; then
+        echo "Found exact match for Bazel UMD version $BAZEL_UMD_VERSION: $EXACT_PKG"
+        apt-get install -y --no-install-recommends "cuda-compat-${COMPAT_VERSION}=${EXACT_PKG}"
+      else
+        echo "Warning: Could not find exact apt package match for $BAZEL_UMD_VERSION. Falling back to default cuda-compat-${COMPAT_VERSION}."
+        apt-get install -y --no-install-recommends "cuda-compat-${COMPAT_VERSION}"
+      fi
+    fi
+
+    COMPAT_DIR="/usr/local/cuda-${COMPAT_VERSION/-/.}/compat"
+    # Clean up any system-wide ldconfig registration to avoid affecting other tasks like Bazel tests.
+    # Instead, we rely on LD_LIBRARY_PATH dynamic configuration below.
+    rm -f /etc/ld.so.conf.d/cuda-compat.conf
+    ldconfig
+    echo "Successfully installed and configured cuda-compat-${COMPAT_VERSION}."
+  fi
+fi
+
+# Move extra wheel files somewhere out of the way. This script
+# expects just one wheel file to exist.
+if [[ "$(ls *.whl | wc -l | tr -d ' ')" != "1" ]]; then
+  echo "More than one wheel file is present: moving the oldest to"
+  echo "$TFCI_OUTPUT_DIR/extra_wheels."
+  # List all .whl files by their modification time (ls -t) and move anything
+  # other than the most recently-modified one (the newest one).
+  mkdir -p $TFCI_OUTPUT_DIR/extra_wheels
+  ls -t *.whl | tail -n +2 | xargs mv -t $TFCI_OUTPUT_DIR/extra_wheels
+fi
+
+# Repair wheels with auditwheel and delete the old one.
+if [[ "$TFCI_WHL_AUDIT_ENABLE" == "1" ]]; then
+  python3 -m auditwheel repair --plat "$TFCI_WHL_AUDIT_PLAT" --wheel-dir . *.whl
+  # if the wheel is already named correctly, auditwheel won't rename it. so we
+  # list all .whl files by their modification time (ls -t) and delete anything
+  # other than the most recently-modified one (the new one).
+  ls -t *.whl | tail -n +2 | xargs rm
+fi
+
+# Check if size is too big. TFCI_WHL_SIZE_LIMIT is in find's format, which can be
+# 'k' for kilobytes, 'M' for megabytes, or 'G' for gigabytes, and the + to indicate
+# "anything greater than" is added by the script.
+if [[ "$TFCI_WHL_SIZE_LIMIT_ENABLE" == "1" ]] && [[ -n "$("$TFCI_FIND_BIN" . -iname "*.whl" -size "+$TFCI_WHL_SIZE_LIMIT")" ]]; then
+  echo "Error: Generated wheel is too big! Limit is $TFCI_WHL_SIZE_LIMIT"
+  echo '(search for TFCI_WHL_SIZE_LIMIT to change it)'
+  ls -sh *.whl
+  exit 2
+fi
+
+# Quick install checks
+venv_dir=$(mktemp -d)
+if [[ $(uname -s) != MSYS_NT* ]]; then
+  "python${TFCI_PYTHON_VERSION}" -m venv "$venv_dir"
+  python="$venv_dir/bin/python3"
+else
+  # When using the Linux-like path, venv creation quietly fails, which is
+  # why it's converted here.
+  venv_dir=$(cygpath -m $venv_dir)
+  "/c/python${TFCI_PYTHON_VERSION}/python.exe" -m venv "$venv_dir"
+  python="$venv_dir/Scripts/python.exe"
+fi
+
+# TODO(b/366266944) Remove the check after tf docker image upgrade for NumPy 2
+# and numpy 1 support is dropped b/361369076.
+if [[ "$TFCI_WHL_NUMPY_VERSION" == 1 ]]; then
+  if [[ "$TFCI_PYTHON_VERSION" == "3.13" ]]; then
+    "$python" -m pip install numpy==1.26.4
+  else
+    "$python" -m pip install numpy==1.26.0
+  fi
+fi
+
+
+if [[ "$TFCI_BAZEL_COMMON_ARGS" =~ gpu|cuda ]]; then
+  echo "Checking to make sure tensorflow[and-cuda] is installable..."
+  "$python" -m pip install "$(echo *.whl)[and-cuda]" $TFCI_PYTHON_VERIFY_PIP_INSTALL_ARGS
+else
+  "$python" -m pip install *.whl $TFCI_PYTHON_VERIFY_PIP_INSTALL_ARGS
+fi
+
+# Detect versioned or unversioned compat library paths for LD_LIBRARY_PATH
+for dir in /usr/local/cuda/compat /usr/local/cuda-*/compat; do
+  if [[ -d "$dir" ]]; then
+    export LD_LIBRARY_PATH="${dir}:${LD_LIBRARY_PATH:-}"
+    break
+  fi
+done
+
+if [[ "$TFCI_WHL_IMPORT_TEST_ENABLE" == "1" ]]; then
+  if [[ "$TFCI_BAZEL_COMMON_ARGS" =~ gpu|cuda ]]; then
+    "$python" -c '
+import tensorflow as tf
+gpus = tf.config.list_physical_devices("GPU")
+if not gpus:
+  raise ValueError("No GPU devices found!")
+print(f"Successfully found GPU devices: {gpus}")
+t1=tf.constant([1,2,3,4])
+t2=tf.constant([5,6,7,8])
+print(tf.add(t1,t2).shape)
+'
+  else
+    "$python" -c '
+import tensorflow as tf
+t1=tf.constant([1,2,3,4])
+t2=tf.constant([5,6,7,8])
+print(tf.add(t1,t2).shape)
+'
+  fi
+  "$python" -c '
+import sys
+import tensorflow as tf
+sys.exit(0 if "keras" in tf.keras.__name__ else 1)
+'
+fi
+
+# Import tf nightly wheel built with numpy2 from PyPI in numpy1 env for testing.
+# This aims to maintain TF compatibility with NumPy 1.x until 2025 b/361369076.
+if [[ "$TFCI_WHL_NUMPY_VERSION" == 1 ]]; then
+  # Uninstall tf nightly wheel built with numpy1.
+  "$python" -m pip uninstall -y tf_nightly_numpy1
+  # Install tf nightly cpu wheel built with numpy2.x from PyPI in numpy1.x env.
+  "$python" -m pip install tf-nightly-cpu
+  if [[ "$TFCI_WHL_IMPORT_TEST_ENABLE" == "1" ]]; then
+    "$python" -c 'import tensorflow as tf; t1=tf.constant([1,2,3,4]); t2=tf.constant([5,6,7,8]); print(tf.add(t1,t2).shape)'
+    "$python" -c 'import sys; import tensorflow as tf; sys.exit(0 if "keras" in tf.keras.__name__ else 1)'
+  fi
+fi

ci/official/utilities/setup.sh

@@ -0,0 +1,155 @@
+#!/bin/bash
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+#
+# Common setup for all TF scripts.
+#
+# Make as FEW changes to this file as possible. It should not contain utility
+# functions (except for tfrun); use dedicated scripts instead and reference them
+# specifically. Use your best judgment to keep the scripts in this directory
+# lean and easy to follow. When in doubt, remember that for CI scripts, "keep it
+# simple" is MUCH more important than "don't repeat yourself."
+
+# -e: abort script if one command fails
+# -u: error if undefined variable used
+# -x: log all commands
+# -o pipefail: entire command fails if pipe fails. watch out for yes | ...
+# -o history: record shell history
+# -o allexport: export all functions and variables to be available to subscripts
+#               (affects 'source $TFCI')
+set -exo pipefail -o history -o allexport
+
+# Set TFCI_GIT_DIR, the root directory for all commands, to two directories
+# above the location of this file (setup.sh). We could also use "git rev-parse
+# --show-toplevel", but that wouldn't work for non-git repos (like if someone
+# downloaded TF as a zip archive).
+export TFCI_GIT_DIR=$(cd $(dirname "$0"); realpath ../../)
+cd "$TFCI_GIT_DIR"
+
+# "TFCI" may optionally be set to the name of an env-type file with TFCI
+# variables in it, OR may be left empty if the user has already exported the
+# relevant variables in their environment. Because of 'set -o allexport' above
+# (which is equivalent to "set -a"), every variable in the file is exported
+# for other files to use.
+#
+# Separately, if TFCI is set *and* there are also additional TFCI_ variables
+# set in the shell environment, those variables will be restored after the
+# TFCI env has been loaded. This is useful for e.g. on-demand "generic" jobs
+# where the user may wish to change just one option.
+if [[ -z "${TFCI:-}" ]]; then
+  echo '==TFCI==: The $TFCI variable is not set. This is fine as long as you'
+  echo 'already sourced a TFCI env file with "set -a; source <path>; set +a".'
+  echo 'If you have not, you will see a lot of undefined variable errors.'
+else
+  FROM_ENV=$(mktemp)
+  # "export -p" prints a list of environment values in a safe-to-source format,
+  # e.g. `declare -x TFCI_BAZEL_COMMON_ARGS="list of args"` for bash.
+  export -p | grep TFCI > "$FROM_ENV"
+
+  # Source the default ci values
+  source ./ci/official/envs/ci_default
+
+  # TODO(angerson) write this documentation
+  # Sources every env, in order, from the comma-separated list "TFCI"
+  # Assumes variables will resolve themselves correctly.
+  set +u
+  for env_file in ${TFCI//,/ }; do
+    source "./ci/official/envs/$env_file"
+  done
+  set -u
+  echo '==TFCI==: Evaluated the following TFCI variables from $TFCI:'
+  export -p | grep TFCI
+
+  # Load those stored pre-existing TFCI_ vars, if any
+  if [[ -s "$FROM_ENV" ]]; then
+    echo '==TFCI==: NOTE: Loading the following env parameters, which were'
+    echo 'already set in the shell environment. If you want to disable this'
+    echo 'behavior, create a new shell.'
+    cat "$FROM_ENV"
+    source "$FROM_ENV"
+    rm "$FROM_ENV"
+  fi
+  set +u
+fi
+
+# If building installer wheels, set the required environment variables that are
+# read by setup.py.
+if [[ "$TFCI_INSTALLER_WHL_ENABLE" == 1 ]]; then
+  export collaborator_build=True
+  # If building nightly installer wheels, set the project name to
+  # nightly equivalent.
+  if [[ "$TFCI_NIGHTLY_UPDATE_VERSION_ENABLE" == 1 ]]; then
+    export TFCI_INSTALLER_WHL_PROJECT_NAME="$TFCI_INSTALLER_WHL_NIGHTLY_PROJECT_NAME"
+  fi
+  export project_name="$TFCI_INSTALLER_WHL_PROJECT_NAME"
+fi
+
+# Mac builds have some specific setup needs. See setup_macos.sh for details
+if [[ "${OSTYPE}" =~ darwin* ]]; then
+  source ./ci/official/utilities/setup_macos.sh
+fi
+
+# Create and expand to the full path of TFCI_OUTPUT_DIR
+export TFCI_OUTPUT_DIR=$(realpath "$TFCI_OUTPUT_DIR")
+mkdir -p "$TFCI_OUTPUT_DIR"
+
+# In addition to dumping all script output to the terminal, place it into
+# $TFCI_OUTPUT_DIR/script.log
+exec > >(tee "$TFCI_OUTPUT_DIR/script.log") 2>&1
+
+# Setup tfrun, a helper function for executing steps that can either be run
+# locally or run under Docker. setup_docker.sh, below, redefines it as "docker
+# exec".
+# Important: "tfrun foo | bar" is "( tfrun foo ) | bar", not "tfrun (foo | bar)".
+# Therefore, "tfrun" commands cannot include pipes -- which is
+# probably for the better. If a pipe is necessary for something, it is probably
+# complex. Write a well-documented script under utilities/ to encapsulate the
+# functionality instead.
+tfrun() { "$@"; }
+
+if [[ $(uname -s) = MSYS_NT* ]]; then
+  source ./ci/official/utilities/windows.sh
+  echo 'Converting MSYS Linux-like paths to Windows paths (for Docker, Python, etc.)'
+  source <(python ./ci/official/utilities/convert_msys_paths_to_win_paths.py --whitelist-prefix TFCI_)
+fi
+
+# Run all "tfrun" commands under Docker. See setup_docker.sh for details
+if [[ "$TFCI_DOCKER_ENABLE" == 1 ]]; then
+  source ./ci/official/utilities/setup_docker.sh
+fi
+
+# Generate an overview page describing the build
+if [[ "$TFCI_INDEX_HTML_ENABLE" == 1 ]]; then
+  ./ci/official/utilities/generate_index_html.sh "$TFCI_OUTPUT_DIR/index.html"
+fi
+
+# Re-try `bazel --version` multiple times to get around
+# Bazel download issues.
+set +e
+MAX_RETRIES=2
+for ((i=1; i <= $MAX_RETRIES; i++)); do
+  tfrun bazel --version
+done
+set -e
+
+# Single handler for all cleanup actions, triggered on an EXIT trap.
+# TODO(angerson) Making this use different scripts may be overkill.
+cleanup() {
+  if [[ "$TFCI_DOCKER_ENABLE" == 1 ]]; then
+    ./ci/official/utilities/cleanup_docker.sh
+  fi
+  ./ci/official/utilities/cleanup_summary.sh
+}
+trap cleanup EXIT

ci/official/utilities/setup_docker.sh

@@ -0,0 +1,74 @@
+#!/bin/bash
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+if [[ "$TFCI_DOCKER_PULL_ENABLE" == 1 ]]; then
+  # Simple retry logic for docker-pull errors. Sleeps if a pull fails.
+  # Pulling an already-pulled container image will finish instantly, so
+  # repeating the command costs nothing.
+  docker pull "$TFCI_DOCKER_IMAGE" || sleep 15
+  docker pull "$TFCI_DOCKER_IMAGE" || sleep 30
+  docker pull "$TFCI_DOCKER_IMAGE" || sleep 60
+  docker pull "$TFCI_DOCKER_IMAGE"
+fi 
+
+if [[ "$TFCI_DOCKER_REBUILD_ENABLE" == 1 ]]; then
+  DOCKER_BUILDKIT=1 docker build --cache-from "$TFCI_DOCKER_IMAGE" -t "$TFCI_DOCKER_IMAGE" $TFCI_DOCKER_REBUILD_ARGS
+  if [[ "$TFCI_DOCKER_REBUILD_UPLOAD_ENABLE" == 1 ]]; then
+    docker push "$TFCI_DOCKER_IMAGE"
+  fi
+fi
+
+# Keep the existing "tf" container if it's already present.
+# The container is not cleaned up automatically! Remove it with:
+# docker rm tf
+if ! docker container inspect tf >/dev/null 2>&1 ; then
+  # Pass all existing TFCI_ variables into the Docker container
+  env_file=$(mktemp)
+  env | grep ^TFCI_ > "$env_file"
+
+  if [[ $(uname -s) == MSYS_NT* ]]; then
+    is_windows=true
+  else
+    is_windows=false
+  fi
+
+  WORKING_DIR="$TFCI_GIT_DIR"
+  if [[ "$is_windows" == true ]]; then
+    env_file=$(cygpath -m $env_file)
+    WORKING_DIR=$(replace_drive_letter_with_prefix "$TFCI_GIT_DIR" "$TFCI_OUTPUT_WIN_DOCKER_DIR")
+    echo "GCE_METADATA_HOST=$IP_ADDR" >> $env_file
+  fi
+
+  docker run $TFCI_DOCKER_ARGS --name tf -w "$WORKING_DIR" -itd --rm \
+      -v "$TFCI_GIT_DIR:$WORKING_DIR" \
+      --env-file "$env_file" \
+      "$TFCI_DOCKER_IMAGE" \
+    bash
+
+  if [[ "$is_windows" == true ]]; then
+    # Allow requests from the container.
+    # Additional setup is contained in ci/official/envs/rbe.
+    CONTAINER_IP_ADDR=$(docker inspect -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' tf)
+    netsh advfirewall firewall add rule name="Allow Metadata Proxy" dir=in action=allow protocol=TCP localport=80 remoteip="$CONTAINER_IP_ADDR"
+
+    # Stop non-essential indexing and link tracking services that
+    # may lock new files or symlinks.
+    # They may be causing sporadic "Permission denied" errors during Bazel builds.
+    # b/461500885
+    docker exec tf powershell -NoProfile -Command 'Stop-Service -Name SysMain,DiagTrack -Force -ErrorAction SilentlyContinue'
+  fi
+
+fi
+tfrun() { docker exec tf "$@"; }

ci/official/utilities/setup_macos.sh

@@ -0,0 +1,138 @@
+#!/bin/bash
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+#
+# macOS specific setup for all TF scripts.
+#
+
+# Mac version of Core utilities differ in usage. Since our scripts are written
+# with the GNU style, we need to set GNU utilities to be default on Mac.
+if [[ -n "$(which grealpath)" ]] &&  [[ -n "$(which gstat)" ]]; then
+  alias realpath=grealpath
+  alias stat=gstat
+  # By default, aliases are only expanded in interactive shells, which means
+  # that they are not substituted for their corresponding commands in shell
+  # scripts. By setting "expand_aliases", we enable alias expansion in
+  # non-interactive shells as well.
+  shopt -s expand_aliases
+else
+  echo '==TFCI==: Error: Cannot find path to grealpath or gstat'
+  echo 'TF CI scripts require GNU core utilties to be installed. Please make'
+  echo 'sure they are present on your system and try again.'
+  exit 1
+fi
+
+# "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
+
+# "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.12 and above which we need
+# for running smoke tests in nightly/release wheel builds.
+if [[ "${TFCI_MACOS_UPGRADE_PYENV_ENABLE}" == 1 ]]; then
+  echo "Upgrading pyenv..."
+  echo "Current pyevn version: $(pyenv --version)"
+
+  # Check if pyenv is managed by homebrew. If so, update and upgrade pyenv.
+  # Otherwise, install the latest pyenv from github.
+  if command -v brew &> /dev/null && brew list pyenv &> /dev/null; then
+    # On "ventura-slcn" VMs, pyenv is managed via Homebrew.
+    echo "pyenv is installed and managed by homebrew."
+    (brew update && brew upgrade pyenv) || true
+  else
+    echo "pyenv is not managed by homebrew. Installing it via github..."
+    # On "ventura" VMs, pyenv is not managed by Homebrew. Install the latest
+    # pyenv from github.
+    rm -rf "$PYENV_ROOT"
+    git clone https://github.com/pyenv/pyenv.git "$PYENV_ROOT"
+  fi
+  echo "Upgraded pyenv version: $(pyenv --version)"
+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
+  # Python 3.12 removed the module `imp` which is needed by gcloud CLI so we set
+  # `CLOUDSDK_PYTHON` to Python 3.11 which is the system Python on TFCI Mac
+  # VMs.
+  pyenv install -s "3.11"
+  if [[ $? -eq 0 ]]; then
+    pyenv local "3.11"
+    export CLOUDSDK_PYTHON=$(pyenv which python3.11)
+  else
+    echo "Python 3.11 not found, falling back to system python"
+    export CLOUDSDK_PYTHON=$(which python3)
+  fi
+  gcloud auth activate-service-account --key-file="${GOOGLE_APPLICATION_CREDENTIALS}"
+fi
+set -x
+
+# "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
+  # Install the necessary Python, unless it's already present
+  pyenv install -s "$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
+
+# TFCI Mac VM images do not have twine installed by default so we need to
+# install it manually. We use Twine in nightly builds to upload Python packages
+# to PyPI.
+if [[ "$TFCI_MACOS_TWINE_INSTALL_ENABLE" == 1 ]]; then
+  pip install twine==3.6.0
+fi
+
+# When cross-compiling with RBE, we need to copy the macOS sysroot to be
+# inside the TensorFlow root directory. We then define them as a filegroup
+# target inside "tensorflow/tools/toolchains/cross_compile/cc" so that Bazel
+# can register it as an input to compile/link actions and send it to the remote
+# VMs when needed.
+# TODO(b/316932689): Avoid copying and replace with a local repository rule.
+if [[ "$TFCI_MACOS_CROSS_COMPILE_ENABLE" == 1 ]]; then
+  mkdir -p "${TFCI_MACOS_CROSS_COMPILE_SDK_DEST}/usr"
+  mkdir -p "${TFCI_MACOS_CROSS_COMPILE_SDK_DEST}/System/Library/Frameworks/"
+  cp -r "${TFCI_MACOS_CROSS_COMPILE_SDK_SOURCE}/usr/include" "${TFCI_MACOS_CROSS_COMPILE_SDK_DEST}/usr/include"
+  cp -r "${TFCI_MACOS_CROSS_COMPILE_SDK_SOURCE}/usr/lib" "${TFCI_MACOS_CROSS_COMPILE_SDK_DEST}/usr/lib"
+  cp -r "${TFCI_MACOS_CROSS_COMPILE_SDK_SOURCE}/System/Library/Frameworks/CoreFoundation.framework" "${TFCI_MACOS_CROSS_COMPILE_SDK_DEST}/System/Library/Frameworks/CoreFoundation.framework"
+  cp -r "${TFCI_MACOS_CROSS_COMPILE_SDK_SOURCE}/System/Library/Frameworks/Security.framework" "${TFCI_MACOS_CROSS_COMPILE_SDK_DEST}/System/Library/Frameworks/Security.framework"
+  cp -r "${TFCI_MACOS_CROSS_COMPILE_SDK_SOURCE}/System/Library/Frameworks/SystemConfiguration.framework" "${TFCI_MACOS_CROSS_COMPILE_SDK_DEST}/System/Library/Frameworks/SystemConfiguration.framework"
+fi

ci/official/utilities/windows.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+#
+# Windows-specific utilities.
+#
+
+# Docker on Windows has difficulty using volumes other than C:\, when it comes
+# to setting up up volume mappings.
+# Thus, the drive letter is replaced with the passed prefix.
+# If no prefix is passed, by default, it's replaced with C:\, in case it's
+# something else (ex. T:), which is a volume used in internal CI.
+function replace_drive_letter_with_prefix () {
+  local path_prefix
+  if [[ -z "$2" ]]; then
+    path_prefix="C:"
+  else
+    path_prefix="$2"
+  fi
+  sed -E "s|^[a-zA-Z]:|${path_prefix}|g" <<< "$1"
+}