Add support for cpu/gpu + check if doctr can be imported

odulcy-mindee · Sep 27, 2023 · 947144a · 947144a
1 parent 437167b
commit 947144a
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 7 deletions.
diff --git a/.github/workflows/public_docker_images.yml b/.github/workflows/public_docker_images.yml
@@ -26,6 +26,7 @@ jobs:
         # Must match version at https://www.python.org/ftp/python/
         python: ["3.8.18", "3.9.18", "3.10.13"]
         framework: ["tf", "torch"]
+        system: ["cpu", "gpu"]
 
     # Sets the permissions granted to the `GITHUB_TOKEN` for the actions in this job.
     permissions:
@@ -50,19 +51,36 @@ jobs:
           images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
           tags: |
             # used on schedule event
-            type=schedule,pattern={{date 'YYYY-MM-DD'}},prefix=${{ matrix.framework }}-py${{ matrix.python }}-
+            type=schedule,pattern={{date 'YYYY-MM-DD'}},prefix=${{ matrix.framework }}-py${{ matrix.python }}-${{ matrix.system }}-
             # set latest tag only if `enable` is True
-            type=raw,value=latest,enable=${{ matrix.framework == 'tf' && matrix.python == '3.8.18' && github.ref == format('refs/heads/{0}', 'main') }}
+            type=raw,value=latest,enable=${{ matrix.framework == 'tf' && matrix.python == '3.8.18' && matrix.system == 'gpu' && github.ref == format('refs/heads/{0}', 'main') }}
             # used on push tag event
-            type=semver,pattern={{raw}},prefix=${{ matrix.framework }}-py${{ matrix.python }}-
+            type=semver,pattern={{raw}},prefix=${{ matrix.framework }}-py${{ matrix.python }}-${{ matrix.system }}-
 
-      - name: Build and push Docker image (latest)
+      - name: Build Docker image
+        id: build
         uses: docker/build-push-action@f2a1d5e99d037542a71f64918e516c093c6f3fc4
         with:
           context: .
           build-args: |
             FRAMEWORK=${{ matrix.framework }}
             PYTHON_VERSION=${{ matrix.python }}
+            SYSTEM=${{ matrix.system }}
+            DOCTR_VERSION=${{ github.sha }}
+          push: false
+          tags: ${{ steps.meta.outputs.tags }}
+
+      - name: Check if `import doctr` works
+        run: docker run ${{ steps.meta.outputs.imageid }} python3 -c 'import doctr'
+
+      - name: Push Docker image
+        uses: docker/build-push-action@f2a1d5e99d037542a71f64918e516c093c6f3fc4
+        with:
+          context: .
+          build-args: |
+            FRAMEWORK=${{ matrix.framework }}
+            PYTHON_VERSION=${{ matrix.python }}
+            SYSTEM=${{ matrix.system }}
             DOCTR_VERSION=${{ github.sha }}
           push: false  # TODO: CHANGE ME
           tags: ${{ steps.meta.outputs.tags }}
diff --git a/Dockerfile b/Dockerfile
@@ -1,10 +1,12 @@
-FROM nvidia/cuda:11.8.0-base-ubuntu22.04 as base
+ARG SYSTEM=cpu
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV LANG=C.UTF-8
 ENV PYTHONUNBUFFERED=1
 ENV PYTHONDONTWRITEBYTECODE=1
 
+FROM nvidia/cuda:11.8.0-base-ubuntu22.04 as gpu
+
 # Enroll NVIDIA GPG public key
 RUN apt-get update && \
     apt-get install -y gnupg ca-certificates wget && \
@@ -30,8 +32,13 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     # - CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
     libcudnn8=8.6.0.163-1+cuda11.8 \
     libnvinfer-plugin8=8.6.1.6-1+cuda11.8 \
-    libnvinfer8=8.6.1.6-1+cuda11.8 \
-    # - Other packages
+    libnvinfer8=8.6.1.6-1+cuda11.8
+
+FROM ubuntu:22.04 as cpu
+
+FROM $SYSTEM
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
     build-essential \
     pkg-config \
     curl \