Write flaky tests to rockset (#70136)
Summary:
Try using Rockset as backend for data instead of RDS
Pull Request resolved: https://github.com/pytorch/pytorch/pull/70136
Reviewed By: suo
Differential Revision: D33242148
Pulled By: janeyx99
fbshipit-source-id: 8935ceb43717fff4922b634165030cca7e934968
diff --git a/tools/stats/print_test_stats.py b/tools/stats/print_test_stats.py
index 8faacbc..3364957 100755
--- a/tools/stats/print_test_stats.py
+++ b/tools/stats/print_test_stats.py
@@ -794,6 +794,7 @@
def assemble_flaky_test_stats(duplicated_tests_by_file: Dict[str, DuplicatedDict]) -> Any:
flaky_tests = []
+ workflow_id = os.environ.get("WORKFLOW_ID", os.environ.get("CIRCLE_WORKFLOW_ID", None))
for file_name, suite_to_dict in duplicated_tests_by_file.items():
for suite_name, testcase_to_runs in suite_to_dict.items():
for testcase_name, list_of_runs in testcase_to_runs.items():
@@ -804,12 +805,21 @@
"suite": suite_name,
"file": file_name,
"num_green": num_green,
- "num_red": num_red
+ "num_red": num_red,
})
if len(flaky_tests) > 0:
+ # write to RDS
register_rds_schema("flaky_tests", schema_from_sample(flaky_tests[0]))
rds_write("flaky_tests", flaky_tests, only_on_master=False)
+ # write to S3 to go to Rockset as well
+ import uuid
+ for flaky_test in flaky_tests:
+ flaky_test["workflow_id"] = workflow_id
+ key = f"flaky_tests/{workflow_id}/{uuid.uuid4()}.json"
+ obj = get_S3_object_from_bucket("ossci-raw-job-status", key)
+ obj.put(Body=json.dumps(flaky_test), ContentType="application/json")
+
def build_info() -> ReportMetaMeta:
return {
