diff --git a/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.cc b/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.cc index 67cd1697547..655003b9a9b 100644 --- a/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.cc +++ b/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.cc @@ -18,6 +18,7 @@ #include #include #include "proto/topology.pb.h" +#include "distributed/rpc/tcp/constants.h" #include "distributed/cluster/topology/utils.h" #include "distributed/cluster/topology/meta_server_node.h" @@ -69,7 +70,7 @@ bool MetaServerNode::InitTCPServer() { return true; } -void MetaServerNode::HandleMessage(const std::shared_ptr &message) { +std::shared_ptr MetaServerNode::HandleMessage(const std::shared_ptr &message) { MS_EXCEPTION_IF_NULL(message); const auto &name = message->Name(); @@ -79,7 +80,7 @@ void MetaServerNode::HandleMessage(const std::shared_ptr &message) const auto &handler = system_msg_handlers_.find(message_name); if (handler == system_msg_handlers_.end()) { MS_LOG(ERROR) << "Unknown system message name: " << message->Name(); - return; + return rpc::NULL_MSG; } system_msg_handlers_[message_name](message); @@ -88,13 +89,14 @@ void MetaServerNode::HandleMessage(const std::shared_ptr &message) const auto &handler = message_handlers_.find(name); if (handler == message_handlers_.end()) { MS_LOG(ERROR) << "Unknown message name: " << name; - return; + return rpc::NULL_MSG; } (*message_handlers_[name])(message->Body()); } + return rpc::NULL_MSG; } -void MetaServerNode::ProcessRegister(const std::shared_ptr &message) { +std::shared_ptr MetaServerNode::ProcessRegister(const std::shared_ptr &message) { MS_EXCEPTION_IF_NULL(message); RegistrationMessage registration; @@ -111,9 +113,10 @@ void MetaServerNode::ProcessRegister(const std::shared_ptr &message } else { MS_LOG(ERROR) << "The node: " << node_id << " have been registered before."; } + return rpc::NULL_MSG; } -void MetaServerNode::ProcessUnregister(const std::shared_ptr &message) { +std::shared_ptr MetaServerNode::ProcessUnregister(const std::shared_ptr &message) { MS_EXCEPTION_IF_NULL(message); UnregistrationMessage unregistration; @@ -124,12 +127,13 @@ void MetaServerNode::ProcessUnregister(const std::shared_ptr &messa std::unique_lock lock(nodes_mutex_); if (nodes_.find(node_id) == nodes_.end()) { MS_LOG(ERROR) << "Received unregistration message from invalid compute graph node: " << node_id; - return; + return rpc::NULL_MSG; } nodes_.erase(node_id); + return rpc::NULL_MSG; } -void MetaServerNode::ProcessHeartbeat(const std::shared_ptr &message) { +std::shared_ptr MetaServerNode::ProcessHeartbeat(const std::shared_ptr &message) { MS_EXCEPTION_IF_NULL(message); HeartbeatMessage heartbeat; @@ -145,6 +149,7 @@ void MetaServerNode::ProcessHeartbeat(const std::shared_ptr &messag } else { MS_LOG(ERROR) << "Invalid node: " << node_id << "."; } + return rpc::NULL_MSG; } void MetaServerNode::UpdateTopoState() { diff --git a/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.h b/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.h index e5378706908..0dd0264c0b6 100644 --- a/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.h +++ b/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.h @@ -80,16 +80,16 @@ class MetaServerNode : public NodeBase { bool InitTCPServer(); // Handle the message received by the tcp server. - void HandleMessage(const std::shared_ptr &message); + std::shared_ptr HandleMessage(const std::shared_ptr &message); // Process the received register message sent from compute graph nodes. - void ProcessRegister(const std::shared_ptr &message); + std::shared_ptr ProcessRegister(const std::shared_ptr &message); // Process the received unregister message sent from compute graph nodes. - void ProcessUnregister(const std::shared_ptr &message); + std::shared_ptr ProcessUnregister(const std::shared_ptr &message); // Process the received heartbeat message sent from compute graph nodes. - void ProcessHeartbeat(const std::shared_ptr &message); + std::shared_ptr ProcessHeartbeat(const std::shared_ptr &message); // Maintain the state which is type of `TopoState` of this cluster topology. void UpdateTopoState(); diff --git a/mindspore/ccsrc/distributed/rpc/tcp/connection.cc b/mindspore/ccsrc/distributed/rpc/tcp/connection.cc index 838dddaa2cd..7f1a43e39b2 100644 --- a/mindspore/ccsrc/distributed/rpc/tcp/connection.cc +++ b/mindspore/ccsrc/distributed/rpc/tcp/connection.cc @@ -422,6 +422,42 @@ void Connection::FillRecvMessage() { recv_message = msg; } +int Connection::Flush() { + int total_send_bytes = 0; + while (!send_message_queue.empty() || total_send_len != 0) { + if (total_send_len == 0) { + FillSendMessage(send_message_queue.front(), source, false); + send_message_queue.pop(); + } + size_t sendLen = 0; + int retval = socket_operation->SendMessage(this, &send_kernel_msg, total_send_len, &sendLen); + if (retval == IO_RW_OK && sendLen > 0) { + total_send_len -= sendLen; + if (total_send_len == 0) { + // update metrics + send_metrics->UpdateError(false); + + output_buffer_size -= send_message->body.size(); + total_send_bytes += send_message->body.size(); + delete send_message; + send_message = nullptr; + break; + } + } else if (retval == IO_RW_OK && sendLen == 0) { + // EAGAIN + MS_LOG(ERROR) << "Failed to send message and update the epoll event"; + (void)recv_event_loop->UpdateEpollEvent(socket_fd, EPOLLOUT | EPOLLIN | EPOLLHUP | EPOLLERR); + continue; + } else { + // update metrics + send_metrics->UpdateError(true, error_code); + state = ConnectionState::kDisconnecting; + break; + } + } + return total_send_bytes; +} + int Connection::AddConnnectEventHandler() { return recv_event_loop->SetEventHandler(socket_fd, EPOLLIN | EPOLLHUP | EPOLLERR, NewConnectEventHandler, reinterpret_cast(this)); diff --git a/mindspore/ccsrc/distributed/rpc/tcp/connection.h b/mindspore/ccsrc/distributed/rpc/tcp/connection.h index ebe9d985ecd..3f032d6954f 100644 --- a/mindspore/ccsrc/distributed/rpc/tcp/connection.h +++ b/mindspore/ccsrc/distributed/rpc/tcp/connection.h @@ -109,6 +109,9 @@ struct Connection { return !(that != nullptr && that->destination == destination && that->is_remote == is_remote); } + // Send all the messages in the message queue. + int Flush(); + // The socket used by this connection. int socket_fd; diff --git a/mindspore/ccsrc/distributed/rpc/tcp/constants.h b/mindspore/ccsrc/distributed/rpc/tcp/constants.h index 432a396b9b2..175892683e8 100644 --- a/mindspore/ccsrc/distributed/rpc/tcp/constants.h +++ b/mindspore/ccsrc/distributed/rpc/tcp/constants.h @@ -30,7 +30,7 @@ namespace mindspore { namespace distributed { namespace rpc { -using MessageHandler = std::function &)>; +using MessageHandler = std::function(const std::shared_ptr &)>; using DeleteCallBack = void (*)(const std::string &from, const std::string &to); using ConnectionCallBack = void (*)(void *conn); @@ -57,6 +57,8 @@ static const int g_httpKmsgEnable = -1; using IntTypeMetrics = std::queue; using StringTypeMetrics = std::queue; +static const std::shared_ptr NULL_MSG = nullptr; + // Server socket listen backlog. static const int SOCKET_LISTEN_BACKLOG = 2048; diff --git a/mindspore/ccsrc/distributed/rpc/tcp/tcp_comm.cc b/mindspore/ccsrc/distributed/rpc/tcp/tcp_comm.cc index 3c770722d08..17355b206e4 100644 --- a/mindspore/ccsrc/distributed/rpc/tcp/tcp_comm.cc +++ b/mindspore/ccsrc/distributed/rpc/tcp/tcp_comm.cc @@ -131,42 +131,6 @@ void OnAccept(int server, uint32_t events, void *arg) { tcpmgr->conn_pool_->AddConnection(conn); } -int DoSend(Connection *conn) { - int total_send_bytes = 0; - while (!conn->send_message_queue.empty() || conn->total_send_len != 0) { - if (conn->total_send_len == 0) { - conn->FillSendMessage(conn->send_message_queue.front(), conn->source, false); - conn->send_message_queue.pop(); - } - size_t sendLen = 0; - int retval = conn->socket_operation->SendMessage(conn, &conn->send_kernel_msg, conn->total_send_len, &sendLen); - if (retval == IO_RW_OK && sendLen > 0) { - conn->total_send_len -= sendLen; - if (conn->total_send_len == 0) { - // update metrics - conn->send_metrics->UpdateError(false); - - conn->output_buffer_size -= conn->send_message->body.size(); - total_send_bytes += conn->send_message->body.size(); - delete conn->send_message; - conn->send_message = nullptr; - break; - } - } else if (retval == IO_RW_OK && sendLen == 0) { - // EAGAIN - MS_LOG(ERROR) << "Failed to send message and update the epoll event"; - (void)conn->recv_event_loop->UpdateEpollEvent(conn->socket_fd, EPOLLOUT | EPOLLIN | EPOLLHUP | EPOLLERR); - continue; - } else { - // update metrics - conn->send_metrics->UpdateError(true, conn->error_code); - conn->state = ConnectionState::kDisconnecting; - break; - } - } - return total_send_bytes; -} - void TCPComm::SetMessageHandler(const MessageHandler &handler) { message_handler_ = handler; } bool TCPComm::Initialize() { @@ -260,7 +224,7 @@ void TCPComm::EventCallBack(void *connection) { if (conn->state == ConnectionState::kConnected) { conn->conn_mutex->lock(); - (void)DoSend(conn); + (void)conn->Flush(); conn->conn_mutex->unlock(); } else if (conn->state == ConnectionState::kDisconnecting) { conn->conn_mutex->lock(); @@ -272,7 +236,7 @@ void TCPComm::WriteCallBack(void *connection) { Connection *conn = reinterpret_cast(connection); if (conn->state == ConnectionState::kConnected) { conn->conn_mutex->lock(); - (void)DoSend(conn); + (void)conn->Flush(); conn->conn_mutex->unlock(); } } @@ -370,7 +334,7 @@ ssize_t TCPComm::Send(MessageBase *msg, bool sync) { } else { (void)conn->send_message_queue.emplace(msg); } - return DoSend(conn); + return conn->Flush(); }; if (sync) { return task(); diff --git a/mindspore/ccsrc/distributed/rpc/tcp/tcp_comm.h b/mindspore/ccsrc/distributed/rpc/tcp/tcp_comm.h index 4f2aaed6d53..be8a72fd399 100644 --- a/mindspore/ccsrc/distributed/rpc/tcp/tcp_comm.h +++ b/mindspore/ccsrc/distributed/rpc/tcp/tcp_comm.h @@ -33,9 +33,6 @@ namespace rpc { // Event handler for new connecting request arrived. void OnAccept(int server, uint32_t events, void *arg); -// Send messages buffered in the connection. -int DoSend(Connection *conn); - void DoDisconnect(int fd, Connection *conn, uint32_t error, int soError); void ConnectedEventHandler(int fd, uint32_t events, void *context); @@ -116,7 +113,6 @@ class TCPComm { std::shared_ptr conn_mutex_; friend void OnAccept(int server, uint32_t events, void *arg); - friend int DoSend(Connection *conn); friend int DoConnect(const std::string &to, Connection *conn, ConnectionCallBack event_callback, ConnectionCallBack write_callback, ConnectionCallBack read_callback); }; diff --git a/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.cc b/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.cc index 10c150b5491..a9574b1de1c 100644 --- a/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.cc +++ b/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.cc @@ -20,6 +20,7 @@ #include #include #include +#include "distributed/rpc/tcp/constants.h" #include "plugin/device/cpu/kernel/rpc/rpc_recv_kernel.h" namespace mindspore { @@ -130,7 +131,7 @@ void RecvActor::EraseInput(const OpContext *context) { } } -void RecvActor::HandleMessage(const std::shared_ptr &msg) { +std::shared_ptr RecvActor::HandleMessage(const std::shared_ptr &msg) { // Block the message handler if the context is invalid. std::unique_lock lock(context_mtx_); context_cv_.wait(lock, [this] { return is_context_valid_; }); @@ -138,9 +139,11 @@ void RecvActor::HandleMessage(const std::shared_ptr &msg) { MS_LOG(INFO) << "Rpc actor recv message for inter-process edge: " << inter_process_edge_name_; - MS_ERROR_IF_NULL_WO_RET_VAL(msg); - MS_ERROR_IF_NULL_WO_RET_VAL(op_context_); + if (msg == nullptr || op_context_ == nullptr) { + return distributed::rpc::NULL_MSG; + } ActorDispatcher::Send(GetAID(), &RecvActor::RunOpInterProcessData, msg, op_context_); + return distributed::rpc::NULL_MSG; } } // namespace runtime } // namespace mindspore diff --git a/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.h b/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.h index 248eb89b70c..76bda913972 100644 --- a/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.h +++ b/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.h @@ -66,7 +66,7 @@ class RecvActor : public RpcActor { private: // The message callback of the tcp server. - void HandleMessage(const std::shared_ptr &msg); + std::shared_ptr HandleMessage(const std::shared_ptr &msg); // The network address of this recv actor. It's generated automatically by rpc module. std::string ip_; diff --git a/tests/ut/cpp/distributed/rpc/tcp/tcp_test.cc b/tests/ut/cpp/distributed/rpc/tcp/tcp_test.cc index 2fa3ce96192..9f800cacaad 100644 --- a/tests/ut/cpp/distributed/rpc/tcp/tcp_test.cc +++ b/tests/ut/cpp/distributed/rpc/tcp/tcp_test.cc @@ -154,7 +154,10 @@ TEST_F(TCPTest, SendOneMessage) { bool ret = server->Initialize(server_url); ASSERT_TRUE(ret); - server->SetMessageHandler([](const std::shared_ptr &message) -> void { IncrDataMsgNum(1); }); + server->SetMessageHandler([](const std::shared_ptr &message) -> std::shared_ptr { + IncrDataMsgNum(1); + return NULL_MSG; + }); // Start the tcp client. auto client_url = "127.0.0.1:1234"; @@ -193,7 +196,10 @@ TEST_F(TCPTest, SendTwoMessages) { bool ret = server->Initialize(server_url); ASSERT_TRUE(ret); - server->SetMessageHandler([](const std::shared_ptr &message) -> void { IncrDataMsgNum(1); }); + server->SetMessageHandler([](const std::shared_ptr &message) -> std::shared_ptr { + IncrDataMsgNum(1); + return NULL_MSG; + }); // Start the tcp client. auto client_url = "127.0.0.1:1234"; @@ -245,7 +251,10 @@ TEST_F(TCPTest, SendSyncMessage) { bool ret = server->Initialize(server_url); ASSERT_TRUE(ret); - server->SetMessageHandler([](const std::shared_ptr &message) -> void { IncrDataMsgNum(1); }); + server->SetMessageHandler([](const std::shared_ptr &message) -> std::shared_ptr { + IncrDataMsgNum(1); + return NULL_MSG; + }); // Start the tcp client. auto client_url = "127.0.0.1:1234"; @@ -283,7 +292,10 @@ TEST_F(TCPTest, SendLargeMessages) { bool ret = server->Initialize(); ASSERT_TRUE(ret); - server->SetMessageHandler([](const std::shared_ptr &message) -> void { IncrDataMsgNum(1); }); + server->SetMessageHandler([](const std::shared_ptr &message) -> std::shared_ptr { + IncrDataMsgNum(1); + return NULL_MSG; + }); // Start the tcp client. auto client_url = "127.0.0.1:1234"; @@ -336,7 +348,10 @@ TEST_F(TCPTest, CreateManyConnectionPairs) { auto port = server->GetPort(); ASSERT_TRUE(ret); - server->SetMessageHandler([](const std::shared_ptr &message) -> void { IncrDataMsgNum(1); }); + server->SetMessageHandler([](const std::shared_ptr &message) -> std::shared_ptr { + IncrDataMsgNum(1); + return NULL_MSG; + }); // Start the tcp client. auto client_url = "127.0.0.1:1234";