Reflect check_labels status as a signal (pytorch#134711)

Fixes the workflow when meta-exported diff (co-dev) doesn't have the required labels, but the signal is suppressed due to job failure (e.g. [see this run](https://github.com/pytorch/pytorch/actions/runs/10590994706/job/29347663526?pr=134484)).

With this change the workflow status correctly reflects the status of the check.

# Testing
* [illegal pr_num](https://github.com/pytorch/pytorch/actions/runs/10603163898/job/29386843591)
* [successful run](https://github.com/pytorch/pytorch/actions/runs/10603279052/job/29387230110) (topic label present)
* no labels: [check fails](https://github.com/pytorch/pytorch/actions/runs/10603310368/job/29387333864)
Pull Request resolved: pytorch#134711
Approved by: https://github.com/clee2000

.github/scripts/check_labels.py

@@ -27,6 +27,12 @@ def parse_args() -> Any:
 
     parser = ArgumentParser("Check PR labels")
     parser.add_argument("pr_num", type=int)
+    # add a flag to return a non-zero exit code if the PR does not have the required labels
+    parser.add_argument(
+        "--exit-non-zero",
+        action="store_true",
+        help="Return a non-zero exit code if the PR does not have the required labels",
+    )
 
     return parser.parse_args()
 
@@ -41,10 +47,13 @@ def main() -> None:
         if not has_required_labels(pr):
             print(LABEL_ERR_MSG)
             add_label_err_comment(pr)
+            if args.exit_non_zero:
+                sys.exit(1)
         else:
             delete_all_label_err_comments(pr)
     except Exception as e:
-        pass
+        if args.exit_non_zero:
+            sys.exit(1)
 
     sys.exit(0)
 

.github/scripts/test_check_labels.py

@@ -18,6 +18,7 @@ def mock_parse_args() -> object:
     class Object:
         def __init__(self) -> None:
             self.pr_num = 76123
+            self.exit_non_zero = False
 
     return Object()
 

.github/workflows/check-labels.yml

@@ -19,6 +19,10 @@ on:
     branches: [gh/**/base]
 
   workflow_dispatch:
+    inputs:
+      pr_number:
+        description: 'PR number to check labels for'
+        required: true
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
@@ -54,7 +58,7 @@ jobs:
       - name: Check labels
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          PR_NUM: ${{ github.event.number }}
+          PR_NUM: ${{ github.event.number || github.event.inputs.pr_number }}
         run: |
           set -ex
-          python3 .github/scripts/check_labels.py "${PR_NUM}"
+          python3 .github/scripts/check_labels.py --exit-non-zero "${PR_NUM}"