[disablebot] Allow reading of disable issue for multiple tests

.github/scripts/test_update_disabled_issues.py

@@ -2,9 +2,9 @@
 
 from update_disabled_issues import (
     condense_disable_jobs,
-    condense_disable_tests,
     filter_disable_issues,
     get_disable_issues,
+    get_disabled_tests,
     OWNER,
     REPO,
     UNSTABLE_PREFIX,
@@ -98,13 +98,13 @@ def test_filter_disable_issues(self, mock_get_disable_issues):
             [item["number"] for item in disabled_jobs], [32132, 42345, 94861]
         )
 
-    def test_condense_disable_tests(self, mock_get_disable_issues):
+    def test_get_disable_tests(self, mock_get_disable_issues):
         mock_get_disable_issues.return_value = MOCK_DATA
 
         disabled_issues = get_disable_issues("dummy token")
 
         disabled_tests, _ = filter_disable_issues(disabled_issues)
-        results = condense_disable_tests(disabled_tests)
+        results = get_disabled_tests(disabled_tests)
 
         self.assertDictEqual(
             {
@@ -125,6 +125,86 @@ def test_condense_disable_tests(self, mock_get_disable_issues):
             results,
         )
 
+    def test_get_disable_tests_aggregate_issue(self, mock_get_disable_issues):
+        # Test that the function can read aggregate issues
+        self.maxDiff = None
+        mock_data = [
+            {
+                "url": "https://github.com/pytorch/pytorch/issues/32644",
+                "number": 32644,
+                "title": "DISABLED MULTIPLE dummy test",
+                "body": "disable the following tests:\n```\ntest_quantized_nn (test_quantization.PostTrainingDynamicQuantTest): mac, win\ntest_zero_redundancy_optimizer (__main__.TestZeroRedundancyOptimizerDistributed)\n```",
+            }
+        ]
+        disabled_tests = get_disabled_tests(mock_data)
+        self.assertDictEqual(
+            {
+                "test_quantized_nn (test_quantization.PostTrainingDynamicQuantTest)": (
+                    str(mock_data[0]["number"]),
+                    mock_data[0]["url"],
+                    ["mac", "win"],
+                ),
+                "test_zero_redundancy_optimizer (__main__.TestZeroRedundancyOptimizerDistributed)": (
+                    str(mock_data[0]["number"]),
+                    mock_data[0]["url"],
+                    [],
+                ),
+            },
+            disabled_tests,
+        )
+
+    def test_get_disable_tests_merge_issues(self, mock_get_disable_issues):
+        # Test that the function can merge multiple issues with the same test
+        # name
+        self.maxDiff = None
+        mock_data = [
+            {
+                "url": "https://github.com/pytorch/pytorch/issues/32644",
+                "number": 32644,
+                "title": "DISABLED MULTIPLE dummy test",
+                "body": "disable the following tests:\n```\ntest_2 (abc.ABC): mac, win\ntest_3 (DEF)\n```",
+            },
+            {
+                "url": "https://github.com/pytorch/pytorch/issues/32645",
+                "number": 32645,
+                "title": "DISABLED MULTIPLE dummy test",
+                "body": "disable the following tests:\n```\ntest_2 (abc.ABC): mac, win, linux\ntest_3 (DEF): mac\n```",
+            },
+            {
+                "url": "https://github.com/pytorch/pytorch/issues/32646",
+                "number": 32646,
+                "title": "DISABLED test_1 (__main__.Test1)",
+                "body": "platforms: linux",
+            },
+            {
+                "url": "https://github.com/pytorch/pytorch/issues/32647",
+                "number": 32647,
+                "title": "DISABLED test_2 (abc.ABC)",
+                "body": "platforms: dynamo",
+            },
+        ]
+        disabled_tests = get_disabled_tests(mock_data)
+        self.assertDictEqual(
+            {
+                "test_2 (abc.ABC)": (
+                    str(mock_data[3]["number"]),
+                    mock_data[3]["url"],
+                    ["dynamo", "linux", "mac", "win"],
+                ),
+                "test_3 (DEF)": (
+                    str(mock_data[1]["number"]),
+                    mock_data[1]["url"],
+                    [],
+                ),
+                "test_1 (__main__.Test1)": (
+                    str(mock_data[2]["number"]),
+                    mock_data[2]["url"],
+                    ["linux"],
+                ),
+            },
+            disabled_tests,
+        )
+
     def test_condense_disable_jobs(self, mock_get_disable_issues):
         mock_get_disable_issues.return_value = MOCK_DATA
 

.github/scripts/update_disabled_issues.py

@@ -19,6 +19,7 @@
 DISABLED_PREFIX = "DISABLED"
 UNSTABLE_PREFIX = "UNSTABLE"
 DISABLED_TEST_ISSUE_TITLE = re.compile(r"DISABLED\s*test_.+\s*\(.+\)")
+DISABLED_TEST_MULTI_ISSUE_TITLE = re.compile(r"DISABLED MULTIPLE")
 JOB_NAME_MAXSPLIT = 2
 
 OWNER = "pytorch"
@@ -122,14 +123,108 @@ def filter_disable_issues(
         if not title or not title.startswith(prefix):
             continue
 
-        if DISABLED_TEST_ISSUE_TITLE.match(title):
+        if DISABLED_TEST_ISSUE_TITLE.match(
+            title
+        ) or DISABLED_TEST_MULTI_ISSUE_TITLE.match(title):
             disable_test_issues.append(issue)
         else:
             disable_job_issues.append(issue)
 
     return disable_test_issues, disable_job_issues
 
 
+def get_disabled_tests(issues: List[Dict[str, Any]]) -> Dict[str, Tuple]:
+    def get_platforms_to_skip(body: str, prefix: str) -> List[str]:
+        # Empty list = all platforms should skip the test
+        platforms_to_skip = []
+        if body is not None:
+            for line in body.splitlines():
+                line = line.lower()
+                if line.startswith(prefix):
+                    platforms_to_skip.extend(
+                        [x.strip() for x in line[len(prefix) :].split(",") if x.strip()]
+                    )
+        return platforms_to_skip
+
+    disabled_tests = {}
+
+    def update_disabled_tests(
+        key: str, number: str, url: str, platforms_to_skip: List[str]
+    ):
+        # merge the list of platforms to skip if the test is disabled by
+        # multiple issues.  This results in some urls being wrong
+        if key not in disabled_tests:
+            disabled_tests[key] = (number, url, platforms_to_skip)
+        else:
+            original_platforms = disabled_tests[key][2]
+            if len(original_platforms) == 0 or len(platforms_to_skip) == 0:
+                platforms = []
+            else:
+                platforms = sorted(set(original_platforms + platforms_to_skip))
+            disabled_tests[key] = (
+                number,
+                url,
+                platforms,
+            )
+
+    test_name_regex = re.compile(r"(test_[a-zA-Z0-9-_\.]+)\s+\(([a-zA-Z0-9-_\.]+)\)")
+
+    def parse_test_name(s: str) -> Optional[str]:
+        test_name_match = test_name_regex.match(s)
+        if test_name_match:
+            return f"{test_name_match.group(1)} ({test_name_match.group(2)})"
+        return None
+
+    for issue in issues:
+        try:
+            url = issue["url"]
+            number = url.split("/")[-1]
+            title = issue["title"].strip()
+            body = issue["body"]
+
+            test_name = parse_test_name(title[len("DISABLED") :].strip())
+            if test_name is not None:
+                update_disabled_tests(
+                    test_name, number, url, get_platforms_to_skip(body, "platforms:")
+                )
+            elif DISABLED_TEST_MULTI_ISSUE_TITLE.match(title):
+                # This is a multi-test issue
+                start = body.lower().find("disable the following tests:")
+                # Format for disabling tests:
+                # Title: DISABLED MULTIPLE anything
+                # disable the following tests:
+                # ```
+                # test_name1 (test_suite1): mac, windows
+                # test_name2 (test_suite2): mac, windows
+                # ```
+                for line in body[start:].splitlines()[2:]:
+                    if "```" in line:
+                        break
+                    split_by_colon = line.split(":")
+
+                    test_name = parse_test_name(split_by_colon[0].strip())
+                    if test_name is None:
+                        continue
+                    update_disabled_tests(
+                        test_name,
+                        number,
+                        url,
+                        get_platforms_to_skip(
+                            split_by_colon[1].strip()
+                            if len(split_by_colon) > 1
+                            else "",
+                            "",
+                        ),
+                    )
+            else:
+                print(f"Unknown disable issue type: {title}")
+        except Exception as e:
+            print(f"Failed to parse issue {issue['url']}: {e}")
+            continue
+
+    return disabled_tests
+
+
 @lru_cache()
 def can_disable_jobs(owner: str, repo: str, username: str, token: str) -> bool:
     url = f"https://api.github.com/repos/{owner}/{repo}/collaborators/{username}/permission"
@@ -146,38 +241,6 @@ def can_disable_jobs(owner: str, repo: str, username: str, token: str) -> bool:
     return perm and perm.get("permission", "").lower() in PERMISSIONS_TO_DISABLE_JOBS
 
 
-def condense_disable_tests(
-    disable_issues: List[Dict[str, Any]],
-) -> Dict[str, Tuple]:
-    disabled_test_from_issues = {}
-    for item in disable_issues:
-        issue_url = item["url"]
-        issue_number = issue_url.split("/")[-1]
-
-        title = item["title"]
-        test_name = title[len(DISABLED_PREFIX) :].strip()
-
-        body = item["body"]
-        platforms_to_skip = []
-        key = "platforms:"
-        # When the issue has no body, it is assumed that all platforms should skip the test
-        if body is not None:
-            for line in body.splitlines():
-                line = line.lower()
-                if line.startswith(key):
-                    platforms_to_skip.extend(
-                        [x.strip() for x in line[len(key) :].split(",") if x.strip()]
-                    )
-
-        disabled_test_from_issues[test_name] = (
-            issue_number,
-            issue_url,
-            platforms_to_skip,
-        )
-
-    return disabled_test_from_issues
-
-
 def condense_disable_jobs(
     disable_issues: List[Any],
     owner: str,
@@ -253,9 +316,7 @@ def main() -> None:
     disable_test_issues, disable_job_issues = filter_disable_issues(disable_issues)
     # Create the list of disabled tests taken into account the list of disabled issues
     # and those that are not flaky anymore
-    dump_json(
-        condense_disable_tests(disable_test_issues), "disabled-tests-condensed.json"
-    )
+    dump_json(get_disabled_tests(disable_test_issues), "disabled-tests-condensed.json")
     dump_json(
         condense_disable_jobs(disable_job_issues, args.owner, args.repo, token),
         "disabled-jobs.json",