[Streaming]Streaming queue support failover (#8161)

This commit is contained in:
wanxing
2020-08-25 14:19:45 +08:00
committed by GitHub
parent 5a787a8253
commit e816e3aefb
16 changed files with 1049 additions and 212 deletions
+169 -10
View File
@@ -257,6 +257,165 @@ class StreamingQueueReaderTestSuite : public StreamingQueueTestSuite {
}
};
class StreamingQueueUpStreamTestSuite : public StreamingQueueTestSuite {
public:
StreamingQueueUpStreamTestSuite(ActorID &peer_actor_id, std::vector<ObjectID> queue_ids,
std::vector<ObjectID> rescale_queue_ids)
: StreamingQueueTestSuite(peer_actor_id, queue_ids, rescale_queue_ids) {
test_func_map_ = {
{"pull_peer_async_test",
std::bind(&StreamingQueueUpStreamTestSuite::PullPeerAsyncTest, this)},
{"get_queue_test",
std::bind(&StreamingQueueUpStreamTestSuite::GetQueueTest, this)}};
}
void GetQueueTest() {
// Sleep 2s, queue shoulde not exist when reader pull
std::this_thread::sleep_for(std::chrono::milliseconds(2000));
auto upstream_handler = ray::streaming::UpstreamQueueMessageHandler::GetService();
ObjectID &queue_id = queue_ids_[0];
RayFunction async_call_func{ray::Language::PYTHON,
ray::FunctionDescriptorBuilder::FromVector(ray::Language::PYTHON, {"", "", "reader_async_call_func", ""})};
RayFunction sync_call_func{ray::Language::PYTHON,
ray::FunctionDescriptorBuilder::FromVector(ray::Language::PYTHON, {"", "", "reader_sync_call_func", ""})};
upstream_handler->SetPeerActorID(queue_id, peer_actor_id_, async_call_func, sync_call_func);
upstream_handler->CreateUpstreamQueue(queue_id, peer_actor_id_, 10240);
STREAMING_LOG(INFO) << "IsQueueExist: "
<< upstream_handler->UpstreamQueueExists(queue_id);
// Sleep 2s, No valid data when reader pull
std::this_thread::sleep_for(std::chrono::milliseconds(2000));
std::this_thread::sleep_for(std::chrono::milliseconds(10 * 1000));
STREAMING_LOG(INFO) << "StreamingQueueUpStreamTestSuite::GetQueueTest done";
status_ = true;
}
void PullPeerAsyncTest() {
// Sleep 2s, queue should not exist when reader pull
std::this_thread::sleep_for(std::chrono::milliseconds(2000));
auto upstream_handler = ray::streaming::UpstreamQueueMessageHandler::GetService();
ObjectID &queue_id = queue_ids_[0];
RayFunction async_call_func{ray::Language::PYTHON,
ray::FunctionDescriptorBuilder::FromVector(ray::Language::PYTHON, {"", "", "reader_async_call_func", ""})};
RayFunction sync_call_func{ray::Language::PYTHON,
ray::FunctionDescriptorBuilder::FromVector(ray::Language::PYTHON, {"", "", "reader_sync_call_func", ""})};
upstream_handler->SetPeerActorID(queue_id, peer_actor_id_, async_call_func, sync_call_func);
std::shared_ptr<WriterQueue> queue =
upstream_handler->CreateUpstreamQueue(queue_id, peer_actor_id_, 10240);
STREAMING_LOG(INFO) << "IsQueueExist: "
<< upstream_handler->UpstreamQueueExists(queue_id);
// Sleep 2s, No valid data when reader pull
std::this_thread::sleep_for(std::chrono::milliseconds(2000));
// message id starts from 1
for (int msg_id = 1; msg_id <= 80; msg_id++) {
uint8_t data[100];
memset(data, msg_id, 100);
STREAMING_LOG(INFO) << "Writer User Push item msg_id: " << msg_id;
ASSERT_TRUE(
queue->Push(msg_id/*seqid*/, data, 100, current_sys_time_ms(), msg_id, msg_id, true).ok());
queue->Send();
}
std::this_thread::sleep_for(std::chrono::milliseconds(5000));
STREAMING_LOG(INFO) << "StreamingQueueUpStreamTestSuite::PullPeerAsyncTest done";
status_ = true;
}
};
class StreamingQueueDownStreamTestSuite : public StreamingQueueTestSuite {
public:
StreamingQueueDownStreamTestSuite(ActorID peer_actor_id, std::vector<ObjectID> queue_ids,
std::vector<ObjectID> rescale_queue_ids)
: StreamingQueueTestSuite(peer_actor_id, queue_ids, rescale_queue_ids) {
test_func_map_ = {
{"pull_peer_async_test",
std::bind(&StreamingQueueDownStreamTestSuite::PullPeerAsyncTest, this)},
{"get_queue_test",
std::bind(&StreamingQueueDownStreamTestSuite::GetQueueTest, this)}};
};
void GetQueueTest() {
auto downstream_handler =
ray::streaming::DownstreamQueueMessageHandler::GetService();
ObjectID &queue_id = queue_ids_[0];
RayFunction async_call_func{ray::Language::PYTHON,
ray::FunctionDescriptorBuilder::FromVector(ray::Language::PYTHON, {"", "", "writer_async_call_func", ""})};
RayFunction sync_call_func{ray::Language::PYTHON,
ray::FunctionDescriptorBuilder::FromVector(ray::Language::PYTHON, {"", "", "writer_sync_call_func", ""})};
downstream_handler->SetPeerActorID(queue_id, peer_actor_id_, async_call_func, sync_call_func);
downstream_handler->CreateDownstreamQueue(queue_id, peer_actor_id_);
bool is_upstream_first_pull_ = false;
downstream_handler->PullQueue(queue_id, 1, is_upstream_first_pull_, 10 * 1000);
ASSERT_TRUE(is_upstream_first_pull_);
downstream_handler->PullQueue(queue_id, 1, is_upstream_first_pull_, 10 * 1000);
ASSERT_FALSE(is_upstream_first_pull_);
STREAMING_LOG(INFO) << "StreamingQueueDownStreamTestSuite::GetQueueTest done";
status_ = true;
}
void PullPeerAsyncTest() {
auto downstream_handler =
ray::streaming::DownstreamQueueMessageHandler::GetService();
ObjectID &queue_id = queue_ids_[0];
RayFunction async_call_func{ray::Language::PYTHON,
ray::FunctionDescriptorBuilder::FromVector(ray::Language::PYTHON, {"", "", "writer_async_call_func", ""})};
RayFunction sync_call_func{ray::Language::PYTHON,
ray::FunctionDescriptorBuilder::FromVector(ray::Language::PYTHON, {"", "", "writer_sync_call_func", ""})};
downstream_handler->SetPeerActorID(queue_id, peer_actor_id_, async_call_func, sync_call_func);
std::shared_ptr<ReaderQueue> queue =
downstream_handler->CreateDownstreamQueue(queue_id, peer_actor_id_);
bool is_first_pull;
downstream_handler->PullQueue(queue_id, 1, is_first_pull, 10 * 1000);
uint64_t count = 0;
uint8_t msg_id = 1;
while (true) {
uint8_t *data = nullptr;
uint32_t data_size = 0;
uint64_t timeout_ms = 1000;
QueueItem item = queue->PopPendingBlockTimeout(timeout_ms * 1000);
if (item.SeqId() == QUEUE_INVALID_SEQ_ID) {
STREAMING_LOG(INFO) << "PopPendingBlockTimeout timeout.";
data = nullptr;
data_size = 0;
} else {
data = item.Buffer()->Data();
data_size = item.Buffer()->Size();
}
STREAMING_LOG(INFO) << "[Reader] count: " << count;
if (data == nullptr) {
STREAMING_LOG(INFO) << "[Reader] data null";
continue;
}
for (uint32_t i = 0; i < data_size; i++) {
ASSERT_EQ(data[i], msg_id);
}
count++;
if (count == 80) {
bool is_upstream_first_pull;
msg_id = 50;
downstream_handler->PullPeerAsync(queue_id, 50, is_upstream_first_pull, 1000);
continue;
}
msg_id++;
STREAMING_LOG(INFO) << "[Reader] count: " << count;
if (count == 110) {
break;
}
}
STREAMING_LOG(INFO) << "StreamingQueueDownStreamTestSuite::PullPeerAsyncTest done";
status_ = true;
}
};
class TestSuiteFactory {
public:
static std::shared_ptr<StreamingQueueTestSuite> CreateTestSuite(
@@ -272,6 +431,9 @@ class TestSuiteFactory {
if (suite_name == "StreamingWriterTest") {
test_suite = std::make_shared<StreamingQueueWriterTestSuite>(
peer_actor_id, queue_ids, rescale_queue_ids);
} else if (suite_name == "StreamingQueueTest") {
test_suite = std::make_shared<StreamingQueueUpStreamTestSuite>(
peer_actor_id, queue_ids, rescale_queue_ids);
} else {
STREAMING_CHECK(false) << "unsurported suite_name: " << suite_name;
}
@@ -279,6 +441,9 @@ class TestSuiteFactory {
if (suite_name == "StreamingWriterTest") {
test_suite = std::make_shared<StreamingQueueReaderTestSuite>(
peer_actor_id, queue_ids, rescale_queue_ids);
} else if (suite_name == "StreamingQueueTest") {
test_suite = std::make_shared<StreamingQueueDownStreamTestSuite>(
peer_actor_id, queue_ids, rescale_queue_ids);
} else {
STREAMING_CHECK(false) << "unsupported suite_name: " << suite_name;
}
@@ -323,9 +488,6 @@ class StreamingWorker {
-1, // metrics_agent_port
};
CoreWorkerProcess::Initialize(options);
reader_client_ = std::make_shared<ReaderClient>();
writer_client_ = std::make_shared<WriterClient>();
STREAMING_LOG(INFO) << "StreamingWorker constructor";
}
@@ -348,7 +510,7 @@ class StreamingWorker {
RAY_CHECK(function_descriptor->Type() ==
ray::FunctionDescriptorType::kPythonFunctionDescriptor);
auto typed_descriptor = function_descriptor->As<ray::PythonFunctionDescriptor>();
STREAMING_LOG(INFO) << "StreamingWorker::ExecuteTask "
STREAMING_LOG(DEBUG) << "StreamingWorker::ExecuteTask "
<< typed_descriptor->ToString();
std::string func_name = typed_descriptor->FunctionName();
@@ -412,6 +574,8 @@ class StreamingWorker {
private:
void HandleInitTask(std::shared_ptr<LocalMemoryBuffer> buffer) {
reader_client_ = std::make_shared<ReaderClient>();
writer_client_ = std::make_shared<WriterClient>();
uint8_t *bytes = buffer->Data();
uint8_t *p_cur = bytes;
uint32_t *magic_num = (uint32_t *)p_cur;
@@ -425,17 +589,12 @@ class StreamingWorker {
queue::protobuf::StreamingQueueMessageType::StreamingQueueTestInitMsgType);
std::shared_ptr<TestInitMessage> message = TestInitMessage::FromBytes(bytes);
STREAMING_LOG(INFO) << "Init message: " << message->ToString();
std::string actor_handle_serialized = message->ActorHandleSerialized();
CoreWorkerProcess::GetCoreWorker().DeserializeAndRegisterActorHandle(
actor_handle_serialized, ObjectID::Nil());
std::shared_ptr<ActorHandle> actor_handle(new ActorHandle(actor_handle_serialized));
STREAMING_CHECK(actor_handle != nullptr);
STREAMING_LOG(INFO) << " actor id from handle: " << actor_handle->GetActorID();
// STREAMING_LOG(INFO) << "actor_handle_serialized: " << actor_handle_serialized;
// peer_actor_handle_ =
// std::make_shared<ActorHandle>(actor_handle_serialized);
STREAMING_LOG(INFO) << "Actor id from handle: " << actor_handle->GetActorID();
STREAMING_LOG(INFO) << "HandleInitTask queues:";
for (auto qid : message->QueueIds()) {
@@ -0,0 +1,29 @@
#include <cstring>
#include <string>
#include "gtest/gtest.h"
#include "queue/message.h"
using namespace ray;
using namespace ray::streaming;
TEST(ProtoBufTest, MessageCommonTest) {
JobID job_id = JobID::FromInt(0);
TaskID task_id = TaskID::ForDriverTask(job_id);
ray::ActorID actor_id = ray::ActorID::Of(job_id, task_id, 0);
ray::ActorID peer_actor_id = ray::ActorID::Of(job_id, task_id, 1);
ObjectID queue_id = ray::ObjectID::FromRandom();
uint8_t data[128];
std::shared_ptr<LocalMemoryBuffer> buffer = std::make_shared<LocalMemoryBuffer>(data, 128, true);
DataMessage msg(actor_id, peer_actor_id, queue_id, 100, 1000, 2000, buffer, true);
std::unique_ptr<LocalMemoryBuffer> serilized_buffer = msg.ToBytes();
std::shared_ptr<DataMessage> msg2 = DataMessage::FromBytes(serilized_buffer->Data());
EXPECT_EQ(msg.ActorId(), msg2->ActorId());
EXPECT_EQ(msg.PeerActorId(), msg2->PeerActorId());
EXPECT_EQ(msg.QueueId(), msg2->QueueId());
}
int main(int argc, char **argv) {
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}
@@ -16,6 +16,12 @@ namespace streaming {
static int node_manager_port;
class StreamingQueueTest : public StreamingQueueTestBase {
public:
StreamingQueueTest()
: StreamingQueueTestBase(1, node_manager_port) {}
};
class StreamingWriterTest : public StreamingQueueTestBase {
public:
StreamingWriterTest() : StreamingQueueTestBase(1, node_manager_port) {}
@@ -26,6 +32,20 @@ class StreamingExactlySameTest : public StreamingQueueTestBase {
StreamingExactlySameTest() : StreamingQueueTestBase(1, node_manager_port) {}
};
TEST_P(StreamingQueueTest, PullPeerAsyncTest) {
STREAMING_LOG(INFO) << "StreamingQueueTest.pull_peer_async_test";
uint32_t queue_num = 1;
SubmitTest(queue_num, "StreamingQueueTest", "pull_peer_async_test", 60 * 1000);
}
TEST_P(StreamingQueueTest, GetQueueTest) {
STREAMING_LOG(INFO) << "StreamingQueueTest.get_queue_test";
uint32_t queue_num = 1;
SubmitTest(queue_num, "StreamingQueueTest", "get_queue_test", 60 * 1000);
}
TEST_P(StreamingWriterTest, streaming_writer_exactly_once_test) {
STREAMING_LOG(INFO) << "StreamingWriterTest.streaming_writer_exactly_once_test";
@@ -36,6 +56,8 @@ TEST_P(StreamingWriterTest, streaming_writer_exactly_once_test) {
60 * 1000);
}
INSTANTIATE_TEST_CASE_P(StreamingTest, StreamingQueueTest, testing::Values(0));
INSTANTIATE_TEST_CASE_P(StreamingTest, StreamingWriterTest, testing::Values(0));
INSTANTIATE_TEST_CASE_P(StreamingTest, StreamingExactlySameTest,