userbenchmark-a100-release.yml

name: Release TorchBench Userbenchmark on A100
on:
  pull_request:
    paths:
      - userbenchmark/release-test/*

jobs:
  run-userbenchmark:
    runs-on: [a100-runner]
    timeout-minutes: 1440 # 24 hours
    environment: docker-s3-upload
    env:
      BASE_CONDA_ENV: "torchbench"
      CONDA_ENV: "userbenchmark-a100"
      PLATFORM_NAME: "gcp_a100"
      SETUP_SCRIPT: "/workspace/setup_instance.sh"
    steps:
      - name: Checkout TorchBench
        uses: actions/checkout@v3
        with:
          path: benchmark
          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
      - name: Tune Nvidia GPU
        run: |
          sudo nvidia-smi -pm 1
          sudo nvidia-smi -ac 1215,1410
          nvidia-smi
      - name: Clone and setup conda env
        run: |
          CONDA_ENV=${BASE_CONDA_ENV} . "${SETUP_SCRIPT}"
          conda create --name "${CONDA_ENV}" --clone "${BASE_CONDA_ENV}"
      - name: Install TorchBench
        run: |
          set -x
          . "${SETUP_SCRIPT}"
          pushd benchmark
          python install.py
      - name: Run user benchmark
        run: |
          set -x
          . "${SETUP_SCRIPT}"
          # remove old results
          if [ -d benchmark-output ]; then rm -Rf benchmark-output; fi
          pushd benchmark

          # Install necessary packages  <-- Added these lines
          pip3 install -e '.[jax_cpu]'
          pip3 install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/cu121'
          pip3 install -e '.[full]'
          
          release_version=$(cat userbenchmark/release-test/version.txt)
          if [ -d .userbenchmark ]; then rm -Rf .userbenchmark; fi
          python run_benchmark.py release-test -c ${release_version}
          cp -r ./.userbenchmark/release-test ../benchmark-output

      - name: Upload artifact
        uses: actions/upload-artifact@v3
        with:
          name: TorchBench result
          path: benchmark-output/
      - name: Clean up Conda env
        if: always()
        run: |
          . "${SETUP_SCRIPT}"
          conda deactivate && conda deactivate
          conda remove -n "${CONDA_ENV}" --all