使用 Bazel 构建¶

Bazel 是一款用于自动化构建和测试软件的免费软件工具。 TensorFlow 和 OpenXLA 都使用它，因此它也非常适合 PyTorch/XLA。

Bazel 依赖项¶

Openxla 是 PyTorch/XLA 的一个 Bazel 外部依赖项，可以在 WORKSPACE 文件中看到。

WORKSPACE

http_archive(
    name = "xla",
    patch_args = [
        "-l",
        "-p1",
    ],
    patch_tool = "patch",
    patches = [
        "//openxla_patches:gpu_nvml.diff",
        "//openxla_patches:gpu_race_condition.diff",
        "//openxla_patches:count_down.diff",
    ],
    strip_prefix = "xla-" + xla_hash,
    urls = [
        "https://github.com/openxla/xla/archive/" + xla_hash + ".tar.gz",
    ],
)

您可以在 WORKSPACE 文件中的 urls 字段中指定您要使用的 OpenXLA 版本。PyTorch/XLA 始终使用确定的 OpenXLA 提交（“OpenXLA pinned”，即 xla_hash）进行构建。可以根据需要添加补丁。Bazel 将解析依赖项，准备代码并将其进行隔离打包。

对于 PyTorch，由于使用了本地的 PyTorch 检出，因此部署了不同的依赖项机制，并且必须从源进行 build，最好是安装到系统中以确保版本兼容性（例如，PyTorch/XLA 中的 codegen 使用 torchgen Python 模块，该模块应安装在系统中）。

本地目录可以在 bazel/dependencies.bzl 中设置，或者在命令行中覆盖。

bazel build --override_repository=org_tensorflow=/path/to/exported/tf_repo //...

bazel build --override_repository=torch=/path/to/exported/and/built/torch_repo //...

请确保覆盖的存储库处于适当的版本，并且对于 torch，已使用 USE_CUDA=0 python setup.py bdist_wheel 构建，以确保所有预期的构建对象都存在；最好是安装到系统中。

WORKSPACE

new_local_repository(
    name = "torch",
    build_file = "//bazel:torch.BUILD",
    path = PYTORCH_LOCAL_DIR,
)

PyTorch 头文件直接来自 torch 依赖项，即本地 PyTorch 检出。共享库（例如 libtorch.so）来自已构建代码的同一本地检出，并且 build/lib/ 包含构建对象。为了使此功能正常工作，需要将 -isystemexternal/torch 传递给编译器，以便它可以找到 system 库并从本地检出中满足它们。其中一些包含为 <system>，一些包含为 "user" 头文件。

Bazel 引入了 pybind11 嵌入式 Python，并链接到它，以通过此机制为插件提供 libpython。Python 头文件也从这里获取，而不是依赖于系统版本。这些是从 "@pybind11//:pybind11_embed" 中获取的，该项设置了链接 libpython 的编译器选项。

如何构建 XLA 库¶

构建库很简单。

bazel build //torch_xla/csrc/runtime/...

Bazel 通过 .bazelrc 进行配置，但也可以通过命令行接收标志。

bazel build --config=remote_cache //torch_xla/csrc/runtime/...

使用 remote_cache 配置通常会使用 gcloud 进行缓存，速度更快，但需要使用 gcloud 进行身份验证。有关配置，请参见 .bazelrc。

使用 Bazel 可以轻松表达复杂的依赖关系，并且通过以相同的方式表达所有内容并构建单个依赖关系图可以获得很多好处。因此，无需像以前那样单独构建 XLA 库，构建整个存储库或链接所有其他内容的插件共享对象就足够了。

如何构建 Torch/XLA 插件¶

可以通过调用标准的 python setup.py bdist_wheel 来实现正常构建，但 C++ 绑定可以通过以下方式轻松构建：

bazel build //:_XLAC.so

这将构建 XLA 客户端和 PyTorch 插件并将其全部链接在一起。这在测试更改时很有用，能够编译 C++ 代码而不构建 Python 插件，从而实现更快的迭代周期。

远程缓存¶

Bazel 内置了远程缓存。可以使用许多缓存后端；我们在 (GCS)[https://bazel.build/remote/caching#cloud-storage] 上部署了缓存。您可以在 .bazelrc 文件中，在配置名称 remote_cache 下查看配置。

远程缓存默认禁用，但由于它极大地加快了增量构建速度，因此几乎总是推荐使用它，并且它在 CI 自动化和 Cloud Build 中默认启用。

要在机器上进行身份验证，请确保您具有以下凭据：

gcloud auth application-default login --no-launch-browser

使用 remote_cache 配置设置配置的远程缓存需要使用 GCP 进行身份验证。有多种方法可以进行 GCP 身份验证。对于可以访问开发 GCP 项目的个人开发者，只需向 Bazel 指定 --config=remote_cache 标志，就会使用默认的 --google_default_credentials，并且如果机器上存在 gcloud 令牌，它将开箱即用，使用登录用户进行身份验证。用户需要具有 GCP 中的远程构建权限（将新开发者添加到 Remote Bazel 角色）。在 CI 中，使用服务帐户密钥进行身份验证，并通过 --config=remote_cache --google_credentials=path/to/service.key 传递给 Bazel。在 Cloud Build 上，使用 docker build --network=cloudbuild 将身份验证从运行 Cloud Build 的服务帐户传递到进行编译的 Docker 镜像：应用程序默认凭据可以完成此操作，并以服务帐户身份进行身份验证。所有帐户（用户帐户和服务帐户）都需要具有远程缓存的读/写权限。

远程缓存使用缓存筒仓。每个唯一的机器和构建都应指定一个唯一的筒仓密钥以受益于一致的缓存。可以使用标志传递筒仓密钥：-remote_default_exec_properties=cache-silo-key=SOME_SILO_KEY'。

使用远程缓存运行构建

BAZEL_REMOTE_CACHE=1 SILO_NAME="cache-silo-YOUR-USER" TPUVM_MODE=1 python setup.py bdist_wheel

添加

GCLOUD_SERVICE_KEY_FILE=~/.config/gcloud/application_default_credentials.json

如果 bazel 找不到身份验证令牌，也可能会有帮助。

这里的 YOUR-USER 可以是作者的用户名或机器名，一个确保良好缓存行为的唯一名称。其他 setup.py 功能也按预期工作（例如 develop）。

首次使用新缓存密钥编译代码时会很慢，因为它将从头开始编译所有内容，但增量编译会非常快。更新 TensorFlow Pin 后，每个密钥第一次编译也会稍微慢一些，然后直到下次更新之前都会很快。

运行测试¶

目前 C++ 代码由 Bazel 构建和测试。Python 代码将在未来迁移。

Bazel 也是一个测试平台，可以轻松运行测试。

bazel test //test/cpp:main

当然，XLA 和 PJRT 配置必须存在于环境中才能运行测试。并非所有环境变量都传递到 Bazel 测试环境中，以确保远程缓存未命中不会太频繁（环境变量是缓存密钥的一部分）。请参阅 .bazelrc 测试配置，查看哪些环境变量已传递，并根据需要添加新的环境变量。

您也可以使用辅助脚本运行测试。

BAZEL_REMOTE_CACHE=1 SILO_NAME="cache-silo-YOUR-USER" ./test/cpp/run_tests.sh -R

xla_client 测试是纯粹的隔离测试，可以轻松执行。torch_xla 插件测试更为复杂：它们需要安装 torch 和 torch_xla，并且它们不能并行运行，因为它们使用同一端口上的 XRT 服务器/客户端，或者因为它们使用 GPU 或 TPU 设备，而当时只有一个可用。因此，torch_xla/csrc/ 下的所有测试都捆绑到一个目标 :main 中，该目标按顺序运行所有测试。

代码覆盖率¶

运行测试时，计算代码覆盖率可能会很有用。

bazel coverage //torch_xla/csrc/runtime/...

可以使用 lcov 可视化覆盖率，如 Bazel 的文档中所述，或者在您选择的编辑器中使用 lcov 插件，例如 VSCode 的 Coverage Gutters。

语言服务器¶

Bazel 可以为语言服务器（如 clangd）提供支持，该服务器为您选择的编辑器带来代码引用、自动补全和对底层代码的语义理解。对于 VSCode，可以使用 Bazel Stack，它可以与 Visual Studio clangd 扩展功能结合使用，提供强大的功能来辅助代码编辑。

构建 PyTorch/XLA¶

一如既往，PyTorch/XLA 可以使用 Python distutils 进行构建。

BAZEL_REMOTE_CACHE=1 SILO_NAME="cache-silo-YOUR-USER" TPUVM_MODE=1 python setup.py bdist_wheel