[原始碼解析] 機器學習引數伺服器ps-lite 之(3) ----- 代理人Customer

羅西的思考發表於2021-08-04

原始碼機器學習伺服器

[原始碼解析] 機器學習引數伺服器ps-lite 之(3) ----- 代理人Customer

[原始碼解析] 機器學習引數伺服器ps-lite 之(3) ----- 代理人Customer

0x00 摘要

本文是引數伺服器第三篇，介紹ps-lite的Customer模組。

目前有了郵局（PostOffice）和通訊模組小推車（Van），接下來就要看看郵局的客戶Customer。

Customer 就是 SimpleApp 在郵局的代理人。因為 worker，server 需要集中精力在演算法上，所以把 worker，server 邏輯上與網路相關的收發訊息功能都總結/轉移到 Customer 之中。

本系列其他文章是：

[原始碼解析] 機器學習引數伺服器ps-lite 之(1) ----- PostOffice

[原始碼解析] 機器學習引數伺服器ps-lite(2) ----- 通訊模組Van

0x01 來源

1.1 目前總體

我們總結一下目前的總體狀態：

PostOffice：一個單例模式的全域性管理類，一個 node 在生命期內具有一個PostOffice，依賴它的類成員對Node進行管理；
Van：通訊模組，負責與其他節點的網路通訊和Message的實際收發工作。PostOffice持有一個Van成員；
SimpleApp：KVServer和KVWorker的父類，它提供了簡單的Request, Wait, Response，Process功能；KVServer和KVWorker分別根據自己的使命重寫了這些功能；
Node ：資訊類，儲存了本節點的對應資訊，每個 Node 可以使用 hostname + port 來唯一標識。
Customer：每個SimpleApp物件持有一個Customer類的成員，且Customer需要在PostOffice進行註冊，該類主要負責：
- 作為一個傳送方，跟蹤由SimpleApp傳送出去的訊息的回覆情況；
- 作為接收方，維護一個Node的訊息佇列，為Node接收訊息；

瞭解一個類的上下文環境可以讓我們更好的理解這個類，所以我們首先需要看看 Customer 在哪裡使用到，我們目前已經分析了兩個類，我們就看看這兩個類中如何使用Customer。

1.2 Postoffice

在 PostOffice 之中，有如下成員變數：

// app_id -> (customer_id -> customer pointer) 
std::unordered_map<int, std::unordered_map<int, Customer*>> customers_;

以及如下成員函式，就是把Customer註冊到customers_：

void Postoffice::AddCustomer(Customer* customer) {
  std::lock_guard<std::mutex> lk(mu_);
  int app_id = CHECK_NOTNULL(customer)->app_id();
  // check if the customer id has existed
  int customer_id = CHECK_NOTNULL(customer)->customer_id();
  customers_[app_id].insert(std::make_pair(customer_id, customer));
  std::unique_lock<std::mutex> ulk(barrier_mu_);
  barrier_done_[app_id].insert(std::make_pair(customer_id, false));
}

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_);
      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;
}

因此，我們可以看出來幾點：

一個 app 例項可以對應多個 Customer；
Customer 需要註冊到 Postoffice 之中；

1.3 Van

在 Van 中，我們可以看到，當處理資料訊息時候，會：

依據訊息中的 app_id 從Postoffice 之中得到 customer_id；
依據 customer_id 從 Postoffice 之中得到 Customer；
呼叫 Customer 的 Accept 方法來處理訊息；

void Van::ProcessDataMsg(Message* msg) {
  // data msg
  int app_id = msg->meta.app_id;
  int customer_id =
      Postoffice::Get()->is_worker() ? msg->meta.customer_id : app_id;
  auto* obj = Postoffice::Get()->GetCustomer(app_id, customer_id, 5);
  obj->Accept(*msg);
}

因此我們知道：

一個 app 例項可以對應多個 Customer；
Customer 需要註冊到 Postoffice 之中；
接受訊息時候會依據訊息中的app id找到 Customer，從而呼叫 Customer 的 Accept 方法來處理具體資料訊息；

1.4 Customer

在 Customer 之中我們可以看到，Accept 的作用就是往 Customer 的 queue 之中插入訊息。

ThreadsafePQueue recv_queue_;
 
inline void Accept(const Message& recved) {
   recv_queue_.Push(recved);
}

Customer物件本身也會啟動一個接受執行緒 recv_thread_，使用 Customer::Receiving()，其中呼叫註冊的recv_handle_函式對訊息進行處理。

std::unique_ptr<std::thread> recv_thread_;

recv_thread_ = std::unique_ptr<std::thread>(new std::thread(&Customer::Receiving, this));

void Customer::Receiving() {
  while (true) {
    Message recv;
    recv_queue_.WaitAndPop(&recv);
    if (!recv.meta.control.empty() &&
        recv.meta.control.cmd == Control::TERMINATE) {
      break;
    }
    recv_handle_(recv);
    if (!recv.meta.request) {
      std::lock_guard<std::mutex> lk(tracker_mu_);
      tracker_[recv.meta.timestamp].second++;
      tracker_cond_.notify_all();
    }
  }
}

1.5 目前邏輯

因此我們可以得出目前邏輯（接受訊息邏輯）如下：

worker節點或者 server節點在程式的最開始會執行Postoffice::start()。
Postoffice::start()會初始化節點資訊，並且呼叫Van::start()。
Van::start() 啟動一個本地執行緒，使用Van::Receiving()來持續監聽收到的message。
Van::Receiving()接收後訊息之後，根據不同命令執行不同動作。針對資料訊息，如果需要下一步處理，會呼叫 ProcessDataMsg：
- 依據訊息中的app id找到 Customer。
- 將訊息傳遞給Customer::Accept函式。
Customer::Accept() 函式將訊息新增到一個佇列recv_queue_；
Customer 物件本身也會啟動一個接受執行緒 recv_thread_，使用 Customer::Receiving()
- 從recv_queue_佇列取訊息。
- 呼叫註冊的recv_handle_函式對訊息進行處理。

簡要版邏輯如下，資料流按照圖上數字順序進行，我們也可以看到， Van，Postoffice，Customer 這三個類彼此之間有些過耦合，可能做一下梳理會更好：

                +--------------------------+
                | Van                      |
                |                          |
DataMessage +----------->  Receiving       |
                |  1           +           |             +---------------------------+
                |              |           |             | Postoffice                |
                |              | 2         |             |                           |
                |              v           | GetCustomer |                           |
                |        ProcessDataMsg <------------------> unordered_map customers_|
                |              +           |      3      |                           |
                |              |           |             +---------------------------+
                +--------------------------+
                               |
                               |
                               | 4
                               |
                +-------------------------+
                | Customer     |          |
                |              |          |
                |              v          |
                |           Accept        |
                |              +          |
                |              |          |
                |              | 5        |
                |              v          |
                |         recv_queue_     |
                |              +          |
                |              | 6        |
                |              |          |
                |              v          |
                |          Receiving      |
                |              +          |
                |              | 7        |
                |              |          |
                |              v          |
                |         recv_handle_    |
                |                         |
                +-------------------------+

下面我們就詳細剖析下具體邏輯。

0x02 基礎類

我們首先要介紹一些基礎類。

2.1 SArray

SArray 有如下特點：

SArray 是共享資料的智慧陣列，提供類似 std::vector 的功能。
SArray 可以從 std::vector 構建出來。
SArray 可以像 C 指標一樣拷貝賦值，當對某個SArray的引用為0時，就自動回收該SArray的記憶體。
可以理解為一個零拷貝的vector，能相容vector的資料結構。

2.2 KVPairs

在ps-lite中，每個server 擁有一段連續的key，以及這些key對應的value。key和value是分開儲存的，每個key可能對應多個value，因此需要記錄每個key的長度，所以就有了 KVPairs。

KVPairs 特點如下：

KVPairs封裝了Key-Value結構，還包含了一個長度選項，擁有keys，values，lens等3個陣列。
KVPairs 包含SArray keys，SArray vals，SArray lens的模板類。Key其實是int64的別名，Val是模板變數。
lens和keys 等長，表示每個key對應的value的個數。
lens可為空，此時values被平分。

舉例而言:

若keys=[1,5]，lens=[2,3]，那麼keys[0] 對應的資料就是：values[0] 和 values[1]，而keys[1] 對應的資料就是 values[2]，values[3]，values[5]。
而如果len為空，則values.size()必須是keys.size()（此處為2）的倍數，key[0]和key[1]各對應一半的values。

定義如下：

struct KVPairs {
  // /** \brief empty constructor */
  // KVPairs() {}
  /** \brief the list of keys */
  SArray<Key> keys;
  /** \brief the according values */
  SArray<Val> vals;
  /** \brief the according value lengths (could be empty) */
  SArray<int> lens; // key對應value的長度vector
  /** \brief priority */
  int priority = 0;
};

2.3 Node

Node封裝了節點資訊，例如角色，ip，埠，是否是恢復節點。

struct Node {
  /** \brief the empty value */
  static const int kEmpty;
  /** \brief default constructor */
  Node() : id(kEmpty), port(kEmpty), is_recovery(false) {}
  /** \brief node roles */
  enum Role { SERVER, WORKER, SCHEDULER };

  /** \brief the role of this node */
  Role role;
  /** \brief node id */
  int id;
  /** \brief customer id */
  int customer_id;
  /** \brief hostname or ip */
  std::string hostname;
  /** \brief the port this node is binding */
  int port;
  /** \brief whether this node is created by failover */
  bool is_recovery;
};

2.4 Control

Control ：封裝了控制訊息的meta資訊，barrier_group（用於標識哪些節點需要同步，當command=BARRIER時使用），node（Node類，用於標識控制命令對哪些節點使用）等，方法簽名。

可以看到，Control 就包含了上面介紹的 Node 型別。

struct Control {
  /** \brief empty constructor */
  Control() : cmd(EMPTY) { }
  /** \brief return true is empty */
  inline bool empty() const { return cmd == EMPTY; }

  /** \brief all commands */
  enum Command { EMPTY, TERMINATE, ADD_NODE, BARRIER, ACK, HEARTBEAT };
  /** \brief the command */
  Command cmd;
  /** \brief node infos */
  std::vector<Node> node;
  /** \brief the node group for a barrier, such as kWorkerGroup */
  int barrier_group;
  /** message signature */
  uint64_t msg_sig;
};

2.5 Meta

Meta ：是訊息的後設資料部分，包括時間戳，傳送者id，接受者id，控制資訊Control，訊息型別等；

struct Meta {
  /** \brief the empty value */
  static const int kEmpty;
  /** \brief default constructor */
  Meta() : head(kEmpty), app_id(kEmpty), customer_id(kEmpty),
           timestamp(kEmpty), sender(kEmpty), recver(kEmpty),
           request(false), push(false), pull(false), simple_app(false) {}
  /** \brief an int head */
  int head;
  /** \brief the unique id of the application of messsage is for*/
  int app_id;
  /** \brief customer id*/
  int customer_id;
  /** \brief the timestamp of this message */
  int timestamp;
  /** \brief the node id of the sender of this message */
  int sender;
  /** \brief the node id of the receiver of this message */
  int recver;
  /** \brief whether or not this is a request message*/
  bool request;
  /** \brief whether or not a push message */
  bool push;
  /** \brief whether or not a pull message */
  bool pull;
  /** \brief whether or not it's for SimpleApp */
  bool simple_app;
  /** \brief an string body */
  std::string body;
  /** \brief data type of message.data[i] */
  std::vector<DataType> data_type;
  /** \brief system control message */
  Control control;
  /** \brief the byte size */
  int data_size = 0;
  /** \brief message priority */
  int priority = 0;
};

2.6 Message

2.6.1 結構

Message 是要傳送的資訊，具體如下：

訊息頭 meta：就是後設資料（使用了Protobuf 進行資料壓縮），包括：
- 控制資訊（Control）表示這個訊息表示的意義（例如終止，確認ACK，同步等），具體包括：
  - 命令型別；
  - 節點列表（vector
    ），節點包括：
    - 節點的角色
    - ip, port
    - id
    - 是否是恢復節點
  - group id表示這個控制命令對誰執行；
  - 方法簽名；
- 傳送者；
- 接受者；
- 時間戳；
- ...
訊息體 body：就是傳送的資料，使用了自定義的 SArray 共享資料，減少資料拷貝；

2.6.2 邏輯關係

幾個類之間的邏輯關係如下：

Message中的某些功能需要依賴Meta來完成，以此類推。

2.6.3 message型別

message 包括如下型別：

ADD_NODE：worker和server向shceduler進行節點新增
BARRIER：節點間的同步阻塞訊息
HEARTBEAT：節點間的心跳訊號，check alive
TERMINATE：節點退出訊號
EMPTY：普通訊息，比如 push or pull

2.6.4 定義

具體定義如下：

struct Message {
  /** \brief the meta info of this message */
  Meta meta;
  /** \brief the large chunk of data of this message */
  std::vector<SArray<char> > data;
  /**
   * \brief push array into data, and add the data type
   */
  template <typename V>
  void AddData(const SArray<V>& val) {
    CHECK_EQ(data.size(), meta.data_type.size());
    meta.data_type.push_back(GetDataType<V>());
    SArray<char> bytes(val);
    meta.data_size += bytes.size();
    data.push_back(bytes);
  }
};

每次傳送訊息時，訊息就按這個格式封裝好，負責傳送訊息的類成員（Customer類）就會按照Meta之中的資訊將訊息送貨上門。

0x03 Customer

3.1 概述

Customer 其實有兩個功能：

作為一個傳送方，用於追蹤SimpleApp傳送出去每個Request對應的Response情況；
作為接收方，因為有自己的接受執行緒和接受訊息佇列，所以Customer實際上是作為一個接受訊息處理引擎（或者說是引擎的一部分）存在；

具體特點如下：

每個SimpleApp物件持有一個Customer類的成員，且Customer需要在PostOffice進行註冊。
因為 Customer 同時又要處理Message 但是其本身並沒有接管網路，因此實際的Response和Message需要外部呼叫者告訴它，所以功能和職責上有點分裂。
每一個連線對應一個Customer例項，每個Customer都與某個node id相繫結，代表當前節點傳送到對應node id節點。連線對方的id和Customer例項的id相同。
新建一次request，會返回一個timestamp，這個timestamp會作為這次request的id，每次請求會自增1，相應的res也會自增1，呼叫wait時會保證後續比如做Wait以此為ID識別。

3.2 定義

3.2.1 成員變數

我們首先看看Customer的成員變數。

需要注意，這裡對於變數功能的理解，我們可以從訊息流程來看，即如果有一個接受訊息，則這個流程資料流如下，所以我們把 Customer 的成員變數也按照這個順序梳理：

Van::ProcessDataMsg ---> Customer::Accept ---> Customer::recv_queue_ ---> Customer::recv_thread_ ---> Customer::recv_handle_

主要成員變數如下：

ThreadsafePQueue recv_queue_ ：執行緒安全的訊息佇列；
std::unique_ptr< std::thread> recv_thread_ ：不斷從 recv_queue 讀取message並呼叫 recv_handle_；
RecvHandle recv_handle_ ：worker 或者 server 的訊息處理函式。
- 繫結Customer接收到request後的處理函式(SimpleApp::Process)；
- Customer會新拉起一個執行緒，用於在customer生命週期內，使用recv_handle_來處理接受的請求，這裡是使用了一個執行緒安全佇列，Accept()用於往佇列中一直髮送訊息，
- 接受到的訊息來自於Van的receiving thread，即每個節點的Van物件收到message後，根據message的不同，推送到不同的customer物件中。
- 對於Worker，比如KVWorker，recv_handle_儲存拉取的msg中的資料，
- 對於Server,需要使用set_request_handle來設定對應的處理函式，如KVServerDefaultHandle，
std::vector<std::pair<int, int>> tracker_ ：request & response 的同步變數。
- tracker_是Customer內用來記錄request和response的狀態的map。記錄了每個 request（使用request id）可能傳送了多少節點以及從多少個節點返回的 response的次數，
- tracker_下標為每個request 的timestamp，即Request編號。
- tracker_[i] . first 表示該請求傳送給了多少節點，即本節點應收到的Response數量。
- tracker_[i] . second 表示目前為止實際收到的Response數量。

3.2.2 具體定義

具體定義如下：

class Customer {
 public:
  /**
   * \brief the handle for a received message
   * \param recved the received message
   */
  using RecvHandle = std::function<void(const Message& recved)>;

  /**
   * \brief constructor
   * \param app_id the globally unique id indicating the application the postoffice
   *               serving for
   * \param customer_id the locally unique id indicating the customer of a postoffice
   * \param recv_handle the functino for processing a received message
   */
  Customer(int app_id, int customer_id, const RecvHandle& recv_handle);

  /**
   * \brief desconstructor
   */
  ~Customer();

  /**
   * \brief return the globally unique application id
   */
  inline int app_id() { return app_id_; }

  /**
   * \brief return the locally unique customer id
   */
  inline int customer_id() { return customer_id_; }

  /**
   * \brief get a timestamp for a new request. threadsafe
   * \param recver the receive node id of this request
   * \return the timestamp of this request
   */
  int NewRequest(int recver);

  /**
   * \brief wait until the request is finished. threadsafe
   * \param timestamp the timestamp of the request
   */
  void WaitRequest(int timestamp);

  /**
   * \brief return the number of responses received for the request. threadsafe
   * \param timestamp the timestamp of the request
   */
  int NumResponse(int timestamp);

  /**
   * \brief add a number of responses to timestamp
   */
  void AddResponse(int timestamp, int num = 1);

  /**
   * \brief accept a received message from \ref Van. threadsafe
   * \param recved the received the message
   */
  inline void Accept(const Message& recved) {
    recv_queue_.Push(recved);
  }

 private:
  /**
   * \brief the thread function
   */
  void Receiving();

  int app_id_;
  int customer_id_;
  RecvHandle recv_handle_;
  ThreadsafePQueue recv_queue_;
  std::unique_ptr<std::thread> recv_thread_;
  std::mutex tracker_mu_;
  std::condition_variable tracker_cond_;
  std::vector<std::pair<int, int>> tracker_;
  DISALLOW_COPY_AND_ASSIGN(Customer);
};

3.3 接受執行緒

在構建函式中，會建立接受執行緒。

recv_thread_ = std::unique_ptr<std::thread>(new std::thread(&Customer::Receiving, this));

執行緒處理函式如下，具體邏輯就是：

在訊息佇列上等待，如果有訊息就取出；
使用 recv_handle_ 處理訊息；
如果 meta.request 為 false，說明是 response，則增加 tracker 之中對應計數。

void Customer::Receiving() {
  while (true) {
    Message recv;
    recv_queue_.WaitAndPop(&recv);
    if (!recv.meta.control.empty() &&
        recv.meta.control.cmd == Control::TERMINATE) {
      break;
    }
    recv_handle_(recv);
    if (!recv.meta.request) {
      std::lock_guard<std::mutex> lk(tracker_mu_);
      tracker_[recv.meta.timestamp].second++;
      tracker_cond_.notify_all();
    }
  }
}

因為是使用 recv_handle_ 來進行具體的業務邏輯，所以我們下面就看看 recv_handle_ 如何設定，其實也就是 Customer 如何構建，使用。

3.4 如何構建

我們需要提前使用下文將要分析的一些類，因為他們是 Customer 的使用者，耦合的太緊密了。

3.4.1 In SimpleApp

首先我們看看SimpleApp，這是具體邏輯功能節點的基類。

每個SimpleApp物件持有一個Customer類的成員，且Customer需要在PostOffice進行註冊，

這裡就是新建一個Custom物件初始化obj_成員。

inline SimpleApp::SimpleApp(int app_id, int customer_id) : SimpleApp() {
  using namespace std::placeholders;
  obj_ = new Customer(app_id, customer_id, std::bind(&SimpleApp::Process, this, _1));
}

我們再看看SimpleApp的兩個子類。

3.4.2 KVServer(app_id)

KVServer類主要用來儲存key-values資料，進行一些業務操作，比如梯度更新。主要方法為：Process() 和Response()。

在其建構函式中會：

新建一個Customer物件來初始化 obj_ 成員；
把 KVServer::Process 傳入Customer建構函式，其實就是把 Process 方法賦予了Customer:: recv_handle_；
對於Server來說，app_id = custom_id = server's id；

建構函式如下：

  /**
   * \brief constructor
   * \param app_id the app id, should match with \ref KVWorker's id
   */
  explicit KVServer(int app_id) : SimpleApp() {
    using namespace std::placeholders;
    obj_ = new Customer(app_id, app_id, std::bind(&KVServer<Val>::Process, this, _1));
  }

3.4.3 KVWorker(app_id, custom_id)

KVWorker類主要用來向Server Push/Pull 自己的 key-value 資料。包括如下方法： Push()，Pull()，Wait()。

在其建構函式中會：

用預設的KVWorker::DefaultSlicer繫結slicer_成員；
新建一個Customer物件初始化obj_ 成員，用KVWorker::Process傳入Customer建構函式，其實就是把 Process 方法賦予了 Customer:: recv_handle_；

  /**
   * \brief constructor
   *
   * \param app_id the app id, should match with \ref KVServer's id
   * \param customer_id the customer id which is unique locally
   */
  explicit KVWorker(int app_id, int customer_id) : SimpleApp() {
    using namespace std::placeholders;
    slicer_ = std::bind(&KVWorker<Val>::DefaultSlicer, this, _1, _2, _3);
    obj_ = new Customer(app_id, customer_id, std::bind(&KVWorker<Val>::Process, this, _1));
  }

3.4.4 Customer

構建函式邏輯如下：

分別用傳入建構函式的引數初始化 app_id_, custom_id_ , recv_handle成員
呼叫PostOffice::AddCustomer將當前Customer註冊到PostOffice；
- PostOffice的customers_成員: 在對應的app_id的元素上新增custom_id;
- PostOffice的barrier_done_成員將該custom_id的同步狀態設為false
新起一個Receiving執行緒recv_thread_；

具體構建函式如下：

Customer::Customer(int app_id, int customer_id, const Customer::RecvHandle& recv_handle)
    : app_id_(app_id), customer_id_(customer_id), recv_handle_(recv_handle) {
  Postoffice::Get()->AddCustomer(this);
  recv_thread_ = std::unique_ptr<std::thread>(new std::thread(&Customer::Receiving, this));
}

3.4.5 梳理

3.4.5.1 示例程式碼

大家可能對 app_id 和 customer_id 有些疑問，比如：

在 KVWorker 構建函式中有：

app_id the app id, should match with KVServer's id
customer_id the customer id which is unique locally

在 KVServer 構建函式中有：

app_id the app id, should match with KVWorker's id

我們使用原始碼自帶的 tests/test_kv_app_multi_workers.cc 來梳理一下 app_id 與 customer_id 的邏輯關係。

我們提前劇透：worker是用 customer_id 來確定自己的身份。customer id 在 worker 程式碼中被用來確定本worker 對應的 key 的範圍。

從指令碼中可以看出來，使用如下做測試：

find test_* -type f -executable -exec ./repeat.sh 4 ./local.sh 2 2 ./{} \;

檔案中啟動了一個 server 和兩個 worker。

server 的 app_id, customer_id 都是 0；
worker 的 app_id 是 0，customer_id 分別是 0，1;
使用 std::thread 來執行 worker，這就是說，在同一個程式內執行兩個 worker 節點。這就可以解釋 KVWorker 構建函式中註釋中的 “the customer id which is unique locally”。
這樣，在 Postoffice 的 std::unordered_map<int, std::unordered_map<int, Customer*>> customers_ 成員變數中有如下：
- [0, [ 0, Customer_0] ]，第一個 0 是 app id, 第二個 0 是customer id
- [0, [ 1, Customer_1] ]，第一個 0 是 app id, 第二個 1 是customer id

因此，我們可以理出來：

app id 用來確定一個應用。
customer id 用來在本應用（app id）內確定一個 local worker，或者一個 server。
所以 KVServer 構建函式中說，app id 需要和 KVWorker's id 一致，其實就是說，大家的 app id 需要一致。
customer id 在下面的 worker 程式碼中被用來確定本worker 對應的 key 的範圍。

具體程式碼如下：

#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(int customer_id) { // 啟動worker
  Start(customer_id);
  if (!IsWorker()) {
    return;
  }
  KVWorker<float> kv(0, customer_id);
  // 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 + customer_id;
    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) {
    kv.Wait(kv.PushPull(keys, vals, &outs));
  }

  float res = 0;
  float res2 = 0;
  for (int i = 0; i < num; ++i) {
    res += fabs(rets[i] - vals[i] * repeat);
    res += 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);
  // stop system
  Finalize(customer_id, true);
}

int main(int argc, char *argv[]) {
  // start system
  bool isWorker = (strcmp(argv[1], "worker") == 0);
  if (!isWorker) {
    Start(0);
    // setup server nodes，啟動server節點
    StartServer();
    Finalize(0, true);
    return 0;
  }
  // run worker nodes，啟動兩個worker節點
  std::thread t0(RunWorker, 0);
  std::thread t1(RunWorker, 1);

  t0.join();
  t1.join();
  return 0;
}

3.4.5.2 確定身份

我們再回憶下 Postoffice 的初始化，可以看到，啟動時候，worker是用 customer_id 來確定自己的身份。於是，customer id 在 worker 程式碼中被用來確定本worker 對應的 key 的範圍。

void Postoffice::Start(int customer_id, const char* argv0, const bool do_barrier) {

    // init node info.
    // 對於所有的worker，進行node設定
    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);
      }
    }
	// 對於所有的server，進行node設定
    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);
      }
    }
	// 設定scheduler的node
    for (int g : {kScheduler, kScheduler + kServerGroup + kWorkerGroup,
                  kScheduler + kWorkerGroup, kScheduler + kServerGroup}) {
      node_ids_[g].push_back(kScheduler);
    }
    init_stage_++;
  }

  // start van
  van_->Start(customer_id); // 這裡有 customer_id

  ......
    
  // do a barrier here，這裡有 customer_id
  if (do_barrier) Barrier(customer_id, kWorkerGroup + kServerGroup + kScheduler);
}

再看看 Van 的初始化，也是用 customer_id 來確定自己的身份。

void Van::Start(int customer_id) {
  if (init_stage == 0) {
    // get my node info
    if (is_scheduler_) {
      my_node_ = scheduler_;
    } else {
      my_node_.hostname = ip;
      my_node_.role = role;
      my_node_.port = port;
      my_node_.id = Node::kEmpty;
      my_node_.customer_id = customer_id; // 這裡有 customer_id
    }
  }
  
  if (!is_scheduler_) {
    // let the scheduler know myself
    Message msg;
    Node customer_specific_node = my_node_;
    customer_specific_node.customer_id = customer_id; // 這裡有 customer_id
    msg.meta.recver = kScheduler;
    msg.meta.control.cmd = Control::ADD_NODE;
    msg.meta.control.node.push_back(customer_specific_node);
    msg.meta.timestamp = timestamp_++;
    Send(msg);
  }
  
  ......
}

所以，也能夠解釋了為什麼在 KVWorker 傳送訊息時候使用 app_id 和 customer_id。

template <typename Val>
void KVWorker<Val>::Send(int timestamp, bool push, bool pull, int cmd, const KVPairs<Val>& kvs) {
  .....
  for (size_t i = 0; i < sliced.size(); ++i) {
    Message msg;
    msg.meta.app_id = obj_->app_id(); // 注意這裡
    msg.meta.customer_id = obj_->customer_id();// 注意這裡
    msg.meta.request     = true;
    ......
    Postoffice::Get()->van()->Send(msg);
  }
}

在 KVServer 之中，也需要在回應訊息時候，使用 app_id 和 customer_id。

template <typename Val>
void KVServer<Val>::Response(const KVMeta& req, const KVPairs<Val>& res) {
  Message msg;
  msg.meta.app_id = obj_->app_id();// 注意這裡
  msg.meta.customer_id = req.customer_id;// 注意這裡
  msg.meta.request     = false;
  msg.meta.push        = req.push;
  msg.meta.pull        = req.pull;
  msg.meta.head        = req.cmd;
  msg.meta.timestamp   = req.timestamp;
  msg.meta.recver      = req.sender;
  ......
  Postoffice::Get()->van()->Send(msg);
}

3.4.5.3 問題

那麼問題來了，為什麼 Server 端，app_id 與 customer_id 相等？

因為目前沒有 ps 的最初程式碼，所以猜測是：

在 ps 程式碼中，Server 端也是有多個 cusomer，但是出於精簡目的，在 ps-lite 之中刪除了這部分功能，因此在 ps-lite 之中，app_id 與 customer_id 相等。

3.5 目前邏輯

因此我們再次梳理流程（接受訊息邏輯）如下：

worker節點或者 server節點在程式的最開始會執行Postoffice::start()。
Postoffice::start()會初始化節點資訊，並且呼叫Van::start()。
Van::start() 啟動一個本地執行緒，使用Van::Receiving()來持續監聽收到的message。
Van::Receiving()接收後訊息之後，根據不同命令執行不同動作。針對資料訊息，如果需要下一步處理，會呼叫 ProcessDataMsg：
- 依據訊息中的app id找到 Customer，即會根據customer id的不同將message發給不同的customer的recv thread。
- 將訊息傳遞給Customer::Accept函式。
Customer::Accept() 函式將訊息新增到一個佇列recv_queue_；
Customer 物件本身也會啟動一個接受執行緒 recv_thread_，使用 Customer::Receiving()
- 從recv_queue_佇列取訊息。
- 如果 (!recv.meta.request) ，就說明是 response，則tracker_[req.timestamp].second++
- 呼叫註冊的recv_handle_函式對訊息進行處理。
對於worker來說，其註冊的recv_handle_是KVWorker::Process()函式。因為worker的recv thread接受到的訊息主要是從server處pull下來的KV對，因此該Process()主要是接收message中的KV對；
而對於Server來說，其註冊的recv_handle_是KVServer::Process()函式。因為server接受的是worker們push上來的KV對，需要對其進行處理，因此該Process()函式中呼叫的使用者通過KVServer::set_request_handle()傳入的函式物件。

目前邏輯如下圖，在第 8 步，recv_handle_ 指向 KVServer::Process 或者 KVWorker::Process。

                +--------------------------+
                | Van                      |
                |                          |
DataMessage +----------->  Receiving       |
                |  1           +           |             +---------------------------+
                |              |           |             | Postoffice                |
                |              | 2         |             |                           |
                |              v           | GetCustomer |                           |
                |        ProcessDataMsg <------------------> unordered_map customers_|
                |              +           |      3      |                           |
                |              |           |             +---------------------------+
                +--------------------------+
                               |
                               |
                               | 4
                               |
                +-------------------------+
                | Customer     |          |
                |              |          |
                |              v          |
                |           Accept        |
                |              +          |
                |              |          |
                |              | 5        |
                |              v          |
                |         recv_queue_     |                +-----------------+
                |              +          |                |KVWorker         |
                |              | 6        |     +--------> |                 |
                |              |          |     |    8     |         Process |
                |              v          |     |          +-----------------+
                |          Receiving      |     |
                |              +          |     |
                |              | 7        |     |
                |              |          |     |          +-----------------+
                |              v          |     |          |KVServer         |
                |         recv_handle_+---------+--------> |                 |
                |                         |          8     |         Process |
                +-------------------------+                +-----------------+

0x04 功能函式

以下這些 Customer 函式都是被其他模組呼叫。

4.1 Customer::NewRequest

4.1.1 實現

此函式的作用是：當傳送一個 request 時候，新增對此 request 的計數。所以，當我們需要給一個Resquest計數的時候，使用此函式。

特點如下：

每次傳送訊息前，先修改此條訊息應收到的 Response數量。
recver表示接收者的node_id，因為ps-lite中一個整數可能對應於多個node_id，所以使用Postoffice解碼獲得所有的真實node_id 的數目。
比如給 kServerGroup 發訊息，kServerGroup 裡面有3 個 server，則 num 為 3，就是應該收到 3 個response。tracker_ 對應的item 就是 [3,0]，表示應該收到 3個，目前收到 0 個。
函式的返回值可以認為是一個時間戳，這個時間戳會作為這次request的id，呼叫wait時會保證後續Wait以此為ID識別。

int Customer::NewRequest(int recver) {
  std::lock_guard<std::mutex> lk(tracker_mu_);
  int num = Postoffice::Get()->GetNodeIDs(recver).size();  // recver 可能會代表一個group。
  tracker_.push_back(std::make_pair(num, 0));
  return tracker_.size() - 1;  // 代表此次請求的時間戳timestamp，後續customer使用這個值代表這個request
}

4.1.2 呼叫

具體呼叫舉例就是在 worker 向 server 推送時候。

  int ZPush(const SArray<Key>& keys,
            const SArray<Val>& vals,
            const SArray<int>& lens = {},
            int cmd = 0,
            const Callback& cb = nullptr,
            int priority = 0) {
    int ts = obj_->NewRequest(kServerGroup); // 這裡會呼叫
    AddCallback(ts, cb);
    KVPairs<Val> kvs;
    kvs.keys = keys;
    kvs.vals = vals;
    kvs.lens = lens;
    kvs.priority = priority;
    Send(ts, true, false, cmd, kvs);
    return ts;
  }

4.2 Customer::AddResponse

4.2.1 實現

作用是：針對request已經返回response進行計數。

特點如下：

當外部呼叫者收到Response時，呼叫AddResponse告訴Customer物件。
主動增加某次請求實際收到的Response數，主要用於客戶端傳送請求時，有時可跳過與某些server的通訊（此次通訊的keys沒有分佈在這些server上），在客戶端就可直接認為已接收到Response。
另外，在Customer::Receiving中，當處理了一條非request請求後，也會增加對應的請求的Response數。 tracker_[recv.meta.timestamp].second++;
這個類有個缺陷，對於過期的以後不會再用到的Request資訊，沒有刪除操作。而這個類的單個物件的生存週期又近乎等於程式的生存週期。因此，基於ps-lite程式跑的時間久了基本都會OOM。

void Customer::AddResponse(int timestamp, int num) {
  std::lock_guard<std::mutex> lk(tracker_mu_);
  tracker_[timestamp].second += num;
}

4.2.2 呼叫

在 KVWorker 的 Send 方法會呼叫，因為某些情況下，（此次通訊的keys沒有分佈在這些server上），在客戶端就可直接認為已接收到Response，所以要跳過。

template <typename Val>
void KVWorker<Val>::Send(int timestamp, bool push, bool pull, int cmd, const KVPairs<Val>& kvs) {
  // slice the message
  SlicedKVs sliced;
  slicer_(kvs, Postoffice::Get()->GetServerKeyRanges(), &sliced);

  // need to add response first, since it will not always trigger the callback
  int skipped = 0;
  for (size_t i = 0; i < sliced.size(); ++i) {
    if (!sliced[i].first) ++skipped;
  }
  
  obj_->AddResponse(timestamp, skipped); // 這裡呼叫
  
  if ((size_t)skipped == sliced.size()) {
    RunCallback(timestamp);
  }

  for (size_t i = 0; i < sliced.size(); ++i) {
    const auto& s = sliced[i];
    if (!s.first) continue;
    Message msg;
    msg.meta.app_id = obj_->app_id();
    msg.meta.customer_id = obj_->customer_id();
    msg.meta.request     = true;
    msg.meta.push        = push;
    msg.meta.pull        = pull;
    msg.meta.head        = cmd;
    msg.meta.timestamp   = timestamp;
    msg.meta.recver      = Postoffice::Get()->ServerRankToID(i);
    msg.meta.priority    = kvs.priority;
    const auto& kvs = s.second;
    if (kvs.keys.size()) {
      msg.AddData(kvs.keys);
      msg.AddData(kvs.vals);
      if (kvs.lens.size()) {
        msg.AddData(kvs.lens);
      }
    }
    Postoffice::Get()->van()->Send(msg);
  }
}

4.3 Customer::WaitRequest

4.3.1 實現

功能是：當我們需要等待某個發出去的Request對應的Response全部收到時，使用此函式會阻塞等待，直到應收到Response數等於實際收到的Response數。

wait操作的過程就是tracker_cond_一直阻塞等待，直到傳送出去的數量和已經返回的數量相等。

void Customer::WaitRequest(int timestamp) {
  std::unique_lock<std::mutex> lk(tracker_mu_);
  tracker_cond_.wait(lk, [this, timestamp]{
      return tracker_[timestamp].first == tracker_[timestamp].second;
    });
}

4.3.2 呼叫

Wait 函式就是使用 WaitRequest 來確保操作完成。

  /**
   * \brief Waits until a push or pull has been finished
   *
   * Sample usage:
   * \code
   *   int ts = w.Pull(keys, &vals);
   *   Wait(ts);
   *   // now vals is ready for use
   * \endcode
   *
   * \param timestamp the timestamp returned by the push or pull
   */
  void Wait(int timestamp) { obj_->WaitRequest(timestamp); }

但是具體如何呼叫，則是使用者自行決定，比如：

  for (int i = 0; i < repeat; ++i) {
    kv.Wait(kv.Push(keys, vals));
  }

於是這就來到了同步策略的問題。

0x05 同步策略

不同的worker同時並行運算的時候，可能因為網路、機器配置等外界原因，導致不同的worker的進度是不一樣的，如何控制worker的同步機制是一個比較重要的課題。

5.1 同步協議

一般來說，有三個級別的非同步控制協議：BSP（Bulk Synchronous Parallel），SSP（Stalness Synchronous Parallel）和ASP（Asynchronous Parallel），它們的同步限制依次放寬。為了追求更快的計算速度，演算法可以選擇更寬鬆的同步協議。

為了解決效能的問題，業界開始探索這裡的一致性模型，最先出來的版本是ASP模式，在ASP之後提出了另一種相對極端的同步協議BSP，後來有人提出將ASP和BSP做一下折中，就是SSP。

這三個協議具體如下：

ASP：task之間完全不用相互等待，完全不顧worker之間的順序，每個worker按照自己的節奏走，跑完一個迭代就update，先完成的task，繼續下一輪的訓練。
- 優點：消除了等待慢task的時間，減少了GPU的空閒時間，因此與BSP相比提高了硬體效率。計算速度快，最大限度利用了叢集的計算能力，所有的worker所在的機器都不用等待
- 缺點：
  - 這個過程可能會導致梯度被計算過時的權重，從而降低統計效率。
  - 適用性差，在一些情況下並不能保證收斂性
BSP：是一般分散式計算採用的同步協議，每一輪迭代中都需要等待所有的task計算完成。每個worker都必須在同一個迭代執行，只有一個迭代任務所有的worker都完成了，才會進行一次worker和server之間的同步和分片更新。
- BSP的模式和單機序列因為僅僅是batch size的區別，所以在模型收斂性上是完全一樣的。同時，因為每個worker在一個週期內是可以平行計算的，所以有了一定的並行能力。spark用的就是這種方式。
- 優點：適用範圍廣；每一輪迭代收斂質量高
- 缺點：每一輪迭代中，，BSP要求每個worker等待或暫停來自其他worker的梯度，這樣就需要等待最慢的task，從而顯著降低了硬體效率，導致整體任務計算時間長。整個worker group的效能由其中最慢的worker決定；這個worker一般稱為straggler。
SSP：允許一定程度的task進度不一致，但這個不一致有一個上限，稱為staleness值，即最快的task最多領先最慢的task staleness輪迭代。
- 就是把將ASP和BSP做一下折中。既然ASP是允許不同worker之間的迭代次數間隔任意大，而BSP則只允許為0，那我就取一個常數s。有了SSP，BSP就可以通過指定s=0而得到。而ASP同樣可以通過制定s=∞來達到。
- 優點：一定程度減少了task之間的等待時間，計算速度較快。
- 缺點：每一輪迭代的收斂質量不如BSP，達到同樣的收斂效果可能需要更多輪的迭代，適用性也不如BSP，部分演算法不適用。

5.2 論文

沐神在論文中提到，parameter server 為使用者提供了多種任務依賴方式：

Sequential：這裡其實是 synchronous task，任務之間是有順序的，只有上一個任務完成，才能開始下一個任務；
Eventual：跟 sequential 相反，所有任務之間沒有順序，各自獨立完成自己的任務，
Bounded Delay：這是sequential 跟 eventual 之間的trade-off，可以設定一個 \(\tau\) 作為最大的延時時間。也就是說，只有 \(>\tau\) 之前的任務都被完成了，才能開始一個新的任務；極端的情況：
- \(\tau\) = 0，情況就是 Sequential；
- \(\tau\) = ∞，情況就是 Eventual；

5.3 ps-lite

ps-lite裡面有幾個涉及到等待同步的地方：

Worker pull 是非同步操作，如果需要等待 pull 完成，則可以呼叫Wait來保證customer裡面的request和response兩者相等，即保證Pull完成後再做其他操作；
在一個worker內，可以存在多個Customer,當第一個傳送barrier後，scheduler接收到request請求，然後根據msg判斷是request，然後，向barrier_group裡的所有node，node接到後， Postoffice::Get()->Manage(*msg)將barrier_done_中的customer_id對應的bool置true，完成同步操作。
當構建節點連線時，也可以進行一個barrier；

更復雜的比如Asp，bsp，ssp可以通過增加相應的Command來完成。

0x06 分散式優化

6.1 問題定義

假設我們要解決以下問題

\[\min_w \sum_{i=1}^n f(x_i, y_i, w) \]

其中 (y_i, x_i) 是一個樣本對，w是模型權重。

我們考慮使用批量大小為b的小批量隨機梯度下降（SGD）來解決上述問題。在步驟 t，該演算法首先隨機選取b個樣本，然後通過下面公式更新權重w

\[ w = w - \eta_t \sum_{i=1}^b \nabla f(x_{k_i}, y_{k_i}, w) \]

我們使用兩個例子來展示在ps-lite之中如何實現一個分散式優化演算法。

6.2 Asynchronous SGD

第一個示例中，我們將SGD擴充套件為非同步SGD。伺服器會維護模型權重w，其中server k 將獲得權重w的第k個階段，由 w_k 表示。一旦Server從worker收到梯度，server k將更新它所維護的權重。

t = 0;
while (Received(&grad)) {
  w_k -= eta(t) * grad;
  t++;
}

對於一個worker來說，每一個步驟會做四件事情

Read(&X, &Y);  // 讀取一個 minibatch 資料
Pull(&w);      // 從伺服器拉去最新的權重
ComputeGrad(X, Y, w, &grad);  // 計算梯度
Push(grad);    // 把權重推送給伺服器

ps-lite將提供push和pull函式，worker 將與具有正確部分資料的server通訊。

請注意：非同步SGD在演算法模式上與單機版本不同。由於worker之間沒有通訊，因此有可能在一個worker計算梯度的時候，其他worker就更新了伺服器上的權重。即，每個worker可能會用到延遲的權重。

6.3 Synchronized SGD

與非同步版本不同，同步版本在語義上與單機演算法相同。就是每一次迭代都要所有的worker計算好梯度，並且同步到server中。

我們使用scheduler 來管理資料同步。

for (t = 0, t < num_iteration; ++t) {
  for (i = 0; i < num_worker; ++i) {
     IssueComputeGrad(i, t);
  }
  for (i = 0; i < num_server; ++i) {
     IssueUpdateWeight(i, t);
  }
  WaitAllFinished();
}

IssueComputeGrad 和 IssueUpdateWeight 會傳送命令給 worker 和 servers，然後 scheduler 會呼叫 WaitAllFinished 等待所有傳送的命令結束。

對於一個worker接受到一個命令，它會做如下：

ExecComputeGrad(i, t) {
   Read(&X, &Y);  // 讀取資料 minibatch = batch / num_workers 個樣本
   Pull(&w);      // 從伺服器拉取最新權重
   ComputeGrad(X, Y, w, &grad);  // 計算梯度
   Push(grad);    // 把權重推送給伺服器
}

這個演算法和ASGD幾乎相同，只是每次步驟中，只有 b/num_workers個樣本被處理。

在 server 節點，與ASGD相比，多了一個聚合步驟。是把所有worker的梯度累計起來之後，再配合學習速率進行迭代。

ExecUpdateWeight(i, t) {
   for (j = 0; j < num_workers; ++j) {
      Receive(&grad);
      aggregated_grad += grad;
   }
   w_i -= eta(t) * aggregated_grad;
}

0x07 總結

PostOffice：一個單例模式的全域性管理類，每一個 node （每個 Node 可以使用 hostname + port 來唯一標識）在生命期內具有一個PostOffice，直接從字面意義可以知道，PostOffice就是郵局；
Van：通訊模組，負責與其他節點的網路通訊和Message的實際收發工作。PostOffice持有一個Van成員，直接從字面意義可以知道，Van就是小推車，用來提供送信的功能；
SimpleApp：KVServer和KVWorker的父類，它提供了簡單的Request, Wait, Response，Process功能；KVServer和KVWorker分別根據自己的使命重寫了這些功能；
Customer：每個SimpleApp物件持有一個Customer類的成員，且Customer需要在PostOffice進行註冊，該類主要負責：
- 作為一個傳送方，跟蹤由SimpleApp傳送出去的訊息的回覆情況；
- 作為接收方，維護一個Node的訊息佇列，為本Node接收訊息；
Customer 由名字就可以知道，是郵局的客戶，就是 SimpleApp 在郵局的代理人。因為需要 worker，server 需要集中精力為演算法上，所以把 worker，server 邏輯上與網路相關的收發訊息功能都總結/轉移到 Customer 之中。

下面給出了邏輯圖。

                +--------------------------+
                | Van                      |
                |                          |
DataMessage +----------->  Receiving       |
                |  1           +           |             +---------------------------+
                |              |           |             | Postoffice                |
                |              | 2         |             |                           |
                |              v           | GetCustomer |                           |
                |        ProcessDataMsg <------------------> unordered_map customers_|
                |              +           |      3      |                           |
                |              |           |             +---------------------------+
                +--------------------------+
                               |
                               |
                               | 4
                               |
                +-------------------------+
                | Customer     |          |
                |              |          |
                |              v          |
                |           Accept        |
                |              +          |
                |              |          |
                |              | 5        |
                |              v          |
                |         recv_queue_     |                +-----------------+
                |              +          |                |KVWorker         |
                |              | 6        |     +--------> |                 |
                |              |          |     |    8     |         Process |
                |              v          |     |          +-----------------+
                |          Receiving      |     |
                |              +          |     |
                |              | 7        |     |
                |              |          |     |          +-----------------+
                |              v          |     |          |KVServer         |
                |         recv_handle_+---------+--------> |                 |
                |                         |          8     |         Process |
                +-------------------------+                +-----------------+