ps-lite是DMLC的子模块,下面介绍其简单使用。
- 环境设置:从https://github.com/dmlc/ps-lite 下载源码编译即可
- 代码:下面的代码是简单的实现Worker向Server push数据,在Server上注册一个事件函数,对推送到Server上的数据进行累加,然后Worker在从Server上pull回数据
#include "ps/ps.h"
#include <iostream>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
using namespace std;
using namespace ps;
template <typename Val>
struct KVServerDefaultHandle1 {
void operator()(
const KVMeta& req_meta, const KVPairs<Val>& req_data, KVServer<Val>* server) {
size_t n = req_data.keys.size();
KVPairs<Val> res;
if (req_meta.push) {
//cout<< "push\n";
CHECK_EQ(n, req_data.vals.size());
} else {
//cout<< "pull\n";
res.keys = req_data.keys; res.vals.resize(n);
}
for (size_t i = 0; i < n; ++i) {
Key key = req_data.keys[i];
if (req_meta.push) {
store[key] += req_data.vals[i];
} else {
res.vals[i] = store[key];
}
}
server->Response(req_meta, res);
}
std::unordered_map<Key, Val> store;
};
void StartServer() {
if (!IsServer()) return;
cout<<"启动Server\n";
auto server = new KVServer<float>(0);
server->set_request_handle(KVServerDefaultHandle1<float>());
RegisterExitCallback([server](){ delete server; });
}
void RunWorker() {
if (!IsWorker()) return;
cout<<"启动Worker rank = " << MyRank() << "\n";
KVWorker<float> kv(0);
// init
int num = 10;
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] = i;
vals[i] = i * 10;
}
// push
int repeat = 1;
std::vector<int> ts;
for (int i = 0; i < repeat; ++i) {
ts.push_back(kv.Push(keys, vals));
}
for (int t : ts) kv.Wait(t);
// pull
std::vector<float> rets;
kv.Wait(kv.Pull(keys, &rets));
for(size_t i = 0; i < rets.size(); i++)
{
cout<< MyRank() <<" rets[" << i << "]" << rets[i] <<endl;
}
cout<<"结束\n";
}
int main(int argc, char *argv[]) {
// setup server nodes
StartServer();
// start system
Start();
// run worker nodes
RunWorker();
// stop system
Finalize();
return 0;
}
假设以上源码存为ps.cpp,用以下指令编译成二进制g++ ps.cpp -std=c++11 -lps -lprotobuf -lzmq -o ps
用如下脚本在本地执行./local.sh 1 3 ./ps,将会启动1个Server和三个Worker
local.sh是ps-lite自带的脚本,位于源码的tests下。
local.sh的脚本如下
#!/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} 1 &
done
# start workers
export DMLC_ROLE='worker'
for ((i=0; i<${DMLC_NUM_WORKER}; ++i)); do
export HEAPPROFILE=./W${i}
${bin} 2 &
done
wait