PS Lite 笔记

本文讲解的 PS Lite 源码版本限定如下：

• GitHub: https://github.com/dmlc/ps-lite/tree/master
• Commit: f45e2e78a7430be09f76264d2f4073fb2b1d54a2

角色

PS Lite 实现了一种轻量级的参数服务器架构，其定义了三种角色：

• Scheduler - 1个，负责协调。
• Server - 若干个，负责存储。
• Worker - 若干个，负责计算。

PostOffice 类

PS Lite 的三种角色虽然职责不同，但其基础功能均由一个全局单例的 PostOffice 类来实现。顾名思义，“邮局”类会维护了一张全局的“地址簿”，记录了所有节点的信息。

除了传递参数的数据消息外，各个节点之间控制信息有 EMPTY、TERMINATE、ADD_NODE、BARRIER、ACK、HEARTBEAT 共 6 种。EMPTY 类型不该出现，ACK 类型只有启用了 Resender 类才会出现。

Van 类

邮局里有了地址簿，就需要有货车来负责拉送物件。PostOffice 类在实例化的时候，会创建一个 Van 类的实例（具体实现是 ZMQVan 类）作为成员变量。该 Van 实例与所属 PostOffice 实例生命周期相同，负责具体的节点间通信。

Resender 类

在分布式系统中，通信也是不可靠的，丢包、延时都是必须考虑的场景。PS Lite 设计了 Resender 类来提高通信的可靠性，它引入了 ACK 机制。即：

• 每一个节点，对于收到的非 ACK/TERMINATE 消息，必须响应一个 ACK 消息。
• 每一个节点，对于发送的每一个非 ACK/TERMINATE 消息，必须在本地缓存下来。存储的数据结构是一个 MAP，根据消息的内容生产唯一的键。
• 每一个节点，对于收到的 ACK 消息，必须根据反馈的键从本地缓存中移除对应的消息。
• 每一个节点运行一个监控线程，每隔 PS_RESEND_TIMEOUT 毫秒检查一下本地缓存。根据每个消息的发送时间戳和当前时间，找出超时的消息进行重发，并累加其重试次数。

线程管理

PS Lite 定义的三种角色采用多线程机制工作，每个线程承担特定的职责，在所属的 Van 实例启动时被创建。具体描述如下：

• Scheduler、Worker 和 Server 的 Van 实例里均持有一个接受数据的线程。
• Worker 和 Server 的 Van 实例里还持有一个间歇地向 Scheduler 发送心跳的线程。
• 如果定义了值不为 0 环境变量 PS_RESEND，那么 Scheduler、Worker 和 Server 还会启动一个监控线程。

心跳机制

为了记录网络的可达性，PS Lite 设计了心跳机制。具体而言：

• 每一个 Worker/Server 节点，每隔 PS_HEARTBEAT_INTERVAL 秒向 Scheduler 发送一条 HEARTBEAT 消息；Scheduler 节点收到后，响应一个 HEARTBEAT 消息。
• 每一个节点的 PostOffice 单例中维护了一个 MAP 结构，存储了心跳关联的节点的活跃信息。键为节点编号，值为上次收到其 HEARTBEAT 消息的时间戳。
• Worker/Server 只记录 Scheduler 的心跳，Scheduler 则记录所有节点的心跳。基于时间戳和心跳超时，可以输出所有的死亡节点。

路由

在多 Server 架构下，一个很重要的问题是如何分布多个参数。换句话说，给定一个参数的键，如何确定其存储在哪一台 Server 上。路由功能直接影响到 Worker 在 Push/Pull 阶段的通信。

PS Lite 将路由逻辑放置在 Worker 端，采用范围划分的策略，即每一个 Server 有自己固定负责的键的范围。这个范围是在 Worker 启动的时候确定的。具体代码参见方法 Postoffice::GetServerKeyRanges()，细节如下：

• 根据编译 PS Lite 时是否设定的宏 USE_KEY32 来决定参数的键的数据类型，要么是 32 位无符号整数，要么是 64 位的。
• 根据键的数据类型，确定其值域的上界。例如 uint32_t 的上界是 4294967295。
• 根据键域的上界和启动时获取的 Server 数量（即环境变量 DMLC_NUM_SERVER 的值）来划分范围。
• 给定上界 MAX 和 Server 数量 N，第 i 个 Server 负责的范围是 [MAX/N*i, MAX/N*(i+1))。

需要注意的是，在不能刚好整除的情况下，键域上界的一小段被丢弃了。

调试

在系统运行中，我们经常希望能打印一些收到的消息来方便定位问题。PS Lite 通过环境变量 PS_DROP_MSG 提供了这一功能，其值代表输出消息的概率（不含百分号）。

例如，我们启动某个 Server 前，配置了环境变量 PS_DROP_MSG=70。那么该 Server 进程会按照 70% 的概率随机打印其收到的消息。