分布式机器学习框架是现代推荐、广告和搜索系统的核心支撑。面对海量训练数据和高维稀疏特征，参数服务器（Parameter Server, PS） 架构应运而生。作为早期经典实现的ps-lite因其简洁性和完整性，成为理解PS原理的绝佳切入点。本文将深入分析ps-lite的核心设计与源码实现。

一、PS架构的必要性

传统分布式计算（如Hadoop/Spark）在推荐场景下面临两大挑战：

参数爆炸：特征空间动辄上亿，单Master节点无法存储全部参数。
带宽瓶颈：每次迭代广播全部参数，网络开销无法承受。

PS架构通过以下设计解决这些问题：

分布式存储：参数分散在多台Server节点存储。
按需通信：Worker只拉取当前Batch所需的特征参数（如Embedding），大幅减少数据传输量。

关键洞察：推荐数据高度稀疏，单个Batch仅涉及少量特征，这是PS高效性的基础。

二、ps-lite核心概念

1. 节点角色

static const int kScheduler = 1;  // 调度器
static const int kServerGroup = 2; // Server组
static const int kWorkerGroup = 4; // Worker组

Scheduler：负责节点管理和跨节点同步
Server：分布式存储部分参数（键值数据库）
Worker：计算梯度，与Server交互

2. 核心通信组件

// 示例：Worker向Server发送请求
KVWorker<float> kv(0, customer_id); 
kv.Push(keys, vals);

PostOffice：单例全局通信管理器（每个进程唯一）
Van：实际网络通信模块（支持ZMQ/IBVerbs等实现）
Customer：消息处理中介（Worker/Server的业务代理）

3. 数据容器

struct KVPairs {SArray<Key> keys;SArray<Val> vals;SArray<int> lens; // 变长数据支持
};

SArray：智能数组，支持零拷贝和自动内存回收
KVPairs：数据传输载体（如feature_id为Key，embedding为Val）

三、通信层深度解析（PostOffice与Van）

1. 节点管理流程

关键代码：

// Van::ProcessAddNodeCommand
if (is_scheduler_) {// 收集所有节点信息nodes->control.node.push_back(my_node_); // 广播全量节点列表for (int r : GetNodeIDs(kWorkerGroup+kServerGroup)) {Send(back);}
}

2. 同步机制(Barrier)

// PostOffice::Barrier
void Barrier(int customer_id, int node_group) {// 发送同步请求给Schedulerreq.meta.control.cmd = Control::BARRIER;van_->Send(req);// 阻塞等待同步完成barrier_cond_.wait(ulk, [this] { return barrier_done_[customer_id]; });
}

关键点：Scheduler统计到达Barrier的节点数，当数量等于目标组节点总数时广播解锁信号。

3. 消息路由

四、业务层实现（Customer/Worker/Server）

1. Worker的Pull/Push流程

// KVWorker::ZPull 零拷贝拉取
int ZPull(const SArray<Key>& keys, SArray<Val>* vals) {int ts = AddPullCB(keys, vals, ...); // 注册回调Send(ts, false, true, cmd, kvs);    // 发送请求return ts;
}// KVWorker::Process 响应处理
void Process(const Message& msg) {recv_kvs_[ts].push_back(kvs); // 缓存数据if (AllResponsesReceived(ts)) {MergeResponses(ts);       // 聚合多Server响应RunCallback(ts);          // 触发回调}
}

2. Server请求处理

// KVServerDefaultHandle示例
void operator()(const KVMeta& meta, const KVPairs<Val>& data) {for (size_t i = 0; i < n; ++i) {if (meta.push) store[key] += data.vals[i]; // 累加梯度if (meta.pull) res.vals[i] = store[key];   // 返回参数}server->Response(meta, res); // 回复Worker
}