Bazel 在 Pytorch/XLA 中¶

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

Bazel 依赖¶

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

WORKSPACE

http_archive(
    name = "org_tensorflow",
    strip_prefix = "tensorflow-f7759359f8420d3ca7b9fd19493f2a01bd47b4ef",
    urls = [
        "https://github.com/tensorflow/tensorflow/archive/f7759359f8420d3ca7b9fd19493f2a01bd47b4ef.tar.gz",
    ],
)

可以通过将此仓库指向不同的修订版来更新 TensorFlow 锁定。可以根据需要添加补丁。Bazel 将会解析依赖项，准备代码并进行密封性修补。

对于 PyTorch，采用了不同的依赖机制，因为使用了本地的 PyTorch 检出，并且此本地检出必须从源代码构建，并且最好安装在系统中以确保版本兼容性（例如，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 找不到 auth 令牌，也可能有所帮助。

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

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

运行测试¶

目前 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