ci: Add automatic GHA job to build + push Docker Container on main

.github/workflows/docker_builder.yml

@@ -0,0 +1,61 @@
+name: 'Torch-TensorRT Docker Build'
+
+# Apply workflow only to main branch
+on:
+  push:
+    branches:
+      - main
+      - nightly
+
+# If pushes to main are made in rapid succession,
+# cancel existing docker builds and use newer commits
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref_name }}
+  cancel-in-progress: true
+
+jobs:
+  build:
+    runs-on: linux.2xlarge
+
+    # Define key environment variables
+    # Container name is of the form torch_tensorrt:<branch_name>
+    env:
+      DOCKER_REGISTRY: ghcr.io/pytorch/tensorrt
+      CONTAINER_NAME: torch_tensorrt:${{ github.ref_name }}
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v3
+
+    - name: Log in to the Container registry
+      uses: docker/login-action@v2
+      with:
+        registry: ${{ env.DOCKER_REGISTRY }}
+        username: ${{ github.actor }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+
+    # Automatically detect TensorRT and cuDNN default versions for Torch-TRT build
+    - name: Build Docker image
+      env:
+        DOCKER_TAG: ${{ env.DOCKER_REGISTRY }}/${{ env.CONTAINER_NAME }}
+      run: |
+        TRT_VERSION=$(python3 -c "import versions; versions.tensorrt_version()")
+        echo "TRT VERSION = ${TRT_VERSION}"
+        CUDNN_VERSION=$(python3 -c "import versions; versions.cudnn_version()")
+        echo "CUDNN VERSION = ${CUDNN_VERSION}"
+
+        DOCKER_BUILDKIT=1 docker build --build-arg TENSORRT_VERSION=$TRT_VERSION --build-arg CUDNN_VERSION=$CUDNN_VERSION -f docker/Dockerfile --tag $DOCKER_TAG .
+
+    - name: Push Docker image
+      env:
+        DOCKER_URL: ${{ env.DOCKER_REGISTRY }}/${{ env.CONTAINER_NAME }}
+      run: docker push $DOCKER_URL
+
+    # Clean up all untagged containers in registry
+    - name: Container Registry Cleanup
+      uses: actions/delete-package-versions@v4
+      with:
+        package-name: "tensorrt/torch_tensorrt"
+        package-type: container
+        min-versions-to-keep: 0
+        delete-only-untagged-versions: True

docker/Dockerfile

@@ -61,6 +61,7 @@ FROM base as torch-tensorrt-builder-base
 ARG ARCH="x86_64"
 ARG TARGETARCH="amd64"
 
+RUN apt-get update
 RUN apt-get install -y python3-setuptools
 RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub