Improve the CI

pytorch-labs · Oct 17, 2024 · f1ad1ad · f1ad1ad
1 parent b8bb8ae
commit f1ad1ad
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
@@ -13,7 +13,6 @@ on:
         description: "PyTorch nightly version"
         required: false
 env:
-  WITH_PUSH: "true"
   CONDA_ENV: "tritonbench"
   DOCKER_IMAGE: "ghcr.io/pytorch-labs/tritonbench:latest"
   SETUP_SCRIPT: "/workspace/setup_instance.sh"
@@ -28,7 +27,7 @@ jobs:
         with:
           path: tritonbench
       - name: Login to GitHub Container Registry
-        if: ${{ env.WITH_PUSH == 'true' }}
+        if: github.event_name != 'pull_request'
         uses: docker/login-action@v2
         with:
           registry: ghcr.io
@@ -39,17 +38,17 @@ jobs:
           set -x
           export NIGHTLY_DATE="${{ github.event.inputs.nightly_date }}"
           cd tritonbench/docker
-          full_ref="${{ github.ref }}"
-          prefix="refs/heads/"
-          branch_name=${full_ref#$prefix}
+          # branch name is github.head_ref when triggered by pull_request
+          # and it is github.ref_name when triggered by workflow_dispatch
+          branch_name=${{ github.head_ref || github.ref_name }}
           docker build . --build-arg TRITONBENCH_BRANCH="${branch_name}" --build-arg FORCE_DATE="${NIGHTLY_DATE}" \
               -f tritonbench-nightly.dockerfile -t ghcr.io/pytorch-labs/tritonbench:latest
           # Extract pytorch version from the docker
           PYTORCH_VERSION=$(docker run -e SETUP_SCRIPT="${SETUP_SCRIPT}" ghcr.io/pytorch-labs/tritonbench:latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')
           export DOCKER_TAG=$(awk '{match($0, /dev[0-9]+/, arr); print arr[0]}' <<< "${PYTORCH_VERSION}")
           docker tag ghcr.io/pytorch-labs/tritonbench:latest ghcr.io/pytorch-labs/tritonbench:${DOCKER_TAG}
       - name: Push docker to remote
-        if: ${{ env.WITH_PUSH == 'true' }}
+        if: github.event_name != 'pull_request'
         run: |
           # Extract pytorch version from the docker
           PYTORCH_VERSION=$(docker run -e SETUP_SCRIPT="${SETUP_SCRIPT}" ghcr.io/pytorch-labs/tritonbench:latest bash -c '. "${SETUP_SCRIPT}"; python -c "import torch; print(torch.__version__)"')