From c683334fc1187a34bc8ae8deaaa545670e2bc7a3 Mon Sep 17 00:00:00 2001
From: gs-olive <113141689+gs-olive@users.noreply.github.com>
Date: Fri, 21 Jul 2023 12:51:20 -0700
Subject: [PATCH] ci: Add automatic CI job to build Docker

- Add GHA automated job to build Docker container with latest `main`
branch for Torch-TensorRT
- Add concurrency control to avoid build clashes
- Add environment variables to reduce code duplication, improve
readability, and make the code easier to modify when versions change
- Add new container registry URL to host built containers
---
 .github/workflows/docker_builder.yml | 61 ++++++++++++++++++++++++++++
 docker/Dockerfile                    |  1 +
 2 files changed, 62 insertions(+)
 create mode 100644 .github/workflows/docker_builder.yml

diff --git a/.github/workflows/docker_builder.yml b/.github/workflows/docker_builder.yml
new file mode 100644
index 0000000000..71e7cb74d4
--- /dev/null
+++ b/.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
diff --git a/docker/Dockerfile b/docker/Dockerfile
index aa3623f32b..ca01bf8e2d 100644
--- a/docker/Dockerfile
+++ b/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