bjutsecurity22
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

!32551 Change the types of input and output of tcp message handler 

Merge pull request !32551 from chengang/add_retrieve_for_cgn
This commit is contained in:
i-robot 2022-04-08 01:25:14 +00:00 committed by Gitee
parent 4636851fc5 18eaf2d9dd
commit efd13aa953
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
9 changed files with 32 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
mindspore/ccsrc/distributed/cluster/topology/meta_server_node.cc
View File

@ -70,8 +70,7 @@ bool MetaServerNode::InitTCPServer() {
return true;
}
std::shared_ptr<MessageBase> MetaServerNode::HandleMessage(const std::shared_ptr<MessageBase> &message) {
MS_EXCEPTION_IF_NULL(message);
MessageBase *const MetaServerNode::HandleMessage(MessageBase *const message) {
const auto &name = message->Name();
// Handle system messages.
@ -83,6 +82,7 @@ std::shared_ptr<MessageBase> MetaServerNode::HandleMessage(const std::shared_ptr
return rpc::NULL_MSG;
}
system_msg_handlers_[message_name](message);
return rpc::NULL_MSG;
// Handle user defined messages.
} else {
@ -92,13 +92,11 @@ std::shared_ptr<MessageBase> MetaServerNode::HandleMessage(const std::shared_ptr
return rpc::NULL_MSG;
}
(*message_handlers_[name])(message->Body());
return rpc::NULL_MSG;
}
return rpc::NULL_MSG;
}
std::shared_ptr<MessageBase> MetaServerNode::ProcessRegister(const std::shared_ptr<MessageBase> &message) {
MS_EXCEPTION_IF_NULL(message);
MessageBase *const MetaServerNode::ProcessRegister(MessageBase *const message) {
RegistrationMessage registration;
const std::string &body = message->Body();
registration.ParseFromArray(body.c_str(), body.length());
@ -116,9 +114,7 @@ std::shared_ptr<MessageBase> MetaServerNode::ProcessRegister(const std::shared_p
return rpc::NULL_MSG;
}
std::shared_ptr<MessageBase> MetaServerNode::ProcessUnregister(const std::shared_ptr<MessageBase> &message) {
MS_EXCEPTION_IF_NULL(message);
MessageBase *const MetaServerNode::ProcessUnregister(MessageBase *const message) {
UnregistrationMessage unregistration;
const std::string &body = message->Body();
unregistration.ParseFromArray(body.c_str(), body.length());
@ -133,9 +129,7 @@ std::shared_ptr<MessageBase> MetaServerNode::ProcessUnregister(const std::shared
return rpc::NULL_MSG;
}
std::shared_ptr<MessageBase> MetaServerNode::ProcessHeartbeat(const std::shared_ptr<MessageBase> &message) {
MS_EXCEPTION_IF_NULL(message);
MessageBase *const MetaServerNode::ProcessHeartbeat(MessageBase *const message) {
HeartbeatMessage heartbeat;
const std::string &body = message->Body();
heartbeat.ParseFromArray(body.c_str(), body.length());
@ -185,7 +179,7 @@ size_t MetaServerNode::GetAliveNodeNum() {
}
bool MetaServerNode::RegisterMessageHandler(const std::string &name,
std::shared_ptr<std::function<void(const std::string &)>> handler) {
std::shared_ptr<std::function<std::string(const std::string &)>> handler) {
if (message_handlers_.find(name) != message_handlers_.end()) {
MS_LOG(ERROR) << "The message name: " << name << " have already been registered";
return false;

12
mindspore/ccsrc/distributed/cluster/topology/meta_server_node.h
View File

@ -73,23 +73,23 @@ class MetaServerNode : public NodeBase {
// Register the message handler for the user defined message which is specified by the `name` parameter.
bool RegisterMessageHandler(const std::string &name,
std::shared_ptr<std::function<void(const std::string &)>> handler);
std::shared_ptr<std::function<std::string(const std::string &)>> handler);
private:
// Create and init the tcp server.
bool InitTCPServer();
// Handle the message received by the tcp server.
std::shared_ptr<MessageBase> HandleMessage(const std::shared_ptr<MessageBase> &message);
MessageBase *const HandleMessage(MessageBase *const message);
// Process the received register message sent from compute graph nodes.
std::shared_ptr<MessageBase> ProcessRegister(const std::shared_ptr<MessageBase> &message);
MessageBase *const ProcessRegister(MessageBase *const message);
// Process the received unregister message sent from compute graph nodes.
std::shared_ptr<MessageBase> ProcessUnregister(const std::shared_ptr<MessageBase> &message);
MessageBase *const ProcessUnregister(MessageBase *const message);
// Process the received heartbeat message sent from compute graph nodes.
std::shared_ptr<MessageBase> ProcessHeartbeat(const std::shared_ptr<MessageBase> &message);
MessageBase *const ProcessHeartbeat(MessageBase *const message);
// Maintain the state which is type of `TopoState` of this cluster topology.
void UpdateTopoState();
@ -106,7 +106,7 @@ class MetaServerNode : public NodeBase {
// All the handlers for compute graph node's user-defined messages processing.
// The `user-defined` means that this kind of message is user defined and has customized message handler.
std::map<std::string, std::shared_ptr<std::function<void(const std::string &)>>> message_handlers_;
std::map<std::string, std::shared_ptr<std::function<std::string(const std::string &)>>> message_handlers_;
// Stores the registered compute graph nodes.
std::map<std::string, std::shared_ptr<ComputeGraphNodeState>> nodes_;

6
mindspore/ccsrc/distributed/rpc/tcp/connection.cc
View File

@ -254,12 +254,9 @@ int Connection::ReceiveMessage() {
return 0;
}
std::shared_ptr<MessageBase> msg(recv_message);
recv_message = nullptr;
// Call msg handler if set
if (message_handler) {
message_handler(msg);
message_handler(recv_message);
} else {
MS_LOG(INFO) << "Message handler was not found";
}
@ -354,7 +351,6 @@ void Connection::FillSendMessage(MessageBase *msg, const std::string &advertiseU
// update metrics
send_metrics->UpdateMax(msg->body.size());
send_metrics->last_send_msg_name = msg->name;
return;
} else {
if (advertise_addr_.empty()) {

4
mindspore/ccsrc/distributed/rpc/tcp/constants.h
View File

@ -30,7 +30,7 @@
namespace mindspore {
namespace distributed {
namespace rpc {
using MessageHandler = std::function<std::shared_ptr<MessageBase>(const std::shared_ptr<MessageBase> &)>;
using MessageHandler = std::function<MessageBase *const(MessageBase *const)>;
using DeleteCallBack = void (*)(const std::string &from, const std::string &to);
using ConnectionCallBack = void (*)(void *conn);
@ -57,7 +57,7 @@ static const int g_httpKmsgEnable = -1;
using IntTypeMetrics = std::queue<int>;
using StringTypeMetrics = std::queue<std::string>;
static const std::shared_ptr<MessageBase> NULL_MSG = nullptr;
static MessageBase *const NULL_MSG = nullptr;
// Server socket listen backlog.
static const int SOCKET_LISTEN_BACKLOG = 2048;

4
mindspore/ccsrc/plugin/device/cpu/kernel/rpc/rpc_kernel.h
View File

@ -37,10 +37,10 @@ class RpcKernelMod : public NativeCpuKernelMod {
void InitKernel(const CNodePtr &kernel_node) override { return; }
// Set remote data as input.
void SetRemoteInput(const std::shared_ptr<MessageBase> &msg) { remote_input_ = msg; }
void SetRemoteInput(MessageBase *const msg) { remote_input_ = msg; }
protected:
std::shared_ptr<MessageBase> remote_input_;
MessageBase *remote_input_;
};
} // namespace kernel
} // namespace mindspore

4
mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.cc
View File

@ -73,7 +73,7 @@ bool RecvActor::StartServer() {
return true;
}
void RecvActor::RunOpInterProcessData(const std::shared_ptr<MessageBase> &msg, OpContext<DeviceTensor> *const context) {
void RecvActor::RunOpInterProcessData(MessageBase *const msg, OpContext<DeviceTensor> *const context) {
// Once recv actor is launched, reset the op_context so that the next step's recv will not be launched in advance.
ResetOpcontext();
@ -132,7 +132,7 @@ void RecvActor::EraseInput(const OpContext<DeviceTensor> *context) {
}
}
std::shared_ptr<MessageBase> RecvActor::HandleMessage(const std::shared_ptr<MessageBase> &msg) {
MessageBase *RecvActor::HandleMessage(MessageBase *const msg) {
// Block the message handler if the context is invalid.
std::unique_lock<std::mutex> lock(context_mtx_);
context_cv_.wait(lock, [this] { return is_context_valid_; });

4
mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.h
View File

@ -59,14 +59,14 @@ class RecvActor : public RpcActor {
bool CheckRunningCondition(const OpContext<DeviceTensor> *context) const override;
// When an inter-process data received, this method is called.
void RunOpInterProcessData(const std::shared_ptr<MessageBase> &msg, OpContext<DeviceTensor> *const context);
void RunOpInterProcessData(MessageBase *const msg, OpContext<DeviceTensor> *const context);
// Besides erasing input data and input controls when finish actor running, inter-process inputs should be erased.
void EraseInput(const OpContext<DeviceTensor> *context) override;
private:
// The message callback of the tcp server.
std::shared_ptr<MessageBase> HandleMessage(const std::shared_ptr<MessageBase> &msg);
MessageBase *HandleMessage(MessageBase *const msg);
// The network address of this recv actor. It's generated automatically by rpc module.
std::string ip_;

7
tests/ut/cpp/distributed/cluster/topology/test_dynamic_networking.cc
View File

@ -83,8 +83,11 @@ TEST_F(TestDynamicNetworking, AddMessageHandler) {
std::string message_name = "route";
static std::string received_message;
auto func = std::make_shared<std::function<void(const std::string &)>>(
[](const std::string &message) { received_message = message; });
auto func =
std::make_shared<std::function<std::string(const std::string &)>>([](const std::string &message) -> std::string {
received_message = message;
return message;
});
msn.RegisterMessageHandler(message_name, func);
ComputeGraphNode cgn("compute_graph_node");

10
tests/ut/cpp/distributed/rpc/tcp/tcp_test.cc
View File

@ -154,7 +154,7 @@ TEST_F(TCPTest, SendOneMessage) {
bool ret = server->Initialize(server_url);
ASSERT_TRUE(ret);
server->SetMessageHandler([](const std::shared_ptr<MessageBase> &message) -> std::shared_ptr<MessageBase> {
server->SetMessageHandler([](MessageBase *const message) -> MessageBase *const {
IncrDataMsgNum(1);
return NULL_MSG;
});
@ -196,7 +196,7 @@ TEST_F(TCPTest, SendTwoMessages) {
bool ret = server->Initialize(server_url);
ASSERT_TRUE(ret);
server->SetMessageHandler([](const std::shared_ptr<MessageBase> &message) -> std::shared_ptr<MessageBase> {
server->SetMessageHandler([](MessageBase *const message) -> MessageBase *const {
IncrDataMsgNum(1);
return NULL_MSG;
});
@ -251,7 +251,7 @@ TEST_F(TCPTest, SendSyncMessage) {
bool ret = server->Initialize(server_url);
ASSERT_TRUE(ret);
server->SetMessageHandler([](const std::shared_ptr<MessageBase> &message) -> std::shared_ptr<MessageBase> {
server->SetMessageHandler([](MessageBase *const message) -> MessageBase *const {
IncrDataMsgNum(1);
return NULL_MSG;
});
@ -292,7 +292,7 @@ TEST_F(TCPTest, SendLargeMessages) {
bool ret = server->Initialize();
ASSERT_TRUE(ret);
server->SetMessageHandler([](const std::shared_ptr<MessageBase> &message) -> std::shared_ptr<MessageBase> {
server->SetMessageHandler([](MessageBase *const message) -> MessageBase *const {
IncrDataMsgNum(1);
return NULL_MSG;
});
@ -348,7 +348,7 @@ TEST_F(TCPTest, CreateManyConnectionPairs) {
auto port = server->GetPort();
ASSERT_TRUE(ret);
server->SetMessageHandler([](const std::shared_ptr<MessageBase> &message) -> std::shared_ptr<MessageBase> {
server->SetMessageHandler([](MessageBase *const message) -> MessageBase *const {
IncrDataMsgNum(1);
return NULL_MSG;
});