add --batch option to docker/build

This builds an image compatible with AWS Batch. Uses CUDA version 11.3.1 instead of 12.1.1 and Ubuntu 20.04 instead of 22.04.
azavea · Aug 9, 2023 · 369fd58 · 369fd58
1 parent ee3fbef
commit 369fd58
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 8 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -1,4 +1,4 @@
-ARG UBUNTU_VERSION=22.04
+ARG UBUNTU_VERSION
 ARG CUDA_VERSION
 FROM nvidia/cuda:${CUDA_VERSION}-cudnn8-runtime-ubuntu${UBUNTU_VERSION}
 

diff --git a/docker/build b/docker/build
@@ -13,7 +13,8 @@ function usage() {
 Build Docker images.
 
 Options:
---arm64 will build image for arm64 architecture
+--arm64   Build image for arm64 architecture.
+--batch   Build an image compatible with AWS Batch. Uses CUDA version 11.3.1 instead of 12.1.1 and Ubuntu 20.04 instead of 22.04.
 "
 }
 
@@ -27,13 +28,33 @@ then
 
     PLATFORM="amd64"
     IMAGE_EXT=""
-    if [ "${1:-}" = "--arm64" ]
-    then
-        PLATFORM="arm64"
-        IMAGE_EXT="-arm64"
-    fi
+    CUDA_VERSION="12.1.1"
+    UBUNTU_VERSION="22.04"
+
+    while [[ $# -gt 0 ]]
+    do
+        case "$1" in
+            --arm64)
+                PLATFORM="arm64"
+                IMAGE_EXT="-arm64"
+                shift
+                ;;
+            --batch)
+                CUDA_VERSION="11.3.1"
+                UBUNTU_VERSION="20.04"
+                shift
+                ;;
+            *)
+                echo "Unknown option: $1"
+                usage
+                exit 1
+                ;;
+        esac
+    done
 
     DOCKER_BUILDKIT=1 docker build \
-        --platform linux/${PLATFORM} --build-arg CUDA_VERSION="11.7.1" \
+        --platform linux/${PLATFORM} \
+        --build-arg CUDA_VERSION="${CUDA_VERSION}" \
+        --build-arg UBUNTU_VERSION="${UBUNTU_VERSION}" \
         -t raster-vision-pytorch${IMAGE_EXT} -f Dockerfile .
 fi