[CUDA] Update Docker to use Ubuntu 24.04, cuda 12.6, cudnn 9.4 and py…

…thon 3.12 (#22339)

### Description
Serve as example to build and run onnxruntime-gpu with latest software
stack.

To build docker image:
```
git clone https://github.com/microsoft/onnxruntime
cd onnxruntime/dockerfiles
docker build -t onnxruntime-cuda -f Dockerfile.cuda ..
```

To launch the docker image built from previous step (and mount the code
directory to run a unit test below):
```
cd ..
docker run --rm -it --gpus all -v $PWD:/code onnxruntime-cuda /bin/bash
```

Then run the following in docker image to verify that the cuda provider
is good:
```
python /code/onnxruntime/test/python/onnxruntime_test_python_cudagraph.py
```

### Motivation and Context
#22335

dockerfiles/Dockerfile.cuda

@@ -2,40 +2,87 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 # --------------------------------------------------------------
-# Build onnxruntime-gpu python package with CUDA 12.5 & CUDNN 9 for Nvidia GPU compute capability 6.1, 7.x, 8.x, 9.x.
+# Build onnxruntime-gpu python package with CUDA 12.6 & CUDNN 9.4 for python 3.12 in Ubuntu 24.04 for Nvidia GPU.
 # If memory is less than 64GB, you may change "--parallel" to "--parallel 4" to avoid out-of-memory error.
 
-FROM nvcr.io/nvidia/cuda:12.5.1-cudnn-devel-ubuntu22.04
+FROM nvcr.io/nvidia/cuda:12.6.1-devel-ubuntu24.04
 
-ARG CMAKE_CUDA_ARCHITECTURES=61;70;75;80;86;90
+# Target CUDA device with compute capability >= 6.1
+ARG CMAKE_CUDA_ARCHITECTURES="61;70;75;80;86;90"
 
-ENV  DEBIAN_FRONTEND=noninteractive
+ENV DEBIAN_FRONTEND=noninteractive
 MAINTAINER Changming Sun "chasun@microsoft.com"
+
+# Add source code to /code
 ADD . /code
 
-ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:${PATH}
-RUN apt-get update \
-&& apt-get install -y --no-install-recommends ca-certificates g++ gcc make git python3-dev python3-pip
+ENV PATH=/usr/local/cuda/bin:${PATH}
+
+# Install required packages
+RUN apt-get update && apt-get install -y --no-install-recommends \
+ ca-certificates \
+ g++ \
+ gcc \
+ make \
+ git \
+ cmake \
+ wget \
+ ninja-build \
+ python3-pip \
+ python3.12-dev \
+ python3.12-venv \
+ && rm -rf /var/lib/apt/lists/*
 
+# Install CUDNN 9.4.0.58 for building ONNX Runtime with CUDA.
+RUN wget https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-9.4.0.58_cuda12-archive.tar.xz \
+ && mkdir -p /code/build/cudnn \
+ && tar -Jxvf cudnn-linux-x86_64-9.4.0.58_cuda12-archive.tar.xz -C /code/build/cudnn --strip=1
+
+# Create a virtual environment and install dependencies, then build ONNX Runtime with CUDA support.
 RUN cd /code \
-&& python3 -m pip install --upgrade pip cmake psutil setuptools wheel packaging ninja \
-&& python3 -m pip install -r tools/ci_build/github/linux/docker/inference/x86_64/python/cpu/scripts/requirements.txt \
-&& /bin/bash ./build.sh --allow_running_as_root --skip_submodule_sync \
---use_cuda --cuda_home /usr/local/cuda --cudnn_home /usr/lib/x86_64-linux-gnu/ \
---build_shared_lib --skip_tests \
---config Release --build_wheel --update --build --parallel \
---cmake_generator Ninja \
---enable_cuda_nhwc_ops \
---cmake_extra_defines ONNXRUNTIME_VERSION=$(cat ./VERSION_NUMBER) \
-onnxruntime_BUILD_UNIT_TESTS=OFF \
-`CMAKE_CUDA_ARCHITECTURES=${CMAKE_CUDA_ARCHITECTURES}`
-
-FROM nvcr.io/nvidia/cuda:12.5.1-cudnn-runtime-ubuntu22.04
-ENV DEBIAN_FRONTEND=noninteractive
-COPY --from=0 /code/build/Linux/Release/dist /root
+ && python3 -m venv /code/env \
+ && . /code/env/bin/activate \
+ && pip install --upgrade psutil setuptools wheel packaging \
+ && pip install -r tools/ci_build/github/linux/docker/inference/x86_64/python/cpu/scripts/requirements.txt \
+ && python /code/tools/ci_build/build.py --build_dir /code/build/Linux \
+ --allow_running_as_root --skip_submodule_sync \
+ --use_cuda --cuda_home /usr/local/cuda \
+ --cudnn_home /code/build/cudnn \
+ --build_shared_lib --skip_tests \
+ --config Release --build_wheel --update --build --parallel \
+ --cmake_generator Ninja \
+ --enable_cuda_nhwc_ops \
+ --cmake_extra_defines ONNXRUNTIME_VERSION=$(cat ./VERSION_NUMBER) "CMAKE_CUDA_ARCHITECTURES=${CMAKE_CUDA_ARCHITECTURES}" onnxruntime_BUILD_UNIT_TESTS=OFF
+
+# Start second stage to copy the build artifacts
+FROM nvcr.io/nvidia/cuda:12.6.1-runtime-ubuntu24.04
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Copy built wheel and license
+COPY --from=0 /code/build/Linux/Release/dist /ort
 COPY --from=0 /code/dockerfiles/LICENSE-IMAGE.txt /code/LICENSE-IMAGE.txt
-RUN apt-get update \
-&& apt-get install -y --no-install-recommends libstdc++6 ca-certificates python3-setuptools python3-wheel python3-pip unattended-upgrades \
-&& unattended-upgrade \
-&& python3 -m pip install /root/*.whl \
-&& rm -rf /root/*.whl
+
+# Set LD_LIBRARY_PATH so that runtime can load CUDA and CUDNN DLLs.
+# CUDNN will be installed by nvidia-cudnn-cu12 python package later.
+# Its location is in the site-packages directory, which can be retrieved like the following:
+# python -c "import sysconfig; print(sysconfig.get_path('purelib'))"
+ENV LD_LIBRARY_PATH="/ort/env/lib/python3.12/site-packages/nvidia/cudnn/lib:/usr/local/cuda/lib64"
+
+# Install runtime dependencies, and run a simple test to verify the installation.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+ libstdc++6 \
+ ca-certificates \
+ python3-pip \
+ python3.12-venv \
+ unattended-upgrades \
+ && unattended-upgrade \
+ && python3 -m venv /ort/env \
+ && . /ort/env/bin/activate \
+ && pip install /ort/*.whl \
+ && pip install nvidia-cudnn-cu12==9.4.0.58 \
+ && python -c 'import onnxruntime; print(onnxruntime.get_available_providers())' \
+ && rm -rf /ort/*.whl \
+ && rm -rf /var/lib/apt/lists/*
+
+# Ensure the virtual environment is always activated when running commands in the container.
+RUN echo ". /ort/env/bin/activate" >> ~/.bashrc