mxnet
Public Types | Public Member Functions | Static Public Member Functions | Protected Attributes | List of all members
mxnet::KVStore Class Referenceabstract

distributed key-value store More...

#include <kvstore.h>

Collaboration diagram for mxnet::KVStore:
Collaboration graph

Public Types

typedef std::function< void(int, const NDArray &, NDArray *)> Updater
 the prototype of user-defined updater More...
 
typedef std::function< void(int, const std::string &)> Controller
 the prototype of a server controller More...
 

Public Member Functions

virtual ~KVStore ()
 virtual destructor More...
 
const std::string & type ()
 return the type More...
 
virtual void Init (const std::vector< int > &keys, const std::vector< NDArray > &values)=0
 Initialize a list of key-value pair to the store. More...
 
virtual void Init (const std::vector< std::string > &str_keys, const std::vector< NDArray > &values)=0
 Initialize a list of key-value pair to the store. More...
 
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 More...
 
virtual void Push (const std::vector< std::string > &str_keys, const std::vector< NDArray > &values, int priority=0)=0
 push a list of key-value pairs into the store More...
 
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 More...
 
virtual void Pull (const std::vector< std::string > &str_keys, const std::vector< NDArray * > &values, int priority=0)=0
 pull a list of key-value pairs from the store More...
 
virtual void set_updater (const Updater &updater)
 set an updater More...
 
void set_barrier_before_exit (const bool barrier_before_exit)
 
virtual int get_rank () const
 
virtual int get_group_size () const
 
virtual int get_num_dead_node (int node_id, int timeout=60) const
 
virtual void Barrier ()
 global barrier among all worker machines More...
 
virtual void SendCommandToServers (int cmd_id, const std::string &cmd_body)
 Send a command to all server nodes. More...
 
virtual void RunServer (const Controller &controller)
 Run as server (or scheduler) More...
 

Static Public Member Functions

static KVStoreCreate (const char *type="local")
 Factory function to create a new KVStore. More...
 
static void InitPSEnv (const std::unordered_map< std::string, std::string > &envs)
 initalize ps-lite environment variables More...
 
static bool IsWorkerNode ()
 
static bool IsServerNode ()
 
static bool IsSchedulerNode ()
 

Protected Attributes

Updater updater_
 the user-defined updater More...
 
std::string type_
 the kvstore type More...
 
std::atomic< bool > barrier_before_exit_ {true}
 whether to do barrier when finalize More...
 

Detailed Description

distributed key-value store

A distributed key-value store for data synchronization over multiple devices/machines. It support user-defined updater.

Member Typedef Documentation

typedef std::function<void(int, const std::string&)> mxnet::KVStore::Controller

the prototype of a server controller

typedef std::function<void(int, const NDArray&, NDArray*)> mxnet::KVStore::Updater

the prototype of user-defined updater

Constructor & Destructor Documentation

virtual mxnet::KVStore::~KVStore ( )
inlinevirtual

virtual destructor

Member Function Documentation

virtual void mxnet::KVStore::Barrier ( )
inlinevirtual

global barrier among all worker machines

But note that, this functions only blocks the main thread of workers until all of them are reached this point. It doesn't guarantee that all operations issued before are actually finished, such as Push and Pull.

static KVStore* mxnet::KVStore::Create ( const char *  type = "local")
static

Factory function to create a new KVStore.

Parameters
typeThe type of the kvstore,
  • 'local' or 'local_update_cpu' or 'local_allreduce_cpu' multi-devices on a single machine. can be also
  • 'device' or 'local_allreduce_device' : same to local but use gpus for kv allreduce
  • 'dist_*' : multi-machines
Returns
a new created KVStore.
virtual int mxnet::KVStore::get_group_size ( ) const
inlinevirtual
Returns
The number of worker nodes
virtual int mxnet::KVStore::get_num_dead_node ( int  node_id,
int  timeout = 60 
) const
inlinevirtual
Returns
the number of dead node(s) specified by {node_id}
Parameters
node_idcan be a node group or a single node
timeouta node fails to send heartbeart in {timeout} seconds will be presumed as 'dead'

Always return 0 when type == "local"

virtual int mxnet::KVStore::get_rank ( ) const
inlinevirtual
Returns
The rank of this node in its group, which is in [0, GroupSize).

Always return 0 when type == "local"

virtual void mxnet::KVStore::Init ( const std::vector< int > &  keys,
const std::vector< NDArray > &  values 
)
pure virtual

Initialize a list of key-value pair to the store.

One must initialize the key before Push and Pull, and a key should be only initialized once

It returns after data have been initialized successfully.

For multiple workers, all workers must call Init. But only worker 0 (get_rank() == 0)'s values are used for initialization. So others' values can be empty (but not keys). This function blocks until all workers are finished. That means, any worker can push and pull on the keys now.

Parameters
keysa list of unique keys
valuesa list of values
virtual void mxnet::KVStore::Init ( const std::vector< std::string > &  str_keys,
const std::vector< NDArray > &  values 
)
pure virtual

Initialize a list of key-value pair to the store.

Parameters
keysa list of unique keys in string format
valuesa list of values
static void mxnet::KVStore::InitPSEnv ( const std::unordered_map< std::string, std::string > &  envs)
inlinestatic

initalize ps-lite environment variables

Parameters
envskey-value environment variables
static bool mxnet::KVStore::IsSchedulerNode ( )
inlinestatic
Returns
whether or not this process is a scheduler node.

Always returns false when type == "local"

static bool mxnet::KVStore::IsServerNode ( )
inlinestatic
Returns
whether or not this process is a server node.

Always returns false when type == "local"

static bool mxnet::KVStore::IsWorkerNode ( )
inlinestatic
Returns
whether or not this process is a worker node.

Always returns true when type == "local"

virtual void mxnet::KVStore::Pull ( const std::vector< int > &  keys,
const std::vector< NDArray * > &  values,
int  priority = 0 
)
pure virtual

pull a list of key-value pairs from the store

One must call Init() on key before. And value should be pre-allocated

This function returns after adding a pull operator to the engine. Any following operator requiring reading value will be blocked until the actual pull is finished. One can wait the pull is finished by

  • when type == "local"
    for (auto& v : values) v.WaitToRead()
  • when type == "dist"
    Wait(keys);
Parameters
keysthe list of keys
valuesthe list of buffers for the pulled data, they should be preallocated
priorityPriority of the action.
virtual void mxnet::KVStore::Pull ( const std::vector< std::string > &  str_keys,
const std::vector< NDArray * > &  values,
int  priority = 0 
)
pure virtual

pull a list of key-value pairs from the store

Parameters
keysthe list of keys in string format
valuesthe list of buffers for the pulled data, they should be preallocated
priorityPriority of the action.
virtual void mxnet::KVStore::Push ( const std::vector< int > &  keys,
const std::vector< NDArray > &  values,
int  priority = 0 
)
pure virtual

push a list of key-value pairs into the store

If a key appears mulitple times in keys, then the according values will be aggregated (summed) before pushing.

The (aggregated) values are merged into the store one by one

updater(key, value, &value_in_store);

One can set a user-defined updater by set_updater. The default updater is Assign.

This function returns after adding a push operator to the engine. Any following operator requiring writing value will be blocked until the actual push is finished. One can wait the push is finished by

  • when type == "local"
    for (auto& v : values) v.WaitToWrite()
  • when type == "dist"
    Wait(keys);

One must call Init() on every key before. And the value NDArray should be always has the same shape as being inited.

Parameters
keysthe list of keys
valuesthe list of values
priorityPriority of the action.
virtual void mxnet::KVStore::Push ( const std::vector< std::string > &  str_keys,
const std::vector< NDArray > &  values,
int  priority = 0 
)
pure virtual

push a list of key-value pairs into the store

Parameters
keysthe list of keys in string format
valuesthe list of values
priorityPriority of the action.
virtual void mxnet::KVStore::RunServer ( const Controller controller)
inlinevirtual

Run as server (or scheduler)

The behavior of a server:

while(receive(x)) {
if (IsCommand(x)) controller(x)
else if (IsKeyValue(x)) updater(x)
}
Parameters
controllerthe user-defined server controller
virtual void mxnet::KVStore::SendCommandToServers ( int  cmd_id,
const std::string &  cmd_body 
)
inlinevirtual

Send a command to all server nodes.

Send a command to all server nodes, which will make each server node run controller

This function returns after the command has been executed in all server nodes

Parameters
cmd_idthe head of the command
cmd_bodythe body of the command
void mxnet::KVStore::set_barrier_before_exit ( const bool  barrier_before_exit)
inline
virtual void mxnet::KVStore::set_updater ( const Updater updater)
inlinevirtual

set an updater

Given a key, assume x is the received (pushed) value and y is the value stored on the store node. The store updates y by h(x, &y). The default h is ASSIGN, namely *y = x.

Parameters
updateruser-defined updater, default is assign
const std::string& mxnet::KVStore::type ( )
inline

return the type

Member Data Documentation

std::atomic<bool> mxnet::KVStore::barrier_before_exit_ {true}
protected

whether to do barrier when finalize

std::string mxnet::KVStore::type_
protected

the kvstore type

Updater mxnet::KVStore::updater_
protected

the user-defined updater


The documentation for this class was generated from the following file: