mxnet
kvstore.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
25 #ifndef MXNET_KVSTORE_H_
26 #define MXNET_KVSTORE_H_
27 #include <dmlc/io.h>
28 #include <vector>
29 #include <utility>
30 #include <unordered_map>
31 #include <string>
32 #include <functional>
33 #include <atomic>
34 #include "../../src/kvstore/gradient_compression.h"
35 #include "./ndarray.h"
36 #if MXNET_USE_DIST_KVSTORE
37 #include "ps/ps.h"
38 #endif // MXNET_USE_DIST_KVSTORE
39 
40 namespace mxnet {
47 class KVStore {
48  public:
50  virtual ~KVStore() {}
51 
62  static KVStore *Create(const char *type = "local");
63 
67  inline const std::string& type() { return type_; }
68 
74  virtual void SetGradientCompression(const std::vector<std::pair<std::string, std::string> >
75  & kwargs) = 0;
76 
93  virtual void Init(const std::vector<int>& keys,
94  const std::vector<NDArray>& values) = 0;
100  virtual void Init(const std::vector<std::string>& str_keys,
101  const std::vector<NDArray>& values) = 0;
138  virtual void Push(const std::vector<int>& keys,
139  const std::vector<NDArray>& values,
140  int priority = 0) = 0;
141 
148  virtual void Push(const std::vector<std::string>& str_keys,
149  const std::vector<NDArray>& values,
150  int priority = 0) = 0;
175  virtual void Pull(const std::vector<int>& keys,
176  const std::vector<NDArray*>& values,
177  int priority = 0, bool ignore_sparse = true) = 0;
185  virtual void Pull(const std::vector<std::string>& str_keys,
186  const std::vector<NDArray*>& values,
187  int priority = 0, bool ignore_sparse = true) = 0;
188 
197  virtual void PullRowSparse(const std::vector<int>& str_keys,
198  const std::vector<std::pair<NDArray*, NDArray>>& val_rowids,
199  int priority = 0) = 0;
200 
209  virtual void PullRowSparse(const std::vector<std::string>& str_keys,
210  const std::vector<std::pair<NDArray*, NDArray>>& val_rowids,
211  int priority = 0) = 0;
212 
216  typedef std::function<void(int, const NDArray&, NDArray*)> Updater;
220  typedef std::function<void(const std::string&, const NDArray&, NDArray*)> StrUpdater;
230  virtual void set_updater(const Updater& updater) {
231  CHECK(updater) << "invalid updater";
232  updater_ = updater;
233  }
234 
244  virtual void set_updater(const StrUpdater& updater) {
245  CHECK(updater) << "invalid updater";
246  str_updater_ = updater;
247  }
248 
249  /******************************************************
250  * the following are used for multi-machines.
251  ******************************************************/
252 
257  static void InitPSEnv(const std::unordered_map<std::string, std::string>& envs) {
258 #if MXNET_USE_DIST_KVSTORE
259  ps::Environment::Init(envs);
260 #else
261  LOG(FATAL) << "compile with USE_DIST_KVSTORE=1 to init parameter server's environment";
262 #endif // MXNET_USE_DIST_KVSTORE
263  }
264 
270  static bool IsWorkerNode() {
271 #if MXNET_USE_DIST_KVSTORE
272  const char* role_str = ps::Environment::Get()->find("DMLC_ROLE");
273  return (role_str == nullptr) || (!strcmp(role_str, "worker"));
274 #else
275  return true;
276 #endif // MXNET_USE_DIST_KVSTORE
277  }
278 
284  static bool IsServerNode() {
285 #if MXNET_USE_DIST_KVSTORE
286  const char* role_str = ps::Environment::Get()->find("DMLC_ROLE");
287  return (role_str != nullptr) && (!strcmp(role_str, "server"));
288 #else
289  return false;
290 #endif // MXNET_USE_DIST_KVSTORE
291  }
292 
293  void set_barrier_before_exit(const bool barrier_before_exit) {
294 #if MXNET_USE_DIST_KVSTORE
295  if (!IsWorkerNode()) LOG(FATAL) << "barrier_before_exit takes effect only on worker nodes";
296  barrier_before_exit_ = barrier_before_exit;
297 #else
298  LOG(FATAL) << "compile with USE_DIST_KVSTORE=1 to enable barrier";
299 #endif
300  }
301 
307  static bool IsSchedulerNode() {
308 #if MXNET_USE_DIST_KVSTORE
309  const char* role_str = ps::Environment::Get()->find("DMLC_ROLE");
310  return (role_str != nullptr) && (!strcmp(role_str, "scheduler"));
311 #else
312  return false;
313 #endif // MXNET_USE_DIST_KVSTORE
314  }
315 
322  virtual int get_rank() const {
323  return 0;
324  }
325 
329  virtual int get_group_size() const {
330  return 1;
331  }
332 
341  virtual int get_num_dead_node(int node_id, int timeout = 60) const {
342  return 0;
343  }
344 
352  virtual void Barrier() { }
353 
365  virtual void SendCommandToServers(int cmd_id, const std::string& cmd_body) { }
366 
370  typedef std::function<void(int, const std::string&)> Controller;
371 
385  virtual void RunServer(const Controller& controller) { }
386 
387  protected:
391  Updater updater_;
392 
396  StrUpdater str_updater_;
397 
401  std::string type_;
402 
407  std::shared_ptr<kvstore::GradientCompression> gradient_compression_;
408 
412  std::atomic<bool> barrier_before_exit_{true};
413 };
414 
415 } // namespace mxnet
416 #endif // MXNET_KVSTORE_H_
distributed key-value store
Definition: kvstore.h:47
std::function< void(int, const NDArray &, NDArray *)> Updater
the prototype of user-defined updater
Definition: kvstore.h:216
virtual void SetGradientCompression(const std::vector< std::pair< std::string, std::string > > &kwargs)=0
Set parameters to use low-bit compressed gradients.
namespace of mxnet
Definition: base.h:118
virtual int get_rank() const
Definition: kvstore.h:322
static KVStore * Create(const char *type="local")
Factory function to create a new KVStore.
virtual void set_updater(const StrUpdater &updater)
set an updater with string keys
Definition: kvstore.h:244
Updater updater_
the user-defined updater
Definition: kvstore.h:391
virtual void PullRowSparse(const std::vector< int > &str_keys, const std::vector< std::pair< NDArray *, NDArray >> &val_rowids, int priority=0)=0
pull a list of key-value pairs from the store. The NDArray pulled back will be in row_sparse storage ...
NDArray interface that handles array arithematics.
const std::string & type()
return the type
Definition: kvstore.h:67
static bool IsSchedulerNode()
Definition: kvstore.h:307
virtual void Barrier()
global barrier among all worker machines
Definition: kvstore.h:352
static void InitPSEnv(const std::unordered_map< std::string, std::string > &envs)
initalize ps-lite environment variables
Definition: kvstore.h:257
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:270
virtual ~KVStore()
virtual destructor
Definition: kvstore.h:50
void set_barrier_before_exit(const bool barrier_before_exit)
Definition: kvstore.h:293
virtual void Pull(const std::vector< int > &keys, const std::vector< NDArray * > &values, int priority=0, bool ignore_sparse=true)=0
pull a list of key-value pairs from the store
StrUpdater str_updater_
the user-defined updater with string keys
Definition: kvstore.h:396
virtual int get_num_dead_node(int node_id, int timeout=60) const
Definition: kvstore.h:341
std::shared_ptr< kvstore::GradientCompression > gradient_compression_
Gradient compression object starts with GC_NONE mode Used if SetGradientCompression sets the type...
Definition: kvstore.h:407
virtual void RunServer(const Controller &controller)
Run as server (or scheduler)
Definition: kvstore.h:385
std::function< void(const std::string &, const NDArray &, NDArray *)> StrUpdater
the prototype of user-defined updater with string keys
Definition: kvstore.h:220
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:365
std::string type_
the kvstore type
Definition: kvstore.h:401
std::function< void(int, const std::string &)> Controller
the prototype of a server controller
Definition: kvstore.h:370
virtual void set_updater(const Updater &updater)
set an updater
Definition: kvstore.h:230
virtual int get_group_size() const
Definition: kvstore.h:329
std::atomic< bool > barrier_before_exit_
whether to do barrier when finalize
Definition: kvstore.h:412
static bool IsServerNode()
Definition: kvstore.h:284