Add modifications to main Dockerfile and build script

nod-ai · Nov 25, 2024 · cd2983d · cd2983d
1 parent d9e142d
commit cd2983d
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 2 deletions.
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -67,6 +67,7 @@ RUN python -m pip install -U pip setuptools
 RUN export USE_CUDA=1
 
 ARG USE_CUDA_VERSION=""
+ARG USE_ROCM_VERSION=""
 
 COPY ./ serve
 
@@ -90,6 +91,16 @@ RUN \
         else \
             python ./ts_scripts/install_dependencies.py;\
         fi; \
+    elif echo "${BASE_IMAGE}" | grep -q "rocm/"; then \
+        # Install ROCm version specific binary when ROCm version is specified as a build arg
+        if [ "$USE_ROCM_VERSION" ]; then \
+            python ./ts_scripts/install_dependencies.py --rocm $USE_ROCM_VERSION\
+            && cd /opt/rocm/share/amd_smi \
+            && python -m pip install . ;\
+        # Install the binary with the latest CPU image on a ROCm base image
+        else \
+            python ./ts_scripts/install_dependencies.py;\
+        fi; \
     # Install the CPU binary
     else \
         python ./ts_scripts/install_dependencies.py; \

diff --git a/docker/build_image.sh b/docker/build_image.sh
@@ -11,6 +11,7 @@ BASE_IMAGE="ubuntu:20.04"
 UPDATE_BASE_IMAGE=false
 USE_CUSTOM_TAG=false
 CUDA_VERSION=""
+ROCM_VERSION=""
 USE_LOCAL_SERVE_FOLDER=false
 BUILD_WITH_IPEX=false
 BUILD_CPP=false
@@ -33,6 +34,7 @@ do
           echo "-bi, --baseimage specify base docker image. Example: nvidia/cuda:11.7.0-cudnn8-runtime-ubuntu20.04 "
           echo "-bt, --buildtype specify for type of created image. Possible values: production, dev, ci."
           echo "-cv, --cudaversion specify to cuda version to use"
+          echo "-rv, --rocmversion spesify to rocm version to use"
           echo "-t, --tag specify tag name for docker image"
           echo "-lf, --use-local-serve-folder specify this option for the benchmark image if the current 'serve' folder should be used during automated benchmarks"
           echo "-ipex, --build-with-ipex specify to build with intel_extension_for_pytorch"
@@ -167,6 +169,24 @@ do
           shift
           shift
           ;;
+        -rv|--rocmversion)
+          ROCM_VERSION="$2"
+          if [ "${ROCM_VERSION}" == "rocm60" ];
+          then
+            BASE_IMAGE="rocm/dev-ubuntu-22.04:6.0.2"
+          elif [ "${ROCM_VERSION}" == "rocm61" ];
+          then
+            BASE_IMAGE="rocm/dev-ubuntu-22.04:6.1.2"
+          elif [ "${ROCM_VERSION}" == "rocm62" ];
+          then
+            BASE_IMAGE="rocm/dev-ubuntu-22.04:6.2.4"
+          else
+            echo "ROCM version not supported"
+            exit 1
+          fi
+          shift
+          shift
+          ;;
     esac
 done
 
@@ -218,6 +238,23 @@ then
       exit 1
     fi
   fi
+
+  if [[ "${MACHINE}" == "gpu" || "${ROCM_VERSION}" != "" ]];
+  then
+    if [ "${ROCM_VERSION}" == "rocm60" ];
+    then
+      BASE_IMAGE="rocm/dev-ubuntu-22.04:6.0.2-complete"
+    elif [ "${ROCM_VERSION}" == "rocm61" ];
+    then
+      BASE_IMAGE="rocm/dev-ubuntu-22.04:6.1.2-complete"
+    elif [ "${ROCM_VERSION}" == "rocm62" ];
+    then
+      BASE_IMAGE="rocm/dev-ubuntu-22.04:6.2.4-complete"
+    else
+      echo "ROCm version $ROCM_VERSION is not supported for CPP"
+      exit 1
+    fi
+  fi
 fi
 
 if [ "${BUILD_TYPE}" == "production" ]; then
@@ -238,10 +275,10 @@ then
 else
   if [ "${BUILD_CPP}" == "true" ]
   then
-    DOCKER_BUILDKIT=1 docker build --file Dockerfile.cpp --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}" --build-arg PYTHON_VERSION="${PYTHON_VERSION}"\
+    DOCKER_BUILDKIT=1 docker build --file Dockerfile.cpp --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}" --build-arg USE_ROCM_VERSION="${ROCM_VERSION}" --build-arg PYTHON_VERSION="${PYTHON_VERSION}"\
     --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg REPO_URL="${REPO_URL}" -t "${DOCKER_TAG}" --target cpp-dev-image .
   else
-    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}"\
+    DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE="${BASE_IMAGE}" --build-arg USE_CUDA_VERSION="${CUDA_VERSION}" --build-arg USE_ROCM_VERSION="${ROCM_VERSION}" --build-arg PYTHON_VERSION="${PYTHON_VERSION}"\
     --build-arg BUILD_NIGHTLY="${BUILD_NIGHTLY}" --build-arg BRANCH_NAME="${BRANCH_NAME}" --build-arg REPO_URL="${REPO_URL}" --build-arg BUILD_FROM_SRC="${BUILD_FROM_SRC}" --build-arg LOCAL_CHANGES="${LOCAL_CHANGES}"\
     --build-arg BUILD_WITH_IPEX="${BUILD_WITH_IPEX}"  -t "${DOCKER_TAG}" --target dev-image  ../
   fi