Update ps lite

MANIFEST.in

@@ -1,7 +1,8 @@
 include */* LICENSE byteps.lds byteps.exp
+recursive-include * *.cc *.h
 prune .git
 prune dist
-recursive-include * *.cc *.h
+prune bin
 prune __pycache__
 prune 3rdparty
 graft 3rdparty/ps-lite

byteps/common/global.cc

@@ -286,8 +286,8 @@ ps::KVWorker<char>* BytePSGlobal::GetOrInitPS() {
   if (!_ps && IsDistributed() &&
       _my_role == BytePSRole::LOCAL_ROOT) {  // only the root needs networking
     // init low-level ps implementation
-    _ps = new ps::KVWorker<char>(0, 0);
-    ps::StartAsync(0, "byteps\0");
+    ps::StartPS(0, ps::Node::WORKER, -1, true, "byteps\0");
+    _ps = new ps::KVWorker<char>(0, 0, 0);
     if (BytePSGlobal::IsResuming() || !ps::Postoffice::Get()->is_recovery()) {
       ps::Postoffice::Get()->Barrier(
           0, ps::kWorkerGroup + ps::kServerGroup + ps::kScheduler);
@@ -344,7 +344,7 @@ void BytePSGlobal::Shutdown() {
 
   if (_ps) {
     // shutdown _ps and wait for the completion acks of other workers/servers
-    ps::Finalize(0, true);
+    ps::Finalize(0, ps::Node::WORKER, true);
     delete _ps;
     _ps = NULL;
   }

byteps/common/shared_memory.cc

@@ -31,7 +31,7 @@ void* BytePSSharedMemory::openSharedMemory(const std::string& prefix,
   std::string shm_name(prefix);
   shm_name += std::to_string(key);
   int shm_fd = shm_open(shm_name.c_str(), O_CREAT | O_RDWR, 0666);
-  BPS_CHECK_GE(shm_fd, 0) << "shm_open failed for " << shm_name;
+  BPS_CHECK_GE(shm_fd, 0) << "shm_open failed for " << shm_name << " " << strerror(errno);
 
   BPS_CHECK_GE(ftruncate(shm_fd, size), 0) << strerror(errno);
 

byteps/server/server.cc

@@ -408,6 +408,18 @@ void init_global_env() {
   // enable to print key profile
   log_key_info_ = GetEnv("PS_KEY_LOG", 0);
 
+  std::string role_str = GetEnv("DMLC_ROLE", "server");
+  role_ = ps::GetRole(role_str);
+  if (role_str == std::string("server")) {
+    is_server_ = true;
+    preferred_rank = -1;
+  } else {
+    is_server_ = false;
+    preferred_rank = 0;
+  }
+
+  LOG(INFO) << "This is a " << role_str << " is_server=" << is_server_;
+
   // enable engine block mode (default disabled)
   is_engine_blocking_ = GetEnv("BYTEPS_SERVER_ENGINE_BLOCKING", 0);
   if (is_engine_blocking_)
@@ -480,16 +492,16 @@ extern "C" void byteps_server() {
   }
 
   // init server instance
-  byteps_server_ = new KVServer<SERVER_DATA_TYPE>(0);
+  ps::StartPS(0, role_, preferred_rank, true, "byteps\0");
+  byteps_server_ = new KVServer<SERVER_DATA_TYPE>(0, false, 0);
   byteps_server_->set_request_handle(BytePSHandler);
-  StartAsync(0, "byteps_server\0");
   if (!Postoffice::Get()->is_recovery()) {
     Postoffice::Get()->Barrier(
         0, ps::kWorkerGroup + ps::kServerGroup + ps::kScheduler);
   }
 
   // clean the server resource
-  Finalize(0, true);
+  Finalize(0, role_, true);
   if (byteps_server_) {
     delete byteps_server_;
     byteps_server_ = nullptr;

byteps/server/server.h

@@ -129,6 +129,10 @@ volatile bool sync_mode_ = true;
 volatile bool debug_mode_ = false;
 volatile bool enable_schedule_ = false;
 
+ps::Node::Role role_;
+int preferred_rank = -1;
+volatile bool is_server_ = true;
+
 // debug
 uint64_t debug_key_;
 std::mutex debug_mu_;

setup.py

@@ -642,7 +642,7 @@ def get_nccl_vals():
         nccl_include_dirs += ['%s/include' % nccl_home]
         nccl_lib_dirs += ['%s/lib' % nccl_home, '%s/lib64' % nccl_home]
 
-    nccl_link_mode = os.environ.get('BYTEPS_NCCL_LINK', 'STATIC')
+    nccl_link_mode = os.environ.get('BYTEPS_NCCL_LINK', 'SHARED')
     if nccl_link_mode.upper() == 'SHARED':
         nccl_libs += ['nccl']
     else: