# build configuration. You must set the following variables to control this
# script:
#
#   TF_BISECT_GOOD: Last known good commit (e.g. commit from the last passing job)
#   TF_BISECT_BAD: First bad commit (e.g. commit from the first failing job)
#   TF_BISECT_SCRIPT: The build script path relative to the TF root dir, e.g.
#     ci/official/wheel.sh

# explicitly, *after* checking if update_version.py needs to be run for a
# nightly job. update_version.py affects TF_VER_SUFFIX, TF_VER_PYTHON, and
# TF_VER_FULL.

# Note: in awk, the command '/search/ {commands}' applies the commands to any line that
# matches the /search/ regular expression. "print $N" prints the Nth "field",
# where fields are strings separated by whitespace.

# limitations under the License.
# ==============================================================================
#
# Usage: get_test_list.sh OUTPUT BAZEL_TEST_COMMAND...
# Writes the list of tests that would be run from BAZEL_TEST_COMMAND to OUTPUT.
# Hides all extra output and always exits with success for now.

OUTPUT=$1
shift

if [[ "$TFCI_DOCKER_PULL_ENABLE" == 1 ]]; then
  # Simple retry logic for docker-pull errors. Sleeps for 15s 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 15
  docker pull "$TFCI_DOCKER_IMAGE"
fi 

if [[ "$TFCI_DOCKER_REBUILD_ENABLE" == 1 ]]; then

# script to run pip-compile for keras, tensorboard, estimator deps.
# if there is a change in requirements.in then all lock files will be updated
# accordingly.

# -e: abort script if one command fails
# -u: error if undefined variable used
# -o pipefail: entire command fails if pipe fails. watch out for yes | ...
# -o history: record shell history
set -euo pipefail -o history

# Check for required arguments
if [ -z "$1" ]; then

# Usage: setup.packages.sh <package_list.txt>
set -e

# Prevent apt install tzinfo from asking our location (assumes UTC)
export DEBIAN_FRONTEND=noninteractive

apt-get update
# Remove commented lines and blank lines
apt-get install -y --no-install-recommends $(sed -e '/^\s*#.*$/d' -e '/^\s*$/d' "$1" | sort -u)

# 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

#
# 1. RUN, TEST, OR BUILD BAZEL TARGET(S) WITHIN A TFCI ENVIRONMENT
#    To use:
#       export TFCI=ci/official/envs/env_goes_here
#       export TF_ANY_TARGETS="quoted list of targets, like on the command line"
#       export TF_ANY_MODE="test" or "build" or "run" (default: "test")
#       ./any.sh
#
# 2. RUN ANY OTHER SCRIPT AND ENV WITH NO SIDE EFFECTS (NO UPLOADS)
#    To use:

export DOCKER_BUILDKIT=1
for target in jax tf; do
  IMAGE="gcr.io/tensorflow-sigs/build-arm64:$target-$TAG"
  docker pull "$IMAGE" || true
  # Due to some flakiness of resources pulled in the build, allow the docker
  # command to reattempt build a few times in the case of failure (b/302558736)
  set +e
  for i in $(seq 1 5)
  do
    docker build \
    --build-arg REQUIREMENTS_FILE=jax.requirements.txt \
    --target=$target \