axlearn/BUILD.bazel at main · apple/axlearn · GitHub

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
load("@REDACTED//tools/buildifier:defs.bzl", "buildifier")
load("@axlearn_pip//:requirements.bzl", "all_whl_requirements")
load("@bazel_gazelle//:def.bzl", "gazelle", "gazelle_binary")
load("@rules_python//python:defs.bzl", "py_library")
load("@rules_python//python/uv:lock.bzl", "lock")
load("@rules_python_gazelle_plugin//manifest:defs.bzl", "gazelle_python_manifest")
load("@rules_python_gazelle_plugin//modules_mapping:def.bzl", "modules_mapping")

# Top-level Gazelle directives.
# gazelle:python_generation_mode file
# gazelle:python_test_file_pattern *_test.py
# gazelle:exclude axlearn/audio
# gazelle:exclude axlearn/cli
# gazelle:exclude axlearn/common/monitoring
# gazelle:exclude axlearn/common/ops
# gazelle:exclude axlearn/common/rattention
# gazelle:exclude axlearn/common/ssm_kernels
# gazelle:exclude axlearn/cloud/gcp/examples
# gazelle:exclude axlearn/cloud/gcp/monitoring
# gazelle:exclude axlearn/experiments/audio
# gazelle:exclude axlearn/experiments/testdata
# gazelle:exclude axlearn/experiments/vision
# gazelle:exclude build
# gazelle:exclude ci

package(default_visibility = ["//visibility:public"])

modules_mapping(
    name = "modules_map",
    wheels = all_whl_requirements,
)

gazelle_python_manifest(
    name = "gazelle_python_manifest",
    modules_mapping = ":modules_map",
    pip_repository_name = "axlearn_pip",
    requirements = "//:requirements_lock.txt",
)

gazelle_binary(
    name = "gazelle_binary",
    languages = ["@rules_python_gazelle_plugin//python"],
)

gazelle(
    name = "gazelle",
    gazelle = ":gazelle_binary",
)

buildifier(
    name = "buildifier",
)

lock(
    name = "requirements",
    srcs = ["requirements.in"],
    out = "requirements_lock.txt",
    args = [
        "--find-links=https://storage.googleapis.com/axlearn-wheels/wheels.html",
        "--find-links=https://storage.googleapis.com/jax-releases/libtpu_releases.html",
    ],
    generate_hashes = False,
)

py_library(
    name = "conftest",
    testonly = True,
    srcs = ["conftest.py"],
    visibility = ["//:__subpackages__"],
    deps = ["@axlearn_pip//absl_py"],
)