xref: /aosp_15_r20/external/pytorch/ios/TestApp/run_on_aws_devicefarm.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1#!/usr/bin/env python3
2
3import datetime
4import os
5import random
6import string
7import sys
8import time
9import warnings
10from typing import Any
11
12import boto3
13import requests
14
15
16POLLING_DELAY_IN_SECOND = 5
17MAX_UPLOAD_WAIT_IN_SECOND = 600
18
19# NB: This is the curated top devices from AWS. We could create our own device
20# pool if we want to
21DEFAULT_DEVICE_POOL_ARN = (
22    "arn:aws:devicefarm:us-west-2::devicepool:082d10e5-d7d7-48a5-ba5c-b33d66efa1f5"
23)
24
25
26def parse_args() -> Any:
27    from argparse import ArgumentParser
28
29    parser = ArgumentParser("Run iOS tests on AWS Device Farm")
30    parser.add_argument(
31        "--project-arn", type=str, required=True, help="the ARN of the project on AWS"
32    )
33    parser.add_argument(
34        "--app-file", type=str, required=True, help="the iOS ipa app archive"
35    )
36    parser.add_argument(
37        "--xctest-file", type=str, required=True, help="the XCTest suite to run"
38    )
39    parser.add_argument(
40        "--name-prefix",
41        type=str,
42        required=True,
43        help="the name prefix of this test run",
44    )
45    parser.add_argument(
46        "--device-pool-arn",
47        type=str,
48        default=DEFAULT_DEVICE_POOL_ARN,
49        help="the name of the device pool to test on",
50    )
51
52    return parser.parse_args()
53
54
55def upload_file(
56    client: Any,
57    project_arn: str,
58    prefix: str,
59    filename: str,
60    filetype: str,
61    mime: str = "application/octet-stream",
62):
63    """
64    Upload the app file and XCTest suite to AWS
65    """
66    r = client.create_upload(
67        projectArn=project_arn,
68        name=f"{prefix}_{os.path.basename(filename)}",
69        type=filetype,
70        contentType=mime,
71    )
72    upload_name = r["upload"]["name"]
73    upload_arn = r["upload"]["arn"]
74    upload_url = r["upload"]["url"]
75
76    with open(filename, "rb") as file_stream:
77        print(f"Uploading {filename} to Device Farm as {upload_name}...")
78        r = requests.put(upload_url, data=file_stream, headers={"content-type": mime})
79        if not r.ok:
80            raise Exception(f"Couldn't upload {filename}: {r.reason}")  # noqa: TRY002
81
82    start_time = datetime.datetime.now()
83    # Polling AWS till the uploaded file is ready
84    while True:
85        waiting_time = datetime.datetime.now() - start_time
86        if waiting_time > datetime.timedelta(seconds=MAX_UPLOAD_WAIT_IN_SECOND):
87            raise Exception(  # noqa: TRY002
88                f"Uploading {filename} is taking longer than {MAX_UPLOAD_WAIT_IN_SECOND} seconds, terminating..."
89            )
90
91        r = client.get_upload(arn=upload_arn)
92        status = r["upload"].get("status", "")
93
94        print(f"{filename} is in state {status} after {waiting_time}")
95
96        if status == "FAILED":
97            raise Exception(f"Couldn't upload {filename}: {r}")  # noqa: TRY002
98        if status == "SUCCEEDED":
99            break
100
101        time.sleep(POLLING_DELAY_IN_SECOND)
102
103    return upload_arn
104
105
106def main() -> None:
107    args = parse_args()
108
109    client = boto3.client("devicefarm")
110    unique_prefix = f"{args.name_prefix}-{datetime.date.today().isoformat()}-{''.join(random.sample(string.ascii_letters, 8))}"
111
112    # Upload the test app
113    appfile_arn = upload_file(
114        client=client,
115        project_arn=args.project_arn,
116        prefix=unique_prefix,
117        filename=args.app_file,
118        filetype="IOS_APP",
119    )
120    print(f"Uploaded app: {appfile_arn}")
121    # Upload the XCTest suite
122    xctest_arn = upload_file(
123        client=client,
124        project_arn=args.project_arn,
125        prefix=unique_prefix,
126        filename=args.xctest_file,
127        filetype="XCTEST_TEST_PACKAGE",
128    )
129    print(f"Uploaded XCTest: {xctest_arn}")
130
131    # Schedule the test
132    r = client.schedule_run(
133        projectArn=args.project_arn,
134        name=unique_prefix,
135        appArn=appfile_arn,
136        devicePoolArn=args.device_pool_arn,
137        test={"type": "XCTEST", "testPackageArn": xctest_arn},
138    )
139    run_arn = r["run"]["arn"]
140
141    start_time = datetime.datetime.now()
142    print(f"Run {unique_prefix} is scheduled as {run_arn}:")
143
144    state = "UNKNOWN"
145    result = ""
146    try:
147        while True:
148            r = client.get_run(arn=run_arn)
149            state = r["run"]["status"]
150
151            if state == "COMPLETED":
152                result = r["run"]["result"]
153                break
154
155            waiting_time = datetime.datetime.now() - start_time
156            print(
157                f"Run {unique_prefix} in state {state} after {datetime.datetime.now() - start_time}"
158            )
159            time.sleep(30)
160    except Exception as error:
161        warnings.warn(f"Failed to run {unique_prefix}: {error}")
162        sys.exit(1)
163
164    if not result or result == "FAILED":
165        print(f"Run {unique_prefix} failed, exiting...")
166        sys.exit(1)
167
168
169if __name__ == "__main__":
170    main()
171