晚上也在思考这篇文章该以怎样的结构来撰写，ps-lite虽然是轻量版，但是很多东西都是在复用，东西都混在一起，所以在解释理清思路的时候有点绕。思来想去还是按照ps程序启动的逻辑来讲，涉及到本文需要讲的点在重点描写。
启动一个完整的 ps 服务，需要启动scheduler、sever和worker,启动的时候都会产生一个实际的物理进程，这个进程里都会包含PostOffice，负责管理全局信息。
ps-lite给出一个简单的demo,启动的顺序是先启动schedule，然后是server 、worker
#!/bin/bash
# set -x
if [ # -lt 3 ]; then
echo "usage:0 num_servers num_workers bin [args..]"
exit -1;
fi
export DMLC_NUM_SERVER=1
shift
export DMLC_NUM_WORKER=1
shift
bin=1
shift
arg="@"
# start the scheduler
export DMLC_PS_ROOT_URI='127.0.0.1'
export DMLC_PS_ROOT_PORT=8000
export DMLC_ROLE='scheduler'
{bin}{arg} &
# start servers
export DMLC_ROLE='server'
for ((i=0; i<{DMLC_NUM_SERVER}; ++i)); do
export HEAPPROFILE=./S{i}
{bin}{arg} &
done
# start workers
export DMLC_ROLE='worker'
for ((i=0; i<{DMLC_NUM_WORKER}; ++i)); do
export HEAPPROFILE=./W{i}
{bin}{arg} &
done
从上面的启动程序可以看到，ps-lite启动任务一些参数是来自环境变量的，你会看到shell脚本里充斥着export。
那我们现在看看它是怎么启动程序的？看一个demo
#include <cmath>
#include "ps/ps.h"
using namespace ps;
void StartServer() {
if (!IsServer()) {
return;
}
auto server = new KVServer<float>(0);
server->set_request_handle(KVServerDefaultHandle<float>());
RegisterExitCallback([server](){ delete server; });
}
void RunWorker() {
if (!IsWorker()) return;
KVWorker<float> kv(0, 0);
// init
int num = 10000;
std::vector<Key> keys(num);
std::vector<float> vals(num);
int rank = MyRank();
srand(rank + 7);
for (int i = 0; i < num; ++i) {
keys[i] = kMaxKey / num * i + rank;
vals[i] = (rand() % 1000);
}
// push
int repeat = 50;
std::vector<int> ts;
for (int i = 0; i < repeat; ++i) {
ts.push_back(kv.Push(keys, vals));
// to avoid too frequency push, which leads huge memory usage
if (i > 10) kv.Wait(ts[ts.size()-10]);
}
for (int t : ts) kv.Wait(t);
// pull
std::vector<float> rets;
kv.Wait(kv.Pull(keys, &rets));
// pushpull
std::vector<float> outs;
for (int i = 0; i < repeat; ++i) {
// PushPull on the same keys should be called serially
kv.Wait(kv.PushPull(keys, vals, &outs));
}
float res = 0;
float res2 = 0;
for (int i = 0; i < num; ++i) {
res += std::fabs(rets[i] - vals[i] * repeat);
res2 += std::fabs(outs[i] - vals[i] * 2 * repeat);
}
CHECK_LT(res / repeat, 1e-5);
CHECK_LT(res2 / (2 * repeat), 1e-5);
LL << "error: " << res / repeat << ", " << res2 / (2 * repeat);
}
int main(int argc, char *argv[]) {
// start system ,这里会根据实际的角色名执行相应的逻辑，比如一开始执行是scheduler的启动任务，
// 会以scheduler的角色启动，后面的StartServer也是不会执行的，当然这里会在是server角色的时候还会在执行一次
Start(0);
// setup server nodes
StartServer();
// run worker nodes
RunWorker();
// stop system
Finalize(0, true);
return 0;
}
ps-lite在这里其实共用了一套代码，也就是说你启动scheduler、sever和worker 这些都会走一套启动代码，根据不同的角色名称去执行相应的代码逻辑，比如只有在角色scheduler的时候才会触发scheduler相关的代码逻辑。
接下来就先以scheduler 启动来介绍
Start(0);
//实际调用的方法
inline void Start(int customer_id, const char* argv0 = nullptr) {
Postoffice::Get()->Start(customer_id, argv0, true);
}
这里会有一个
Postoffice::Get()
调用Get方法是去获取PostOffice全局单例对象，这里想要强调的一点就是PostOffice是单例，即一个进程内只有这一个对象，全局变量。
无论scheduler还是worker 都会调用，那么你可以理解这里PostOffice单例是相对而言的，scheduler进程下有一个，sever进程下也有一个。
接下来再来看看 Start 函数做了哪些事情？
void Postoffice::Start(int customer_id, const char* argv0, const bool do_barrier) {
start_mu_.lock();
if (init_stage_ == 0) {
// 初始化环境变量，主要是从 shell 执行脚本中获取相关参数比如 role 角色变量、server 数量和worker 数量
InitEnvironment();
// init glog
if (argv0) {
dmlc::InitLogging(argv0);
} else {
dmlc::InitLogging("ps-lite\0");
}
// init node info.
for (int i = 0; i < num_workers_; ++i) {
int id = WorkerRankToID(i);
for (int g : {id, kWorkerGroup, kWorkerGroup + kServerGroup,
kWorkerGroup + kScheduler,
kWorkerGroup + kServerGroup + kScheduler}) {
node_ids_[g].push_back(id);
}
}
for (int i = 0; i < num_servers_; ++i) {
int id = ServerRankToID(i);
for (int g : {id, kServerGroup, kWorkerGroup + kServerGroup,
kServerGroup + kScheduler,
kWorkerGroup + kServerGroup + kScheduler}) {
node_ids_[g].push_back(id);
}
}
for (int g : {kScheduler, kScheduler + kServerGroup + kWorkerGroup,
kScheduler + kWorkerGroup, kScheduler + kServerGroup}) {
node_ids_[g].push_back(kScheduler);
}
init_stage_++;
}
start_mu_.unlock();
// start van
van_->Start(customer_id);
start_mu_.lock();
if (init_stage_ == 1) {
// record start time
start_time_ = time(NULL);
init_stage_++;
}
start_mu_.unlock();
// do a barrier here
if (do_barrier) Barrier(customer_id, kWorkerGroup + kServerGroup + kScheduler);
}
现在来看看这个初始化环境的函数做了哪些事情？
void Postoffice::InitEnvironment() {
const char* val = NULL;
std::string van_type = GetEnv("DMLC_PS_VAN_TYPE", "zmq");
// 核心的一个点就是创建了 Van，至于Van 是什么后续也会做相应的详细介绍
van_ = Van::Create(van_type);
//接下来都是在解析环境变量，对于我们而言就是判断这次启动的是哪个？
val = CHECK_NOTNULL(Environment::Get()->find("DMLC_NUM_WORKER"));
num_workers_ = atoi(val);
val = CHECK_NOTNULL(Environment::Get()->find("DMLC_NUM_SERVER"));
num_servers_ = atoi(val);
val = CHECK_NOTNULL(Environment::Get()->find("DMLC_ROLE"));
std::string role(val);
is_worker_ = role == "worker";
is_server_ = role == "server";
is_scheduler_ = role == "scheduler";
verbose_ = GetEnv("PS_VERBOSE", 0);
}
ok，让我们再回到PostOffice的start函数里，假设我们启动的是scheduler角色下的程序
for (int g : {kScheduler, kScheduler + kServerGroup + kWorkerGroup,
kScheduler + kWorkerGroup, kScheduler + kServerGroup}) {
node_ids_[g].push_back(kScheduler);
}
init_stage_++;
那这个node_ids_ 是干啥的？这里就需要讨论下 node_id 的概念
其实是可以分为两个部分：node group 和 single node_id
首先我们介绍下 node id 映射功能，就是如何在逻辑节点和物理节点之间做映射，如何把物理节点划分成各个逻辑组，如何用简便的方法做到给组内物理节点统一发消息。
三个逻辑组的定义如下：
/** \brief node ID for the scheduler */
static const int kScheduler = 1;
/**
* \brief the server node group ID
*
* group id can be combined:
* - kServerGroup + kScheduler means all server nodes and the scheuduler
* - kServerGroup + kWorkerGroup means all server and worker nodes
*/
static const int kServerGroup = 2;
/** \brief the worker node group ID */
static const int kWorkerGroup = 4;
for (int i = 0; i < num_workers_; ++i) {
int id = WorkerRankToID(i);
for (int g : {id, kWorkerGroup, kWorkerGroup + kServerGroup,
kWorkerGroup + kScheduler,
kWorkerGroup + kServerGroup + kScheduler}) {
node_ids_[g].push_back(id);
}
如下面代码所示，如果配置了 3 个worker，则 worker 的 rank 从 0 ～ 2，那么这几个 worker 实际对应的 物理 node ID 就会使用 WorkerRankToID 来计算出来。
node id 是物理节点的唯一标示，rank 是每一个逻辑概念（scheduler，work，server）内部的唯一标示。这两个标示由一个算法来确定。
node id 是物理节点的唯一标示，rank 是每一个逻辑概念（scheduler，work，server）内部的唯一标示。这两个标示由一个算法来确定。
如下面代码所示，如果配置了 3 个worker，则 worker 的 rank 从 0 ～ 2，那么这几个 worker 实际对应的 物理 node ID 就会使用 WorkerRankToID 来计算出来。
for (int i = 0; i < num_workers_; ++i) {
int id = WorkerRankToID(i);
for (int g : {id, kWorkerGroup, kWorkerGroup + kServerGroup,
kWorkerGroup + kScheduler,
kWorkerGroup + kServerGroup + kScheduler}) {
node_ids_[g].push_back(id);
}
}
具体计算规则如下：
/**
* \brief convert from a worker rank into a node id
* \param rank the worker rank
*/
static inline int WorkerRankToID(int rank) {
return rank * 2 + 9;
}
/**
* \brief convert from a server rank into a node id
* \param rank the server rank
*/
static inline int ServerRankToID(int rank) {
return rank * 2 + 8;
}
/**
* \brief convert from a node id into a server or worker rank
* \param id the node id
*/
static inline int IDtoRank(int id) {
#ifdef _MSC_VER
#undef max
#endif
return std::max((id - 8) / 2, 0);
}
而且这个算法保证server id为偶数，node id为奇数。
这样我们可以知道，1-7 的id表示的是node group，单个节点的id 就从 8 开始。
具体计算规则如下：
因为有时请求要发送给多个节点，所以ps-lite用了一个 map 来存储每个 node group / single node 对应的实际的node节点集合，即 确定每个id值对应的节点id集。
std::unordered_map<int, std::vector<int>> node_ids_
for (int i = 0; i < num_workers_; ++i) {
int id = WorkerRankToID(i);
for (int g : {id, kWorkerGroup, kWorkerGroup + kServerGroup,
kWorkerGroup + kScheduler,
kWorkerGroup + kServerGroup + kScheduler}) {
node_ids_[g].push_back(id);
}
}
这 5 个id 相对应，即需要在 node_ids_ 这个映射表中对应的 4, 4 + 1, 4 + 2, 4 +1 + 2, 12 这五个 item 之中添加。就是上面代码中的内部 for 循环条件。即，node_ids_ [4], node_ids_ [5],node_ids_ [6],node_ids_ [7] ,node_ids_ [12] 之中，都需要把 12 添加到 vector 最后。
所以，为了实现 “设置 1-7 内任意一个数字 可以发送给其对应的 所有node” 这个功能，对于每一个新节点，需要将其对应多个id（node，node group）上，这些id组就是本节点可以与之通讯的节点。例如对于 worker 2 来说，其 node id 是 2 * 2 + 8 = 12，所以需要将它与
还是花了不少的功夫在讲解node，那么这个node 的标记是用来干啥的？
这些node的标记实际上与我们的worker还有server都是对应的关心，所以通过这些node标记就可以快速找打，这样通信同步一些数据就方便。
在记录完node_id之后，开始调用Van的启动程序。Van其实是一个通信模块。Van的东西还是蛮多的，打算放在下一篇文章里讲了。
在继续就是讲到 Barrier
if (do_barrier) Barrier(customer_id, kWorkerGroup + kServerGroup + kScheduler);
总的来讲，schedular节点通过计数的方式实现各个节点的同步。具体来说就是：
ps-lite 使用 Barrier 来控制系统的初始化，就是大家都准备好了再一起前进。这是一个可选项。具体如下：
Node会调用 Barrier 函数 告知Scheduler，随即自己进入等待状态。
注意，调用时候是
if (do_barrier) Barrier(customer_id, kWorkerGroup + kServerGroup + kScheduler);
复制代码
void Postoffice::Barrier(int customer_id, int node_group) {
if (GetNodeIDs(node_group).size() <= 1) return;
auto role = van_->my_node().role;
if (role == Node::SCHEDULER) {
CHECK(node_group & kScheduler);
} else if (role == Node::WORKER) {
CHECK(node_group & kWorkerGroup);
} else if (role == Node::SERVER) {
CHECK(node_group & kServerGroup);
}
std::unique_lock<std::mutex> ulk(barrier_mu_);
barrier_done_[0][customer_id] = false;
Message req;
req.meta.recver = kScheduler;
req.meta.request = true;
req.meta.control.cmd = Control::BARRIER;
req.meta.app_id = 0;
req.meta.customer_id = customer_id;
req.meta.control.barrier_group = node_group; // 记录了等待哪些
req.meta.timestamp = van_->GetTimestamp();
van_->Send(req); // 给 scheduler 发给 BARRIER
barrier_cond_.wait(ulk, [this, customer_id] { // 然后等待
return barrier_done_[0][customer_id];
});
}
这就是说，等待所有的 group，即 scheduler 节点也要给自己发送消息。
处理等待的动作在 Van 类之中，我们提前放出来。
具体ProcessBarrierCommand逻辑如下：
request==false的BARRIER消息。barrier_done_置为true，然后通知所有等待条件变量barrier_cond_.notify_all()。void Van::ProcessBarrierCommand(Message* msg) {
auto& ctrl = msg->meta.control;
if (msg->meta.request) { // scheduler收到了消息，因为 Postoffice::Barrier函数 会在发送时候做设置为true。
if (barrier_count_.empty()) {
barrier_count_.resize(8, 0);
}
int group = ctrl.barrier_group;
++barrier_count_[group]; // Scheduler会对Barrier请求进行计数
if (barrier_count_[group] ==
static_cast<int>(Postoffice::Get()->GetNodeIDs(group).size())) { // 如果相等，说明已经收到了最后一个请求，所以发送解除 barrier 消息。
barrier_count_[group] = 0;
Message res;
res.meta.request = false; // 回复时候，这里就是false
res.meta.app_id = msg->meta.app_id;
res.meta.customer_id = msg->meta.customer_id;
res.meta.control.cmd = Control::BARRIER;
for (int r : Postoffice::Get()->GetNodeIDs(group)) {
int recver_id = r;
if (shared_node_mapping_.find(r) == shared_node_mapping_.end()) {
res.meta.recver = recver_id;
res.meta.timestamp = timestamp_++;
Send(res);
}
}
}
} else { // 说明这里收到了 barrier respones，可以解除 barrier了。具体见上面的设置为false处。
Postoffice::Get()->Manage(*msg);
}
}
Manage 函数就是解除了 barrier。
void Postoffice::Manage(const Message& recv) {
CHECK(!recv.meta.control.empty());
const auto& ctrl = recv.meta.control;
if (ctrl.cmd == Control::BARRIER && !recv.meta.request) {
barrier_mu_.lock();
auto size = barrier_done_[recv.meta.app_id].size();
for (size_t customer_id = 0; customer_id < size; customer_id++) {
barrier_done_[recv.meta.app_id][customer_id] = true;
}
barrier_mu_.unlock();
barrier_cond_.notify_all(); // 这里解除了barrier
}
}
在上面的启动程序中可能没见到下面两个函数的调用，但是这也是 Postoffice 重要的成员组成
到现在为止，邮车和customer都有了，信件本身无非就是embedding这些参数，但是这些参数的存放也是有讲究的，这也是在上一篇文章中提到的分布式存储，这个分布式是如何体现的？
const std::vector<Range>& Postoffice::GetServerKeyRanges() {
server_key_ranges_mu_.lock();
//循环遍历所有的server，配置server key 的范围
//本质上就是根据server的数量均匀划分而已，就是这么简单
if (server_key_ranges_.empty()) {
for (int i = 0; i < num_servers_; ++i) {
server_key_ranges_.push_back(Range(
kMaxKey / num_servers_ * i,
kMaxKey / num_servers_ * (i+1)));
}
}
server_key_ranges_mu_.unlock();
return server_key_ranges_;
}
通过以上的操作的确解决了数据分布式存储，而且可以明确在worker向server端拉取数据的时候要去哪个server拉数据的问题。
现在大概知道了邮车，那么怎么知道要给哪些customer送信件呢？邮局需要管理一份用户的名单。
Customer* Postoffice::GetCustomer(int app_id, int customer_id, int timeout) const {
Customer* obj = nullptr;
for (int i = 0; i < timeout * 1000 + 1; ++i) {
{
std::lock_guard<std::mutex> lk(mu_);
// app_id 是对应 kv存储的id，举个例子FM 里存在一阶weight app_id=0
// 通过app_id 去寻找customer，一般 worker 会有多个thread 对应不同的customer
//但是消费的都是同一个 kv，所以根据app_id可以找到对应的 customer
const auto it = customers_.find(app_id);
if (it != customers_.end()) {
std::unordered_map<int, Customer*> customers_in_app = it->second;
obj = customers_in_app[customer_id];
break;
}
}
std::this_thread::sleep_for(std::chrono::milliseconds(1));
}
return obj;
}
这个 GetCustomer 的操作主要是在Van中的 ProcessDataMsg 调用，这里就是Van要把传递的信件交给customer，然后通过 GetCustomer 这个方式来获取相应的customer。
上面的函数列的是读取，还有 AddCustomer 和 RemoveCustomer 负责添加和删除。