[testing] upload test stats: Add info to the invoking file summary and some other changes (#164016)

* Changes some internal logic for grouping so hopefully it's slightly less annoying write code for
* Changes the invoking file summary to just use file, which I think is correct most of the time
* Adds some fields to the file summary, like skips, errors, etc so I can reuse it for file report regression things

Output should be the same, maybe with slightly more fields since I got rid of some of the pops

Pull Request resolved: https://github.com/pytorch/pytorch/pull/164016
Approved by: https://github.com/huydhn

tools/stats/test_dashboard.py

@@ -89,74 +89,101 @@ def get_td_exclusions(
         return grouped_tests
 
 
-def group_test_cases(test_cases: list[dict[str, Any]]) -> dict[str, Any]:
+def group_test_cases(test_cases: list[dict[str, Any]]) -> list[list[dict[str, Any]]]:
+    # Returns a list of lists. Each inner list contains test cases with the same
+    # build name, test config, file name, class name, and test name (ex if it was run multiple times)
     start = time.time()
-    grouped_tests: dict[str, Any] = defaultdict(
-        lambda: defaultdict(
-            lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
-        )
-    )
+    test_case_with_job_info = defaultdict(list)
+
     for test_case in test_cases:
         job_name = get_job_name(test_case["job_id"])
         build_name = get_build_name(job_name)
         if "bazel" in build_name:
             continue
         test_config = get_test_config(job_name)
-        class_name = test_case.pop("classname", "NoClass")
-        name = test_case.pop("name", "NoName")
-        invoking_file = test_case.pop("invoking_file", "NoFile")
-        invoking_file = invoking_file.replace(".", "/")
-        test_case.pop("workflow_id")
-        test_case.pop("workflow_run_attempt")
-        grouped_tests[build_name][test_config][invoking_file][class_name][name].append(
-            test_case
+
+        test_case["job_name"] = job_name
+        test_case["build_name"] = build_name
+        test_case["test_config"] = test_config
+
+        key = (
+            build_name,
+            test_config,
+            test_case.get("file", "NoFile"),
+            test_case.get("classname", "NoClass"),
+            test_case.get("name", "NoName"),
         )
+        test_case_with_job_info[key].append(test_case)
 
     print(f"Time taken to group tests: {time.time() - start}")
-    return grouped_tests
+    return list(test_case_with_job_info.values())
 
 
-def get_reruns(grouped_tests: dict[str, Any]) -> dict[str, Any]:
+def get_reruns(grouped_tests: list[list[dict[str, Any]]]) -> dict[str, Any]:
     reruns: dict[str, Any] = defaultdict(
         lambda: defaultdict(
             lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
         )
     )
-    for build_name, build in grouped_tests.items():
-        for test_config, test_config_data in build.items():
-            for invoking_file, invoking_file_data in test_config_data.items():
-                for class_name, class_data in invoking_file_data.items():
-                    for test_name, test_data in class_data.items():
-                        if len(test_data) > 1:
-                            if invoking_file in (
-                                "distributed/test_distributed_spawn",
-                                "onnx/test_fx_to_onnx_with_onnxruntime",
-                                "distributed/algorithms/quantization/test_quantization",
-                            ):
-                                continue
-                            reruns[build_name][test_config][invoking_file][class_name][
-                                test_name
-                            ] = test_data
+
+    for tests in grouped_tests:
+        if len(tests) > 1:
+            build_name = tests[0]["build_name"]
+            test_config = tests[0]["test_config"]
+            file = tests[0].get("file", "NoFile")
+            class_name = tests[0].get("classname", "NoClass")
+            test_name = tests[0].get("name", "NoName")
+            if file in (
+                "distributed/test_distributed_spawn.py",
+                "onnx/test_fx_to_onnx_with_onnxruntime.py",
+                "distributed/algorithms/quantization/test_quantization.py",
+            ):
+                continue
+            reruns[build_name][test_config][file][class_name][test_name] = tests
+
     return reruns
 
 
-def get_invoking_file_summary(grouped_tests: dict[str, Any]) -> dict[str, Any]:
+def get_invoking_file_summary(
+    grouped_tests: list[list[dict[str, Any]]],
+) -> dict[str, Any]:
+    summary_flat = {}
+    for tests in grouped_tests:
+        build_name = tests[0]["build_name"]
+        test_config = tests[0]["test_config"]
+        short_job_name = f"{build_name} / ({test_config})"
+        file = tests[0].get("file", "NoFile")
+
+        key = (build_name, test_config, file)
+        if key not in summary_flat:
+            summary_flat[key] = {
+                "count": 0,
+                "time": 0.0,
+                "skipped": 0,
+                "failures": 0,
+                "errors": 0,
+                "successes": 0,
+                "short_job_name": short_job_name,
+                "file": file,
+            }
+        summary_flat[key]["count"] += 1
+        status = "successes"
+        for test in tests:
+            summary_flat[key]["time"] += test["time"]
+            if "skipped" in test:
+                status = "skipped"
+            elif "failure" in test:
+                status = "failures"
+            elif "error" in test:
+                status = "errors"
+        summary_flat[key][status] += 1
+
     invoking_file_summary: dict[str, Any] = defaultdict(
         lambda: defaultdict(lambda: defaultdict(lambda: {"count": 0, "time": 0.0}))
     )
-    for build_name, build in grouped_tests.items():
-        for test_config, test_config_data in build.items():
-            for invoking_file, invoking_file_data in test_config_data.items():
-                for class_data in invoking_file_data.values():
-                    for test_data in class_data.values():
-                        invoking_file_summary[build_name][test_config][invoking_file][
-                            "count"
-                        ] += 1
-                        for i in test_data:
-                            invoking_file_summary[build_name][test_config][
-                                invoking_file
-                            ]["time"] += i["time"]
 
+    for (build_name, test_config, file), data in summary_flat.items():
+        invoking_file_summary[build_name][test_config][file] = data
     return invoking_file_summary