Fix v2 nightly and bisection workflows (#2366)

Summary:
Fix the broken v2 nightly workflow.

Pull Request resolved: #2366

Test Plan: https://github.com/pytorch/benchmark/actions/runs/10118537615/job/27985526154

Reviewed By: aaronenyeshi

Differential Revision: D60320533

Pulled By: xuzhao9

fbshipit-source-id: 0bc65c00cd000416855ea033c6e2b075f14212e6

.github/workflows/v2-bisection.yml

@@ -11,15 +11,10 @@ jobs:
   bisection:
     environment: docker-s3-upload
     env:
-      BISECT_CONDA_ENV: "bisection-ci-v2"
+      CONDA_ENV_NAME: "bisection-ci-v2"
       BISECT_DIR: ".torchbench/v2-bisection-ci"
       BISECT_BRANCH: "v2.0"
-      PYTHON_VER: "3.10"
-      CUDA_VER: "12.1"
-      NUMPY_VER: "1.21.2"
-      CMAKE_VER: "3.26"
-      MKL_VER: "2021.2.0"
-      MAGMA_VER: "magma-cuda121"
+      SETUP_SCRIPT: "/data/nvme/bin/setup_instance.sh"
     if: ${{ github.repository_owner == 'pytorch' }}
     runs-on: [self-hosted, bm-runner]
     timeout-minutes: 2880 # 48 hours
@@ -31,32 +26,27 @@ jobs:
       - name: Create conda environment
         run: |
           set -x
-          conda create -y -n "${BISECT_CONDA_ENV}" python="${PYTHON_VER}"
-          . activate "${BISECT_CONDA_ENV}"
-          . /data/nvme/bin/setup_instance.sh
-          conda install -y numpy="${NUMPY_VER}"  mkl="${MKL_VER}" mkl-include="${MKL_VER}" \
-                           requests ninja pyyaml setuptools cffi sympy ffmpeg \
-                           typing_extensions future six dataclasses tabulate gitpython git-lfs tqdm regex
-          # install cmake 3.26 from conda-forge, cmake > 3.24 is required by torchaudio
-          conda install -y cmake="${CMAKE_VER}" -c conda-forge
-          # Install magma
-          conda install -y -c pytorch "${MAGMA_VER}"
+          python3 ./utils/python_utils.py --create-conda-env ${CONDA_ENV_NAME}
+          . activate "${CONDA_ENV_NAME}"
+          . "${SETUP_SCRIPT}"
+          python utils/cuda_utils.py --install-torch-build-deps
+          python utils/cuda_utils.py --install-torchbench-deps
       - name: Bisection
         run: |
+          . activate "${CONDA_ENV_NAME}"
+          . "${SETUP_SCRIPT}"
           export BISECT_ISSUE="${{ github.event.inputs.issue_name }}"
           export BISECT_BASE="${HOME}/${BISECT_DIR}/${BISECT_ISSUE}"
           export TORCHBENCH_SRC_DIR="${PWD}"
-          . activate "$BISECT_CONDA_ENV"
-          . /data/nvme/bin/setup_instance.sh
           bash ./.github/scripts/run-bisection.sh
           # Update the result json symbolic link
           ln -sf "${BISECT_BASE}/gh${GITHUB_RUN_ID}/result.json" "${BISECT_BASE}/result.json"
       - name: Analyze bisection result
         run: |
+          . activate "${CONDA_ENV_NAME}"
+          . "${SETUP_SCRIPT}"
           export BISECT_ISSUE="${{ github.event.inputs.issue_name }}"
           export BISECT_BASE="${HOME}/${BISECT_DIR}/${BISECT_ISSUE}"
-          . activate "$BISECT_CONDA_ENV"
-          . /data/nvme/bin/setup_instance.sh
           python ./.github/scripts/bmutils/analyze-bisection-result.py --bisection-root "${BISECT_BASE}" --gh-workflow-id "${GITHUB_RUN_ID}"
           cp -r "${BISECT_BASE}" ./bisection-result
       - name: Create the github issue

.github/workflows/v2-nightly.yml

@@ -10,19 +10,16 @@ jobs:
     env:
       TORCHBENCH_VER: "v2"
       CONFIG_VER: "v2"
-      PYTHON_VER: "3.10"
-      CUDA_VER: "12.1"
-      MAGMA_VERSION: "magma-cuda121"
       CONDA_ENV_NAME:  "torchbench-v2-nightly-ci"
       OUTPUT_DIR: ".torchbench/v2-nightly-ci"
       BISECTION_ROOT: ".torchbench/v2-bisection-ci"
-      CUDA_VERSION: "cu121"
       SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }}
       AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
       AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
       IS_GHA: 1
       AWS_DEFAULT_REGION: us-east-1
       BUILD_ENVIRONMENT: benchmark-nightly
+      SETUP_SCRIPT: "/data/nvme/bin/setup_instance.sh"
     if: ${{ github.repository_owner == 'pytorch' }}
     runs-on: [self-hosted, bm-runner]
     steps:
@@ -32,41 +29,29 @@ jobs:
           ref: v2.0
       - name: Create conda env
         run: |
-          conda create -y -q --name "${CONDA_ENV_NAME}" python=${{ env.PYTHON_VER }}
+          python3 ./utils/python_utils.py --create-conda-env ${CONDA_ENV_NAME}
       - name: Install PyTorch nightly
         run: |
           . activate "${CONDA_ENV_NAME}"
-          . /data/nvme/bin/setup_instance.sh
-          # Install dependencies
-          pip install requests bs4 argparse gitpython boto3 regex
-          # Check if nightly builds are available
-          NIGHTLIES=$(python torchbenchmark/util/torch_nightly.py --packages torch)
-          # If failed, the script will generate empty result
-          if [ -z $NIGHTLIES ]; then
-              echo "Torch nightly build failed. Cancel the workflow."
-              exit 1
-          fi
-          # Install magma
-          conda install -y -c pytorch "${MAGMA_VERSION}"
-          # Install PyTorch nightly from pip
-          pip install --no-cache-dir --pre torch torchvision torchaudio --index-url \
-            https://download.pytorch.org/whl/nightly/${CUDA_VERSION}
-      - name: Install other TorchBench dependencies
+          . "${SETUP_SCRIPT}"
+          python utils/cuda_utils.py --install-torch-deps
+          python utils/cuda_utils.py --install-torch-nightly
+      - name: Install Torchbench models
         run: |
           . activate "${CONDA_ENV_NAME}"
-          . /data/nvme/bin/setup_instance.sh
-          conda install -y git-lfs
+          . "${SETUP_SCRIPT}"
           python install.py
       - name: Run benchmark
         run: |
           . activate "${CONDA_ENV_NAME}"
-          . /data/nvme/bin/setup_instance.sh
+          . "${SETUP_SCRIPT}"
           WORKFLOW_HOME="${HOME}/${{ env.OUTPUT_DIR }}/gh${GITHUB_RUN_ID}"
           bash ./.github/scripts/run.sh "${WORKFLOW_HOME}"
       - name: Generate the bisection config
         run: |
           set -x
           . activate "${CONDA_ENV_NAME}"
+          . "${SETUP_SCRIPT}"
           WORKFLOW_HOME="${HOME}/${{ env.OUTPUT_DIR }}/gh${GITHUB_RUN_ID}"
           mkdir -p benchmark-output/
           # Update the self-hosted pytorch version