Enable Python 3.11 as the default version.

.github/workflows/build-nightly-docker.yml

@@ -28,7 +28,10 @@ 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}
+          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' }}

docker/gcp-a100-runner-dind.dockerfile

@@ -14,7 +14,8 @@ RUN sudo apt-get install -y git jq \
                             vim wget curl ninja-build cmake \
                             libgl1-mesa-glx libsndfile1-dev kmod libxml2-dev libxslt1-dev \
                             fontconfig libfontconfig1-dev \
-                            libpango-1.0-0 libpangoft2-1.0-0
+                            libpango-1.0-0 libpangoft2-1.0-0 \
+                            libsdl2-dev libsdl2-2.0-0
 
 # get switch-cuda utility
 RUN sudo wget -q https://raw.githubusercontent.com/phohenecker/switch-cuda/master/switch-cuda.sh -O /usr/bin/switch-cuda.sh
@@ -85,11 +86,10 @@ RUN wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.s
     chmod +x Miniconda3-latest-Linux-x86_64.sh && \
     bash ./Miniconda3-latest-Linux-x86_64.sh -b -u
 
-# Use Python 3.10 as default
+# Test activate miniconda
 RUN . ${HOME}/miniconda3/etc/profile.d/conda.sh && \
     conda activate base && \
-    conda init && \
-    conda install -y python=3.10
+    conda init
 
 RUN echo "\
 . \${HOME}/miniconda3/etc/profile.d/conda.sh\n\

docker/torchbench-nightly.dockerfile

@@ -1,12 +1,19 @@
 # default base image: xzhao9/gcp-a100-runner-dind:latest
 ARG BASE_IMAGE=xzhao9/gcp-a100-runner-dind:latest
+
 FROM ${BASE_IMAGE}
 
 ENV CONDA_ENV=torchbench
 ENV SETUP_SCRIPT=/workspace/setup_instance.sh
+ARG TORCHBENCH_BRANCH=${TORCHBENCH_BRANCH:-main}
+
+# Setup dependencies
+RUN sudo apt install -y libsdl2-dev libsdl2-2.0-0
 
 # 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} && \

requirements.txt

@@ -15,9 +15,8 @@ MonkeyType
 psutil
 pyyaml
 # pytorch build script pins numpy version
-# https://github.com/pytorch/builder/blob/ae5c82e65cb3d8bac6df50e742a195019af91ad3/wheel/build_wheel.sh#L145
-numpy==1.21.2
-# https://github.com/opencv/opencv-python/issues/885
-opencv-python==4.7.0.72
+# https://github.com/pytorch/builder/blob/main/wheel/build_wheel.sh
+numpy==1.23.5
+opencv-python
 submitit
 pynvml

torchbenchmark/models/soft_actor_critic/requirements.txt

@@ -1,3 +1,2 @@
-gym==0.25.2
-# pygame 2.1.3 requires sdl2-config, which doesn't exist in our benchmark environment
-pygame==2.1.2
+gym
+pygame

utils/python_utils.py

@@ -1,7 +1,7 @@
 import argparse
 import subprocess
 
-DEFAULT_PYTHON_VERSION = "3.10"
+DEFAULT_PYTHON_VERSION = "3.11"
 
 PYTHON_VERSION_MAP = {
     "3.8": {