xref: /aosp_15_r20/external/pytorch/tools/stats/upload_test_stats_intermediate.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1import argparse
2import sys
3
4from tools.stats.test_dashboard import upload_additional_info
5from tools.stats.upload_test_stats import get_tests
6
7
8if __name__ == "__main__":
9    parser = argparse.ArgumentParser(description="Upload test stats to Rockset")
10    parser.add_argument(
11        "--workflow-run-id",
12        required=True,
13        help="id of the workflow to get artifacts from",
14    )
15    parser.add_argument(
16        "--workflow-run-attempt",
17        type=int,
18        required=True,
19        help="which retry of the workflow this is",
20    )
21    args = parser.parse_args()
22
23    print(f"Workflow id is: {args.workflow_run_id}")
24
25    test_cases = get_tests(args.workflow_run_id, args.workflow_run_attempt)
26
27    # Flush stdout so that any errors in Rockset upload show up last in the logs.
28    sys.stdout.flush()
29
30    upload_additional_info(args.workflow_run_id, args.workflow_run_attempt, test_cases)
31