From 2e1d7c62ed68b3a37580005032d5cbae9aad7183 Mon Sep 17 00:00:00 2001
From: rsbowman-striveworks
 <148816332+rsbowman-striveworks@users.noreply.github.com>
Date: Wed, 22 Nov 2023 17:33:36 -0600
Subject: [PATCH] Fix CUDA_VERSION env var conflict (#2807)

CUDA_VERSION may be set in the container to the dot separated CUDA
version number.  This conflicts with the docker build arg CUDA_VERSION
used to select which CUDA version to use.  Rename the latter to
USE_CUDA_VERSION to avoid the conflict.

Co-authored-by: Sean Bowman <r.bowman@striveworks.us>
---
 docker/Dockerfile     | 6 +++---
 docker/build_image.sh | 8 ++++----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 340045f14b..bcca48c056 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -62,7 +62,7 @@ RUN python -m pip install -U pip setuptools
 # This is only useful for cuda env
 RUN export USE_CUDA=1
 
-ARG CUDA_VERSION=""
+ARG USE_CUDA_VERSION=""
 
 RUN git clone --depth 1 --recursive https://github.com/pytorch/serve.git
 
@@ -71,8 +71,8 @@ WORKDIR "serve"
 RUN \
     if echo "$BASE_IMAGE" | grep -q "cuda:"; then \
         # Install CUDA version specific binary when CUDA version is specified as a build arg
-        if [ "$CUDA_VERSION" ]; then \
-            python ./ts_scripts/install_dependencies.py --cuda $CUDA_VERSION; \
+        if [ "USE_CUDA_VERSION" ]; then \
+            python ./ts_scripts/install_dependencies.py --cuda $USE_CUDA_VERSION; \
         # Install the binary with the latest CPU image on a CUDA base image
         else \
             python ./ts_scripts/install_dependencies.py; \
diff --git a/docker/build_image.sh b/docker/build_image.sh
index 6e615d7984..75c2d549fc 100755
--- a/docker/build_image.sh
+++ b/docker/build_image.sh
@@ -155,10 +155,10 @@ fi
 
 if [ "${BUILD_TYPE}" == "production" ]
 then
-  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" -t "${DOCKER_TAG}" --target production-image  .
+  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" -t "${DOCKER_TAG}" --target production-image  .
 elif [ "${BUILD_TYPE}" == "ci" ]
 then
-  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}"  -t "${DOCKER_TAG}" --target ci-image  .
+  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}"  -t "${DOCKER_TAG}" --target ci-image  .
 else
-  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg BUILD_WITH_IPEX="${BUILD_WITH_IPEX}"  -t "${DOCKER_TAG}" --target dev-image  .
-fi
\ No newline at end of file
+  DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}"  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg BUILD_WITH_IPEX="${BUILD_WITH_IPEX}"  -t "${DOCKER_TAG}" --target dev-image  .
+fi