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: + 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