From 2c243ed112b4061e8e287be9c10fc22a463053d9 Mon Sep 17 00:00:00 2001
From: Xu Zhao <xzhao9@fb.com>
Date: Thu, 9 Sep 2021 12:35:36 -0700
Subject: [PATCH] Enable the on-demand performance PR testing to run on a
 specified TB branch (#64701)

Summary:
This is to enable performance testing of experimental features such as LazyTensor.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/64701

Test Plan:
TorchBench CI

RUN_TORCHBENCH: BERT_pytorch, mobilenet_v3_large
TORCHBENCH_BRANCH: v1.0

Reviewed By: seemethere

Differential Revision: D30847389

Pulled By: xuzhao9

fbshipit-source-id: 6853b368fa6f1ba8ffde517805c74bf318dcb35b
---
 .github/scripts/run_torchbench.py    | 30 +++++++++++++++++++++++++++++-
 .github/workflows/run_torchbench.yml | 11 ++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/.github/scripts/run_torchbench.py b/.github/scripts/run_torchbench.py
index b3c7ee2..9037308 100644
--- a/.github/scripts/run_torchbench.py
+++ b/.github/scripts/run_torchbench.py
@@ -13,6 +13,7 @@ Testing environment:
 # 1. Does not reuse the build artifact in other CI workflows
 # 2. CI jobs are serialized because there is only one worker
 import os
+import git  # type: ignore[import]
 import pathlib
 import argparse
 import subprocess
@@ -23,6 +24,7 @@ CUDA_VERSION = "cu102"
 PYTHON_VERSION = "3.7"
 TORCHBENCH_CONFIG_NAME = "config.yaml"
 MAGIC_PREFIX = "RUN_TORCHBENCH:"
+MAGIC_TORCHBENCH_PREFIX = "TORCHBENCH_BRANCH:"
 ABTEST_CONFIG_TEMPLATE = """# This config is automatically generated by run_torchbench.py
 start: {control}
 end: {treatment}
@@ -57,7 +59,7 @@ def extract_models_from_pr(torchbench_path: str, prbody_file: str) -> List[str]:
         lines = map(lambda x: x.strip(), pf.read().splitlines())
         magic_lines = list(filter(lambda x: x.startswith(MAGIC_PREFIX), lines))
         if magic_lines:
-            # Only the first magic line will be respected.
+            # Only the first magic line will be recognized.
             model_list = list(map(lambda x: x.strip(), magic_lines[0][len(MAGIC_PREFIX):].split(",")))
     # Shortcut: if model_list is ["ALL"], run all the tests
     if model_list == ["ALL"]:
@@ -71,6 +73,26 @@ def extract_models_from_pr(torchbench_path: str, prbody_file: str) -> List[str]:
             return []
     return model_list
 
+def identify_torchbench_branch(torchbench_path: str, prbody_file: str) -> None:
+    branch_name: str
+    with open(prbody_file, "r") as pf:
+        lines = map(lambda x: x.strip(), pf.read().splitlines())
+        magic_lines = list(filter(lambda x: x.startswith(MAGIC_TORCHBENCH_PREFIX), lines))
+        if magic_lines:
+            # Only the first magic line will be recognized.
+            branch_name = magic_lines[0][len(MAGIC_TORCHBENCH_PREFIX):].strip()
+    # If not specified, directly return without the branch checkout
+    if not branch_name:
+        return
+    try:
+        print(f"Checking out the TorchBench branch: {branch_name} ...")
+        repo = git.Repo(torchbench_path)
+        origin = repo.remotes.origin
+        origin.fetch(branch_name)
+        repo.create_head(branch_name, origin.refs[branch_name]).checkout()
+    except git.exc.GitCommandError:
+        raise RuntimeError(f'{branch_name} doesn\'t exist in the pytorch/benchmark repository. Please double check.')
+
 def run_torchbench(pytorch_path: str, torchbench_path: str, output_dir: str) -> None:
     # Copy system environment so that we will not override
     env = dict(os.environ)
@@ -96,6 +118,12 @@ if __name__ == "__main__":
     if not models:
         print("Can't parse the model filter from the pr body. Currently we only support allow-list.")
         exit(1)
+    # Identify the specified TorchBench branch, verify the branch exists, and checkout the branch
+    try:
+        identify_torchbench_branch(args.torchbench_path, args.pr_body)
+    except RuntimeError as e:
+        print(f"Identify TorchBench branch failed: {str(e)}")
+        exit(1)
     print(f"Ready to run TorchBench with benchmark. Result will be saved in the directory: {output_dir}.")
     # Run TorchBench with the generated config
     torchbench_config = gen_abtest_config(args.pr_base_sha, args.pr_head_sha, models)
diff --git a/.github/workflows/run_torchbench.yml b/.github/workflows/run_torchbench.yml
index cee27e1..5b0eb5e 100644
--- a/.github/workflows/run_torchbench.yml
+++ b/.github/workflows/run_torchbench.yml
@@ -32,12 +32,21 @@ jobs:
           # shellcheck disable=SC1091
           . "${HOME}"/anaconda3/etc/profile.d/conda.sh
           conda activate pr-ci
-          conda install -y numpy=1.17 requests=2.22 ninja pyyaml mkl mkl-include setuptools cmake cffi typing_extensions future six dataclasses pillow pytest tabulate
+          conda install -y numpy=1.17 requests=2.22 ninja pyyaml mkl mkl-include setuptools cmake cffi typing_extensions future six dataclasses pillow pytest tabulate gitpython
+          conda install -y -c pytorch-nightly torchtext torchvision
       - name: Update self-hosted PyTorch
         run: |
           pushd "${HOME}"/pytorch
+          git remote prune origin
           git fetch
           popd
+      - name: Install TorchBench dependencies
+        run: |
+          # shellcheck disable=SC1091
+          . "${HOME}"/anaconda3/etc/profile.d/conda.sh
+          conda activate pr-ci
+          pushd "${PWD}"/benchmark
+          python install.py
       - name: Run TorchBench
         run: |
           pushd "${HOME}"/pytorch
-- 
2.7.4