From d88a8c41d5e71e438bbc008178133bd01e4a7232 Mon Sep 17 00:00:00 2001
From: Benjamin Glass <bglass@quansight.com>
Date: Mon, 30 Dec 2024 17:15:31 +0000
Subject: [PATCH] Fix flaky "Upload test stats" job (#143991)

Test stat uploading was intermittently failing due to certain XML strings being opportunistically converted to numbers, when string output was expected. This PR makes the conversion behavior optional, which should fix the stat uploads.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/143991
Approved by: https://github.com/clee2000, https://github.com/huydhn
---
 tools/stats/check_disabled_tests.py |  3 ++-
 tools/stats/upload_test_stats.py    | 22 ++++++++++++----------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/tools/stats/check_disabled_tests.py b/tools/stats/check_disabled_tests.py
index 024795b6e31..d49d9764322 100644
--- a/tools/stats/check_disabled_tests.py
+++ b/tools/stats/check_disabled_tests.py
@@ -42,7 +42,8 @@ def process_report(
     all_tests: dict[str, dict[str, int]] = {}
 
     for test_case in root.iter(TESTCASE_TAG):
-        parsed_test_case = process_xml_element(test_case)
+        # Parse the test case as string values only.
+        parsed_test_case = process_xml_element(test_case, output_numbers=False)
 
         # Under --rerun-disabled-tests mode, a test is skipped when:
         # * it's skipped explicitly inside PyTorch code
diff --git a/tools/stats/upload_test_stats.py b/tools/stats/upload_test_stats.py
index d436f69b739..21644476972 100644
--- a/tools/stats/upload_test_stats.py
+++ b/tools/stats/upload_test_stats.py
@@ -56,7 +56,9 @@ def parse_xml_report(
     return test_cases
 
 
-def process_xml_element(element: ET.Element) -> dict[str, Any]:
+def process_xml_element(
+    element: ET.Element, output_numbers: bool = True
+) -> dict[str, Any]:
     """Convert a test suite element into a JSON-serializable dict."""
     ret: dict[str, Any] = {}
 
@@ -69,15 +71,15 @@ def process_xml_element(element: ET.Element) -> dict[str, Any]:
 
     # The XML format encodes all values as strings. Convert to ints/floats if
     # possible to make aggregation possible in SQL.
-    for k, v in ret.items():
-        try:
-            ret[k] = int(v)
-        except ValueError:
-            pass
-        try:
-            ret[k] = float(v)
-        except ValueError:
-            pass
+    if output_numbers:
+        for k, v in ret.items():
+            try:
+                ret[k] = int(v)
+            except ValueError:
+                try:
+                    ret[k] = float(v)
+                except ValueError:
+                    pass
 
     # Convert inner and outer text into special dict elements.
     # e.g.