diff --git a/docker/Dockerfile b/docker/Dockerfile index 94f4a1ba99..5fff586895 100644 --- 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 index b5b9f8e87e..f998150a66 100755 --- 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