Fix for Cuda 10.1 /Cuda 10.2 related torch package installation issue in Docker.

docker/Dockerfile

@@ -19,7 +19,7 @@
 ARG BASE_IMAGE=ubuntu:18.04
 
 FROM ${BASE_IMAGE} AS compile-image
-
+ARG BASE_IMAGE=ubuntu:18.04
 ENV PYTHONUNBUFFERED TRUE
 
 RUN --mount=type=cache,id=apt-dev,target=/var/cache/apt \
@@ -41,15 +41,29 @@ RUN python3 -m venv /home/venv
 
 ENV PATH="/home/venv/bin:$PATH"
 
-RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1
-RUN update-alternatives --install /usr/local/bin/pip pip /usr/local/bin/pip3 1
+RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1 \
+    && update-alternatives --install /usr/local/bin/pip pip /usr/local/bin/pip3 1
 
-# This is only useful for cuda env 
+# This is only useful for cuda env
 RUN export USE_CUDA=1
 
-RUN pip install -U pip setuptools
-
-RUN pip install --no-cache-dir torch torchvision torchtext torchserve torch-model-archiver
+ARG CUDA_VERSION=""
+
+RUN TORCH_VER=$(curl --silent --location https://pypi.org/pypi/torch/json | python -c "import sys, json, pkg_resources; releases = json.load(sys.stdin)['releases']; print(sorted(releases, key=pkg_resources.parse_version)[-1])") && \
+    TORCH_VISION_VER=$(curl --silent --location https://pypi.org/pypi/torchvision/json | python -c "import sys, json, pkg_resources; releases = json.load(sys.stdin)['releases']; print(sorted(releases, key=pkg_resources.parse_version)[-1])") && \
+    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 \
+            pip install --no-cache-dir torch==$TORCH_VER+$CUDA_VERSION torchvision==$TORCH_VISION_VER+$CUDA_VERSION -f https://download.pytorch.org/whl/torch_stable.html; \
+        # Install the binary with the latest CUDA version support
+        else \
+            pip install --no-cache-dir torch torchvision; \
+        fi \
+    # Install the CPU binary
+    else \
+        pip install --no-cache-dir torch==$TORCH_VER+cpu torchvision==$TORCH_VISION_VER+cpu -f https://download.pytorch.org/whl/torch_stable.html; \
+    fi
+RUN pip install --no-cache-dir torchtext torchserve torch-model-archiver
 
 # Final image for production
 FROM ${BASE_IMAGE} AS runtime-image

docker/README.md

@@ -32,10 +32,16 @@ For creating CPU based image :
 DOCKER_BUILDKIT=1 docker build --file Dockerfile -t torchserve:latest .
 ```
 
-For creating GPU based image :
+For creating GPU based image with the latest CUDA version PyTorch supports (ex. CUDA 10.2 as of Oct 2020):
 
 ```bash
-DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE=nvidia/cuda:10.1-cudnn7-runtime-ubuntu18.04 -t torchserve:latest .
+DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE=nvidia/cuda:10.2-cudnn7-runtime-ubuntu18.04 -t torchserve:latest .
+```
+
+For creating GPU based image with older CUDA versions (ex. CUDA 10.1), make sure that the `--build-arg CUDA_VERSION=<version>` is specified. The version is in the format "cuda92", "cuda101":
+
+```bash
+DOCKER_BUILDKIT=1 docker build --file Dockerfile --build-arg BASE_IMAGE=nvidia/cuda:10.1-cudnn7-runtime-ubuntu18.04 --build-arg CUDA_VERSION=cu101 -t torchserve:latest .
 ```
 
 ## Start a container with a TorchServe image