wandb

1
#!/usr/bin/env python
2
"""Tool to interact with CircleCI jobs using API key.
3

4
Get the current status of a circleci pipeline based on branch/commit
5
    ```
6
    $ ./circleci-tool status
7
    $ ./circleci-tool status --wait
8
    ```
9

10
Trigger (re)execution of a branch
11
    ```
12
    $ ./circleci-tool.py trigger
13
    $ ./circleci-tool.py trigger --wait
14
    $ ./circleci-tool.py trigger --platform mac
15
    $ ./circleci-tool.py trigger --platform mac --test-file tests/test.py
16
    $ ./circleci-tool.py trigger --platform win --test-file tests/test.py --test-name test_this
17
    $ ./circleci-tool.py trigger --platform win --test-file tests/test.py --test-name test_this --test-repeat 4
18
    $ ./circleci-tool.py trigger --toxenv py36,py37 --loop 3
19
    $ ./circleci-tool.py trigger --wait --platform win --test-file tests/test_notebooks.py --parallelism 5 --xdist 2
20
    $ ./circleci-tool.py trigger --toxenv func-s_service-py37 --loop 3
21

22
    ```
23

24
Trigger nightly run
25
    ```
26
    $ ./circleci-tool.py trigger-nightly --slack-notify
27
    ```
28

29
Download artifacts from an executed workflow
30
    ```
31
    $ ./circleci-tool download
32
    ```
33

34
"""
35

36

37
import argparse
38
import os
39
import subprocess
40
import sys
41
import time
42

43
import requests
44

45
CIRCLECI_API_TOKEN = "CIRCLECI_TOKEN"
46

47
NIGHTLY_SHARDS = (
48
    "standalone-cpu",
49
    "standalone-gpu",
50
    "kfp",
51
    "standalone-gpu-win",
52
    "regression",
53
)
54

55
platforms_dict = dict(linux="test", lin="test", mac="mac", win="win")
56
platforms_short_dict = dict(linux="lin", lin="lin", mac="mac", win="win")
57
py_name_dict = dict(
58
    py37="py37",
59
    py38="py38",
60
    py39="py39",
61
    py310="py310",
62
    py311="py311",
63
)
64
py_image_dict = dict(
65
    py37="python:3.7",
66
    py38="python:3.8",
67
    py39="python:3.9",
68
    py310="python:3.10",
69
    py311="python:3.11",
70
)
71

72

73
def poll(args, pipeline_id=None, workflow_ids=None):
74
    print(f"Waiting for pipeline to complete (Branch: {args.branch})...")
75
    while True:
76
        num = 0
77
        done = 0
78
        if pipeline_id:
79
            url = f"https://circleci.com/api/v2/pipeline/{pipeline_id}/workflow"
80
            r = requests.get(url, auth=(args.api_token, ""))
81
            assert r.status_code == 200, f"Error making api request: {r}"
82
            d = r.json()
83
            workflow_ids = [item["id"] for item in d["items"]]
84
        num = len(workflow_ids)
85
        for work_id in workflow_ids:
86
            work_status_url = f"https://circleci.com/api/v2/workflow/{work_id}"
87
            r = requests.get(work_status_url, auth=(args.api_token, ""))
88
            # print("STATUS", work_status_url)
89
            assert r.status_code == 200, f"Error making api work request: {r}"
90
            w = r.json()
91
            status = w["status"]
92
            print("Status:", status)
93
            if status not in ("running", "failing"):
94
                done += 1
95
        if num and done == num:
96
            print("Finished")
97
            return
98
        time.sleep(20)
99

100

101
def trigger(args):
102
    url = "https://circleci.com/api/v2/project/gh/wandb/wandb/pipeline"
103
    payload = {
104
        "branch": args.branch,
105
    }
106
    manual: bool = any(
107
        [args.platform, args.toxenv, args.test_file, args.test_name, args.test_repeat]
108
    )
109
    if manual:
110
        parameters = {"manual": True}
111
        platforms = args.platform.split(",") if args.platform else ["linux"]
112
        toxenv = args.toxenv or "py37"
113
        toxcmd = toxenv
114
        if args.test_file or args.test_repeat:
115
            toxcmd += " --"
116
        if args.test_file:
117
            toxcmd += " " + args.test_file
118
            if args.test_name:
119
                toxcmd += " -k " + args.test_name
120
        if args.test_repeat:
121
            toxcmd += f" --flake-finder --flake-runs={args.test_repeat}"
122
        # get last token split by hyphen as python version
123
        pyver = toxenv.split("-")[-1]
124
        pyname = py_name_dict.get(pyver)
125
        assert pyname, f"unknown pyver: {pyver}"
126
        # handle more complex pyenv (func tests)
127
        if pyver != toxenv:
128
            toxsplit = toxenv.split("-")
129
            assert len(toxsplit) == 3
130
            tsttyp, tstshard, tstver = toxsplit
131
            prefix = "s_"
132
            if tstshard.startswith(prefix):
133
                tstshard = tstshard[len(prefix) :]
134
            pyname = f"{pyname}-{tsttyp}-{tstshard}"
135
        pyimage = py_image_dict.get(pyver)
136
        assert pyimage, f"unknown pyver: {pyver}"
137
        for p in platforms:
138
            job = platforms_dict.get(p)
139
            assert job, f"unknown platform: {p}"
140
            pshort = platforms_short_dict.get(p)
141
            jobname = f"{pshort}-{pyname}"
142
            parameters["manual_" + job] = True
143
            parameters["manual_" + job + "_name"] = jobname
144
            if job == "test":
145
                parameters["manual_" + job + "_image"] = pyimage
146
            parameters["manual_" + job + "_toxenv"] = toxcmd
147
            if args.parallelism:
148
                parameters["manual_parallelism"] = args.parallelism
149
            if args.xdist:
150
                parameters["manual_xdist"] = args.xdist
151
        payload["parameters"] = parameters
152
    print("Sending to CircleCI:", payload)
153
    if args.dryrun:
154
        return
155
    r = requests.post(url, json=payload, auth=(args.api_token, ""))
156
    assert r.status_code == 201, "Error making api request"
157
    d = r.json()
158
    uuid = d["id"]
159
    print("CircleCI workflow started:", uuid)
160
    if args.wait or args.loop:
161
        poll(args, pipeline_id=uuid)
162

163

164
def trigger_nightly(args):
165
    url = "https://circleci.com/api/v2/project/gh/wandb/wandb/pipeline"
166

167
    default_shards = set(NIGHTLY_SHARDS)
168
    shards = {
169
        f"nightly_execute_{shard.replace('-', '_')}": False for shard in default_shards
170
    }
171

172
    requested_shards = set(args.shards.split(",")) if args.shards else default_shards
173

174
    # check that all requested shards are valid and that there is at least one
175
    if not requested_shards.issubset(default_shards):
176
        raise ValueError(
177
            f"Requested invalid shards: {requested_shards}. "
178
            f"Valid shards are: {default_shards}"
179
        )
180
    # flip the requested shards to True
181
    for shard in requested_shards:
182
        shards[f"nightly_execute_{shard.replace('-', '_')}"] = True
183

184
    payload = {
185
        "branch": args.branch,
186
        "parameters": {
187
            **{
188
                "manual": True,
189
                "manual_nightly": True,
190
                "nightly_git_branch": args.branch,
191
                "nightly_slack_notify": args.slack_notify,
192
            },
193
            **shards,
194
        },
195
    }
196

197
    print("Sending to CircleCI:", payload)
198
    if args.dryrun:
199
        return
200
    r = requests.post(url, json=payload, auth=(args.api_token, ""))
201
    assert r.status_code == 201, "Error making api request"
202
    d = r.json()
203
    uuid = d["id"]
204
    number = d["number"]
205
    print("CircleCI workflow started.")
206
    print(f"UUID: {uuid}")
207
    print(f"Number: {number}")
208
    if args.wait:
209
        poll(args, pipeline_id=uuid)
210

211

212
def get_ci_builds(args, completed=True):
213
    bname = args.branch
214
    # TODO: extend pagination if not done
215
    url = "https://circleci.com/api/v1.1/project/gh/wandb/wandb?shallow=true&limit=100"
216
    if completed:
217
        url = url + "&filter=completed"
218
    # print("SEND", url)
219
    r = requests.get(url, auth=(args.api_token, ""))
220
    assert r.status_code == 200, f"Error making api request: {r}"
221
    lst = r.json()
222
    cfirst = None
223
    ret = []
224
    done = False
225
    for d in lst:
226
        b = d.get("branch")
227
        if b != bname:
228
            continue
229
        v = d.get("vcs_revision")
230
        n = d.get("build_num")
231
        j = d.get("workflows", {}).get("job_name")
232
        w = d.get("workflows", {}).get("workflow_id")
233
        # print("DDD", d)
234
        cfirst = cfirst or v
235
        if cfirst != v:
236
            done = True
237
            break
238
        ret.append((v, n, j, w))
239
    if not done:
240
        return
241
    return ret
242

243

244
def grab(args, vhash, bnum):
245
    # curl -H "Circle-Token: $CIRCLECI_TOKEN" https://circleci.com/api/v1.1/project/github/wandb/wandb/61238/artifacts
246
    # curl -L  -o out.dat -H "Circle-Token: $CIRCLECI_TOKEN" https://61238-86031674-gh.circle-artifacts.com/0/cover-results/.coverage
247
    cachedir = ".circle_cache"
248
    cfbase = f"cover-{vhash}-{bnum}.xml"
249
    cfname = os.path.join(cachedir, cfbase)
250
    if not os.path.exists(cachedir):
251
        os.mkdir(cachedir)
252
    if os.path.exists(cfname):
253
        return
254
    url = f"https://circleci.com/api/v1.1/project/github/wandb/wandb/{bnum}/artifacts"
255
    r = requests.get(url, auth=(args.api_token, ""))
256
    assert r.status_code == 200, f"Error making api request: {r}"
257
    lst = r.json()
258
    if not lst:
259
        return
260
    for item in lst:
261
        p = item.get("path")
262
        u = item.get("url")
263
        # print("got", p)
264
        if p != "cover-results/coverage.xml":
265
            continue
266
        # print("GRAB", p, u)
267
        # TODO: use tempfile
268
        print("Downloading circle artifacts...")
269
        s, o = subprocess.getstatusoutput(
270
            f'curl -L  -o out.dat -H "Circle-Token: {args.api_token}" {u!r}'
271
        )
272
        assert s == 0
273
        os.rename("out.dat", cfname)
274

275

276
def status(args):
277
    # TODO: check for current git hash only
278
    got = get_ci_builds(args, completed=False)
279
    if not got:
280
        print("ERROR: couldn't find job, maybe we should poll?")
281
        sys.exit(1)
282
    work_ids = [workid for _, _, _, workid in got]
283
    poll(args, workflow_ids=[work_ids[0]])
284

285

286
def download(args):
287
    print(f"Checking for circle artifacts (Branch: {args.branch})...")
288
    got = get_ci_builds(args)
289
    assert got
290
    for v, n, _, _ in got:
291
        grab(args, v, n)
292

293

294
def process_args():
295
    parser = argparse.ArgumentParser()
296

297
    subparsers = parser.add_subparsers(
298
        dest="action", title="action", description="Action to perform"
299
    )
300
    parser.add_argument("--api_token", help=argparse.SUPPRESS)
301
    parser.add_argument("--branch", help="git branch (autodetected)")
302
    parser.add_argument("--dryrun", action="store_true", help="Don't do anything")
303

304
    parse_trigger = subparsers.add_parser("trigger")
305
    parse_trigger.add_argument(
306
        "--platform", help="comma-separated platform (linux,mac,win)"
307
    )
308
    parse_trigger.add_argument("--toxenv", help="single toxenv (py36,py37,py38,py39)")
309
    parse_trigger.add_argument("--test-file", help="test file (ex: tests/test.py)")
310
    parse_trigger.add_argument("--test-name", help="test name (ex: test_dummy)")
311
    parse_trigger.add_argument("--test-repeat", type=int, help="repeat N times (ex: 3)")
312
    parse_trigger.add_argument("--parallelism", type=int, help="CircleCI parallelism")
313
    parse_trigger.add_argument("--xdist", type=int, help="pytest xdist parallelism")
314
    parse_trigger.add_argument("--loop", type=int, help="Outer loop (implies wait)")
315
    parse_trigger.add_argument(
316
        "--wait", action="store_true", help="Wait for finish or error"
317
    )
318

319
    parse_trigger_nightly = subparsers.add_parser("trigger-nightly")
320
    parse_trigger_nightly.add_argument(
321
        "--slack-notify", action="store_true", help="post notifications to slack"
322
    )
323
    parse_trigger_nightly.add_argument(
324
        "--shards",
325
        default=",".join(NIGHTLY_SHARDS),
326
        help="comma-separated shards (standalone-{cpu,gpu,gpu-win},kfp,regression)",
327
    )
328
    parse_trigger_nightly.add_argument(
329
        "--wait", action="store_true", help="Wait for finish or error"
330
    )
331

332
    parse_status = subparsers.add_parser("status")
333
    parse_status.add_argument(
334
        "--wait", action="store_true", help="Wait for finish or error"
335
    )
336

337
    parse_download = subparsers.add_parser("download")
338
    parse_download.add_argument(
339
        "--wait", action="store_true", help="Wait for finish or error"
340
    )
341

342
    args = parser.parse_args()
343
    return parser, args
344

345

346
def process_environment(args):
347
    api_token = os.environ.get(CIRCLECI_API_TOKEN)
348
    assert api_token, f"Set environment variable: {CIRCLECI_API_TOKEN}"
349
    args.api_token = api_token
350

351

352
def process_workspace(args):
353
    branch = args.branch
354
    if not branch:
355
        code, branch = subprocess.getstatusoutput("git branch --show-current")
356
        assert code == 0, "failed git command"
357
        args.branch = branch
358

359

360
def main():
361
    parser, args = process_args()
362
    process_environment(args)
363
    process_workspace(args)
364

365
    if args.action == "trigger":
366
        for i in range(args.loop or 1):
367
            if args.loop:
368
                print(f"Loop: {i + 1} of {args.loop}")
369
            trigger(args)
370
    elif args.action == "trigger-nightly":
371
        trigger_nightly(args)
372
    elif args.action == "status":
373
        # find my workflow report status, wait on it (if specified)
374
        status(args)
375
    elif args.action == "download":
376
        download(args)
377
    else:
378
        parser.print_help()
379

380

381
if __name__ == "__main__":
382
    main()
383