mxnet
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros
kvstore.h
Go to the documentation of this file.
1 
6 #ifndef MXNET_KVSTORE_H_
7 #define MXNET_KVSTORE_H_
8 #include <dmlc/io.h>
9 #include <vector>
10 #include <unordered_map>
11 #include <string>
12 #include <functional>
13 #include "./ndarray.h"
14 #if MXNET_USE_DIST_KVSTORE
15 #include "ps/ps.h"
16 #endif // MXNET_USE_DIST_KVSTORE
17 
18 namespace mxnet {
25 class KVStore {
26  public:
28  virtual ~KVStore() {}
29 
40  static KVStore *Create(const char *type = "local");
41 
45  inline const std::string& type() { return type_; }
46 
63  virtual void Init(const std::vector<int>& keys,
64  const std::vector<NDArray>& values) = 0;
101  virtual void Push(const std::vector<int>& keys,
102  const std::vector<NDArray>& values,
103  int priority = 0) = 0;
127  virtual void Pull(const std::vector<int>& keys,
128  const std::vector<NDArray*>& values,
129  int priority = 0) = 0;
130 
134  typedef std::function<void(int, const NDArray&, NDArray*)> Updater;
144  virtual void set_updater(const Updater& updater) {
145  CHECK(updater) << "invalid updater";
146  updater_ = updater;
147  }
148 
149  /******************************************************
150  * the following are used for multi-machines.
151  ******************************************************/
152 
157  static void InitPSEnv(const std::unordered_map<std::string, std::string>& envs) {
158 #if MXNET_USE_DIST_KVSTORE
159  ps::Environment::Init(envs);
160 #else
161  LOG(FATAL) << "compile with USE_DIST_KVSTORE=1 to init parameter server's environment";
162 #endif // MXNET_USE_DIST_KVSTORE
163  }
164 
170  static bool IsWorkerNode() {
171 #if MXNET_USE_DIST_KVSTORE
172  const char* role_str = ps::Environment::Get()->find("DMLC_ROLE");
173  return (role_str == nullptr) || (!strcmp(role_str, "worker"));
174 #else
175  return true;
176 #endif // MXNET_USE_DIST_KVSTORE
177  }
178 
184  static bool IsServerNode() {
185 #if MXNET_USE_DIST_KVSTORE
186  const char* role_str = ps::Environment::Get()->find("DMLC_ROLE");
187  return (role_str != nullptr) && (!strcmp(role_str, "server"));
188 #else
189  return false;
190 #endif // MXNET_USE_DIST_KVSTORE
191  }
192 
193 
199  static bool IsSchedulerNode() {
200 #if MXNET_USE_DIST_KVSTORE
201  const char* role_str = ps::Environment::Get()->find("DMLC_ROLE");
202  return (role_str != nullptr) && (!strcmp(role_str, "scheduler"));
203 #else
204  return false;
205 #endif // MXNET_USE_DIST_KVSTORE
206  }
207 
214  virtual int get_rank() const {
215  return 0;
216  }
217 
221  virtual int get_group_size() const {
222  return 1;
223  }
224 
232  virtual void Barrier() { }
233 
245  virtual void SendCommandToServers(int cmd_id, const std::string& cmd_body) { }
246 
250  typedef std::function<void(int, const std::string&)> Controller;
251 
265  virtual void RunServer(const Controller& controller) { }
266 
267  protected:
272 
276  std::string type_;
277 };
278 
279 } // namespace mxnet
280 #endif // MXNET_KVSTORE_H_
distributed key-value store
Definition: kvstore.h:25
std::function< void(int, const NDArray &, NDArray *)> Updater
the prototype of user-defined updater
Definition: kvstore.h:134
virtual int get_rank() const
Definition: kvstore.h:214
static KVStore * Create(const char *type="local")
Factory function to create a new KVStore.
Updater updater_
the user-defined updater
Definition: kvstore.h:271
NDArray interface that handles array arithematics.
const std::string & type()
return the type
Definition: kvstore.h:45
virtual void Pull(const std::vector< int > &keys, const std::vector< NDArray * > &values, int priority=0)=0
pull a list of key-value pairs from the store
static bool IsSchedulerNode()
Definition: kvstore.h:199
virtual void Barrier()
global barrier among all worker machines
Definition: kvstore.h:232
static void InitPSEnv(const std::unordered_map< std::string, std::string > &envs)
initalize ps-lite environment variables
Definition: kvstore.h:157
virtual void Init(const std::vector< int > &keys, const std::vector< NDArray > &values)=0
Initialize a list of key-value pair to the store.
static bool IsWorkerNode()
Definition: kvstore.h:170
virtual ~KVStore()
virtual destructor
Definition: kvstore.h:28
virtual void RunServer(const Controller &controller)
Run as server (or scheduler)
Definition: kvstore.h:265
virtual void Push(const std::vector< int > &keys, const std::vector< NDArray > &values, int priority=0)=0
push a list of key-value pairs into the store
virtual void SendCommandToServers(int cmd_id, const std::string &cmd_body)
Send a command to all server nodes.
Definition: kvstore.h:245
std::string type_
the kvstore type
Definition: kvstore.h:276
std::function< void(int, const std::string &)> Controller
the prototype of a server controller
Definition: kvstore.h:250
virtual void set_updater(const Updater &updater)
set an updater
Definition: kvstore.h:144
virtual int get_group_size() const
Definition: kvstore.h:221
static bool IsServerNode()
Definition: kvstore.h:184