# The RBE machine itself has older kernel mode driver, and it requires
# nvidia driver to be installed.
nvidia-driver-580-open
# TODO(b/445248346): The Docker image shouldn't have cuda-compat installed.
# However, hermetic CUDA forward-compatibility mode is still missing some
# libraries.

nvidia-cublas-cu12>=12.5.3.2,<13.0
nvidia-cuda-cupti-cu12>=12.5.82,<13.0
nvidia-cuda-nvcc-cu12>=12.5.82,<13.0
nvidia-cuda-nvrtc-cu12>=12.5.82,<13.0
nvidia-cuda-runtime-cu12>=12.5.82,<13.0
# The upper bound is set for the CUDNN API compatibility.
# See
# https://docs.nvidia.com/deeplearning/cudnn/backend/latest/developer/forward-compatibility.html#cudnn-api-compatibility
nvidia-cudnn-cu12>=9.3.0.75,<10.0
nvidia-cufft-cu12>=11.2.3.61,<12.0

   
   **1. Installing **TensorFlow-GPU** (TF) prebuilt binaries**
   
   
   Make sure you are using compatible TF and CUDA versions.
   Please refer following TF version and CUDA version compatibility table.
   
   | TF  | CUDA |
   
   | :-------------: | :-------------: |
   
   | 2.5.0  | 11.2 |
   
   | 2.4.0  | 11.0 |
   
   | 2.1.0 - 2.3.0  | 10.1 |
   

export DEBIAN_FRONTEND=noninteractive

# Fetch the NVIDIA key.
apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub;

# Set up sources for NVIDIA CUDNN.
cat >/etc/apt/sources.list.d/nvidia.list <<SOURCES
# NVIDIA
deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /

RUN ln -sf /usr/bin/python3.12 /usr/bin/python
RUN ln -sf /usr/lib/python3.12 /usr/lib/tf_python

# Link the compat driver to the location if available.
RUN if [ -e "/usr/local/cuda/compat/libcuda.so.1" ]; then ln -s /usr/local/cuda/compat/libcuda.so.1 /usr/lib/x86_64-linux-gnu/libcuda.so.1; fi

# Install various tools.
# - bats: bash unit testing framework
# - bazelisk: always use the correct bazel version

WIP ML Build Docker container for ML repositories (Tensorflow, JAX and XLA).

This container branches off from
/tensorflow/tools/tf_sig_build_dockerfiles/. However, since
hermetic CUDA and hermetic Python is now available for Tensorflow, a lot of the
requirements installed on the original container can be removed to reduce the
footprint of the container and make it more reusable across different ML

    "$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

      description: If compiling from source
  - type: input
    id: Compiler
    attributes:
      label: GCC/compiler version
      description: If compiling from source
  - type: input
    id: Cuda
    attributes:
      label: CUDA/cuDNN version
  - type: input
    id: Gpu
    attributes:
      label: GPU model and memory
      description: If compiling from source
  - type: textarea
    id: what-happened

-   type: input id: Compiler attributes: label: GCC/Compiler version
    description: if compiling from source placeholder: validations: required:
    false

-   type: input id: Cuda attributes: label: CUDA/cuDNN version description:
    placeholder: validations: required: false

-   type: input id: Gpu attributes: label: GPU model and memory description: if

TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64
# Note: this is not set to "--cpu", because that changes the package name
# to tensorflow_cpu. These ARM builds are supposed to have the name "tensorflow"
# despite lacking Nvidia CUDA support.
TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow"
TFCI_DOCKER_ENABLE=1
TFCI_DOCKER_IMAGE=us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest