Skip to content

Commit c683334

Browse files
committed
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
1 parent f7b03f4 commit c683334

File tree

2 files changed

+62
-0
lines changed

2 files changed

+62
-0
lines changed

.github/workflows/docker_builder.yml

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
name: 'Torch-TensorRT Docker Build'
2+
3+
# Apply workflow only to main branch
4+
on:
5+
push:
6+
branches:
7+
- main
8+
- nightly
9+
10+
# If pushes to main are made in rapid succession,
11+
# cancel existing docker builds and use newer commits
12+
concurrency:
13+
group: ${{ github.workflow }}-${{ github.ref_name }}
14+
cancel-in-progress: true
15+
16+
jobs:
17+
build:
18+
runs-on: linux.2xlarge
19+
20+
# Define key environment variables
21+
# Container name is of the form torch_tensorrt:<branch_name>
22+
env:
23+
DOCKER_REGISTRY: ghcr.io/pytorch/tensorrt
24+
CONTAINER_NAME: torch_tensorrt:${{ github.ref_name }}
25+
26+
steps:
27+
- name: Checkout repository
28+
uses: actions/checkout@v3
29+
30+
- name: Log in to the Container registry
31+
uses: docker/login-action@v2
32+
with:
33+
registry: ${{ env.DOCKER_REGISTRY }}
34+
username: ${{ github.actor }}
35+
password: ${{ secrets.GITHUB_TOKEN }}
36+
37+
# Automatically detect TensorRT and cuDNN default versions for Torch-TRT build
38+
- name: Build Docker image
39+
env:
40+
DOCKER_TAG: ${{ env.DOCKER_REGISTRY }}/${{ env.CONTAINER_NAME }}
41+
run: |
42+
TRT_VERSION=$(python3 -c "import versions; versions.tensorrt_version()")
43+
echo "TRT VERSION = ${TRT_VERSION}"
44+
CUDNN_VERSION=$(python3 -c "import versions; versions.cudnn_version()")
45+
echo "CUDNN VERSION = ${CUDNN_VERSION}"
46+
47+
DOCKER_BUILDKIT=1 docker build --build-arg TENSORRT_VERSION=$TRT_VERSION --build-arg CUDNN_VERSION=$CUDNN_VERSION -f docker/Dockerfile --tag $DOCKER_TAG .
48+
49+
- name: Push Docker image
50+
env:
51+
DOCKER_URL: ${{ env.DOCKER_REGISTRY }}/${{ env.CONTAINER_NAME }}
52+
run: docker push $DOCKER_URL
53+
54+
# Clean up all untagged containers in registry
55+
- name: Container Registry Cleanup
56+
uses: actions/delete-package-versions@v4
57+
with:
58+
package-name: "tensorrt/torch_tensorrt"
59+
package-type: container
60+
min-versions-to-keep: 0
61+
delete-only-untagged-versions: True

docker/Dockerfile

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ FROM base as torch-tensorrt-builder-base
6161
ARG ARCH="x86_64"
6262
ARG TARGETARCH="amd64"
6363

64+
RUN apt-get update
6465
RUN apt-get install -y python3-setuptools
6566
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub
6667

0 commit comments

Comments
 (0)