Build with ref

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

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/.github/workflows/build-nightly-docker.yml b/.github/workflows/build-nightly-docker.yml
@@ -28,7 +28,8 @@ jobs:
           export TODAY=$(date +'%Y%m%d')
           export DOCKER_TAG=dev${TODAY}
           cd benchmark/docker
-          docker build . -f torchbench-nightly.dockerfile -t ghcr.io/pytorch/torchbench:${DOCKER_TAG}
+          echo "Github Ref: ${{ github.ref }}"
+          docker build . --build-arg TORCHBENCH_BRANCH="${{ github.ref }}" -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,12 +1,16 @@
 # 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
 
 # Setup Conda env and CUDA
-RUN git clone https://github.com/pytorch/benchmark /workspace/benchmark
+RUN git clone -b ${TORCHBENCH_BRANCH} --single-branch \
+ https://github.com/pytorch/benchmark /workspace/benchmark
+
 RUN cd /workspace/benchmark && \
     . ${SETUP_SCRIPT} && \
     python ./utils/python_utils.py --create-conda-env ${CONDA_ENV} && \