本文讲解的 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% 的概率随机打印其收到的消息。