[source code analysis] machine learning parameter server PS Lite ----- application node implementation

Posted by M4F on Tue, 04 Jan 2022 09:24:14 +0100

[source code analysis] machine learning parameter server PS Lite (4) -- application node implementation

0x00 summary

This is the fourth article on parameter server, which introduces kvworker and kvserver.

Other articles in this series are:

[ Source code analysis] machine learning parameter server PS Lite (1) -- postoffice

[ Source code analysis] machine learning parameter server PS Lite (2) -- communication module Van

[ Source code analysis] machine learning parameter server PS Lite (3) -- agent Customer

Kvworker and kvserver are abstractions of Server / Worker nodes, which are defined by Van - > Customer - > recv_ handle_ To start as part of the engine.

This article will first introduce some basic support classes, then introduce the base class SimpleApp of Server / Worker, and finally introduce the specific implementation of Server / Worker.

The overall flow chart is as follows:

0x01 basic class

We first need to introduce some basic classes.

1.1 Range

The Range class is used to determine which server the parameters to be pulled are on and the Range of the key corresponding to a server according to this Range.

The Range class provides the following functions:

  • The position of uint64 in begin() and end();
  • size() gets the size of this range, i.e. end_- begin_;
class Range {
 public:
  Range() : Range(0, 0) {}
  Range(uint64_t begin, uint64_t end) : begin_(begin), end_(end) { }

  uint64_t begin() const { return begin_; }
  uint64_t end() const { return end_; }
  uint64_t size() const { return end_ - begin_; }
 private:
  uint64_t begin_;
  uint64_t end_;
};

1.2 TreadsafeQueue

TreadsafeQueue is a queue that can be read by multiple threads. It achieves thread safety through lock and condition quantity cooperation. It is used as a message queue.

/**
 * \brief thread-safe queue allowing push and waited pop
 */
class ThreadsafePQueue {
 public:
  ThreadsafePQueue() { }
  ~ThreadsafePQueue() { }

  /**
   * \brief push an value into the end. threadsafe.
   * \param new_value the value
   */
  void Push(Message new_value) {
    mu_.lock();
    queue_.push(std::move(new_value));
    mu_.unlock();
    cond_.notify_all();
  }

  /**
   * \brief wait until pop an element from the beginning, threadsafe
   * \param value the poped value
   */
  void WaitAndPop(Message* value) { // The waiting queue is not empty. Pop messages according to priority
    std::unique_lock<std::mutex> lk(mu_);
    cond_.wait(lk, [this]{return !queue_.empty();});
    *value = std::move(queue_.top());
    queue_.pop();
  }

 private:
  class Compare {
   public:
    bool operator()(const Message &l, const Message &r) {
      return l.meta.priority <= r.meta.priority;
    }
  };
  mutable std::mutex mu_; //Data synchronization mutex
  std::priority_queue<Message, std::vector<Message>, Compare> queue_; // message priority queue
  std::condition_variable cond_; //Queue is not empty condition variable
};

0x02 SimpleApp

2.1 general

SimpleApp is a base class that abstracts the functions of application nodes.

  • It provides basic sending functions and simple message processing functions (Request, Wait, Response).
  • The message types are: int type head and string type body.
  • It has 2 derived classes. KVServer and KVWorker.

2.2 definitions

2.2. 1 support

SimpleData defines the basic format of Request and Response.

struct SimpleData {
  /** \brief the int head */
  int head;
  /** \brief the string body */
  std::string body;
  /** \brief sender's node id */
  int sender;
  /** \brief the associated timestamp */
  int timestamp;
  /** \brief sender's customer id */
  int customer_id;
};

2.2. 2 member variables

SimpleApp mainly has the following member variables:

  • Customer* obj_ : The customer of this App controls the connection request;
  • Handle request_handle_ : Request processing function;
  • Handle response_handle_ : Response processing function;
  • set_request_handle,set_response_handle: set member request_handle_, response_handle_. When the client calls SimpleApp::Process, according to message The indicator variable in meta determines whether it is request or response, and calls the corresponding handle for processing;
class SimpleApp {
 public:
  /**
   * \brief constructor
   * @param app_id the app id, should match with the remote node app with which this app
   * @param customer_id the customer_id, should be node-locally unique
   * is communicated
   */
  explicit SimpleApp(int app_id, int customer_id);

  /** \brief deconstructor */
  virtual ~SimpleApp() { delete obj_; obj_ = nullptr; }

  /**
   * \brief send a request to a remote node
   *
   * \param req_head request head
   * \param req_body request body
   * \param recv_id remote node id
   *
   * @return the timestamp of this request
   */
  virtual inline int Request(int req_head, const std::string& req_body, int recv_id);

  /**
   * \brief wait until a request is finished
   *
   * \param timestamp
   */
  virtual inline void Wait(int timestamp) { obj_->WaitRequest(timestamp); }


  /**
   * \brief send back a response for a request
   * \param recv_req the received request
   * \param the response body
   */
  virtual inline void Response(const SimpleData& recv_req, const std::string& res_body = "");

  /**
   * \brief the handle to proces a received request/respoonse
   *
   * \param recved the received request or response
   * \param app this pointer
   */
  using Handle = std::function<void(const SimpleData& recved, SimpleApp* app)>;

  /**
   * \brief set the request handle
   * \param request_handle the request handle
   */
  virtual inline void set_request_handle(const Handle& request_handle) {
    CHECK(request_handle) << "invalid request handle";
    request_handle_ = request_handle;
  }

  /**
   * \brief set the response handle
   * \param response_handle the response handle
   */
  virtual inline void set_response_handle(const Handle& response_handle) {
    CHECK(response_handle) << "invalid response handle";
    response_handle_ = response_handle;
  }

  /**
   * \brief returns the customer
   */
  virtual inline Customer* get_customer() { return obj_; }

 protected:
  /** \brief empty construct */
  inline SimpleApp() : obj_(nullptr) {
    request_handle_ = [](const SimpleData& recved, SimpleApp* app) {
      app->Response(recved);
    };
    response_handle_ = [](const SimpleData& recved, SimpleApp* app) { };
  }

  /** \brief process a received message */
  virtual inline void Process(const Message& msg);

  /** \brief ps internal object */
  Customer* obj_;

 private:
  /** \brief request handle */
  Handle request_handle_;
  /** \brief request handle */
  Handle response_handle_;
};

2.3 function

Three simple functions are as follows:

Request is to call Van to send a message.

inline int SimpleApp::Request(int req_head, const std::string& req_body, int recv_id) {
  // setup message
  Message msg;
  msg.meta.head = req_head;
  if (req_body.size()) msg.meta.body = req_body;
  int ts = obj_->NewRequest(recv_id);
  msg.meta.timestamp = ts;
  msg.meta.request = true;
  msg.meta.simple_app = true;
  msg.meta.app_id = obj_->app_id();
  msg.meta.customer_id = obj_->customer_id();

  // send
  for (int r : Postoffice::Get()->GetNodeIDs(recv_id)) {
    msg.meta.recver = r;
    Postoffice::Get()->van()->Send(msg);
  }
  return ts;
}

Response is to call Van to reply to the message.

inline void SimpleApp::Response(const SimpleData& req, const std::string& res_body) {
  // setup message
  Message msg;
  msg.meta.head = req.head;
  if (res_body.size()) msg.meta.body = res_body;
  msg.meta.timestamp = req.timestamp;
  msg.meta.request = false;
  msg.meta.simple_app = true;
  msg.meta.app_id = obj_->app_id();
  msg.meta.customer_id = req.customer_id;
  msg.meta.recver = req.sender;

  // send
  Postoffice::Get()->van()->Send(msg);
}

The Process function is based on message The indicator variable in meta determines whether it is request or response, and calls the corresponding handle for processing.

inline void SimpleApp::Process(const Message& msg) {
  SimpleData recv;
  recv.sender    = msg.meta.sender;
  recv.head      = msg.meta.head;
  recv.body      = msg.meta.body;
  recv.timestamp = msg.meta.timestamp;
  recv.customer_id = msg.meta.customer_id;
  if (msg.meta.request) { // Judge whether it is request or response, and call the corresponding handle for processing
    CHECK(request_handle_);
    request_handle_(recv, this);
  } else {
    CHECK(response_handle_);
    response_handle_(recv, this);
  }
}

0x03 KVServer

KVServer is an abstraction of the Server node. Its function is to receive information, process information and return results. Its main functions are:

  • Maintain key value pairs data;
  • Handle & respond to push & pull requests from clients;
    • Function request_handle_ Processing request:
      • Request is called when KVServer::Process is called_ handle_ .
      • request_handle_ The default is KVServerDefaultHandle.
    • The Response function is used to return data;

3.1 definitions

request_handle_ It is a request processing function and needs to be customized.

  • In this callback function, the user needs to implement the model weight gradient update algorithm and model weight return operation of various optimizers.
  • You can directly refer to the default version of KVServerDefaultHandle implemented by PS Lite.
/**
 * \brief A server node for maintaining key-value pairs
 */
template <typename Val>
class KVServer : public SimpleApp {
 public:
  /**
   * \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));
  }

  /** \brief deconstructor */
  virtual ~KVServer() { delete obj_; obj_ = nullptr; }

  /**
   * \brief the handle to process a push/pull request from a worker
   * \param req_meta meta-info of this request
   * \param req_data kv pairs of this request
   * \param server this pointer
   */
  using ReqHandle = std::function<void(const KVMeta& req_meta,
                                       const KVPairs<Val>& req_data,
                                       KVServer* server)>;
  void set_request_handle(const ReqHandle& request_handle) {
    CHECK(request_handle) << "invalid request handle";
    request_handle_ = request_handle;
  }

  /**
   * \brief response to the push/pull request
   * \param req the meta-info of the request
   * \param res the kv pairs that will send back to the worker
   */
  void Response(const KVMeta& req, const KVPairs<Val>& res = KVPairs<Val>());

 private:
  /** \brief internal receive handle */
  void Process(const Message& msg);
  /** \brief request handle */
  ReqHandle request_handle_; // You need to implement it yourself
};

3.2 function

3.2.1 Response

Response() is to send response information to the called worker. Compared with SimpleApp, KVServer has new processing for head and body.

Note that the response function should be a user-defined request_handle_ Call, i.e. request_handle_ Process the received message, then call Response to reply to worker.

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;
  if (res.keys.size()) {
    msg.AddData(res.keys);
    msg.AddData(res.vals);
    if (res.lens.size()) {
      msg.AddData(res.lens);
    }
  }
  Postoffice::Get()->van()->Send(msg);
}

3.2.2 Process

Process() is registered in the Customer object. When the receiving thread of the Customer object receives a message, Process() is called to process the data.

The internal logic of Process() is:

  • Extract the meta information of the message and build a KVMeta.
  • As you can see, there is no maintenance of KV data in Process.
  • Process calls a request implemented by the user_ handle_ (std::function function object) process the data.
  • In the callback function request_handle_ Users need to implement the model weight gradient update algorithm and model weight return operation of various optimizers.
template <typename Val>
void KVServer<Val>::Process(const Message& msg) {
  if (msg.meta.simple_app) {
    SimpleApp::Process(msg); return;
  }
  KVMeta meta;
  meta.cmd       = msg.meta.head;
  meta.push      = msg.meta.push;
  meta.pull      = msg.meta.pull;
  meta.sender    = msg.meta.sender;
  meta.timestamp = msg.meta.timestamp;
  meta.customer_id = msg.meta.customer_id;
  KVPairs<Val> data;
  int n = msg.data.size();
  if (n) {
    CHECK_GE(n, 2);
    data.keys = msg.data[0];
    data.vals = msg.data[1];
    if (n > 2) {
      CHECK_EQ(n, 3);
      data.lens = msg.data[2];
      CHECK_EQ(data.lens.size(), data.keys.size());
    }
  }
  CHECK(request_handle_);
  request_handle_(meta, data, this);
}

3.2. 3 example function

KVServerDefaultHandle is an example provided by PS Lite to demonstrate how to maintain KV, process messages and return requests.

A hash table unordered is maintained here_ Map, record key and value, and respond to push and pull requests.

Use std::unordered_map store stores the parameters of the server. When the request is push, the store parameters are updated. When the request is pull, the parameters are pulled;

/**
 * \brief an example handle adding pushed kv into store
 */
template <typename Val>
struct KVServerDefaultHandle {
  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.pull) {
      CHECK_EQ(n, req_data.vals.size());
    } else {
      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];
      }
      if (req_meta.pull) {
        res.vals[i] = store[key];
      }
    }
    server->Response(req_meta, res);
  }
  std::unordered_map<Key, Val> store;
};

3.2. 4 process

We will continue to sort out and refine the process above.

  • The worker node or server node executes Postoffice::start() at the beginning of the program.

  • Postoffice::start() initializes the node information and calls Van::start().

  • Each node listens to a local port; The connected node was already connected at startup.

  • Van::start() starts a local thread to receive socket information, and uses Van::Receiving() to continuously listen for received message s.

    • receiver_thread_ = std::unique_ptr<std::thread>(new std::thread(&Van::Receiving, this));
      
  • After receiving the message, Van::Receiving() performs different actions according to different commands. For data messages, if the next step is required, ProcessDataMsg will be called:

    • Find the customer according to the app id in the message (each app task will bind a custom class), that is, the message will be sent to the recv thread of different customers according to different customer IDs.
    • Pass the message to the Customer::Accept function.
  • The Customer::Accept() function adds the message to a queue recv_queue_;

  • The Customer object itself starts an accept thread recv_thread_, Using Customer::Receiving():

    • Constantly from recv_queue_ Queue fetch message.
    • If (! recv.meta.request), it means response, and the tracker_[req.timestamp].second++
    • Call the registered user-defined recv_handle_ Function to process the message.
  • For workers, their registered recv_handle_ Is the KVWorker::Process() function. Because the message received by the worker's recv thread is mainly the KV pair pull ed from the server, the process () mainly receives the KV pair in the message;

  • For the Server, its registered recv_handle_ Is the KVServer::Process() function.

  • Because we are KVServer here, and server accepts the KV pairs of worker's push, which need to be processed, so the users invoked in the Process() function are KVServer:: set_. request_ The function object passed in by handle().

  • In user-defined request_handle_ In the function, if you need to send a response to the worker, call KVServer::Response.

The current logic is shown in the figure below. In step 8, recv_handle_ Point to KVServer::Process or KVWorker::Process (this section is server, so the corresponding is KVServer::Process). In step 10, return the response to the worker.

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

0x04 KVWorker

4.1 General

KVWorker is used to push and pull key value pairs to the server node, which are various parameters that need to be processed in parallel during the algorithm process.

  • Both push and pull operations in Worker return an ID asynchronously, and then use the ID to wait for blocking, that is, synchronous operation.
  • Or, when calling asynchronously, a Callback is passed in for subsequent operations.

4.2 definitions

The main variables of KVWorker are:

  • std::unordered_ map<int, std::vector<KVPairs>> recv_ KVs: pull result received: kv value;
  • std::unordered_ Map < int, callback > callbacks: the callback function executed after receiving all the response s of the request;
  • Slicer slicer_ : The default slice function variable, which slices KVPairs according to the Range of each server when calling the Send function;

The main functions are:

  • ZPush zero copy push function

  • ZPull zero copy pull function

  • AddPullCB key reorganization function

  • Process message processing function

  • DefaultSlicer default slice handler

  • set_slicer: set slicer_ Member, which slices KVPairs according to the Range of each server when calling the Send function;

/**
 * \brief A worker node that can \ref Push (\ref Pull) key-value pairs to (from) server
 * nodes
 *
 * \tparam Val the type of value, which should be primitive types such as
 * int32_t and float
 */
template<typename Val>
class KVWorker : public SimpleApp {
 public:
  /** avoid too many this-> */
  using SimpleApp::obj_; // Customer object
  /**
   * \brief callback function for \ref Push and \ref Pull
   *
   * It is called by the data receiving thread of this instance when the push or
   * pull is actually finished. Namely the kv pairs have already written into
   * servers' data structure or the kv pairs have already pulled back.
   */
  using Callback = std::function<void()>;

  /**
   * \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));
  }

  /** \brief deconstructor */
  virtual ~KVWorker() { delete obj_; obj_ = nullptr; }

  using SlicedKVs = std::vector<std::pair<bool, KVPairs<Val>>>;
  /**
   * \brief a slicer partitions a key-value list according to the key ranges
   * \param send the kv list for partitioning
   * \param ranges the key ranges, ranges[i] is the key range of server i
   * \param sliced the sliced lists. slices[i] should only contains keys in
   * ranges[i] and the according values
   */
  using Slicer = std::function<void(
      const KVPairs<Val>& send, const std::vector<Range>& ranges,
      SlicedKVs* sliced)>;

  /**
   * \brief set a user-defined slicer
   */
  void set_slicer(const Slicer& slicer) {
    CHECK(slicer); slicer_ = slicer;
  }

 private:
  /**
   * \brief add a callback for a request. threadsafe.
   * @param cb callback
   * @param timestamp the timestamp of the request
   */
  void AddCallback(int timestamp, const Callback& cb) {
    if (!cb) return;
    std::lock_guard<std::mutex> lk(mu_);
    callbacks_[timestamp] = cb;
  }

  /** \brief data buffer for received kvs for each timestamp */
  std::unordered_map<int, std::vector<KVPairs<Val>>> recv_kvs_; // kv value received 
  /** \brief callbacks for each timestamp */
  std::unordered_map<int, Callback> callbacks_; // Callback function executed after receiving all response s of request
  /** \brief lock */
  std::mutex mu_;
  /** \brief kv list slicer */
  Slicer slicer_; // Default slice function variable
};

4.3 function

4.3.1 Push & ZPush

Because Push calls ZPush, we will introduce it together.

The main methods of Push are:

  • Send the data (KV list) to the corresponding server node;
  • The KV list is sent by zone according to the Key range maintained by each server;
  • Push is an asynchronous direct return. If you want to know the return result, you can:
    • Use Wait to Wait, that is, use tracker_ To record the sent request amount and the corresponding response request amount. When the sent amount is equal to the received amount, it means that each request has been successfully sent, so as to achieve the purpose of synchronization;
    • Use callback so that you can call back to at the end.

ZPush method is:

  • Using obj_ (Customer type) to record the sent request quantity and the corresponding response request quantity, and return a timestamp;
  • Set the callback corresponding to the timestamp;
  • Construct the KVPair object with the passed in parameters, and call Send to Send the object;
  int Push(const std::vector<Key>& keys,
           const std::vector<Val>& vals,
           const std::vector<int>& lens = {},
           int cmd = 0,
           const Callback& cb = nullptr,
           int priority = 0) {
    return ZPush(
        SArray<Key>(keys), SArray<Val>(vals), SArray<int>(lens), cmd, cb,
        priority);
  }
  
  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;
  }  

How to call can refer to its notes:

   * Sample usage: the following codes push two KV pairs `{1, (1.1, 1.2)}` and `{3,
   * (3.1,3.2)}` to server nodes, where the value is a length-2 float vector
   * \code
   *   KVWorker<float> w;
   *   std::vector<Key> keys = {1, 3};
   *   std::vector<float> vals = {1.1, 1.2, 3.1, 3.2};
   *   w.Push(keys, vals);
   * \endcode

4.3.2 Pull

The logic of the pull method is similar to that of the push method:

  • Bind a callback function to copy data and get a timestamp.
  • According to key_vector pulls Val from the Server_ vector,
  • Finally, the timestamp is returned,
  • This function is not blocked. You can use worker Wait (timestamp);
  int Pull(const std::vector<Key>& keys,
           std::vector<Val>* vals,
           std::vector<int>* lens = nullptr,
           int cmd = 0,
           const Callback& cb = nullptr,
           int priority = 0) {
    SArray<Key> skeys(keys);
    int ts = AddPullCB(skeys, vals, lens, cmd, cb);
    KVPairs<Val> kvs;
    kvs.keys = skeys;
    kvs.priority = priority;
    Send(ts, false, true, cmd, kvs);
    return ts;
  }

4.3.3 ZPull

The logic is consistent with Pull, but the process of copying to the system is omitted. Therefore, it is necessary to ensure that the caller does not change the key_ before ZPull completes. vector;

  int ZPull(const SArray<Key>& keys,
            SArray<Val>* vals,
            SArray<int>* lens = nullptr,
            int cmd = 0,
            const Callback& cb = nullptr,
            int priority = 0) {
    int ts = AddPullCB(keys, vals, lens, cmd, cb);
    KVPairs<Val> kvs;
    kvs.keys = keys;
    kvs.priority = priority;
    Send(ts, false, true, cmd, kvs);
    return ts;
  }

4.3.4 Send

Both Push() and Pull() will call the send() function at last. Send() will segment KVPairs. Because each Server only retains some parameters, SlicedKVpairs will be sent to different servers.

If it is skipped, callback will be called directly.

Otherwise, the traversal is sent.

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.5 DefaultSlicer

The slicing function can be rewritten by the user. The default is DefaultSlicer. Each SlicedKVPairs is wrapped into a Message object and sent with van::send().

Partition the data to be sent according to the partition range information of STD:: vector & ranges. Currently, Postoffice::GetServerKeyRanges is used by default to divide the partition range.

template <typename Val>
void KVWorker<Val>::DefaultSlicer(
    const KVPairs<Val>& send, const std::vector<Range>& ranges,
    typename KVWorker<Val>::SlicedKVs* sliced) {
  sliced->resize(ranges.size());

  // find the positions in msg.key
  size_t n = ranges.size();
  std::vector<size_t> pos(n+1);
  const Key* begin = send.keys.begin();
  const Key* end = send.keys.end();
  for (size_t i = 0; i < n; ++i) {
    if (i == 0) {
      pos[0] = std::lower_bound(begin, end, ranges[0].begin()) - begin;
      begin += pos[0];
    } else {
      CHECK_EQ(ranges[i-1].end(), ranges[i].begin());
    }
    size_t len = std::lower_bound(begin, end, ranges[i].end()) - begin;
    begin += len;
    pos[i+1] = pos[i] + len;

    // don't send it to servers for empty kv
    sliced->at(i).first = (len != 0);
  }
  CHECK_EQ(pos[n], send.keys.size());
  if (send.keys.empty()) return;

  // the length of value
  size_t k = 0, val_begin = 0, val_end = 0;
  if (send.lens.empty()) {
    k = send.vals.size() / send.keys.size();
    CHECK_EQ(k * send.keys.size(), send.vals.size());
  } else {
    CHECK_EQ(send.keys.size(), send.lens.size());
  }

  // slice
  for (size_t i = 0; i < n; ++i) {
    if (pos[i+1] == pos[i]) {
      sliced->at(i).first = false;
      continue;
    }
    sliced->at(i).first = true;
    auto& kv = sliced->at(i).second;
    kv.keys = send.keys.segment(pos[i], pos[i+1]);
    if (send.lens.size()) {
      kv.lens = send.lens.segment(pos[i], pos[i+1]);
      for (int l : kv.lens) val_end += l;
      kv.vals = send.vals.segment(val_begin, val_end);
      val_begin = val_end;
    } else {
      kv.vals = send.vals.segment(pos[i]*k, pos[i+1]*k);
    }
  }
}


4.3.6 PushPull & ZPushPull

Is to aggregate push and pull together.

  int PushPull(const std::vector<Key>& keys,
               const std::vector<Val>& vals,
               std::vector<Val>* outs,
               std::vector<int>* lens = nullptr,
               int cmd = 0,
               const Callback& cb = nullptr,
               int priority = 0) {
    CHECK_NOTNULL(outs);
    if (outs->empty())
      outs->resize(vals.size());
    else
      CHECK_EQ(vals.size(), outs->size());

    SArray<Key> skeys(keys);
    SArray<Val> svals(vals);
    auto souts = new SArray<Val>(outs->data(), outs->size());
    SArray<int>* slens = lens ?
        new SArray<int>(lens->data(), lens->size()) : nullptr;
    int ts = ZPushPull(skeys, svals, souts, slens, cmd,
        [this, cb, souts, slens]() {
          delete souts;
          delete slens;
          if (cb) cb();
        }, priority);
    return ts;
  }

  int ZPushPull(const SArray<Key>& keys,
                const SArray<Val>& vals,
                SArray<Val>* outs,
                SArray<int>* lens = nullptr,
                int cmd = 0,
                const Callback& cb = nullptr,
                int priority = 0) {
    int ts = AddPullCB(keys, outs, lens, cmd, cb);
    KVPairs<Val> kvs;
    kvs.keys = keys;
    kvs.vals = vals;
    kvs.priority = priority;
    if (lens)
      kvs.lens = *lens;
    Send(ts, true, true, cmd, kvs);
    re

4.3.7 Callback related

We mentioned the settings of some callback functions. Let's see how to use them.

4.3. 7.1 setting

We can see that for each timestamp, a callback function is set to form a list of callback functions.

Each time a request is sent, a callback function is registered in this list.

  using Callback = std::function<void()>;
  
  /** \brief callbacks for each timestamp */
  std::unordered_map<int, Callback> callbacks_;  // Callback function list
  
  void AddCallback(int timestamp, const Callback& cb) {
    if (!cb) return;
    std::lock_guard<std::mutex> lk(mu_);
    callbacks_[timestamp] = cb; // Add callback function
  }

4.3.7.2 AddPullCB

This is the callback function to get the response after pull, which is used to copy the returned data.

However, if multiple servers should have returns, how should they be handled? No matter push or pull, the value pulled from each server will be filled into the local vals only after receiving all the responses.

template <typename Val>
template <typename C, typename D>
int KVWorker<Val>::AddPullCB(
    const SArray<Key>& keys, C* vals, D* lens, int cmd,
    const Callback& cb) {
  int ts = obj_->NewRequest(kServerGroup);
  AddCallback(ts, [this, ts, keys, vals, lens, cb]() mutable {
      mu_.lock();
      auto& kvs = recv_kvs_[ts];
      mu_.unlock();

      // do check
      size_t total_key = 0, total_val = 0;
      for (const auto& s : kvs) { // Conduct effectiveness verification
        Range range = FindRange(keys, s.keys.front(), s.keys.back()+1);
        CHECK_EQ(range.size(), s.keys.size())
            << "unmatched keys size from one server";
        if (lens) CHECK_EQ(s.lens.size(), s.keys.size());
        total_key += s.keys.size();
        total_val += s.vals.size();
      }
      CHECK_EQ(total_key, keys.size()) << "lost some servers?";

      // fill vals and lens
      std::sort(kvs.begin(), kvs.end(), [](
          const KVPairs<Val>& a, const KVPairs<Val>& b) {
                  return a.keys.front() < b.keys.front();
        });
      CHECK_NOTNULL(vals);
      if (vals->empty()) {
        vals->resize(total_val);
      } else {
        CHECK_EQ(vals->size(), total_val);
      }
      Val* p_vals = vals->data();
      int *p_lens = nullptr;
      if (lens) {
        if (lens->empty()) {
          lens->resize(keys.size());
        } else {
          CHECK_EQ(lens->size(), keys.size());
        }
        p_lens = lens->data();
      }
      for (const auto& s : kvs) { // Copy returned data
        memcpy(p_vals, s.vals.data(), s.vals.size() * sizeof(Val));
        p_vals += s.vals.size();
        if (p_lens) {
          memcpy(p_lens, s.lens.data(), s.lens.size() * sizeof(int));
          p_lens += s.lens.size();
        }
      }

      mu_.lock();
      recv_kvs_.erase(ts);
      mu_.unlock();
      if (cb) cb();
    });

  return ts;
}

4.3. 7.3 operation

Find the callback function according to the timestamp, run it, and then delete it.

When to call, that is, it will be called in Process. We will introduce it right away.

template <typename Val>
void KVWorker<Val>::RunCallback(int timestamp) {
  mu_.lock();
  auto it = callbacks_.find(timestamp);
  if (it != callbacks_.end()) {
    mu_.unlock();

    CHECK(it->second);
    it->second();

    mu_.lock();
    callbacks_.erase(it);
  }
  mu_.unlock();
}

4.3.8 Process

If it is a Pull Response, the recv will be saved first in the values returned by each received Response_ kvs_ Inside, recv_kvs_[ts].push_back(kvs);

No matter push or pull, the value pulled from each server will be filled into the local vals only after receiving all the responses.

template <typename Val>
void KVWorker<Val>::Process(const Message& msg) {
  if (msg.meta.simple_app) {
    SimpleApp::Process(msg); return;
  }
  // store the data for pulling
  int ts = msg.meta.timestamp;
  if (msg.meta.pull) {
    CHECK_GE(msg.data.size(), (size_t)2);
    KVPairs<Val> kvs;
    kvs.keys = msg.data[0];
    kvs.vals = msg.data[1];
    if (msg.data.size() > (size_t)2) {
      kvs.lens = msg.data[2];
    }
    mu_.lock();
    recv_kvs_[ts].push_back(kvs);
    mu_.unlock();
  }

  // Finished, run calls, only after all responses are received
  if (obj_->NumResponse(ts) == Postoffice::Get()->num_servers() - 1)  {
    RunCallback(ts); // RunCallback is called here.
  }
}

0x05 summary

Finally, let's make a summary with a messaging process to see how each part is used. The overall flow chart is as follows:

  1. The worker node wants to Send a message, so the Send method is called.
  2. The Send method will call the Customer's NewRequest to create a new request.
  3. Postoffice::start() initializes the node information and calls Van::start().
  4. The send method will call Van's send method for network interaction.
  5. After the network transmission, the process comes to the Server. For the Server, this is a Request and calls Van's Receiving. After Receiving the message, Van::Receiving() performs different actions according to different commands. For data messages, ProcessDataMsg will be called if further processing is required.
  6. Continue calling to ProcessDataMsg of Van, and then call GetCustomer.
  7. GetCustomer will call Postoffice. For customers_ Handle accordingly.
  8. The Customer will use Accept to process the message.
  9. The Customer::Accept() function adds the message to a queue recv_queue_.
  10. The Customer object itself starts an accept thread recv_thread_, Using Customer::Receiving():
    1. Constantly from recv_queue_ Queue fetch message.
    2. If (! recv.meta.request), it means response, and the tracker_[req.timestamp].second++
    3. Call the registered user-defined recv_handle_ Function to process the message.
  11. Van::Receiving() calls the registered user-defined recv_handle_ Function to process the message.
  12. For Server, its registered recv_handle_ Is the KVServer::Process() function.
  13. Process function call request_handle_ Continue processing, generate a Response and return it to the Worker.
  14. The Response is passed to the Worker through the network.
  15. The operation returned to the Worker and came to the Van of the Worker. For the Worker, this is a Request, which calls Van's Receiving. (the following sequence of operations is similar to Server)
  16. After receiving the message, Van::Receiving() performs different actions according to different commands. For data messages, ProcessDataMsg will be called if further processing is required.
  17. The Customer will use Accept to process the message.
  18. The Customer::Accept() function adds the message to a queue recv_queue_.
  19. Here's a decoupling by a new thread recv_thread_ handle.
  20. The Customer object itself has started a new thread recv_thread_, Use Customer::Receiving().
  21. For workers, their registered recv_handle_ Is the KVWorker::Process() function.
  22. Call the KVWorker::Process() function to process the Response message Response.
+---------------------+       +------------------------+   Worker   +  Server            +--------------------------+
| KVWorker            |  1    |  Van                   |      3     |                    | Van                      |
|          Send  +--------+--------------->  send +-----------------+----->  Request +----------->  Receiving       |
|                     |   |   |                        |                                 |              +           |
|                     |   |   |       Receiving   <---------+       |           4        |              |           |             +---------------------------+
|                     |   |   |           +            |    |       |                    |              |           |             | Postoffice                |
|    Process          |   |   |           | 16         |    |       |                    |              | 5         |             |                           |
|      ^              |   |   |           v            |    | 15    |                    |              v           | GetCustomer |                           |
|      |              |   |   |     ProcessDataMsg     |    |       |                    |        ProcessDataMsg <------------------> unordered_map customers_|
|      |              |   |   |           +            |    |       |                    |              +           |      6      |                           |
|      |              |   |   |           |            |    |       |                    |              |           |             +---------------------------+
+---------------------+   |   +------------------------+    |       |                    +--------------------------+
       |                  |               |                 |       |                                   |
       |                  |2              |  17             |       |                                   | 7
       |                  |               |                 |       |                                   |
       |     +---------------------------------------+      |       |         +------------------------------------+
       |     | Customer   |               |          |      |       |         | Customer                |          |
       |     |            |               v          |      |       |         |                         |          |
       |     |            v                          |      |       |         |                         v          |
       |     |        NewRequest        Accept       |      |       |         |                      Accept        |
       |     |                            +          |      |       |         |                         +          |
       |     |                            |  18      |      |       |         |                         |          |
       |     |                            |          |      |       |         |                         | 8        |
       |     |                            v          |      |       |         |                         v          |
       |     |                      revc_queue_      |      |       |         |                    recv_queue_     |
       |     |                            +          |      |       |         |                         +          |
    22 |     |                            |  19      |      |       |         |                         | 9        |
       |     |                            |          |      |       |         |                         |          |
       |     |                  20        v          |      |       |         |                10       v          |
       |     | recv_thread_ +-------> Receving       |      |       |         |  recv_thread_ +---> Receiving      |
       |     |                            |          |      |       |         |                         +          |
       |     |                            |  21      |      |       |         |                         | 11       |
       |     |                            |          |      |       |         |                         |          |                +------------------+
       |     |                            v          |      |       |         |                         v          |                |KVServer          |
       +---------------------------+ recv_handle     |      |       |         |                    recv_handle_+------------------> |                  |
             |                                       |      |       |         |                                    |          12    |         Process  |
             +---------------------------------------+      |       |         +------------------------------------+                |           +      |
                                                            |       |                                                               +------------------+
                                                            |       |                                                                           |
                                                            |       |                                                                           |  13
                                                            |       |                                                                           v
                                                            |       |                                                               +-----------+-------+
                                                            |       |                                                               | request_handle_   |
                                                            |       |                                            14                 |                   |
                                                            +<-----------+   Response <----------------------------------------------------+ Response   |
                                                                    |                                                               |                   |
                                                                    |                                                               +-------------------+
                                                                    +

Mobile phones are as follows:

So far, the introduction of PS Lite has been completed. Let's start with the introduction of Douban parameter server Paracel. Please look forward to it.

0xEE personal information

★★★★★★★ thinking about life and technology ★★★★★★

Wechat public account: Rossi's thinking

If you want to get the news push of personal articles in time, or want to see the technical materials recommended by yourself, please pay attention.

0xFF reference

The most comprehensive understanding of PS Lite in history

Implementation of machine learning parameter server framework from zero (2)

Topics: Machine Learning Deep Learning