Specify the head branch when upload perf stats to Rockset (pytorch#97643)

Before this, my assumption was that the workflow was only run on the main branch. This is not correct anymore as it could also now be run as part of the PR, i.e. https://hud.pytorch.org/pr/91316.  So this change does two things:

* Always upload inductor-A100-perf-nightly artifacts to S3 once completed by removing the main branch gating.
* Add `head_branch` to Rockset records, so that the [dashboard](https://torchci-git-fork-huydhn-add-compilers-bench-74abf8-fbopensource.vercel.app/benchmark/compilers) knows if the records come from the daily schedule on the main branch or from experimental PR.  The `head_branch` would be set to `master` in the former.

Pull Request resolved: pytorch#97643
Approved by: https://github.com/desertfire

Author: huydhn

.github/workflows/upload-torch-dynamo-perf-stats.yml

@@ -5,9 +5,6 @@ on:
     workflows: [inductor-A100-perf-nightly]
     types:
       - completed
-    branches:
-      - master
-      - main
 
 jobs:
   get-conclusion:
@@ -59,5 +56,6 @@ jobs:
           WORKFLOW_RUN_ID: ${{ github.event.workflow_run.id }}
           WORKFLOW_RUN_ATTEMPT: ${{ github.event.workflow_run.run_attempt }}
           REPO_FULLNAME: ${{ github.event.workflow_run.repository.full_name }}
+          HEAD_BRANCH: ${{ github.event.workflow_run.head_branch }}
         run: |
-          python3 -m tools.stats.upload_dynamo_perf_stats --workflow-run-id "${WORKFLOW_RUN_ID}" --workflow-run-attempt "${WORKFLOW_RUN_ATTEMPT}" --repo "${REPO_FULLNAME}"
+          python3 -m tools.stats.upload_dynamo_perf_stats --workflow-run-id "${WORKFLOW_RUN_ID}" --workflow-run-attempt "${WORKFLOW_RUN_ATTEMPT}" --repo "${REPO_FULLNAME}" --head-branch "${HEAD_BRANCH}"

tools/stats/upload_dynamo_perf_stats.py

@@ -18,7 +18,10 @@
 
 
 def upload_dynamo_perf_stats_to_rockset(
-    repo: str, workflow_run_id: int, workflow_run_attempt: int
+    repo: str,
+    workflow_run_id: int,
+    workflow_run_attempt: int,
+    head_branch: str,
 ) -> List[Dict[str, Any]]:
     perf_stats = []
     with TemporaryDirectory() as temp_dir:
@@ -65,6 +68,7 @@ def upload_dynamo_perf_stats_to_rockset(
                                     "runner": runner,
                                     "job_id": job_id,
                                     "filename": filename,
+                                    "head_branch": head_branch,
                                 }
                             )
                             perf_stats.append(row)
@@ -97,9 +101,15 @@ def upload_dynamo_perf_stats_to_rockset(
         required=True,
         help="which GitHub repo this workflow run belongs to",
     )
+    parser.add_argument(
+        "--head-branch",
+        type=str,
+        required=True,
+        help="Head branch of the workflow",
+    )
     args = parser.parse_args()
     perf_stats = upload_dynamo_perf_stats_to_rockset(
-        args.repo, args.workflow_run_id, args.workflow_run_attempt
+        args.repo, args.workflow_run_id, args.workflow_run_attempt, args.head_branch
     )
     upload_to_rockset(
         collection="torch_dynamo_perf_stats",