wandb

1
import os
2
import platform
3
from typing import Callable, List
4

5
import nox
6

7
CORE_VERSION = "0.17.0b9"
8

9

10
@nox.session(python=False, name="build-rust")
11
def build_rust(session: nox.Session) -> None:
12
    """Builds the wandb-core wheel with maturin."""
13
    with session.chdir("client"):
14
        session.run(
15
            "maturin",
16
            "build",
17
            "--release",
18
            "--strip",
19
            external=True,
20
        )
21

22

23
@nox.session(python=False, name="install")
24
def install(session: nox.Session) -> None:
25
    # find latest wheel file in ./target/wheels/:
26
    wheel_file = [
27
        f
28
        for f in os.listdir("./client/target/wheels/")
29
        if f.startswith(f"wandb_core-{CORE_VERSION}") and f.endswith(".whl")
30
    ][0]
31
    session.run(
32
        "pip",
33
        "install",
34
        "--force-reinstall",
35
        f"./client/target/wheels/{wheel_file}",
36
        external=True,
37
    )
38

39

40
@nox.session(python=False, name="develop")
41
def develop(session: nox.Session) -> None:
42
    with session.chdir("client"):
43
        session.run(
44
            "maturin",
45
            "develop",
46
            "--release",
47
            "--strip",
48
            external=True,
49
        )
50

51

52
@nox.session(python=False, name="list-failing-tests-wandb-core")
53
def list_failing_tests_wandb_core(session: nox.Session) -> None:
54
    """Lists the core failing tests grouped by feature."""
55
    import pandas as pd
56
    import pytest
57

58
    class MyPlugin:
59
        def __init__(self):
60
            self.collected = []
61
            self.features = []
62

63
        def pytest_collection_modifyitems(self, items):
64
            for item in items:
65
                marks = item.own_markers
66
                for mark in marks:
67
                    if mark.name == "wandb_core_failure":
68
                        self.collected.append(item.nodeid)
69
                        self.features.append(
70
                            {
71
                                "name": item.nodeid,
72
                                "feature": mark.kwargs.get("feature", "unspecified"),
73
                            }
74
                        )
75

76
        def pytest_collection_finish(self):
77
            session.log("\n\nFailing tests grouped by feature:")
78
            df = pd.DataFrame(self.features)
79
            for feature, group in df.groupby("feature"):
80
                session.log(f"\n{feature}:")
81
                for name in group["name"]:
82
                    session.log(f"  {name}")
83

84
    my_plugin = MyPlugin()
85
    pytest.main(
86
        [
87
            "-m",
88
            "wandb_core_failure",
89
            "tests/pytest_tests/system_tests/test_core",
90
            "--collect-only",
91
        ],
92
        plugins=[my_plugin],
93
    )
94

95

96
@nox.session(python=False, name="build-apple-stats-monitor")
97
def build_apple_stats_monitor(session: nox.Session) -> None:
98
    """Builds the apple stats monitor binary for the current platform.
99

100
    The binary will be located in
101
    core/pkg/monitor/apple/.build/<arch>-apple-macosx/release/AppleStats
102
    """
103
    with session.chdir("core/pkg/monitor/apple"):
104
        session.run(
105
            "swift",
106
            "build",
107
            "--configuration",
108
            "release",
109
            "-Xswiftc",
110
            "-cross-module-optimization",
111
            external=True,
112
        )
113
        # copy the binary to core/pkg/monitor/apple/AppleStats
114
        session.run(
115
            "cp",
116
            f".build/{platform.machine().lower()}-apple-macosx/release/AppleStats",
117
            "AppleStats",
118
        )
119

120

121
@nox.session(python=False, name="graphql-codegen-schema-change")
122
def graphql_codegen_schema_change(session: nox.Session) -> None:
123
    """Runs the GraphQL codegen script and saves the previous api version.
124

125
    This will save the current generated go graphql code gql_gen.go
126
    in core/internal/gql/v[n+1]/gql_gen.go, run the graphql codegen script,
127
    and save the new generated go graphql code as core/internal/gql/gql_gen.go.
128
    The latter will always point to the latest api version, while the versioned
129
    gql_gen.go files can be used in versioning your GraphQL API requests,
130
    for example when communicating with an older server.
131

132
    Should use whenether there is a schema change on the Server side that
133
    affects your GraphQL API. Do not use this if there is no schema change
134
    and you are e.g. just adding a new query or mutation
135
    against the schema that already supports it.
136
    """
137
    session.run(
138
        "./core/scripts/generate-graphql.sh",
139
        "--schema-change",
140
        external=True,
141
    )
142

143

144
@nox.session(python=False, name="local-testcontainer-registry")
145
def local_testcontainer_registry(session: nox.Session) -> None:
146
    """Ensure we collect and store the latest local-testcontainer in the registry.
147

148
    This will find the latest released version (tag) of wandb/core,
149
    find associated commit hash, and then pull the local-testcontainer
150
    image with the same commit hash from
151
    us-central1-docker.pkg.dev/wandb-production/images/local-testcontainer
152
    and push it to the SDK's registry with the release tag,
153
    if it doesn't already exist there.
154

155
    To run locally, you must have the following environment variables set:
156
    - GITHUB_ACCESS_TOKEN: a GitHub personal access token with the repo scope
157
    - GOOGLE_APPLICATION_CREDENTIALS: path to a service account key file
158
      or a JSON string containing the key file contents
159

160
    To run this for a specific release tag, use:
161
    nox -s local-testcontainer-registry -- <release_tag>
162
    """
163
    tags: List[str] = session.posargs or []
164

165
    import subprocess
166

167
    def query_github(payload: dict[str, str]) -> dict[str, str]:
168
        import json
169

170
        import requests
171

172
        headers = {
173
            "Authorization": f"bearer {os.environ['GITHUB_ACCESS_TOKEN']}",
174
            "Content-Type": "application/json",
175
        }
176

177
        url = "https://api.github.com/graphql"
178
        response = requests.post(url, headers=headers, data=json.dumps(payload))
179
        data = response.json()
180

181
        return data
182

183
    def get_release_tag_and_commit_hash(tags: List[str]):
184
        if not tags:
185
            # Get the latest release tag and commit hash
186
            query = """
187
            {
188
            repository(owner: "wandb", name: "core") {
189
                latestRelease {
190
                tagName
191
                tagCommit {
192
                    oid
193
                }
194
                }
195
            }
196
            }
197
            """
198
            payload = {"query": query}
199

200
            data = query_github(payload)
201

202
            return (
203
                data["data"]["repository"]["latestRelease"]["tagName"],
204
                data["data"]["repository"]["latestRelease"]["tagCommit"]["oid"],
205
            )
206
        else:
207
            # Get the commit hash for the given release tag
208
            query = """
209
            query($owner: String!, $repo: String!, $tag: String!) {
210
            repository(owner: $owner, name: $repo) {
211
                ref(qualifiedName: $tag) {
212
                target {
213
                    oid
214
                }
215
                }
216
            }
217
            }
218
            """
219
            # TODO: allow passing multiple tags?
220
            variables = {"owner": "wandb", "repo": "core", "tag": tags[0]}
221
            payload = {"query": query, "variables": variables}
222

223
            data = query_github(payload)
224

225
            return tags[0], data["data"]["repository"]["ref"]["target"]["oid"]
226

227
    local_release_tag, commit_hash = get_release_tag_and_commit_hash(tags)
228

229
    release_tag = local_release_tag.removeprefix("local/v")
230
    session.log(f"Release tag: {release_tag}")
231
    session.log(f"Commit hash: {commit_hash}")
232

233
    if not release_tag or not commit_hash:
234
        session.error("Failed to get release tag or commit hash.")
235

236
    subprocess.check_call(["gcloud", "config", "set", "project", "wandb-client-cicd"])
237

238
    # Check if image with tag already exists in the SDK's Artifact registry
239
    images = (
240
        subprocess.Popen(
241
            [
242
                "gcloud",
243
                "artifacts",
244
                "docker",
245
                "tags",
246
                "list",
247
                "us-central1-docker.pkg.dev/wandb-client-cicd/images/local-testcontainer",
248
            ],
249
            stdout=subprocess.PIPE,
250
        )
251
        .communicate()[0]
252
        .decode()
253
        .split("\n")
254
    )
255
    images = [img for img in images if img]
256

257
    if any(release_tag in img for img in images):
258
        session.warn(f"Image with tag {release_tag} already exists.")
259
        return
260

261
    source_image = f"us-central1-docker.pkg.dev/wandb-production/images/local-testcontainer:{commit_hash}"
262
    target_image = f"us-central1-docker.pkg.dev/wandb-client-cicd/images/local-testcontainer:{release_tag}"
263

264
    # install gcrane: `go install github.com/google/go-containerregistry/cmd/gcrane@latest`
265
    subprocess.check_call(["gcrane", "cp", source_image, target_image])
266

267
    session.log(f"Successfully copied image {target_image}")
268

269

270
@nox.session(python=False, name="bump-core-version")
271
def bump_core_version(session: nox.Session) -> None:
272
    args = session.posargs
273
    if not args:
274
        session.log("Usage: nox -s bump-core-version -- <args>\n")
275
        # Examples:
276
        session.log(
277
            "For example, to bump from 0.17.0b8/0.17.0-beta.8 to 0.17.0b9/0.17.0-beta.9:"
278
        )
279
        session.log("nox -s bump-core-version -- pre")
280
        return
281

282
    for cfg in (".bumpversion.core.cfg", ".bumpversion.cargo.cfg"):
283
        session.run(
284
            "bump2version",
285
            "--config-file",
286
            cfg,
287
            *args,
288
        )
289

290

291
@nox.session(python=False, name="proto-go")
292
def proto_go(session: nox.Session) -> None:
293
    """Generate Go bindings for protobufs."""
294
    _generate_proto_go(session)
295

296

297
def _generate_proto_go(session: nox.Session) -> None:
298
    session.run("./core/scripts/generate-proto.sh", external=True)
299

300

301
@nox.session(name="proto-python")
302
@nox.parametrize("pb", [3, 4])
303
def proto_python(session: nox.Session, pb: int) -> None:
304
    """Generate Python bindings for protobufs.
305

306
    The pb argument is the major version of the protobuf package to use.
307

308
    Tested with Python 3.10 on a Mac with an M1 chip.
309
    Absolutely does not work with Python 3.7.
310
    """
311
    _generate_proto_python(session, pb=pb)
312

313

314
def _generate_proto_python(session: nox.Session, pb: int) -> None:
315
    if pb == 3:
316
        session.install("protobuf~=3.20.3")
317
        session.install("mypy-protobuf~=3.3.0")
318
        session.install("grpcio~=1.48.0")
319
        session.install("grpcio-tools~=1.48.0")
320
    elif pb == 4:
321
        session.install("protobuf~=4.23.4")
322
        session.install("mypy-protobuf~=3.5.0")
323
        session.install("grpcio~=1.50.0")
324
        session.install("grpcio-tools~=1.50.0")
325
    else:
326
        session.error("Invalid protobuf version given. `pb` must be 3 or 4.")
327

328
    session.install("-r", "requirements_build.txt")
329
    session.install(".")
330

331
    with session.chdir("wandb/proto"):
332
        session.run("python", "wandb_internal_codegen.py")
333

334

335
def _ensure_no_diff(
336
    session: nox.Session,
337
    after: Callable[[], None],
338
    in_directory: str,
339
) -> None:
340
    """Fails if the callable modifies the directory."""
341
    saved = session.create_tmp()
342
    session.run("cp", "-r", in_directory, saved, external=True)
343
    after()
344
    session.run("diff", in_directory, saved, external=True)
345
    session.run("rm", "-rf", saved, external=True)
346

347

348
@nox.session(name="proto-check-python", tags=["proto-check"])
349
@nox.parametrize("pb", [3, 4])
350
def proto_check_python(session: nox.Session, pb: int) -> None:
351
    """Regenerates Python protobuf files and ensures nothing changed."""
352
    _ensure_no_diff(
353
        session,
354
        after=lambda: _generate_proto_python(session, pb=pb),
355
        in_directory=f"wandb/proto/v{pb}/.",
356
    )
357

358

359
@nox.session(name="proto-check-go", tags=["proto-check"])
360
def proto_check_go(session: nox.Session) -> None:
361
    """Regenerates Go protobuf files and ensures nothing changed."""
362
    _ensure_no_diff(
363
        session,
364
        after=lambda: _generate_proto_go(session),
365
        in_directory="core/pkg/service/.",
366
    )
367