pytorch

1
import argparse
2
import os
3
import re
4
from tempfile import TemporaryDirectory
5

6
from tools.stats.upload_stats_lib import download_gha_artifacts, upload_file_to_s3
7

8
ARTIFACTS = [
9
    "sccache-stats",
10
    "test-jsons",
11
    "test-reports",
12
    "usage-log",
13
]
14
BUCKET_NAME = "gha-artifacts"
15
FILENAME_REGEX = r"-runattempt\d+"
16

17

18
def get_artifacts(repo: str, workflow_run_id: int, workflow_run_attempt: int) -> None:
19
    with TemporaryDirectory() as temp_dir:
20
        print("Using temporary directory:", temp_dir)
21
        os.chdir(temp_dir)
22

23
        for artifact in ARTIFACTS:
24
            artifact_paths = download_gha_artifacts(
25
                artifact, workflow_run_id, workflow_run_attempt
26
            )
27

28
            for artifact_path in artifact_paths:
29
                # GHA artifact is named as follows: NAME-runattempt${{ github.run_attempt }}-SUFFIX.zip
30
                # and we want remove the run_attempt to conform with the naming convention on S3, i.e.
31
                # pytorch/pytorch/WORKFLOW_ID/RUN_ATTEMPT/artifact/NAME-SUFFIX.zip
32
                s3_filename = re.sub(FILENAME_REGEX, "", artifact_path.name)
33
                upload_file_to_s3(
34
                    file_name=str(artifact_path.resolve()),
35
                    bucket=BUCKET_NAME,
36
                    key=f"{repo}/{workflow_run_id}/{workflow_run_attempt}/artifact/{s3_filename}",
37
                )
38

39

40
if __name__ == "__main__":
41
    parser = argparse.ArgumentParser(description="Upload test artifacts from GHA to S3")
42
    parser.add_argument(
43
        "--workflow-run-id",
44
        type=int,
45
        required=True,
46
        help="id of the workflow to get artifacts from",
47
    )
48
    parser.add_argument(
49
        "--workflow-run-attempt",
50
        type=int,
51
        required=True,
52
        help="which retry of the workflow this is",
53
    )
54
    parser.add_argument(
55
        "--repo",
56
        type=str,
57
        required=True,
58
        help="which GitHub repo this workflow run belongs to",
59
    )
60
    args = parser.parse_args()
61
    get_artifacts(args.repo, args.workflow_run_id, args.workflow_run_attempt)
62