!20478 Node id is configured through environment variables

Merge pull request !20478 from anancds/node

mindspore/ccsrc/ps/constants.h

@@ -37,6 +37,7 @@ constexpr char kEnvWorkerNum[] = "MS_WORKER_NUM";
 constexpr char kEnvSchedulerHost[] = "MS_SCHED_HOST";
 constexpr char kEnvSchedulerPort[] = "MS_SCHED_PORT";
 constexpr char kEnvSchedulerManagePort[] = "MS_SCHED_MANAGE_PORT";
+constexpr char kEnvNodeId[] = "MS_NODE_ID";
 
 constexpr char kCommTypeOfIBVerbs[] = "ibverbs";
 constexpr char kRoleOfPServer[] = "server";
@@ -95,6 +96,7 @@ constexpr uint32_t kMaxMessageSize = static_cast<uint32_t>(100 * (uint32_t(1) <<
 constexpr char kServerNum[] = "server_num";
 constexpr char kWorkerNum[] = "worker_num";
 constexpr char kNodesIds[] = "node_ids";
+constexpr char kNodeId[] = "node_id";
 
 constexpr int64_t kSubmitTaskIntervalInMs = 1;
 constexpr int64_t kMaxTaskNum = 10240;

mindspore/ccsrc/ps/core/abstract_node.cc

@@ -567,11 +567,13 @@ void AbstractNode::ProcessSendMetadata(const std::shared_ptr<TcpConnection> &con
                << ", cluster state is:" << CommUtil::ClusterStateToString(current_cluster_state_)
                << ", the rank id:" << node_info_.rank_id_;
 
+  client_mutex_.lock();
   nodes_address_.clear();
   for (const auto &it : send_meta_message.servers_meta()) {
     nodes_address_[std::make_pair(NodeRole::SERVER, it.rank_id())] = std::make_pair(it.ip(), it.port());
     MS_LOG(INFO) << "The server ip is:" << it.ip() << ", the port is:" << it.port() << ", the rank id:" << it.rank_id();
   }
+  client_mutex_.unlock();
   server_->SendMessage(conn, meta, Protos::RAW, data, size);
   is_ready_ = true;
   wait_start_cond_.notify_all();
@@ -916,7 +918,15 @@ void AbstractNode::InitServerHandler() {
 }
 
 void AbstractNode::InitNodeInfo(const NodeRole &role) {
-  node_info_.node_id_ = CommUtil::GenerateUUID();
+  if (PSContext::instance()->node_id().empty() && config_->Exists(kNodeId)) {
+    node_info_.node_id_ = config_->Get(kNodeId, "");
+  } else {
+    node_info_.node_id_ = PSContext::instance()->node_id();
+  }
+
+  if (node_info_.node_id_.empty()) {
+    node_info_.node_id_ = CommUtil::GenerateUUID();
+  }
   node_info_.node_role_ = role;
   node_info_.ip_ = server_->BoundIp();
   node_info_.port_ = server_->BoundPort();

mindspore/ccsrc/ps/core/abstract_node.h

@@ -192,7 +192,6 @@ class AbstractNode : public Node {
 
   // the key is: <node_role,rank_id>, the value is: <ip, port>
   std::map<std::pair<NodeRole, uint32_t>, std::pair<std::string, uint16_t>> nodes_address_;
-  std::mutex client_mutex_;
   // the map's key is: rank_id
   std::unordered_map<uint32_t, std::shared_ptr<TcpClient>> connected_nodes_;
 

mindspore/ccsrc/ps/core/communicator/tcp_server.cc

@@ -32,6 +32,7 @@
 namespace mindspore {
 namespace ps {
 namespace core {
+TcpConnection::~TcpConnection() { bufferevent_free(buffer_event_); }
 void TcpConnection::InitConnection(const messageReceive &callback) { tcp_message_handler_.SetCallback(callback); }
 
 void TcpConnection::OnReadHandler(const void *buffer, size_t num) { tcp_message_handler_.ReceiveMessage(buffer, num); }
@@ -432,12 +433,10 @@ void TcpServer::EventCallback(struct bufferevent *bev, std::int16_t events, void
     }
     // Free connection structures
     srv->RemoveConnection(conn->GetFd());
-    bufferevent_free(bev);
   } else if (events & BEV_EVENT_ERROR) {
     MS_LOG(WARNING) << "Event buffer remain data: " << remain;
     // Free connection structures
     srv->RemoveConnection(conn->GetFd());
-    bufferevent_free(bev);
 
     // Notify about disconnection
     if (srv->client_disconnection_) {

mindspore/ccsrc/ps/core/communicator/tcp_server.h

@@ -53,7 +53,7 @@ class TcpConnection {
   explicit TcpConnection(struct bufferevent *bev, const evutil_socket_t &fd, TcpServer *server)
       : buffer_event_(bev), fd_(fd), server_(server) {}
   TcpConnection(const TcpConnection &);
-  virtual ~TcpConnection() = default;
+  virtual ~TcpConnection();
 
   using Callback = std::function<void(const std::shared_ptr<CommMessage>)>;
 

mindspore/ccsrc/ps/core/node.h

@@ -117,6 +117,8 @@ class Node {
   //  }
   // }
   std::unique_ptr<Configuration> config_;
+  // Used to synchronize the connected nodes
+  std::mutex client_mutex_;
 };
 }  // namespace core
 }  // namespace ps

mindspore/ccsrc/ps/core/node_manager.cc

@@ -33,7 +33,12 @@ uint32_t NodeManager::NextRankId(const RegisterMessage &register_message, const
 
   const std::string &node_id = register_message.node_id();
   if (registered_nodes_info_.find(node_id) != registered_nodes_info_.end()) {
+    const std::string &new_ip = register_message.ip();
+    uint32_t new_port = register_message.port();
     rank_id = registered_nodes_info_[node_id].rank_id_;
+    registered_nodes_info_[node_id].is_alive = true;
+    registered_nodes_info_[node_id].ip_ = new_ip;
+    registered_nodes_info_[node_id].port_ = new_port;
     MS_LOG(INFO) << "The node id: " << node_id << " is already assigned!";
     return rank_id;
   }

mindspore/ccsrc/ps/core/scheduler_node.cc

@@ -83,7 +83,15 @@ void SchedulerNode::Initialize() {
   InitCommandHandler();
   CreateTcpServer();
   is_already_stopped_ = false;
-  node_info_.node_id_ = CommUtil::GenerateUUID();
+  if (PSContext::instance()->node_id().empty() && config_->Exists(kNodeId)) {
+    node_info_.node_id_ = config_->Get(kNodeId, "");
+  } else {
+    node_info_.node_id_ = PSContext::instance()->node_id();
+  }
+
+  if (node_info_.node_id_.empty()) {
+    node_info_.node_id_ = CommUtil::GenerateUUID();
+  }
   node_info_.node_role_ = NodeRole::SCHEDULER;
   leader_scaler_ = std::make_unique<LeaderScaler>(this);
   MS_LOG(INFO) << "[Scheduler start]: 2. The node role is:" << CommUtil::NodeRoleToString(node_info_.node_role_)
@@ -142,6 +150,12 @@ void SchedulerNode::ProcessRegister(const std::shared_ptr<TcpServer> &server,
   const std::string &node_id = register_message.node_id();
   node_manager_.UpdateHeartbeat(node_id);
 
+  client_mutex_.lock();
+  if (connected_nodes_.count(node_id)) {
+    connected_nodes_.erase(node_id);
+  }
+  client_mutex_.unlock();
+
   RegisterRespMessage register_resp_message;
   register_resp_message.set_node_id(node_id);
 
@@ -414,6 +428,7 @@ void SchedulerNode::StartUpdateClusterStateTimer() {
 }
 
 const std::shared_ptr<TcpClient> &SchedulerNode::GetOrCreateClient(const NodeInfo &node_info) {
+  std::lock_guard<std::mutex> lock(client_mutex_);
   if (connected_nodes_.count(node_info.node_id_)) {
     return connected_nodes_[node_info.node_id_];
   } else {
@@ -423,8 +438,9 @@ const std::shared_ptr<TcpClient> &SchedulerNode::GetOrCreateClient(const NodeInf
     std::string ip = node_info.ip_;
     uint16_t port = node_info.port_;
     auto client = std::make_shared<TcpClient>(ip, port, config_.get());
-    client->SetMessageCallback([&](std::shared_ptr<MessageMeta> meta, const Protos &protos, const void *data,
-                                   size_t size) { NotifyMessageArrival(meta); });
+    client->SetMessageCallback([&](const std::shared_ptr<MessageMeta> &meta, const Protos &, const void *, size_t) {
+      NotifyMessageArrival(meta);
+    });
     client->Init();
     if (is_client_started_ == false) {
       is_client_started_ = true;

mindspore/ccsrc/ps/core/server_node.cc

@@ -27,14 +27,14 @@ bool ServerNode::Start(const uint32_t &timeout) {
   MS_LOG(INFO) << "[Server start]: 4. The node role:" << CommUtil::NodeRoleToString(node_info_.node_role_)
                << " the node id:" << node_info_.node_id_ << " successfully registered to the scheduler!";
 
-  StartHeartbeatTimer(client_to_scheduler_);
-  MS_LOG(INFO) << "[Server start]: 5. Server start heartbeat timer!";
-
   if (!WaitForStart(timeout)) {
     MS_LOG(ERROR) << "Start server node timeout!";
     return false;
   }
 
+  StartHeartbeatTimer(client_to_scheduler_);
+  MS_LOG(INFO) << "[Server start]: 5. Server start heartbeat timer!";
+
   MsException::Instance().CheckException();
   MS_LOG(INFO) << "[Server start]: 6. Successfully start server node!";
   return true;

mindspore/ccsrc/ps/core/worker_node.cc

@@ -27,14 +27,14 @@ bool WorkerNode::Start(const uint32_t &timeout) {
   MS_LOG(INFO) << "[Worker start]: 4. The node role:" << CommUtil::NodeRoleToString(node_info_.node_role_)
                << " the node id:" << node_info_.node_id_ << " successfully registered to the scheduler!";
 
-  StartHeartbeatTimer(client_to_scheduler_);
-  MS_LOG(INFO) << "[Worker start]: 5. Worker start heartbeat timer!";
-
   if (!WaitForStart(timeout)) {
     MS_LOG(ERROR) << "Start Worker node timeout!";
     return false;
   }
 
+  StartHeartbeatTimer(client_to_scheduler_);
+  MS_LOG(INFO) << "[Worker start]: 5. Worker start heartbeat timer!";
+
   MsException::Instance().CheckException();
   MS_LOG(INFO) << "[Worker start]: 6. Successfully start worker node!";
   return true;
@@ -69,7 +69,7 @@ void WorkerNode::CreateTcpServer() {
   std::string server_ip;
   CommUtil::GetAvailableInterfaceAndIP(&interface, &server_ip);
   server_ = std::make_shared<TcpServer>(server_ip, 0, config_.get());
-  server_->SetMessageCallback([&](std::shared_ptr<TcpConnection> conn, std::shared_ptr<MessageMeta> meta,
+  server_->SetMessageCallback([&](const std::shared_ptr<TcpConnection> &conn, const std::shared_ptr<MessageMeta> &meta,
                                   const Protos &protos, const void *data, size_t size) {
     if (server_handler_.count(meta->cmd()) == 0) {
       MS_LOG(EXCEPTION) << "The cmd:" << meta->cmd() << " is not supported!";

mindspore/ccsrc/ps/ps_context.cc

@@ -54,6 +54,7 @@ void PSContext::SetPSEnable(bool enabled) {
     scheduler_port_ = std::strtol(common::GetEnv(kEnvSchedulerPort).c_str(), nullptr, kBase);
     scheduler_manage_port_ = std::strtol(common::GetEnv(kEnvSchedulerManagePort).c_str(), nullptr, kBase);
     cluster_config_ = std::make_unique<core::ClusterConfig>(worker_num_, server_num_, scheduler_host_, scheduler_port_);
+    node_id_ = common::GetEnv(kEnvNodeId);
   } else {
     MS_LOG(INFO) << "PS mode is disabled.";
     is_worker_ = false;
@@ -407,5 +408,9 @@ uint16_t PSContext::scheduler_manage_port() const { return scheduler_manage_port
 void PSContext::set_config_file_path(const std::string &path) { config_file_path_ = path; }
 
 std::string PSContext::config_file_path() const { return config_file_path_; }
+
+void PSContext::set_node_id(const std::string &node_id) { node_id_ = node_id; }
+
+const std::string &PSContext::node_id() const { return node_id_; }
 }  // namespace ps
 }  // namespace mindspore

mindspore/ccsrc/ps/ps_context.h

@@ -178,6 +178,9 @@ class PSContext {
   void set_encrypt_type(const std::string &encrypt_type);
   const std::string &encrypt_type() const;
 
+  void set_node_id(const std::string &node_id);
+  const std::string &node_id() const;
+
  private:
   PSContext()
       : ps_enabled_(false),
@@ -215,7 +218,8 @@ class PSContext {
         dp_eps_(50),
         dp_delta_(0.01),
         dp_norm_clip_(1.0),
-        encrypt_type_(kNotEncryptType) {}
+        encrypt_type_(kNotEncryptType),
+        node_id_("") {}
   bool ps_enabled_;
   bool is_worker_;
   bool is_pserver_;
@@ -304,6 +308,9 @@ class PSContext {
 
   // Secure mechanism for federated learning. Used in federated learning for now.
   std::string encrypt_type_;
+
+  // Unique id of the node
+  std::string node_id_;
 };
 }  // namespace ps
 }  // namespace mindspore

tests/st/fl/albert/run_hybrid_train_sched.py

@@ -15,6 +15,7 @@
 
 import argparse
 import subprocess
+import os
 
 parser = argparse.ArgumentParser(description="Run train_cloud.py case")
 parser.add_argument("--device_target", type=str, default="CPU")
@@ -34,6 +35,7 @@ scheduler_ip = args.scheduler_ip
 scheduler_port = args.scheduler_port
 scheduler_manage_port = args.scheduler_manage_port
 
+os.environ['MS_NODE_ID'] = "20"
 cmd_sched = "execute_path=$(pwd) && self_path=$(dirname \"${script_self}\") && rm -rf ${execute_path}/scheduler/ &&"
 cmd_sched += "mkdir ${execute_path}/scheduler/ &&"
 cmd_sched += "cd ${execute_path}/scheduler/ || exit && export GLOG_v=1 &&"

tests/st/fl/albert/run_hybrid_train_server.py

@@ -15,6 +15,7 @@
 
 import argparse
 import subprocess
+import os
 
 parser = argparse.ArgumentParser(description="Run train_cloud.py case")
 parser.add_argument("--device_target", type=str, default="CPU")
@@ -77,6 +78,7 @@ if local_server_num == -1:
 assert local_server_num <= server_num, "The local server number should not be bigger than total server number."
 
 for i in range(local_server_num):
+    os.environ['MS_NODE_ID'] = str(i)
     cmd_server = "execute_path=$(pwd) && self_path=$(dirname \"${script_self}\") && "
     cmd_server += "rm -rf ${execute_path}/server_" + str(i) + "/ &&"
     cmd_server += "mkdir ${execute_path}/server_" + str(i) + "/ &&"

tests/st/fl/hybrid_lenet/config.json

@@ -9,5 +9,6 @@
     "client_cert_path": "client.crt",
     "client_password": "client_password",
     "ca_cert_path": "ca.crt",
-    "Key_encrypt_decrypt_algorithm": ""
+    "Key_encrypt_decrypt_algorithm": "",
+    "node_id":1
 }

tests/st/fl/hybrid_lenet/run_hybrid_train_sched.py

@@ -15,6 +15,7 @@
 
 import argparse
 import subprocess
+import os
 
 parser = argparse.ArgumentParser(description="Run test_hybrid_train_lenet.py case")
 parser.add_argument("--device_target", type=str, default="CPU")
@@ -36,6 +37,7 @@ scheduler_port = args.scheduler_port
 scheduler_manage_port = args.scheduler_manage_port
 config_file_path = args.config_file_path
 
+os.environ['MS_NODE_ID'] = "20"
 cmd_sched = "execute_path=$(pwd) && self_path=$(dirname \"${script_self}\") && rm -rf ${execute_path}/scheduler/ &&"
 cmd_sched += "mkdir ${execute_path}/scheduler/ &&"
 cmd_sched += "cd ${execute_path}/scheduler/ || exit && export GLOG_v=1 &&"

tests/st/fl/hybrid_lenet/run_hybrid_train_server.py

@@ -15,6 +15,7 @@
 
 import argparse
 import subprocess
+import os
 
 parser = argparse.ArgumentParser(description="Run test_hybrid_train_lenet.py case")
 parser.add_argument("--device_target", type=str, default="CPU")
@@ -78,6 +79,7 @@ if local_server_num == -1:
 assert local_server_num <= server_num, "The local server number should not be bigger than total server number."
 
 for i in range(local_server_num):
+    os.environ['MS_NODE_ID'] = str(i)
     cmd_server = "execute_path=$(pwd) && self_path=$(dirname \"${script_self}\") && "
     cmd_server += "rm -rf ${execute_path}/server_" + str(i) + "/ &&"
     cmd_server += "mkdir ${execute_path}/server_" + str(i) + "/ &&"

tests/st/fl/hybrid_lenet/run_hybrid_train_worker.py

@@ -15,6 +15,7 @@
 
 import argparse
 import subprocess
+import os
 
 parser = argparse.ArgumentParser(description="Run test_hybrid_train_lenet.py case")
 parser.add_argument("--device_target", type=str, default="CPU")
@@ -46,6 +47,7 @@ if local_worker_num == -1:
 assert local_worker_num <= worker_num, "The local worker number should not be bigger than total worker number."
 
 for i in range(local_worker_num):
+    os.environ['MS_NODE_ID'] = str(10 + i)
     cmd_worker = "execute_path=$(pwd) && self_path=$(dirname \"${script_self}\") && "
     cmd_worker += "rm -rf ${execute_path}/worker_" + str(i) + "/ &&"
     cmd_worker += "mkdir ${execute_path}/worker_" + str(i) + "/ &&"

tests/st/fl/mobile/config.json

@@ -9,5 +9,6 @@
     "client_cert_path": "client.crt",
     "client_password": "client_password",
     "ca_cert_path": "ca.crt",
-    "Key_encrypt_decrypt_algorithm": ""
+    "Key_encrypt_decrypt_algorithm": "",
+    "node_id": 1
 }

tests/st/fl/mobile/run_mobile_sched.py

@@ -15,6 +15,7 @@
 
 import argparse
 import subprocess
+import os
 
 parser = argparse.ArgumentParser(description="Run test_mobile_lenet.py case")
 parser.add_argument("--device_target", type=str, default="CPU")
@@ -39,6 +40,7 @@ if __name__ == "__main__":
     scheduler_manage_port = args.scheduler_manage_port
     config_file_path = args.config_file_path
 
+    os.environ['MS_NODE_ID'] = "20"
     cmd_sched = "execute_path=$(pwd) && self_path=$(dirname \"${script_self}\") && rm -rf ${execute_path}/scheduler/ &&"
     cmd_sched += "mkdir ${execute_path}/scheduler/ &&"
     cmd_sched += "cd ${execute_path}/scheduler/ || exit && export GLOG_v=1 &&"

tests/st/fl/mobile/run_mobile_server.py

@@ -15,6 +15,7 @@
 
 import argparse
 import subprocess
+import os
 
 parser = argparse.ArgumentParser(description="Run test_mobile_lenet.py case")
 parser.add_argument("--device_target", type=str, default="CPU")
@@ -79,6 +80,7 @@ if __name__ == "__main__":
     assert local_server_num <= server_num, "The local server number should not be bigger than total server number."
 
     for i in range(local_server_num):
+        os.environ['MS_NODE_ID'] = str(i)
         cmd_server = "execute_path=$(pwd) && self_path=$(dirname \"${script_self}\") && "
         cmd_server += "rm -rf ${execute_path}/server_" + str(i) + "/ &&"
         cmd_server += "mkdir ${execute_path}/server_" + str(i) + "/ &&"