PyTorch/XLA 中的 Bazel¶
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 Pin。可以根据需要添加补丁。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 插件并将它们全部链接在一起。这在测试更改时很有用,可以在不构建 Python 插件的情况下编译 C++ 代码,从而加快迭代周期。
远程缓存¶
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 镜像中:Application Default Credentials 在此处完成工作并以服务帐户进行身份验证。所有帐户,包括用户和服务帐户,都需要具有远程缓存读/写权限。
远程缓存使用缓存筒仓(cache silos)。每个唯一的机器和构建都应指定一个唯一的筒仓密钥,以受益于一致的缓存。可以使用标志传递筒仓密钥:-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¶
与往常一样,可以使用 Python distutils
构建 PyTorch/XLA
BAZEL_REMOTE_CACHE=1 SILO_NAME="cache-silo-YOUR-USER" TPUVM_MODE=1 python setup.py bdist_wheel