ps-lite本地模式使用简介

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