pytorch

1
#!/usr/bin/env python3
2

3
import datetime
4
import os
5
import random
6
import string
7
import sys
8
import time
9
import warnings
10
from typing import Any
11

12
import boto3
13
import requests
14

15
POLLING_DELAY_IN_SECOND = 5
16
MAX_UPLOAD_WAIT_IN_SECOND = 600
17

18
# NB: This is the curated top devices from AWS. We could create our own device
19
# pool if we want to
20
DEFAULT_DEVICE_POOL_ARN = (
21
    "arn:aws:devicefarm:us-west-2::devicepool:082d10e5-d7d7-48a5-ba5c-b33d66efa1f5"
22
)
23

24

25
def parse_args() -> Any:
26
    from argparse import ArgumentParser
27

28
    parser = ArgumentParser("Run iOS tests on AWS Device Farm")
29
    parser.add_argument(
30
        "--project-arn", type=str, required=True, help="the ARN of the project on AWS"
31
    )
32
    parser.add_argument(
33
        "--app-file", type=str, required=True, help="the iOS ipa app archive"
34
    )
35
    parser.add_argument(
36
        "--xctest-file",
37
        type=str,
38
        required=True,
39
        help="the XCTest suite to run",
40
    )
41
    parser.add_argument(
42
        "--name-prefix",
43
        type=str,
44
        required=True,
45
        help="the name prefix of this test run",
46
    )
47
    parser.add_argument(
48
        "--device-pool-arn",
49
        type=str,
50
        default=DEFAULT_DEVICE_POOL_ARN,
51
        help="the name of the device pool to test on",
52
    )
53

54
    return parser.parse_args()
55

56

57
def upload_file(
58
    client: Any,
59
    project_arn: str,
60
    prefix: str,
61
    filename: str,
62
    filetype: str,
63
    mime: str = "application/octet-stream",
64
):
65
    """
66
    Upload the app file and XCTest suite to AWS
67
    """
68
    r = client.create_upload(
69
        projectArn=project_arn,
70
        name=f"{prefix}_{os.path.basename(filename)}",
71
        type=filetype,
72
        contentType=mime,
73
    )
74
    upload_name = r["upload"]["name"]
75
    upload_arn = r["upload"]["arn"]
76
    upload_url = r["upload"]["url"]
77

78
    with open(filename, "rb") as file_stream:
79
        print(f"Uploading {filename} to Device Farm as {upload_name}...")
80
        r = requests.put(upload_url, data=file_stream, headers={"content-type": mime})
81
        if not r.ok:
82
            raise Exception(f"Couldn't upload {filename}: {r.reason}")
83

84
    start_time = datetime.datetime.now()
85
    # Polling AWS till the uploaded file is ready
86
    while True:
87
        waiting_time = datetime.datetime.now() - start_time
88
        if waiting_time > datetime.timedelta(seconds=MAX_UPLOAD_WAIT_IN_SECOND):
89
            raise Exception(
90
                f"Uploading {filename} is taking longer than {MAX_UPLOAD_WAIT_IN_SECOND} seconds, terminating..."
91
            )
92

93
        r = client.get_upload(arn=upload_arn)
94
        status = r["upload"].get("status", "")
95

96
        print(f"{filename} is in state {status} after {waiting_time}")
97

98
        if status == "FAILED":
99
            raise Exception(f"Couldn't upload {filename}: {r}")
100
        if status == "SUCCEEDED":
101
            break
102

103
        time.sleep(POLLING_DELAY_IN_SECOND)
104

105
    return upload_arn
106

107

108
def main() -> None:
109
    args = parse_args()
110

111
    client = boto3.client("devicefarm")
112
    unique_prefix = f"{args.name_prefix}-{datetime.date.today().isoformat()}-{''.join(random.sample(string.ascii_letters, 8))}"
113

114
    # Upload the test app
115
    appfile_arn = upload_file(
116
        client=client,
117
        project_arn=args.project_arn,
118
        prefix=unique_prefix,
119
        filename=args.app_file,
120
        filetype="IOS_APP",
121
    )
122
    print(f"Uploaded app: {appfile_arn}")
123
    # Upload the XCTest suite
124
    xctest_arn = upload_file(
125
        client=client,
126
        project_arn=args.project_arn,
127
        prefix=unique_prefix,
128
        filename=args.xctest_file,
129
        filetype="XCTEST_TEST_PACKAGE",
130
    )
131
    print(f"Uploaded XCTest: {xctest_arn}")
132

133
    # Schedule the test
134
    r = client.schedule_run(
135
        projectArn=args.project_arn,
136
        name=unique_prefix,
137
        appArn=appfile_arn,
138
        devicePoolArn=args.device_pool_arn,
139
        test={"type": "XCTEST", "testPackageArn": xctest_arn},
140
    )
141
    run_arn = r["run"]["arn"]
142

143
    start_time = datetime.datetime.now()
144
    print(f"Run {unique_prefix} is scheduled as {run_arn}:")
145

146
    state = "UNKNOWN"
147
    result = ""
148
    try:
149
        while True:
150
            r = client.get_run(arn=run_arn)
151
            state = r["run"]["status"]
152

153
            if state == "COMPLETED":
154
                result = r["run"]["result"]
155
                break
156

157
            waiting_time = datetime.datetime.now() - start_time
158
            print(
159
                f"Run {unique_prefix} in state {state} after {datetime.datetime.now() - start_time}"
160
            )
161
            time.sleep(30)
162
    except Exception as error:
163
        warnings.warn(f"Failed to run {unique_prefix}: {error}")
164
        sys.exit(1)
165

166
    if not result or result == "FAILED":
167
        print(f"Run {unique_prefix} failed, exiting...")
168
        sys.exit(1)
169

170

171
if __name__ == "__main__":
172
    main()
173