Support different CUDA versions for the slim Python Predictor image (#1263)

Author: RobertLucian

build/build-image.sh

@@ -41,14 +41,27 @@ dir=$1
 image=$2
 
 docker build "$ROOT" \
-    -f $dir/Dockerfile \
-    -t cortexlabs/${image} \
-    -t cortexlabs/${image}:${CORTEX_VERSION}
+  -f $dir/Dockerfile \
+  -t cortexlabs/${image} \
+  -t cortexlabs/${image}:${CORTEX_VERSION}
 
 if [ "$slim" == "true" ]; then
-    docker build "$ROOT" \
+  if [ "${image}" == "python-predictor-gpu" ]; then
+    cuda=("10.0" "10.1" "10.2" "11.0")
+    cudnn=("7" "7" "7" "8")
+
+    for i in ${!cudnn[@]}; do
+      docker build "$ROOT" \
         -f $dir/Dockerfile \
         --build-arg SLIM=true \
-        -t cortexlabs/${image}-slim \
-        -t cortexlabs/${image}-slim:${CORTEX_VERSION}
+        --build-arg CUDA_VERSION=${cuda[$i]} \
+        --build-arg CUDNN=${cudnn[$i]} \
+        -t cortexlabs/${image}-slim:${CORTEX_VERSION}-cuda${cuda[$i]}
+    done
+  else
+    docker build "$ROOT" \
+      -f $dir/Dockerfile \
+      --build-arg SLIM=true \
+      -t cortexlabs/${image}-slim:${CORTEX_VERSION}
+  fi
 fi

build/push-image.sh

@@ -42,5 +42,11 @@ echo "$DOCKER_PASSWORD" | docker login -u "$DOCKER_USERNAME" --password-stdin
 docker push cortexlabs/${image}:${CORTEX_VERSION}
 
 if [ "$slim" == "true" ]; then
-  docker push cortexlabs/${image}-slim:${CORTEX_VERSION}
+  if [ "$image" == "python-predictor-gpu" ]; then
+    for cuda in 10.0 10.1 10.2 11.0; do
+      docker push cortexlabs/${image}-slim:${CORTEX_VERSION}-cuda${cuda}
+    done
+  else
+    docker push cortexlabs/${image}-slim:${CORTEX_VERSION}
+  fi
 fi

docs/deployments/system-packages.md

@@ -50,7 +50,7 @@ Cortex's base Docker images are listed below. Depending on the Cortex Predictor
 
 <!-- CORTEX_VERSION_BRANCH_STABLE x6 -->
 * Python Predictor (CPU): `cortexlabs/python-predictor-cpu-slim:master`
-* Python Predictor (GPU): `cortexlabs/python-predictor-gpu-slim:master`
+* Python Predictor (GPU): `cortexlabs/python-predictor-gpu-slim:master-cuda10.1` (also available in cuda10.0, cuda10.2, and cuda11.0)
 * Python Predictor (Inferentia): `cortexlabs/python-predictor-inf-slim:master`
 * TensorFlow Predictor (CPU, GPU, Inferentia): `cortexlabs/tensorflow-predictor-slim:master`
 * ONNX Predictor (CPU): `cortexlabs/onnx-predictor-cpu-slim:master`

images/python-predictor-gpu/Dockerfile

@@ -1,4 +1,6 @@
-FROM nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04
+ARG CUDA_VERSION=10.1
+ARG CUDNN=7
+FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN-devel-ubuntu18.04
 
 RUN apt-get update -qq && apt-get install -y -q \
         build-essential \