Skip to content

feat: Implement Dynamic shapes + fallback support for export path (#2… #195

feat: Implement Dynamic shapes + fallback support for export path (#2…

feat: Implement Dynamic shapes + fallback support for export path (#2… #195

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: |
python3 -m pip install pyyaml
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