[CI][AMD] Run the benchmark on MI350 (#502)

Author: xuzhao9

.ci/test_infra/oss_ci_benchmark_v3.py

@@ -5,7 +5,10 @@
 
 import argparse
 import json
+import os
 import re
+import sys
+from os.path import abspath, exists
 from pathlib import Path
 
 from typing import Any, Dict, List, Tuple
@@ -15,10 +18,37 @@
         "name": "gcp-h100-runner",
         "gpu_count": 1,
         "avail_gpu_mem_in_gb": 80,
-    }
+    },
+    "amd-mi350-runner": {
+        "name": "amd-mi350-runner",
+        "gpu_count": 1,
+        "avail_gpu_mem_in_gb": 288,
+    },
 }
 
 
+def setup_tritonbench_cwd():
+    original_dir = abspath(os.getcwd())
+
+    for tritonbench_dir in (
+        ".",
+        "../../tritonbench",
+    ):
+        if exists(tritonbench_dir):
+            break
+
+    if exists(tritonbench_dir):
+        tritonbench_dir = abspath(tritonbench_dir)
+        os.chdir(tritonbench_dir)
+        sys.path.append(tritonbench_dir)
+    return original_dir
+
+
+setup_tritonbench_cwd()
+
+from tritonbench.utils.scuba_utils import get_github_env
+
+
 def parse_runners(
     runner_name: str, runner_type: str, envs: Dict[str, str]
 ) -> List[Dict[str, Any]]:
@@ -130,6 +160,11 @@ def v3_json_to_str(v3_json: List[Dict[str, Any]], to_lines: bool = True) -> str:
         required=True,
         help="Upload benchmark result json file.",
     )
+    parser.add_argument(
+        "--add-github-env",
+        action="store_true",
+        help="Add github env to the result json file.",
+    )
     parser.add_argument("--output", required=True, help="output json.")
     args = parser.parse_args()
     upload_file_path = Path(args.json)
@@ -138,8 +173,13 @@ def v3_json_to_str(v3_json: List[Dict[str, Any]], to_lines: bool = True) -> str:
     ), f"Specified result json path {args.json} does not exist."
     with open(upload_file_path, "r") as fp:
         benchmark_result = json.load(fp)
-    oss_ci_v3_json = generate_oss_ci_benchmark_v3_json(benchmark_result)
-    out_str = v3_json_to_str(oss_ci_v3_json)
+    if args.add_github_env:
+        github_env = get_github_env()
+        benchmark_result["github"] = github_env
+        out_str = v3_json_to_str(benchmark_result, to_lines=False)
+    else:
+        oss_ci_v3_json = generate_oss_ci_benchmark_v3_json(benchmark_result)
+        out_str = v3_json_to_str(oss_ci_v3_json)
     output_dir = Path(args.output).parent
     output_dir.mkdir(parents=True, exist_ok=True)
     with open(args.output, "w") as fp:

.ci/tritonbench/run-benchmark.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-set -x
+set -xeuo pipefail
 
 if [ -z "${SETUP_SCRIPT}" ]; then
   echo "ERROR: SETUP_SCRIPT is not set"
@@ -31,6 +31,10 @@ fi
 tritonbench_dir=$(dirname "$(readlink -f "$0")")/../..
 cd "${tritonbench_dir}"
 
+# check if the current repo has "dubious ownership" issue
+git config --global --add safe.directory '*'
+git rev-parse --verify HEAD
+
 echo "Running ${BENCHMARK_NAME} benchmark under conda env ${CONDA_ENV}"
 
 . "${SETUP_SCRIPT}"

.ci/upload/requirements.txt

@@ -1 +1,3 @@
 boto3
+pyyaml
+requests

.github/workflows/_linux-benchmark-mi350.yml

@@ -58,6 +58,10 @@ jobs:
           # The max duration enforced by the server side
           role-duration-seconds: 18000
           aws-region: us-east-1
+      - name: Setup uploader dependencies
+        run: |
+          sudo apt-get install -y python3-pip
+          pip3 install -r .ci/upload/requirements.txt
       - name: Pull docker image
         uses: pytorch/test-infra/.github/actions/pull-docker-image@main
         with:
@@ -77,21 +81,24 @@ jobs:
             --security-opt seccomp=unconfined \
             --shm-size=32g \
             --cap-add=SYS_PTRACE \
+            --user root \
             -v "${GITHUB_WORKSPACE}:/tmp/workspace" \
             -w /tmp/workspace \
             "${DOCKER_IMAGE}"
           )
+
           # write container id to env
           echo "TRITONBENCH_CONTAINER_ID=${container_name}" >> $GITHUB_ENV
       - name: Compile Triton (On Demand)
         if: ${{ inputs.side_a_triton && inputs.side_a_commit }}
         run: |
           docker exec -t -w /tmp/workspace "${TRITONBENCH_CONTAINER_ID}" bash -c "
             set -eux
-            bash ./.ci/triton/compile.sh --repo ${{ inputs.side_a_triton }} --commit ${{ inputs.side_a_commit }} --side a
+            bash ./.ci/triton/install.sh --repo ${{ inputs.side_a_triton }} --commit ${{ inputs.side_a_commit }} --side a
           "
       - name: Benchmarking
         run: |
+          set -eux
           if [ -n "${{ inputs.side_a_triton }}" ] && [ -n "${{ inputs.side_a_commit }}" ]; then
             docker exec -t -w /tmp/workspace "${TRITONBENCH_CONTAINER_ID}" bash -c "
               set -eux
@@ -104,25 +111,23 @@ jobs:
             "
           fi
           cp -r ".benchmarks/${{ inputs.benchmark_name }}" benchmark-output
+          # post-process result.json
+          latest_result_json=$(find ./benchmark-output -name "result.json"  | sort -r | head -n 1)
+          python3 ./.ci/test_infra/oss_ci_benchmark_v3.py --json ${latest_result_json} \
+            --add-github-env --output ${latest_result_json}
       - name: Upload result to GH Actions Artifact
         uses: actions/upload-artifact@v4
         with:
           name: ${{ env.JOB_NAME }}
           path: benchmark-output/
-      - name: Setup uploader dependencies
-        run: |
-          sudo apt-get install -y python3-pip
-          pip3 install -y pyyaml
       - name: Upload result to Scribe
         run: |
-          . "${SETUP_SCRIPT}"
           latest_result_json=$(find ./benchmark-output -name "result.json"  | sort -r | head -n 1)
-          python ./.ci/upload/scribe.py --json ${latest_result_json}
+          python3 ./.ci/upload/scribe.py --json ${latest_result_json}
       - name: Rewrite Tritonbench json to ClickHouse style
         run: |
-          . "${SETUP_SCRIPT}"
           latest_result_json=$(find ./benchmark-output -name "result.json"  | sort -r | head -n 1)
-          python ./.ci/test_infra/oss_ci_benchmark_v3.py --json ${latest_result_json} \
+          python3 ./.ci/test_infra/oss_ci_benchmark_v3.py --json ${latest_result_json} \
                  --output benchmark-output/results/result.json
       - name: Upload result to ClickHouse
         uses: pytorch/test-infra/.github/actions/upload-benchmark-results@main
@@ -131,3 +136,7 @@ jobs:
           dry-run: false
           schema-version: v3
           github-token: ${{ secrets.GITHUB_TOKEN }}
+      - name: Kill the container
+        if: always()
+        run: |
+          docker kill "${TRITONBENCH_CONTAINER_ID}" || true

.github/workflows/nightly.yml

@@ -28,7 +28,7 @@ on:
   pull_request:
     paths:
       - benchmarks/nightly/**
-      - .github/workflows/_linux-benchmark-h100.yml
+      - .github/workflows/_linux-benchmark-*.yml
       - .github/workflows/nightly.yml
 
 jobs:
@@ -42,6 +42,16 @@ jobs:
       side_a_commit: ${{ inputs.side_a_commit }}
     secrets:
       TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN }}
+  mi350-triton-main-nightly-periodic:
+    uses: ./.github/workflows/_linux-benchmark-mi350.yml
+    if: ${{ inputs.test_type != 'abtest' }}
+    with:
+      conda_env: "triton-main"
+      benchmark_name: "nightly"
+      side_a_triton: ${{ inputs.side_a_triton }}
+      side_a_commit: ${{ inputs.side_a_commit }}
+    secrets:
+      TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN }}
   h100-triton-nightly-abtest:
     uses: ./.github/workflows/_linux-benchmark-abtest-h100.yml
     if: ${{ inputs.test_type == 'abtest' }}

.github/workflows/pr.yaml

@@ -1,6 +1,10 @@
 name: TritonBench PR Test
 on:
   pull_request:
+    paths:
+      - .github/workflows/pr.yml
+      - .github/workflows/_linux-test-*.yml
+      - tritonbench/**
   push:
     branches:
       - main

tritonbench/utils/gpu_utils.py

@@ -229,5 +229,5 @@ def has_nvidia_smi() -> bool:
     try:
         subprocess.check_output("nvidia-smi")
         return True
-    except subprocess.SubprocessError:
+    except (subprocess.SubprocessError, FileNotFoundError):
         return False

tritonbench/utils/run_utils.py

@@ -10,6 +10,8 @@
 
 from typing import Dict, List, Optional
 
+import torch
+
 import yaml
 
 from tritonbench.utils.env_utils import is_fbcode
@@ -44,8 +46,6 @@ def get_run_env(
     Gather environment of the benchmark.
     repo_locs: Git repository dict of the repositories.
     """
-    import torch
-
     run_env = {}
     run_env["benchmark_date"] = run_timestamp
     run_env["cuda_version"] = torch.version.cuda if torch.version.cuda else "unknown"
@@ -73,31 +73,6 @@ def get_run_env(
     return run_env
 
 
-def get_github_env() -> Dict[str, str]:
-    assert (
-        "GITHUB_RUN_ID" in os.environ
-    ), "GITHUB_RUN_ID environ must exist to obtain GitHub env"
-    out = {}
-    out["GITHUB_ACTION"] = os.environ["GITHUB_ACTION"]
-    out["GITHUB_ACTOR"] = os.environ["GITHUB_ACTOR"]
-    out["GITHUB_BASE_REF"] = os.environ["GITHUB_BASE_REF"]
-    out["GITHUB_REF"] = os.environ["GITHUB_REF"]
-    out["GITHUB_REF_PROTECTED"] = os.environ["GITHUB_REF_PROTECTED"]
-    out["GITHUB_REPOSITORY"] = os.environ["GITHUB_REPOSITORY"]
-    out["GITHUB_RUN_ATTEMPT"] = os.environ["GITHUB_RUN_ATTEMPT"]
-    out["GITHUB_RUN_ID"] = os.environ["GITHUB_RUN_ID"]
-    out["GITHUB_RUN_NUMBER"] = os.environ["GITHUB_RUN_NUMBER"]
-    out["GITHUB_WORKFLOW"] = os.environ["GITHUB_WORKFLOW"]
-    out["GITHUB_WORKFLOW_REF"] = os.environ["GITHUB_WORKFLOW_REF"]
-    out["GITHUB_WORKFLOW_SHA"] = os.environ["GITHUB_WORKFLOW_SHA"]
-    out["JOB_NAME"] = os.environ["JOB_NAME"]
-    out["RUNNER_ARCH"] = os.environ["RUNNER_ARCH"]
-    out["RUNNER_TYPE"] = os.environ["RUNNER_TYPE"]
-    out["RUNNER_NAME"] = os.environ["RUNNER_NAME"]
-    out["RUNNER_OS"] = os.environ["RUNNER_OS"]
-    return out
-
-
 def run_config(config_file: str, args: List[str]):
     assert Path(config_file).exists(), f"Config file {config_file} must exist."
     with open(config_file, "r") as fp:

tritonbench/utils/scuba_utils.py

@@ -13,7 +13,6 @@
 
 from tritonbench.utils.gpu_utils import get_nvidia_gpu_states, has_nvidia_smi
 from tritonbench.utils.path_utils import REPO_PATH
-from tritonbench.utils.run_utils import get_github_env, get_run_env
 
 CATEGORY_NAME = "perfpipe_pytorch_user_benchmarks"
 
@@ -58,6 +57,30 @@
 }
 
 
+def get_github_env() -> Dict[str, str]:
+    if "GITHUB_RUN_ID" not in os.environ:
+        return {}
+    out = {}
+    out["GITHUB_ACTION"] = os.environ["GITHUB_ACTION"]
+    out["GITHUB_ACTOR"] = os.environ["GITHUB_ACTOR"]
+    out["GITHUB_BASE_REF"] = os.environ["GITHUB_BASE_REF"]
+    out["GITHUB_REF"] = os.environ["GITHUB_REF"]
+    out["GITHUB_REF_PROTECTED"] = os.environ["GITHUB_REF_PROTECTED"]
+    out["GITHUB_REPOSITORY"] = os.environ["GITHUB_REPOSITORY"]
+    out["GITHUB_RUN_ATTEMPT"] = os.environ["GITHUB_RUN_ATTEMPT"]
+    out["GITHUB_RUN_ID"] = os.environ["GITHUB_RUN_ID"]
+    out["GITHUB_RUN_NUMBER"] = os.environ["GITHUB_RUN_NUMBER"]
+    out["GITHUB_WORKFLOW"] = os.environ["GITHUB_WORKFLOW"]
+    out["GITHUB_WORKFLOW_REF"] = os.environ["GITHUB_WORKFLOW_REF"]
+    out["GITHUB_WORKFLOW_SHA"] = os.environ["GITHUB_WORKFLOW_SHA"]
+    out["JOB_NAME"] = os.environ["JOB_NAME"]
+    out["RUNNER_ARCH"] = os.environ["RUNNER_ARCH"]
+    out["RUNNER_TYPE"] = os.environ["RUNNER_TYPE"]
+    out["RUNNER_NAME"] = os.environ["RUNNER_NAME"]
+    out["RUNNER_OS"] = os.environ["RUNNER_OS"]
+    return out
+
+
 class ScribeUploader:
     def __init__(self, category, schema):
         self.category = category
@@ -133,6 +156,7 @@ def decorate_benchmark_data(
     name, run_timestamp, ci: bool, benchmark_data: List[Dict[str, Any]]
 ):
     """aggregate benchmark_data into a single object"""
+    from tritonbench.utils.run_utils import get_run_env
 
     repo_locs = {
         "tritonbench": REPO_PATH,