mirror of
https://github.com/wassname/ray.git
synced 2026-09-10 12:38:43 +08:00
[Streaming]Streaming queue support failover (#8161)
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user