# All required CUDA packages
cuda-compat-13-0
cuda-command-line-tools-13-0
cuda-cudart-dev-13-0
cuda-nvcc-13-0
cuda-cupti-13-0
cuda-nvprune-13-0
cuda-libraries-13-0
cuda-libraries-dev-13-0
cuda-nvml-dev-13-0
libcufft-13-0
libcurand-13-0
libcusolver-dev-13-0
libcusparse-dev-13-0
libcublas-13-0
libcublas-dev-13-0
libnccl-dev=2.27.7-1+cuda13.0
libnccl2=2.27.7-1+cuda13.0

	var (
		counter int
		mv      *MetaValue
		body    = `{ "bio": "bada bing bada boom", "id": 1, "name": "Charles" }
{ "bio": "bada bing bada boom", "id": 2, "name": "Charles" }
{ "bio": "bada bing bada boom", "id": 3, "name": "Charles" }
{ "bio": "bada bing bada boom", "id": 4, "name": "Charles" }
{ "bio": "bada bing bada boom", "id": 5, "name": "Charles" }
`
	)

	decoder := NewDecoder(mkReader(body), 0)

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

# 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.

    write_repo_env_to_bazelrc('cuda', env_var, local_path)


def set_other_cuda_vars(environ_cp):
  """Set other CUDA related variables."""
  # If CUDA is enabled, always use GPU during build and test.
  if environ_cp.get('TF_CUDA_CLANG') == '1':
    write_to_bazelrc('build --config=cuda_clang')
  else:
    write_to_bazelrc('build --config=cuda')

#     release_cpu_linux:               Toolchain and CUDA options for Linux CPU builds.
#     release_gpu_linux:               Toolchain and CUDA options for Linux GPU builds.
#     release_cpu_macos:               Toolchain and CUDA options for MacOS CPU builds.
#     release_cpu_windows:             Toolchain and CUDA options for Windows CPU builds.
# LINT.IfChange

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

    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 \

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

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/ /