pytorch

1
import argparse
2
import json
3
import os
4
from pathlib import Path
5
from tempfile import TemporaryDirectory
6
from typing import Any, Dict, List
7

8
from tools.stats.upload_stats_lib import (
9
    download_gha_artifacts,
10
    download_s3_artifacts,
11
    unzip,
12
    upload_to_rockset,
13
)
14

15

16
def get_sccache_stats(
17
    workflow_run_id: int, workflow_run_attempt: int
18
) -> List[Dict[str, Any]]:
19
    with TemporaryDirectory() as temp_dir:
20
        print("Using temporary directory:", temp_dir)
21
        os.chdir(temp_dir)
22

23
        # Download and extract all the reports (both GHA and S3)
24
        download_s3_artifacts("sccache-stats", workflow_run_id, workflow_run_attempt)
25

26
        artifact_paths = download_gha_artifacts(
27
            "sccache-stats", workflow_run_id, workflow_run_attempt
28
        )
29
        for path in artifact_paths:
30
            unzip(path)
31

32
        stats_jsons = []
33
        for json_file in Path(".").glob("**/*.json"):
34
            with open(json_file) as f:
35
                stats_jsons.append(json.load(f))
36
        return stats_jsons
37

38

39
if __name__ == "__main__":
40
    parser = argparse.ArgumentParser(description="Upload test stats to Rockset")
41
    parser.add_argument(
42
        "--workflow-run-id",
43
        type=int,
44
        required=True,
45
        help="id of the workflow to get artifacts from",
46
    )
47
    parser.add_argument(
48
        "--workflow-run-attempt",
49
        type=int,
50
        required=True,
51
        help="which retry of the workflow this is",
52
    )
53
    args = parser.parse_args()
54
    stats = get_sccache_stats(args.workflow_run_id, args.workflow_run_attempt)
55
    upload_to_rockset("sccache_stats", stats)
56