pytorch

1
import argparse
2
import sys
3
from pathlib import Path
4

5
from pytest_caching_utils import (
6
    download_pytest_cache,
7
    GithubRepo,
8
    PRIdentifier,
9
    upload_pytest_cache,
10
)
11

12
TEMP_DIR = "./tmp"  # a backup location in case one isn't provided
13

14

15
def parse_args() -> argparse.Namespace:
16
    parser = argparse.ArgumentParser(
17
        description="Upload this job's the pytest cache to S3"
18
    )
19

20
    mode = parser.add_mutually_exclusive_group(required=True)
21
    mode.add_argument(
22
        "--upload", action="store_true", help="Upload the pytest cache to S3"
23
    )
24
    mode.add_argument(
25
        "--download",
26
        action="store_true",
27
        help="Download the pytest cache from S3, merging it with any local cache",
28
    )
29

30
    parser.add_argument(
31
        "--cache_dir",
32
        required=True,
33
        help="Path to the folder pytest uses for its cache",
34
    )
35
    parser.add_argument("--pr_identifier", required=True, help="A unique PR identifier")
36
    parser.add_argument(
37
        "--job_identifier",
38
        required=True,
39
        help="A unique job identifier that should be the same for all runs of job",
40
    )
41
    parser.add_argument(
42
        "--sha", required="--upload" in sys.argv, help="SHA of the commit"
43
    )  # Only required for upload
44
    parser.add_argument(
45
        "--test_config", required="--upload" in sys.argv, help="The test config"
46
    )  # Only required for upload
47
    parser.add_argument(
48
        "--shard", required="--upload" in sys.argv, help="The shard id"
49
    )  # Only required for upload
50

51
    parser.add_argument(
52
        "--repo",
53
        required=False,
54
        help="The github repository we're running in, in the format 'owner/repo-name'",
55
    )
56
    parser.add_argument(
57
        "--temp_dir", required=False, help="Directory to store temp files"
58
    )
59
    parser.add_argument(
60
        "--bucket", required=False, help="The S3 bucket to upload the cache to"
61
    )
62

63
    args = parser.parse_args()
64

65
    return args
66

67

68
def main() -> None:
69
    args = parse_args()
70

71
    pr_identifier = PRIdentifier(args.pr_identifier)
72
    print(f"PR identifier for `{args.pr_identifier}` is `{pr_identifier}`")
73

74
    repo = GithubRepo.from_string(args.repo)
75
    cache_dir = Path(args.cache_dir)
76
    if args.temp_dir:
77
        temp_dir = Path(args.temp_dir)
78
    else:
79
        temp_dir = Path(TEMP_DIR)
80

81
    if args.upload:
82
        print(f"Uploading cache with args {args}")
83

84
        # verify the cache dir exists
85
        if not cache_dir.exists():
86
            print(f"The pytest cache dir `{cache_dir}` does not exist. Skipping upload")
87
            return
88

89
        upload_pytest_cache(
90
            pr_identifier=pr_identifier,
91
            repo=repo,
92
            job_identifier=args.job_identifier,
93
            sha=args.sha,
94
            test_config=args.test_config,
95
            shard=args.shard,
96
            cache_dir=cache_dir,
97
            bucket=args.bucket,
98
            temp_dir=temp_dir,
99
        )
100

101
    if args.download:
102
        print(f"Downloading cache with args {args}")
103
        download_pytest_cache(
104
            pr_identifier=pr_identifier,
105
            repo=repo,
106
            job_identifier=args.job_identifier,
107
            dest_cache_dir=cache_dir,
108
            bucket=args.bucket,
109
            temp_dir=temp_dir,
110
        )
111

112

113
if __name__ == "__main__":
114
    main()
115