forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathpytest_cache.py
More file actions
115 lines (95 loc) · 3.21 KB
/
Copy pathpytest_cache.py
File metadata and controls
115 lines (95 loc) · 3.21 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import argparse
import sys
from pathlib import Path
from pytest_caching_utils import (
download_pytest_cache,
GithubRepo,
PRIdentifier,
upload_pytest_cache,
)
TEMP_DIR = "./tmp" # a backup location in case one isn't provided
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Upload this job's the pytest cache to S3"
)
mode = parser.add_mutually_exclusive_group(required=True)
mode.add_argument(
"--upload", action="store_true", help="Upload the pytest cache to S3"
)
mode.add_argument(
"--download",
action="store_true",
help="Download the pytest cache from S3, merging it with any local cache",
)
parser.add_argument(
"--cache_dir",
required=True,
help="Path to the folder pytest uses for its cache",
)
parser.add_argument("--pr_identifier", required=True, help="A unique PR identifier")
parser.add_argument(
"--job_identifier",
required=True,
help="A unique job identifier that should be the same for all runs of job",
)
parser.add_argument(
"--sha", required="--upload" in sys.argv, help="SHA of the commit"
) # Only required for upload
parser.add_argument(
"--test_config", required="--upload" in sys.argv, help="The test config"
) # Only required for upload
parser.add_argument(
"--shard", required="--upload" in sys.argv, help="The shard id"
) # Only required for upload
parser.add_argument(
"--repo",
required=False,
help="The github repository we're running in, in the format 'owner/repo-name'",
)
parser.add_argument(
"--temp_dir", required=False, help="Directory to store temp files"
)
parser.add_argument(
"--bucket", required=False, help="The S3 bucket to upload the cache to"
)
args = parser.parse_args()
return args
def main() -> None:
args = parse_args()
pr_identifier = PRIdentifier(args.pr_identifier)
print(f"PR identifier for `{args.pr_identifier}` is `{pr_identifier}`")
repo = GithubRepo.from_string(args.repo)
cache_dir = Path(args.cache_dir)
if args.temp_dir:
temp_dir = Path(args.temp_dir)
else:
temp_dir = Path(TEMP_DIR)
if args.upload:
print(f"Uploading cache with args {args}")
# verify the cache dir exists
if not cache_dir.exists():
print(f"The pytest cache dir `{cache_dir}` does not exist. Skipping upload")
return
upload_pytest_cache(
pr_identifier=pr_identifier,
repo=repo,
job_identifier=args.job_identifier,
sha=args.sha,
test_config=args.test_config,
shard=args.shard,
cache_dir=cache_dir,
bucket=args.bucket,
temp_dir=temp_dir,
)
if args.download:
print(f"Downloading cache with args {args}")
download_pytest_cache(
pr_identifier=pr_identifier,
repo=repo,
job_identifier=args.job_identifier,
dest_cache_dir=cache_dir,
bucket=args.bucket,
temp_dir=temp_dir,
)
if __name__ == "__main__":
main()