llama.cpp 分布式推理介绍(5) RPC 通信协议

Chapter 5: RPC 通信协议

在上一章中，我们探索了 RPC 服务器 (rpc_server)，它是我们远程计算架构中强大的"执行者"。我们现在已经认识了客户端的"遥控器"和服务器端的"执行者"。但是，它们之间是如何沟通的呢？它们说的是同一种"语言"吗？如果客户端想让服务器分配内存，它应该如何表达？

这就引出了我们本章的主题------RPC 通信协议。这套协议就是客户端和服务器之间进行沟通所共同遵守的一套"语言"规则。

1. 问题的提出：我们如何避免"鸡同鸭讲"？

想象一下，两位来自不同国家的外交官需要进行一场至关重要的谈判。如果他们语言不通，一个说中文，一个说法语，那结果必然是一场灾难。为了确保沟通顺畅无误，他们必须约定一种共同的语言（比如英语），并且使用标准、无歧义的外交辞令。

ggml-rpc 中的客户端和服务器也面临同样的问题。客户端不能随意地向服务器发送一串含糊不清的字节流，期望服务器能"猜到"它的意图。它们必须严格遵守一套预先定义好的规则，这套规则详细规定了：

• 可以执行哪些操作（命令的"词汇表"）。
• 每个操作的请求应该长什么样子（请求的"语法"）。
• 服务器的回复又应该是什么格式（响应的"语法"）。

这套规则就是 RPC 通信协议。它确保了客户端的每一个请求都能被服务器准确无误地理解和执行，反之亦然。

2. 核心概念：一套严格的"外交辞令"

RPC 通信协议 定义了客户端和服务器之间交换的所有消息的命令类型和数据格式。它就像是双方的"通信法典"，确保信息交换的精确性和一致性。

这个"法典"主要包含两个核心部分：

1. 命令集 (Command Set) ：这是一个预定义的、有限的操作列表，规定了客户端可以向服务器请求的所有动作。这就像外交辞令中的标准短语，如"表示赞同"、"提出抗议"等。在 ggml-rpc 中，这些命令以枚举类型 rpc_cmd 的形式存在。
2. 消息结构 (Message Structure)：它为每个命令定义了对应的请求消息和响应消息的二进制格式。这就像规定了"提出抗议"这份外交照会的标准格式：开头写什么，中间包含哪些要素，结尾如何落款。

命令集：我们能做什么？

ggml-rpc 定义了一系列命令，覆盖了从内存管理到计算执行的各种操作。让我们来看几个关键的命令，它们定义在 ggml-rpc.cpp 文件中：

// 文件: ggml-rpc.cpp

enum rpc_cmd {
    RPC_CMD_ALLOC_BUFFER = 0,    // 请求分配一个缓冲区
    RPC_CMD_SET_TENSOR,          // 设置一个张量的数据
    RPC_CMD_GET_TENSOR,          // 获取一个张量的数据
    RPC_CMD_GRAPH_COMPUTE,       // 请求执行一个计算图
    RPC_CMD_GET_DEVICE_MEMORY,   // 获取远程设备的内存信息
    RPC_CMD_HELLO,               // 客户端和服务器"握手"，检查版本
    // ... 还有其他命令
};

这个 enum 就是我们协议的"词汇表"。客户端每次发送请求时，都会在消息的开头附上其中一个命令值，来告诉服务器它想干什么。

消息结构：我们该怎么说？

光有词汇还不够，我们还需要语法。ggml-rpc 的消息结构非常简洁高效。

一个典型的请求 (Request) 消息流看起来是这样的：

| 命令 (1字节) | 请求体大小 (8字节) | 请求体数据 (可变长度) |

• 命令 ：就是上面 rpc_cmd 中的一个值，告诉服务器这是什么操作。
• 请求体大小 ：一个 64 位整数，告诉服务器接下来需要读取多少字节的数据作为请求体。如果某个命令不需要额外数据（比如 GET_DEVICE_MEMORY），这个值就是 0。
• 请求体数据：实际的请求内容，比如要分配多大的内存，或者序列化后的计算图数据。

一个典型的响应 (Response) 消息流看起来是这样的：

| 响应体大小 (8字节) | 响应体数据 (可变长度) |

• 响应体大小：告诉客户端接下来需要读取多少字节的数据作为响应。
• 响应体数据：服务器返回的结果，比如设备的内存大小，或者计算图的执行状态。

有些简单的命令可能没有请求体或响应体，但整体结构保持一致。这种设计清晰明了，使得双方解析数据都非常容易。

3. 实例解析：一次简单的"问候"

让我们通过一个最简单的命令 RPC_CMD_GET_DEVICE_MEMORY，来看看一次完整的请求-响应周期是如何运作的。这个命令用于查询服务器的可用内存和总内存。

第一步：客户端发送请求

客户端想要查询内存，于是它构建了一个 GET_DEVICE_MEMORY 请求。这个请求非常简单，因为它不需要携带任何额外的数据。

• 命令 : RPC_CMD_GET_DEVICE_MEMORY (假设其值为 11)
• 请求体大小 : 0
• 请求体数据: (无)

所以，客户端通过网络发送的二进制数据流是：

| 0x0B | 0x0000000000000000 |

(1 字节的命令 + 8 字节的大小)

第二步：服务器处理并响应

服务器接收到数据，首先读取第 1 个字节，发现是 GET_DEVICE_MEMORY 命令。然后读取后面 8 个字节，发现请求体大小为 0。服务器明白了客户端的意图，于是它调用本地的 GGML 函数查询 GPU 内存，得到了两个值：free_mem 和 total_mem。

服务器需要将这两个值返回给客户端。协议为此定义了一个专门的响应结构体：

// 文件: ggml-rpc.cpp

// #pragma pack(push, 1) 确保结构体成员之间没有填充字节，
// 使得它在内存中的布局和网络传输的字节流完全一致。
#pragma pack(push, 1)
struct rpc_msg_get_device_memory_rsp {
    uint64_t free_mem;   // 8 字节
    uint64_t total_mem;  // 8 字节
};
#pragma pack(pop)

假设服务器查询到可用内存为 22 GB，总内存为 24 GB。它会填充这个结构体，然后将其作为响应体发送出去。

• 响应体大小 : 16 (因为结构体大小是 8 + 8 = 16 字节)
• 响应体数据 : 填充了内存数据的 rpc_msg_get_device_memory_rsp 结构体的二进制内容。

所以，服务器发送的二进制数据流是：

| 0x0000000000000010 | { ... free_mem ... } | { ... total_mem ... } |

(8 字节的大小 + 16 字节的数据)

第三步：客户端接收并解析

客户端接收到服务器的响应，先读取 8 字节的大小，得知后面有 16 字节的数据。然后它读取这 16 个字节，并直接将它们"映射"或"转换"回一个 rpc_msg_get_device_memory_rsp 结构体，从而轻松地取出了 free_mem 和 total_mem 的值。

至此，一次简单而完整的 RPC 通信就完成了！

4. 幕后探秘：协议的实现

让我们通过一个简化的时序图和代码片段，来更深入地理解这个过程。

高层流程图

下图展示了 ggml_backend_dev_get_memory 调用背后的协议交互：

sequenceDiagram participant 应用程序 participant ggml-rpc 客户端 participant 网络 participant ggml-rpc 服务端 应用程序->>ggml-rpc 客户端: 调用 ggml_backend_dev_get_memory() activate ggml-rpc 客户端 ggml-rpc 客户端->>网络: 发送请求: [CMD=GET_MEM, SIZE=0] 网络->>ggml-rpc 服务端: 接收请求 activate ggml-rpc 服务端 ggml-rpc 服务端->>ggml-rpc 服务端: 解析命令，查询本地 GPU 内存 ggml-rpc 服务端->>网络: 发送响应: [SIZE=16, DATA={...}] deactivate ggml-rpc 服务端 网络->>ggml-rpc 客户端: 接收响应 ggml-rpc 客户端->>ggml-rpc 客户端: 解析响应数据 ggml-rpc 客户端-->>应用程序: 返回 free_mem, total_mem deactivate ggml-rpc 客户端

深入代码

这一切是如何在代码中实现的呢？

在客户端，有一个核心函数 send_rpc_cmd，它负责封装和发送请求，并接收响应。

// 文件: ggml-rpc.cpp (客户端侧，简化版)

static bool send_rpc_cmd(
    const std::shared_ptr<socket_t> & sock, // 网络连接
    enum rpc_cmd cmd,                        // 要发送的命令
    const void * input, size_t input_size,   // 请求体数据和大小
    void * output, size_t output_size        // 用于接收响应的缓冲区
) {
    // 1. 发送 1 字节的命令
    if (!send_data(sock->fd, &cmd, sizeof(cmd))) { return false; }
    // 2. 发送 8 字节的请求体大小
    if (!send_data(sock->fd, &input_size, sizeof(input_size))) { return false; }
    // 3. 发送请求体数据
    if (!send_data(sock->fd, input, input_size)) { return false; }

    // 4. 接收 8 字节的响应体大小
    uint64_t out_size;
    if (!recv_data(sock->fd, &out_size, sizeof(out_size))) { return false; }
    // 5. 接收响应体数据
    if (!recv_data(sock->fd, output, output_size)) { return false; }

    return true;
}

这段代码清晰地反映了我们之前描述的协议格式：先发命令，再发大小，再发数据；然后接收大小，再接收数据。

在服务器端，rpc_serve_client 函数中有一个巨大的 switch 语句，它就像一个"接线员"，根据收到的命令字节，将请求转交给相应的处理函数。

// 文件: ggml-rpc.cpp (服务器侧，简化版)

static void rpc_serve_client(sockfd_t sockfd, ...) {
    while (true) {
        // 读取 1 字节的命令
        uint8_t cmd;
        if (!recv_data(sockfd, &cmd, 1)) { break; }

        // 根据命令进行分发
        switch (cmd) {
            case RPC_CMD_GET_DEVICE_MEMORY: {
                // 读取请求体 (大小为0)
                if (!recv_msg(sockfd, nullptr, 0)) { return; }

                // 准备响应
                rpc_msg_get_device_memory_rsp response;
                response.free_mem = ...; // 获取内存信息
                response.total_mem = ...;

                // 发送响应
                if (!send_msg(sockfd, &response, sizeof(response))) { return; }
                break;
            }
            // ... 其他命令的处理 ...
        }
    }
}

这个 switch 结构是服务器的核心逻辑，确保每个命令都得到正确的处理。

5. 总结与展望

在本章中，我们揭开了 ggml-rpc 客户端和服务器之间沟通的秘密------RPC 通信协议。

• 我们理解了协议就像一套严格的"外交辞令"，定义了命令集 和消息结构，以确保双方能够准确无误地沟通。
• 我们学习了请求和响应消息的基本二进制格式，它们由命令、大小和数据体组成。
• 通过 GET_DEVICE_MEMORY 的例子，我们完整地走了一遍从请求构建、发送、服务器处理到响应返回的全过程。

我们已经知道，对于像 GRAPH_COMPUTE 这样的复杂命令，客户端需要将整个计算图"打包"成二进制数据流发送给服务器。这个"打包"的过程称为序列化 。特别是，计算图中的核心元素 ggml_tensor 是如何被转换成适合网络传输的格式的呢？

这正是我们下一章要深入探讨的主题。我们将详细了解 rpc_tensor 这个特殊的数据结构，它是实现张量远程传输的关键。