bump: ⬆️ CUDA Version (#1769)

mindee · Nov 18, 2024 · 99842ba · 99842ba
1 parent c56bf41
commit 99842ba
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 51 deletions.
diff --git a/.github/workflows/public_docker_images.yml b/.github/workflows/public_docker_images.yml
@@ -24,7 +24,6 @@ jobs:
         # Must match version at https://www.python.org/ftp/python/
         python: ["3.10.13", "3.11.8", "3.12.7"]
         framework: ["tf", "torch", "tf,viz,html,contrib", "torch,viz,html,contrib"]
-        system: ["cpu", "gpu"]
 
     # Sets the permissions granted to the `GITHUB_TOKEN` for the actions in this job.
     permissions:
@@ -42,16 +41,22 @@ jobs:
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
+      - name: Sanitize docker tag
+        run:
+          PREFIX_DOCKER_TAG="${{ matrix.framework }}-py${{ matrix.python }}-"
+          PREFIX_DOCKER_TAG=$(echo PREFIX_DOCKER_TAG|sed 's/,/-/g')
+          echo PREFIX_DOCKER_TAG=${PREFIX_DOCKER_TAG} >> $GITHUB_ENV
+
       - name: Extract metadata (tags, labels) for Docker
         id: meta
         uses: docker/metadata-action@9ec57ed1fcdbf14dcef7dfbe97b2010124a938b7
         with:
           images: ${{ env.REGISTRY }}/${{ github.repository }}
           tags: |
             # used only on schedule event
-            type=schedule,pattern={{date 'YYYY-MM'}},prefix=${{ matrix.framework }}-py${{ matrix.python }}-${{ matrix.system }}-
+            type=schedule,pattern={{date 'YYYY-MM'}},prefix=${{ env.PREFIX_DOCKER_TAG }}
             # used only if a tag following semver is published
-            type=semver,pattern={{raw}},prefix=${{ matrix.framework }}-py${{ matrix.python }}-${{ matrix.system }}-
+            type=semver,pattern={{raw}},prefix=${{ env.PREFIX_DOCKER_TAG }}
 
       - name: Build Docker image
         id: build
@@ -61,7 +66,6 @@ jobs:
           build-args: |
             FRAMEWORK=${{ matrix.framework }}
             PYTHON_VERSION=${{ matrix.python }}
-            SYSTEM=${{ matrix.system }}
             DOCTR_REPO=${{ github.repository }}
             DOCTR_VERSION=${{ github.sha }}
           push: false  # push only if `import doctr` works
@@ -72,14 +76,13 @@ jobs:
 
       - name: Push Docker image
         # Push only if the CI is not triggered by "PR on main"
-        if: github.ref == 'refs/heads/main' && github.event_name != 'pull_request'
+        if: ${{ (github.ref == 'refs/heads/main' && github.event_name != 'pull_request') || (startsWith(github.ref, 'refs/tags') && github.event_name == 'push') }}
         uses: docker/build-push-action@f2a1d5e99d037542a71f64918e516c093c6f3fc4
         with:
           context: .
           build-args: |
             FRAMEWORK=${{ matrix.framework }}
             PYTHON_VERSION=${{ matrix.python }}
-            SYSTEM=${{ matrix.system }}
             DOCTR_REPO=${{ github.repository }}
             DOCTR_VERSION=${{ github.sha }}
           push: true

diff --git a/Dockerfile b/Dockerfile
@@ -1,38 +1,10 @@
-FROM ubuntu:22.04
+FROM nvidia/cuda:12.2.0-base-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV LANG=C.UTF-8
 ENV PYTHONUNBUFFERED=1
 ENV PYTHONDONTWRITEBYTECODE=1
 
-ARG SYSTEM=gpu
-
-# Enroll NVIDIA GPG public key and install CUDA
-RUN if [ "$SYSTEM" = "gpu" ]; then \
-    apt-get update && \
-    apt-get install -y gnupg ca-certificates wget && \
-    # - Install Nvidia repo keys
-    # - See: https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#network-repo-installation-for-ubuntu
-    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb && \
-    dpkg -i cuda-keyring_1.1-1_all.deb && \
-    apt-get update && apt-get install -y --no-install-recommends \
-    cuda-command-line-tools-11-8 \
-    cuda-cudart-dev-11-8 \
-    cuda-nvcc-11-8 \
-    cuda-cupti-11-8 \
-    cuda-nvprune-11-8 \
-    cuda-libraries-11-8 \
-    cuda-nvrtc-11-8 \
-    libcufft-11-8 \
-    libcurand-11-8 \
-    libcusolver-11-8 \
-    libcusparse-11-8 \
-    libcublas-11-8 \
-    # - 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; \
-fi
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
     # - Other packages
@@ -48,8 +20,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     # - Packages for docTR
     libgl1-mesa-dev libsm6 libxext6 libxrender-dev libpangocairo-1.0-0 \
     && apt-get clean \
-    && rm -rf /var/lib/apt/lists/* \
-fi
+    && rm -rf /var/lib/apt/lists/*
 
 # Install Python
 ARG PYTHON_VERSION=3.10.13

diff --git a/README.md b/README.md
@@ -254,37 +254,37 @@ Check out our [TensorFlow.js demo](https://github.com/mindee/doctr-tfjs-demo) to
 
 ### Docker container
 
-[We offer Docker container support for easy testing and deployment](https://github.com/mindee/doctr/pkgs/container/doctr).
+We offer Docker container support for easy testing and deployment. [Here are the available docker tags.](https://github.com/mindee/doctr/pkgs/container/doctr).
 
 #### Using GPU with docTR Docker Images
 
-The docTR Docker images are GPU-ready and based on CUDA `11.8`.
-However, to use GPU support with these Docker images, please ensure that Docker is configured to use your GPU.
+The docTR Docker images are GPU-ready and based on CUDA `12.2`. Make sure your host is **at least `12.2`**, otherwise Torch or TensorFlow won't be able to initialize the GPU.
+Please ensure that Docker is configured to use your GPU.
 
 To verify and configure GPU support for Docker, please follow the instructions provided in the [NVIDIA Container Toolkit Installation Guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html).
 
 Once Docker is configured to use GPUs, you can run docTR Docker containers with GPU support:
 
 ```shell
-docker run -it --gpus all ghcr.io/mindee/doctr:tf-py3.8.18-gpu-2023-09 bash
+docker run -it --gpus all ghcr.io/mindee/doctr:torch-py3.9.18-2024-10 bash
 ```
 
 #### Available Tags
 
-The Docker images for docTR follow a specific tag nomenclature: `<framework>-py<python_version>-<system>-<doctr_version|YYYY-MM>`. Here's a breakdown of the tag structure:
+The Docker images for docTR follow a specific tag nomenclature: `<deps>-py<python_version>-<doctr_version|YYYY-MM>`. Here's a breakdown of the tag structure:
 
-- `<framework>`: `tf` (TensorFlow) or `torch` (PyTorch).
-- `<python_version>`: `3.8.18`, `3.9.18`, or `3.10.13`.
-- `<system>`: `cpu` or `gpu`
-- `<doctr_version>`: a tag >= `v0.7.1`
-- `<YYYY-MM>`: e.g. `2023-09`
+- `<deps>`: `tf`, `torch`, `tf-viz-html-contrib` or `torch-viz-html-contrib`.
+- `<python_version>`: `3.9.18`, `3.10.13` or `3.11.8`.
+- `<doctr_version>`: a tag >= `v0.11.0`
+- `<YYYY-MM>`: e.g. `2014-10`
 
 Here are examples of different image tags:
 
 | Tag                        | Description                                       |
 |----------------------------|---------------------------------------------------|
-| `tf-py3.8.18-cpu-v0.7.1`       | TensorFlow version `3.8.18` with docTR `v0.7.1`. |
-| `torch-py3.9.18-gpu-2023-09`| PyTorch version `3.9.18` with GPU support and a monthly build from `2023-09`. |
+| `tf-py3.10.13-v0.11.0`       | TensorFlow version `3.10.13` with docTR `v0.11.0`. |
+| `torch-viz-html-contrib-py3.11.8-2024-10`       | Torch with extra dependencies version `3.11.8` from latest commit on `main` in `2024-10`. |
+| `torch-py3.11.8-2024-10`| PyTorch version `3.11.8` from latest commit on `main` in `2024-10`. |
 
 #### Building Docker Images Locally
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -53,7 +53,8 @@ dependencies = [
 [project.optional-dependencies]
 tf = [
     # cf. https://github.com/mindee/doctr/pull/1461
-    "tensorflow>=2.15.0,<3.0.0",
+    "tensorflow[and-cuda]>=2.15.0,<3.0.0; sys_platform != 'darwin'",
+    "tensorflow>=2.15.0,<3.0.0; sys_platform == 'darwin'", # TensorFlow compatible for macOS
     "tf-keras>=2.15.0,<3.0.0",  # Keep keras 2 compatibility
     "tf2onnx>=1.16.0,<2.0.0",  # cf. https://github.com/onnx/tensorflow-onnx/releases/tag/v1.16.0
 ]
@@ -97,7 +98,8 @@ docs = [
 dev = [
     # Tensorflow
     # cf. https://github.com/mindee/doctr/pull/1461
-    "tensorflow>=2.15.0,<3.0.0",
+    "tensorflow[and-cuda]>=2.15.0,<3.0.0; sys_platform != 'darwin'",
+    "tensorflow>=2.15.0,<3.0.0; sys_platform == 'darwin'", # TensorFlow compatible for macOS
     "tf-keras>=2.15.0,<3.0.0",  # Keep keras 2 compatibility
     "tf2onnx>=1.16.0,<2.0.0",  # cf. https://github.com/onnx/tensorflow-onnx/releases/tag/v1.16.0
     # PyTorch