Internal Wiki

主题

PJRT Collectives 与分布式协调接口 ​

本页所有地址均适用于 libtpu-0.0.40-cp314 wheel 中的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，build libtpu_lts_20260413_b_RC00；.text VMA == file offset 0xe63c000）。其他版本会有所不同。

摘要 ​

多主机 TPU 作业有两个协作问题，PJRT 在两个不同位置解决它们。第一个是稳态 collective，即在已经跨芯片分片的张量上执行 all-reduce、all-gather、reduce-scatter；它由编译器降低为 ICI ring 流量，并且从不以数据形式穿过 PJRT C-ABI。第二个是bootstrap rendezvous：在任何芯片能与其他芯片通信之前，每台主机上的每个 worker 进程必须就集群拓扑达成一致、交换网络端点，并交接一个 barrier。第二个问题正是客户端提供的 key/value store 所解决的问题，也是本页主题。

本页负责 PJRT 级分布式协调 C-ABI 接口：宿主框架（JAX、TF）用于跨进程 rendezvous 的 key/value get/put 原语、插件从中 bootstrap 的协调 agent，以及客户端创建如何把该 agent 传递到 megascale/ICI 启动流程。它不是设备上的 collective 算法；那些内容在 On-Pod Collectives 中；也不是进程内 Collectives PJRT 扩展（type 21），后者是一个由 CPU executor 支撑、与 TPU 多主机协调无关的接口，位于 PJRT Extension Chain。它也不同于 MegaScaleTransport 运行的芯片 fabric rendezvous，后者见 Megascale Bootstrap。

核心发现先说明，以免重新实现者被上游 PJRT C 头文件误导：此构建不公开 PJRT_KeyValueStore 结构体，也不公开 PJRT_KeyValueGetCallback / PJRT_KeyValuePutCallback 函数指针 typedef。 规范 OpenXLA PJRT_Client_Create 接收的 KV-store 接口在这里不是一个离散扩展。跨进程协调改为通过 tsl::CoordinationServiceAgent 抽象接口（XLA 分布式运行时 agent）到达；它通过 TF_CoordinationService*KeyValue* 调用跨越扁平 C ABI，这些调用在 libtpu 中只作为 WEAK UND 导入 stub 存在，并在运行时绑定到宿主 TF runtime；此外还通过两个二进制内的 plugin-agent 类到达。KV 语义与上游 xla::KeyValueStoreInterface 完全相同；绑定形态不同。

陷阱 — 如果重新实现者按上游 pjrt_c_api.h 推导，会在 PJRT_Client_Create_Args 中寻找一对 kv_get_callback / kv_put_callback。它们不存在于 libtpu 的主表，也不存在于全部 17 个扩展中。协调连线通过安装在进程上的 coordination agent 带外提供给插件，而不是作为创建时 callback。不要为了匹配头文件而合成 KV-store 扩展；应记录实际存在的 agent 接口。

对重新实现而言，契约是：

  • KV 原语集合及其字节级签名：get（阻塞、带超时、try）、insert（带覆盖标志）、delete、dir-list，以及异步 callback 形态。这是每个多主机后端都需要的 rendezvous 词汇。
  • 三个绑定层，它们最终都落到同一个 CoordinationServiceAgent：C-ABI TF_CoordinationService* thunk、CPluginCoordinationServiceAgent / DirectPluginCoordinationServiceAgent plugin agent，以及 tsl::CoordinationServiceAgent 基类。
  • 客户端创建如何到达它：GetPjrtApi → GetTpuPjrtApi 构建 PJRT_Api；实际的 PjRtClient 由 PjrtClientFactoryRegistry::GetPjrtClient(DeviceType, PjrtClientFactoryOptions) 产生，而 coordination agent 是在 megascale 启动前交换分区拓扑的通道。
  • 关注点分离：这个 Python 级 rendezvous 通道（xla.coordination）与 C++ 级芯片 fabric rendezvous（xla.megascale.runtime）分离，它们是两个从不共享 barrier ID 的 protobuf 命名空间。
PJRT 入口GetPjrtApi @0xe6a83a0 → pjrt::tpu_plugin::GetTpuPjrtApi @0xe6aa440
客户端工厂xla::PjrtClientFactoryRegistry::GetPjrtClient(DeviceType, PjrtClientFactoryOptions) @0x10849c60
KV 接口tsl::CoordinationServiceAgent（抽象）；plugin agents CPlugin… / DirectPlugin…CoordinationServiceAgent
C-ABI shimTF_CoordinationService{Insert,Get,GetWithTimeout,TryGet,Delete}KeyValue, TF_GetCoordinationServiceAgent（thunks @0x213f0a10–0x213f0bc0）
KV protostensorflow::{Get,Insert,Delete,TryGet}KeyValueRequest/Response, KeyValueEntry, GetKeyValueDirRequest（xla.coordination）
PJRT_KeyValueStore 结构体此构建中不存在（HIGH）
PJRT_KeyValueGet/PutCallback此构建中不存在（HIGH）
下层（链接，不重复）On-Pod Collectives, Megascale Bootstrap, ICI fabric

1. KV-store 接口位于何处 ​

上游 PJRT 允许客户端在 PJRT_Client_Create 时把 key/value store 交给插件：两个 callback（kv_get、kv_put）加一个用户上下文，插件将其包装为 xla::KeyValueStoreInterface，并用它 bootstrap 分布式运行时。规范结构体是 PJRT_KeyValueGetCallback / PJRT_KeyValuePutCallback。

这个结构体不在此二进制中。扫描反编译符号表中的 KeyValueStoreInterface、PjRtKeyValue、KeyValueGetCallback、KeyValuePutCallback 和 CApiKeyValueStore 都没有结果。17 个 PJRT 扩展（extension chain）中没有任何一个是 KV-store 扩展，PJRT_Client_Create 选项摄取（client-and-device）也不接收 get/put callback 对。实际存在的是完整的 tsl::CoordinationServiceAgent 系列，也就是上游 PjRtClient 原本会通过 KeyValueStoreInterface 驱动的同一种 agent 类型，在这里被直接暴露。功能上这就是 KV store；结构上它作为进程驻留 coordination agent 到达，而不是作为创建时 callback 到达。

特性 — 上游设计分层为 PjRtClient → KeyValueStoreInterface → CoordinationServiceAgent。libtpu 折叠了中间层：agent 就是接口。重新实现者应提供该 agent（或其 C-ABI shim），并跳过 KeyValueStoreInterface wrapper，因为此二进制中没有任何内容消费它。

为什么是 agent，而不是 callback ​

Coordination agent 是比双 callback KV store 更丰富的契约：除了 get/put，它还携带 barrier、heartbeat、task-state 和异步路径。插件为了容错多主机训练需要所有这些能力（见 extension chain §Megascale 中的 megascale error-aggregator）。只传递 kv_get/kv_put 会丢失 barrier 和 heartbeat 通道，因此绑定暴露整个 agent。代价是 rendezvous 词汇比上游 KV 最小集更宽，但下面每个原语都是真实的导出入口。

2. KV 原语集合 ​

所有 key/value 原语都是 tsl::CoordinationServiceAgent 的方法。键和值是 std::string_view / std::string（UTF-8 字节串，不以 null 结尾）；没有类型系统，store 是一个集群范围共享的不透明字节映射。下面集合对本构建是穷尽的。

方法映射 ​

原语Agent 符号（地址）签名形态
阻塞 getCoordinationServiceAgent::GetKeyValue @0x1dafd700(string_view key) → StatusOr<string>
带超时 get…::GetKeyValue @0x1dafd720(string_view key, absl::Duration timeout) → StatusOr<string>
Try get…::TryGetKeyValue @0x1dafe0a0(string_view key) → StatusOr<string>（非阻塞；不存在时 NotFound）
异步 get…::GetKeyValueAsync @0x1dafdc20(string_view key, function<void(const StatusOr<string>&)> done)
Insert…::InsertKeyValue @0x1dafe660(string_view key, string_view value) → Status
Insert（覆盖）…::InsertKeyValue @0x1dafe680(string_view key, string_view value, bool allow_overwrite) → Status
Delete…::DeleteKeyValue @0x1dafe9c0(string_view key) → Status（前缀删除）
Dir list…::GetKeyValueDir @0x1dafe2c0(string_view dir) → StatusOr<vector<KeyValueEntry>>
Dir list async…::GetKeyValueDirAsync @0x1dafe3a0(string_view dir, function<void(const StatusOr<vector<KeyValueEntry>>&)> done)
Init check…::IsInitialized @0x1dafd5a0() → bool
Own task…::GetOwnTask @0x1dafd5e0() → StatusOr<CoordinatedTask>

两个 InsertKeyValue 重载（@0x1dafe660 和 @0x1dafe680，mangled 后缀 …S5_ 与 …S5_b）只差末尾的 bool allow_overwrite。普通重载会在键重复时返回 AlreadyExists；覆盖形式会静默替换。每个 worker 的端点只发布一次的 rendezvous 协议会使用普通形式，这样重复发布会暴露重启。

异步 callback 形态 ​

GetKeyValueAsync（及其 dir 对应项）是非阻塞 rendezvous 使用的加载路径。continuation 是 std::function<void(const absl::StatusOr<std::string>&)>，由 0x1dafdc20 处的 mangled 类型 …NS1_8functionIFvRKN4absl8StatusOrINS1_12basic_stringIcS4_NS1_9allocatorIcEEEEEEEE 确认。反编译的 prologue 接收 (this, key_ptr, key_len, function*)，并把 callable 克隆进 agent 的 pending 表（GetKeyValueAsync lambda 的 __policy::__large_clone thunk 位于 0x1dafee80（$_0）和 0x1daff180（$_1）），因此 callback 的生命周期超过调用栈帧，并会在 coordinator 返回值时触发。这是异步 kv_get 的对应物。

c
// tsl::CoordinationServiceAgent::GetKeyValueAsync   @0x1dafdc20
function GetKeyValueAsync(key /*string_view*/, done /*function<void(const StatusOr<string>&)>*/):
    clone done into the agent's pending-request table   // __large_clone @0x1dafee80
    issue GetKeyValueRequest{ key } over the coordination channel
    // on response: invoke done(StatusOr<string>) from the agent's callback thread
```text

> **说明 —** 没有单独的 “put callback” 对。Insert 默认是同步的；异步接口只存在于 *get* 侧，因为 rendezvous get 是等待 peer 的阻塞操作。重新实现者若要匹配上游 `kv_put`，只需提供阻塞 insert。

---

## 3. 三个绑定层

§2 中每个原语都通过三个层之一到达，它们最终都落在同一个 `CoordinationServiceAgent` 上。这些层用于服务不同 ABI 边界上的调用者；它们不是替代实现。

### Layer A — C-ABI thunks（`TF_CoordinationService*`）

一个扁平 C ABI 为无法使用 C++ 名字修饰符号的调用者（插件自身的 C 边界，以及任何 out-of-tree consumer）镜像 agent。此二进制中的每个 `TF_CoordinationService*` 名称都是 **`WEAK` `UND` 导入**（`.dynsym` 中 value `0x0`，nm class `w`）；libtpu *不*定义这些符号，它只为每个名称携带一个 `.plt` stub。`0x213f0…` 处的 stub 通过 `.got.plt` 槽位做间接 `jmp *[GOT]`，动态加载器会在运行时把它绑定到*宿主*框架 TF runtime 提供的定义（同一个安装 coordination agent 的进程），如果框架不存在则保持未绑定。thunk 及其 GOT 槽位：

| C-ABI 入口 | `.plt` thunk | `.got.plt` 槽位 | 映射到 |
|---|---|---|---|
| `TF_GetCoordinationServiceAgent` | `0x213f0a10` | `0x224c2a80` | 获取进程 agent 句柄 |
| `TF_CoordinationServiceInsertKeyValue` | `0x213f0b70` | `0x224c2b30` | `InsertKeyValue` |
| `TF_CoordinationServiceGetKeyValue` | `0x213f0b80` | `0x224c2b38` | `GetKeyValue`（阻塞） |
| `TF_CoordinationServiceGetKeyValueWithTimeout` | `0x213f0b90` | `0x224c2b40` | `GetKeyValue(key, timeout)` |
| `TF_CoordinationServiceTryGetKeyValue` | `0x213f0ba0` | `0x224c2b48` | `TryGetKeyValue` |
| `TF_CoordinationServiceDeleteKeyValue` | `0x213f0bb0` | `0x224c2b50` | `DeleteKeyValue` |
| `TF_CoordinationServiceIsInitialized` | `0x213f0bc0` | `0x224c2b58` | `IsInitialized` |

没有可反汇编的**二进制内实现**；`.plt` stub 就是 libtpu 对这些调用的全部足迹。重新实现者必须带外提供定义（它们属于 TF 的 `c_api_coordination` 接口），方式与加载 libtpu 的运行时完全一致。调用*形态*由下面的 Layer B 固定，后者是二进制内调用者。

### Layer B — plugin agents

两个具体的 `CoordinationServiceAgent` 子类把 agent 适配到 TF plugin op-kernel 边界：

- **`tensorflow::CPluginCoordinationServiceAgent`** — C-ABI 支撑的 agent。它的 `InsertKeyValue`（`0xe71e260`）已完整反编译，是这些层能够组合的证据：它分配一个 `TF_Status`（`TF_NewStatus`），调用 `TF_CoordinationServiceInsertKeyValue(key_ptr, key_len, val_ptr, val_len, self+8, status)`，用 `tsl::StatusFromTF_Status` 转换结果，然后释放 status。`self+8` 槽位是底层 C 句柄。Get/TryGet/Delete（`0xe71e2e0` / `0xe71e520` / `0xe71e5a0`）以及带超时 get（`0xe71e480`）遵循相同的 C-bounce 模板。
- **`tensorflow::DirectPluginCoordinationServiceAgent`** — 进程内 agent（无 C ABI），方法位于 `0xe71c5c0`–`0xe71c640`。它直接调用 C++ agent，用于 coordination service 与插件同驻一个地址空间的情况。

二者都通过 `GetPluginCoordinationServiceAgent` 从 op-kernel context 获得（`CPluginOpKernelContext` @`0xe71b520`，`DirectPluginOpKernelContext` @`0xe71dbc0`）。

```c
// tensorflow::CPluginCoordinationServiceAgent::InsertKeyValue   @0xe71e260
function InsertKeyValue(key /*string_view*/, value /*string_view*/) -> Status:
    status = TF_NewStatus()
    TF_CoordinationServiceInsertKeyValue(                 // .plt 0x213f0b70 → host TF runtime
        key.data, key.size, value.data, value.size,
        self->c_agent /* self+8 */, status)
    result = StatusFromTF_Status(status)                 // tsl::StatusFromTF_Status @0x10900bc0
    if (status) TF_DeleteStatus(status)
    return result

Layer C — 抽象基类 ​

tsl::CoordinationServiceAgent 是抽象接口（vtable 分发；0x1daf… 处的公开符号是虚分发 trampoline，这就是为什么若干项反编译成自尾调用）。真正的 coordinator client 逻辑位于 coordination_service_agent.cc translation unit（静态初始化器 GLOBAL__sub_I_coordination_service_agent.cc @0x2121b570 / 0x21362f10），与 coordination_service.cc（GLOBAL__sub_I @0x2137a6d0）一起构建。基类拥有 §4 中的 request/response proto，以及到 coordinator 的 gRPC channel。

4. 传输 protos ​

KV 原语序列化为 xla.coordination protobuf 消息，也就是 megascale bootstrap 指出的独立 Python 级 rendezvous 命名空间。每个 request/response 对都是标准 proto2::MessageLite，符号表中存在常规的 Clear / MergeImpl / ctor / dtor 四件套：

Message符号（代表地址）作用
tensorflow::KeyValueEntryctor 0x20811080, Clear 0x208113a0一个 {key, value} 对；dir-list 元素类型
GetKeyValueRequest / Response0x20811a20 / 0x20811dc0阻塞 + 带超时 get
InsertKeyValueRequest（配对 Response）（与 KeyValueEntry 配对）publish
DeleteKeyValueRequestctor 0x20813640, Clear 0x20813820前缀删除
TryGetKeyValueRequest（存在于符号表）非阻塞 get
GetKeyValueDirRequestClear 0x20812ee0目录枚举

这些 proto 上的 tensorflow:: C++ 命名空间（相对于 xla.coordination protobuf package）是历史 TF→XLA 迁移遗留；它们就是 coordination-service 消息。C 侧的响应处理经过 tensorflow::(anonymous namespace)::ProcessGetKeyValueResult(TF_Buffer*, TSL_Status*) @0xe71e360。

说明 — 这些 xla.coordination.*KeyValue* 消息与芯片 fabric rendezvous 使用的 xla.megascale.runtime.* 消息是不同的 protobuf 类型树。两者共享名称（例如都定义 BarrierRequest），但从不共享 wire schema 或 barrier ID。关于这种拆分，见 Megascale Bootstrap §Cross-References。

5. 客户端创建如何到达 coordination ​

PJRT C-ABI 通过单个导出入口 GetPjrtApi @0xe6a83a0 获得；它是直接跳到 pjrt::tpu_plugin::GetTpuPjrtApi @0xe6aa440 的 thunk，后者用 __cxa_guard 初始化扩展链并返回静态 PJRT_Api。该表的 PJRT_Client_Create 槽位摄取客户端选项（client-and-device）；随后具体的 PjRtClient 由平台工厂构建：

text
GetPjrtApi  @0xe6a83a0
  └─ GetTpuPjrtApi  @0xe6aa440          ── builds PJRT_Api + 17-node extension chain
        └─ PJRT_Client_Create (main table slot)
              └─ xla::PjrtClientFactoryRegistry::GetPjrtClient(DeviceType, PjrtClientFactoryOptions)
                   @0x10849c60   (registry singleton: PjrtClientFactoryRegistry::Get @0x10849a60)
```text

`PjrtClientFactoryOptions` 是把每进程分布式身份（process index、process count 和 coordination handle）带入工厂的结构体。工厂注册表模式中，`Get()` 返回 singleton，`GetPjrtClient(device_type, options)` 分发到已注册的 TPU 工厂；这就是多主机构建附着 coordination agent 的位置。TF 侧路径（`tensorflow::GetPjRtClient` @`0x10848a40`、`PjRtState::GetPjRtClient` @`0x10848ca0`、`GetPjRtClientWrapper` @`0xf79caa0`）为 TF runtime 包装同一个工厂。

### Bootstrap 交接

一旦客户端带着 coordination agent 存在，该 agent 的 KV store 就是 **Python 级** rendezvous（process index 分配、run id、shard layout）的介质；它必须在 C++ 级芯片 fabric rendezvous 之前完成。依赖链自上而下如下：

```text
[host framework]  JAX / TF distributed init
      │  publishes per-process state via InsertKeyValue / reads peers via GetKeyValue
      ▼
PJRT distributed CoordinationService   (xla.coordination)   ── THIS PAGE
      │  agent.IsInitialized() gates progress; barrier on the agent releases all workers
      ▼
xla::megascale::runtime CommunicationBackend   (xla.megascale.runtime)  ── Megascale Bootstrap
      │  GetMultiSliceTopology cross-slice rendezvous → multi-slice address table
      ▼
tpunetd  ── intra-slice ICI fabric bring-up
      ▼
steady-state on-device collectives over ICI   ── On-Pod Collectives

箭头是单向的：coordination KV rendezvous 必须收敛，megascale 拓扑发现才能运行；megascale 拓扑发现必须收敛，ICI collectives 才能发出。本文记录的 KV store 是第一个环节，也是宿主框架唯一直接接触的环节。

特性 — KV store 不携带张量数据，也从不处于 collective 热路径。它在启动时每进程移动几百字节（端点、id），之后除 heartbeat 外基本空闲。重新实现者不应把它的吞吐需求与 collective 层混淆：它是控制面 rendezvous，不是数据面。

6. 存在与缺失 ​

接口此构建中的状态
tsl::CoordinationServiceAgent KV 方法（get/insert/delete/dir/async）PRESENT
TF_CoordinationService*KeyValue* C-ABI thunksPRESENT（作为 WEAK UND 导入）
CPlugin / DirectPlugin coordination agentsPRESENT
xla.coordination KV protosPRESENT
PjrtClientFactoryRegistry + PjrtClientFactoryOptionsPRESENT
PJRT_KeyValueStore C 结构体ABSENT（HIGH）
PJRT_KeyValueGetCallback / PutCallback typedefsABSENT（HIGH）
规范 KeyValueStoreInterface wrapperABSENT（HIGH）

说明 — CreateCommunicators / 跨主机 communicator-handle 接口是进程内 Collectives extension（type 21）：一个由 CPU executor 支撑的 XLA 接口，不是 TPU 多主机协调路径，也不携带 KV store。它记录在 扩展链 中。本页覆盖的是不同的 PJRT 分布式协调 KV 接口，它会 bootstrap 多主机执行。

总结：PJRT 级collective communicator 接口，即为 communicator factory 提供输入的离散 KV-store 扩展，在此构建中很薄甚至不存在。实际存在并承担负载的是 CoordinationServiceAgent KV rendezvous，以及把它传递到 megascale 的客户端工厂。本页记录现有内容。

交叉引用 ​