#
# setup.python.sh: Install a specific Python version and packages for it.
# Usage: setup.python.sh <pyversion> <requirements.txt>

# Sets up custom apt sources for our TF images.

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

# Set up shared custom sources
apt-get update
apt-get install -y gnupg ca-certificates

          # the end of the line.
          # See https://cloud.google.com/architecture/using-container-images#exploring_image_manifests_digests_and_tags
          echo -n "Trying to map name $1 to tag $2... "
          digest=$(curl -s --head "https://gcr.io/v2/tensorflow-sigs/build/manifests/$2" | egrep -o "sha256:[[:alnum:]]*")
          # Find the line matching the regex "sigbuild-r2.9" (with quotes) and

TFCI_MACOS_BAZEL_TEST_DIR_PATH="/Volumes/BuildData/bazel_output"
TFCI_OUTPUT_DIR=build_output
TFCI_WHL_BAZEL_TEST_ENABLE=1
TFCI_WHL_SIZE_LIMIT=240M
TFCI_WHL_SIZE_LIMIT_ENABLE=1

# 3.11 is the system python on our images
case $TFCI_PYTHON_VERSION in
3.11)
  TFCI_MACOS_PYENV_INSTALL_ENABLE=0
  ;;
*)
  TFCI_MACOS_PYENV_INSTALL_ENABLE=1
  ;;

  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

    --build-arg REQUIREMENTS_FILE=jax.requirements.txt \
    --target=$target \
    --cache-from "$IMAGE" \
    -t "$IMAGE"  . && break
  done
  final=$?
  if [ $final -ne 0 ]; then
    exit $final
  fi
  set -e

  if [[ -n "$KOKORO_BUILD_ID" ]]; then
    gcloud auth configure-docker
    docker push "$IMAGE"
  fi

TFCI_MACOS_TWINE_INSTALL_ENABLE=1
TFCI_OUTPUT_DIR=build_output
TFCI_WHL_BAZEL_TEST_ENABLE=1
TFCI_WHL_SIZE_LIMIT=255M
TFCI_WHL_SIZE_LIMIT_ENABLE=1

# 3.11 is the system python on our images
case $TFCI_PYTHON_VERSION in
3.11)
  TFCI_MACOS_PYENV_INSTALL_ENABLE=0
  TFCI_MACOS_UPGRADE_PYENV_ENABLE=0
  ;;
*)
  TFCI_MACOS_PYENV_INSTALL_ENABLE=1
  TFCI_MACOS_UPGRADE_PYENV_ENABLE=1
  ;;

running the `tf.train.Server`.

## Untrusted inputs during training and prediction

TensorFlow supports a wide range of input data formats. For example it can
process images, audio, videos, and text. There are several modules specialized
in taking those formats, modifying them, and/or converting them to intermediate
formats that can be processed by TensorFlow.

# different compilation methods. E.g. for a PR to test a new CUDA version, set
# the CACHEBUSTER to the PR number.
build --action_env=CACHEBUSTER=20220325

# Use Python 3.X as installed in container image
build --action_env PYTHON_BIN_PATH="/usr/local/bin/python3"
build --python_path="/usr/local/bin/python3"

# Build TensorFlow v2
build --define=tf_api_version=2 --action_env=TF2_BEHAVIOR=1

/usr/bin/$VERSION get-pip.py
/usr/bin/$VERSION -m pip install --no-cache-dir --upgrade pip

# Disable the cache dir to save image space, and install packages

# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
libcudnn8=8.9.6.50-1+cuda12.2
libcudnn8-dev=8.9.6.50-1+cuda12.2

# This can be removed once NVIDIA publishes a cuda-12.3.2 Docker image.
# For now it ensures that we install at least version 12.3.107 of PTXAS,
# since 12.3.103 has a bug.