From 9b24c0a9b6e0c186c8dbb9af21e8a43f3f790b26 Mon Sep 17 00:00:00 2001
From: Huy Do <huydhn@gmail.com>
Date: Mon, 16 Dec 2024 21:41:51 -0800
Subject: [PATCH] Try schema v3

---
 .github/workflows/_linux-benchmark-cuda.yml | 92 ++++++++-------------
 .github/workflows/torchao.yml               |  2 +
 userbenchmark/torchao/run.py                |  1 +
 3 files changed, 38 insertions(+), 57 deletions(-)

diff --git a/.github/workflows/_linux-benchmark-cuda.yml b/.github/workflows/_linux-benchmark-cuda.yml
index 9d5e01ad0..a85ff252e 100644
--- a/.github/workflows/_linux-benchmark-cuda.yml
+++ b/.github/workflows/_linux-benchmark-cuda.yml
@@ -15,81 +15,59 @@ on:
         required: false
         description: |
           HF auth token to avoid rate limits when downloading models or datasets from hub
-      AWS_ACCESS_KEY_ID:
-        required: true
-        description: |
-          AWS access token for S3 uploading
-      AWS_SECRET_ACCESS_KEY:
-        required: true
-        description: |
-          AWS secret access key for S3 uploading
 
 jobs:
-  # Run a specific userbenchmark with given arguments
-  # Need to pass in userbenchmark name and arguments
   benchmark:
     # Don't run on forked repos
     if: github.repository_owner == 'pytorch'
     runs-on: linux.aws.a100
-    timeout-minutes: 1440 # 24 hours
+    timeout-minutes: 1440
     environment: docker-s3-upload
     env:
-      BASE_CONDA_ENV: "torchbench"
-      CONDA_ENV: "userbenchmark"
-      SETUP_SCRIPT: "/workspace/setup_instance.sh"
+      OUTPUT_DIR: '.userbenchmark'
       HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
-      AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
-      AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
     steps:
       - name: Checkout TorchBench
         uses: actions/checkout@v3
         with:
           path: benchmark
+
       - name: Remove result if it already exists
-        if: always()
-        run: |
-          # remove old results if exists
-          if [ -d benchmark-output ]; then rm -Rf benchmark-output; fi
-          pushd benchmark
-          if [ -d .userbenchmark ]; then rm -Rf .userbenchmark; fi
-      - name: Clone and setup conda env
+        shell: bash
+        working-directory: benchmark
         run: |
-          CONDA_ENV=${BASE_CONDA_ENV} . "${SETUP_SCRIPT}"
-          conda create --name "${CONDA_ENV}" --clone "${BASE_CONDA_ENV}"
+          set -eux
+
+          if [[ -d "${OUTPUT_DIR}" ]]; then
+            rm -rf "${OUTPUT_DIR}"
+          fi
+
+      - name: Setup miniconda
+        uses: pytorch/test-infra/.github/actions/setup-miniconda@main
+        with:
+          python-version: "3.9"
+
       - name: Install benchmark
+        shell: bash
+        working-directory: benchmark
         run: |
-          . "${SETUP_SCRIPT}"
-          pushd benchmark
-          python install.py
+          set -eux
+          ${CONDA_RUN} python install.py
+
       - name: Run benchmark
+        shell: bash
+        working-directory: benchmark
         run: |
-          . "${SETUP_SCRIPT}"
-          pushd benchmark
-          python run_benchmark.py ${{ inputs.userbenchmark }} ${{ inputs.userbenchmark-run-args }}
-      - name: Copy benchmark logs
-        if: always()
-        run: |
-          pushd benchmark
-          cp -r ./.userbenchmark/${{ inputs.userbenchmark }} ../benchmark-output
-      - name: Upload benchmark result to GH Actions Artifact
-        uses: actions/upload-artifact@v3
-        if: always()
+          set -eux
+          ${CONDA_RUN} python run_benchmark.py ${{ inputs.userbenchmark }} ${{ inputs.userbenchmark-run-args }}
+
+          # DEBUG
+          ls -laR "${OUTPUT_DIR}"
+
+      - name: Upload the benchmark results to OSS benchmark database for the dashboard
+        uses: pytorch/test-infra/.github/actions/upload-benchmark-results@main
         with:
-          name: ${{ inputs.userbenchmark }} benchmarking result
-          path: benchmark-output/
-      - name: Copy artifact and upload to Amazon S3
-        env:
-          WORKFLOW_RUN_ID: ${{ github.run_id }}
-          WORKFLOW_RUN_ATTEMPT: ${{ github.run_attempt }}
-        run: |
-          . "${SETUP_SCRIPT}"
-          pushd benchmark
-          # Upload the result json to Amazon S3
-          python ./scripts/userbenchmark/upload_s3_csv.py --s3-prefix torchbench-csv --userbenchmark ${{ inputs.userbenchmark }} \
-                                                          --upload-path ../benchmark-output --match-filename "^${{ inputs.userbenchmark }}.*\.csv"
-      - name: Clean up Conda env
-        if: always()
-        run: |
-          . "${SETUP_SCRIPT}"
-          conda deactivate && conda deactivate
-          conda remove -n "${CONDA_ENV}" --all
+          benchmark-results-dir: ${{ env.OUTPUT_DIR }}
+          dry-run: true  # DEBUG: TO BE REMOVED
+          schema-version: v3
+          github-token: ${{ secrets.GITHUB_TOKEN }}
diff --git a/.github/workflows/torchao.yml b/.github/workflows/torchao.yml
index 99c5d95ff..a570647f8 100644
--- a/.github/workflows/torchao.yml
+++ b/.github/workflows/torchao.yml
@@ -1,5 +1,7 @@
 name: Torchao nightly workflow (A100)
 on:
+  # DEBUG
+  pull_request:
   workflow_dispatch:
   schedule:
     - cron: '00 18 * * *' # run at 6:00 PM UTC, K8s containers will roll out at 12PM EST
diff --git a/userbenchmark/torchao/run.py b/userbenchmark/torchao/run.py
index 626ae55cb..b2f67f684 100644
--- a/userbenchmark/torchao/run.py
+++ b/userbenchmark/torchao/run.py
@@ -29,6 +29,7 @@ def _get_ci_args(
         f"--{experiment}",
         "--output",
         f"{str(OUTPUT_DIR.joinpath(output_file_name).resolve())}",
+        "--only hf_T5",  # DEBUG: TO BE REMOVED
     ]
     return ci_args