Build docker from git ref

pytorch · Sep 28, 2023 · cfc129c · cfc129c
1 parent 480353c
commit cfc129c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/.github/workflows/build-nightly-docker.yml b/.github/workflows/build-nightly-docker.yml
@@ -28,8 +28,10 @@ jobs:
           export TODAY=$(date +'%Y%m%d')
           export DOCKER_TAG=dev${TODAY}
           cd benchmark/docker
-          echo "Github Ref: ${{ github.ref }}"
-          docker build . --build-arg TORCHBENCH_BRANCH="${{ github.ref }}" -f torchbench-nightly.dockerfile -t ghcr.io/pytorch/torchbench:${DOCKER_TAG}
+          full_ref="${{ github.ref }}"
+          prefix="refs/heads/"
+          branch_name=${full_ref#$prefix}
+          docker build . --build-arg TORCHBENCH_BRANCH="${branch_name}" -f torchbench-nightly.dockerfile -t ghcr.io/pytorch/torchbench:${DOCKER_TAG}
           docker tag ghcr.io/pytorch/torchbench:${DOCKER_TAG} ghcr.io/pytorch/torchbench:latest
       - name: Push docker to remote
         if: ${{ env.WITH_PUSH == 'true' }}

diff --git a/docker/torchbench-nightly.dockerfile b/docker/torchbench-nightly.dockerfile
@@ -1,14 +1,14 @@
 # default base image: xzhao9/gcp-a100-runner-dind:latest
 ARG BASE_IMAGE=xzhao9/gcp-a100-runner-dind:latest
-ARG TORCHBENCH_BRANCH=main
 
 FROM ${BASE_IMAGE}
 
 ENV CONDA_ENV=torchbench
 ENV SETUP_SCRIPT=/workspace/setup_instance.sh
+ARG TORCHBENCH_BRANCH=${TORCHBENCH_BRANCH:-main}
 
 # Setup Conda env and CUDA
-RUN git clone -b ${TORCHBENCH_BRANCH} --single-branch \
+RUN git clone -b "${TORCHBENCH_BRANCH}" --single-branch \
  https://github.com/pytorch/benchmark /workspace/benchmark
 
 RUN cd /workspace/benchmark && \