Skip to content

Commit

Permalink
[fleet_executor] Add message bus test with brpc (#37533)
Browse files Browse the repository at this point in the history
  • Loading branch information
FeixLiu authored Nov 26, 2021
1 parent 6b7c061 commit 206db39
Show file tree
Hide file tree
Showing 4 changed files with 135 additions and 7 deletions.
6 changes: 1 addition & 5 deletions paddle/fluid/distributed/fleet_executor/message_bus.cc
Original file line number Diff line number Diff line change
Expand Up @@ -51,15 +51,11 @@ void MessageBus::Init(
#endif

ListenPort();

std::call_once(once_flag_, []() {
std::atexit([]() { MessageBus::Instance().Release(); });
});
}

bool MessageBus::IsInit() const { return is_init_; }

void MessageBus::Release() {
MessageBus::~MessageBus() {
VLOG(3) << "Message bus releases resource.";
#if defined(PADDLE_WITH_DISTRIBUTE) && defined(PADDLE_WITH_PSCORE) && \
!defined(PADDLE_WITH_ASCEND_CL)
Expand Down
4 changes: 2 additions & 2 deletions paddle/fluid/distributed/fleet_executor/message_bus.h
Original file line number Diff line number Diff line change
Expand Up @@ -50,11 +50,11 @@ class MessageBus final {

bool IsInit() const;

void Release();

// called by Interceptor, send InterceptorMessage to dst
bool Send(const InterceptorMessage& interceptor_message);

~MessageBus();

DISABLE_COPY_AND_ASSIGN(MessageBus);

private:
Expand Down
4 changes: 4 additions & 0 deletions paddle/fluid/distributed/fleet_executor/test/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2,3 +2,7 @@ set_source_files_properties(interceptor_ping_pong_test.cc PROPERTIES COMPILE_FLA
set_source_files_properties(compute_interceptor_test.cc PROPERTIES COMPILE_FLAGS ${DISTRIBUTE_COMPILE_FLAGS})
cc_test(interceptor_ping_pong_test SRCS interceptor_ping_pong_test.cc DEPS fleet_executor ${BRPC_DEPS})
cc_test(compute_interceptor_test SRCS compute_interceptor_test.cc DEPS fleet_executor ${BRPC_DEPS})
if(WITH_DISTRIBUTE AND WITH_PSCORE AND NOT (WITH_ASCEND OR WITH_ASCEND_CL))
set_source_files_properties(interceptor_ping_pong_with_brpc_test.cc PROPERTIES COMPILE_FLAGS ${DISTRIBUTE_COMPILE_FLAGS})
cc_test(interceptor_ping_pong_with_brpc_test SRCS interceptor_ping_pong_with_brpc_test.cc DEPS fleet_executor ${BRPC_DEPS})
endif()
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#include <sys/socket.h>
#include <time.h>
#include <iostream>
#include <unordered_map>

#include "gtest/gtest.h"

#include "paddle/fluid/distributed/fleet_executor/carrier.h"
#include "paddle/fluid/distributed/fleet_executor/interceptor.h"
#include "paddle/fluid/distributed/fleet_executor/message_bus.h"

namespace paddle {
namespace distributed {

class PingPongInterceptor : public Interceptor {
public:
PingPongInterceptor(int64_t interceptor_id, TaskNode* node)
: Interceptor(interceptor_id, node) {
RegisterMsgHandle([this](const InterceptorMessage& msg) { PingPong(msg); });
}

void PingPong(const InterceptorMessage& msg) {
std::cout << GetInterceptorId() << " recv msg, count=" << count_
<< std::endl;
++count_;
if (count_ == 20 && GetInterceptorId() == 0) {
InterceptorMessage stop;
stop.set_message_type(STOP);
Send(0, stop);
Send(1, stop);
return;
}

InterceptorMessage resp;
int64_t dst = GetInterceptorId() == 0 ? 1 : 0;
Send(dst, resp);
}

private:
int count_{0};
};

REGISTER_INTERCEPTOR(PingPong, PingPongInterceptor);

TEST(InterceptorTest, PingPong) {
std::cout << "Ping pong test through brpc" << std::endl;
unsigned int seed = time(0);
// random generated two ports in from 6000 to 9000
int port0 = 6000 + rand_r(&seed) % 3000;
int port1 = port0 + 1;

// using socket to check the availability of the port
int server_fd = -1;
server_fd = socket(AF_INET, SOCK_STREAM, 0);
int opt = 1;
linger ling;
ling.l_onoff = 1;
ling.l_linger = 0;
setsockopt(server_fd, SOL_SOCKET, SO_LINGER, &ling, sizeof(ling));
setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
struct sockaddr_in address;
address.sin_family = AF_INET;
address.sin_addr.s_addr = INADDR_ANY;
address.sin_port = htons(port0);
while (bind(server_fd, (struct sockaddr*)&address, sizeof(address)) == -1) {
port0++;
address.sin_port = htons(port0);
}
close(server_fd);

// use another socket to check another port
server_fd = socket(AF_INET, SOCK_STREAM, 0);
setsockopt(server_fd, SOL_SOCKET, SO_LINGER, &ling, sizeof(ling));
setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
port1 = port0 + 1;
address.sin_port = htons(port1);
while (bind(server_fd, (struct sockaddr*)&address, sizeof(address)) == -1) {
port1++;
address.sin_port = htons(port1);
}
close(server_fd);

std::string ip0 = "127.0.0.1:" + std::to_string(port0);
std::string ip1 = "127.0.0.1:" + std::to_string(port1);
std::cout << "ip0: " << ip0 << std::endl;
std::cout << "ip1: " << ip1 << std::endl;

int pid = fork();
if (pid == 0) {
MessageBus& msg_bus = MessageBus::Instance();
msg_bus.Init({{0, 0}, {1, 1}}, {{0, ip0}, {1, ip1}}, ip0);

Carrier& carrier = Carrier::Instance();

Interceptor* a = carrier.SetInterceptor(
0, InterceptorFactory::Create("PingPong", 0, nullptr));
carrier.SetCreatingFlag(false);

InterceptorMessage msg;
a->Send(1, msg);
} else {
MessageBus& msg_bus = MessageBus::Instance();
msg_bus.Init({{0, 0}, {1, 1}}, {{0, ip0}, {1, ip1}}, ip1);

Carrier& carrier = Carrier::Instance();

carrier.SetInterceptor(1,
InterceptorFactory::Create("PingPong", 1, nullptr));
carrier.SetCreatingFlag(false);
}
}

} // namespace distributed
} // namespace paddle

0 comments on commit 206db39

Please sign in to comment.