6 #ifndef MXNET_KVSTORE_H_
7 #define MXNET_KVSTORE_H_
10 #include <unordered_map>
14 #if MXNET_USE_DIST_KVSTORE
16 #endif // MXNET_USE_DIST_KVSTORE
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;
134 typedef std::function<void(int, const NDArray&, NDArray*)>
Updater;
145 CHECK(updater) <<
"invalid updater";
157 static void InitPSEnv(
const std::unordered_map<std::string, std::string>& envs) {
158 #if MXNET_USE_DIST_KVSTORE
159 ps::Environment::Init(envs);
161 LOG(FATAL) <<
"compile with USE_DIST_KVSTORE=1 to init parameter server's environment";
162 #endif // MXNET_USE_DIST_KVSTORE
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"));
176 #endif // MXNET_USE_DIST_KVSTORE
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"));
190 #endif // MXNET_USE_DIST_KVSTORE
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"));
205 #endif // MXNET_USE_DIST_KVSTORE
250 typedef std::function<void(int, const std::string&)>
Controller;
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