diff --git a/src/ray/core_worker/experimental_mutable_object_provider.cc b/src/ray/core_worker/experimental_mutable_object_provider.cc index c61efb06bcf0..ecf33a89b8e7 100644 --- a/src/ray/core_worker/experimental_mutable_object_provider.cc +++ b/src/ray/core_worker/experimental_mutable_object_provider.cc @@ -120,8 +120,9 @@ void MutableObjectProvider::HandleRegisterMutableObject( RegisterReaderChannel(reader_object_id); } -void MutableObjectProvider::HandlePushMutableObject( +Status MutableObjectProvider::HandlePushMutableObject( const rpc::PushMutableObjectRequest &request, rpc::PushMutableObjectReply *reply) { + reply->set_done(false); LocalReaderInfo info; const ObjectID writer_object_id = ObjectID::FromBinary(request.writer_object_id()); { @@ -130,6 +131,17 @@ void MutableObjectProvider::HandlePushMutableObject( RAY_CHECK(it != remote_writer_object_to_local_reader_.end()); info = it->second; } + auto clear_written_so_far = [this, &writer_object_id]() { + absl::MutexLock guard(&written_so_far_lock_); + written_so_far_.erase(writer_object_id); + }; + Status channel_status = + object_manager_->GetChannelStatus(info.local_object_id, /*is_reader=*/true); + if (!channel_status.ok()) { + clear_written_so_far(); + return channel_status; + } + size_t total_data_size = request.total_data_size(); size_t total_metadata_size = request.total_metadata_size(); @@ -151,17 +163,23 @@ void MutableObjectProvider::HandlePushMutableObject( if (tmp_written_so_far == 0u) { // We set `metadata` to nullptr since the metadata is at the end of the object, which // we will not have until the last chunk is received. - RAY_CHECK_OK(object_manager_->WriteAcquire(info.local_object_id, - total_data_size, - /*metadata=*/nullptr, - total_metadata_size, - info.num_readers, - object_backing_store)); + Status status = object_manager_->WriteAcquire(info.local_object_id, + total_data_size, + /*metadata=*/nullptr, + total_metadata_size, + info.num_readers, + object_backing_store); + if (!status.ok()) { + clear_written_so_far(); + return status; + } } else { - RAY_CHECK_OK(object_manager_->GetObjectBackingStore(info.local_object_id, - total_data_size, - total_metadata_size, - object_backing_store)); + Status status = object_manager_->GetObjectBackingStore( + info.local_object_id, total_data_size, total_metadata_size, object_backing_store); + if (!status.ok()) { + clear_written_so_far(); + return status; + } } RAY_CHECK(object_backing_store); @@ -174,11 +192,16 @@ void MutableObjectProvider::HandlePushMutableObject( request.metadata().data(), total_metadata_size); // The entire object has been written, so call `WriteRelease()`. - RAY_CHECK_OK(object_manager_->WriteRelease(info.local_object_id)); + Status status = object_manager_->WriteRelease(info.local_object_id); + if (!status.ok()) { + clear_written_so_far(); + return status; + } reply->set_done(true); } else { reply->set_done(false); } + return Status::OK(); } Status MutableObjectProvider::WriteAcquire(const ObjectID &object_id, @@ -248,7 +271,12 @@ void MutableObjectProvider::PollWriterClosure( RAY_CHECK(object->GetData()); RAY_CHECK(object->GetMetadata()); - std::shared_ptr num_replied = std::make_shared(0); + struct PushCompletionState { + absl::Mutex mutex; + size_t num_replied ABSL_GUARDED_BY(mutex) = 0; + bool failed ABSL_GUARDED_BY(mutex) = false; + }; + auto completion_state = std::make_shared(); for (const auto &reader : *remote_readers) { reader->PushMutableObject( writer_object_id, @@ -256,16 +284,35 @@ void MutableObjectProvider::PollWriterClosure( object->GetMetadata()->Size(), object->GetData()->Data(), object->GetMetadata()->Data(), - [this, &io_context, writer_object_id, remote_readers, num_replied]( + [this, &io_context, writer_object_id, remote_readers, completion_state]( const Status &push_object_status, const rpc::PushMutableObjectReply &reply) { - *num_replied += 1; if (!push_object_status.ok()) { RAY_LOG(ERROR) << "Failed to transfer object to a remote node for an object id " - << writer_object_id << ". It can cause hang."; + << writer_object_id << ": " << push_object_status; } - if (*num_replied == remote_readers->size()) { + bool close_writer = false; + bool continue_polling = false; + { + absl::MutexLock guard(&completion_state->mutex); + ++completion_state->num_replied; + if (!push_object_status.ok() && !completion_state->failed) { + completion_state->failed = true; + close_writer = true; + } + continue_polling = completion_state->num_replied == remote_readers->size() && + !completion_state->failed; + } + if (close_writer) { + // TODO(pzhu): Add failure recovery and retries for remote push failures. + Status close_status = object_manager_->SetError(writer_object_id); + if (!close_status.ok()) { + RAY_LOG(ERROR) << "Failed to close mutable object channel " + << writer_object_id + << " after a remote push failure: " << close_status; + } + } else if (continue_polling) { io_context.post( [this, &io_context, writer_object_id, remote_readers]() { PollWriterClosure(io_context, writer_object_id, remote_readers); diff --git a/src/ray/core_worker/experimental_mutable_object_provider.h b/src/ray/core_worker/experimental_mutable_object_provider.h index 3883c22a5d35..0294efe04e3c 100644 --- a/src/ray/core_worker/experimental_mutable_object_provider.h +++ b/src/ray/core_worker/experimental_mutable_object_provider.h @@ -50,8 +50,8 @@ class MutableObjectProvider : public MutableObjectProviderInterface { int64_t num_readers, const ObjectID &reader_object_id) override; - void HandlePushMutableObject(const rpc::PushMutableObjectRequest &request, - rpc::PushMutableObjectReply *reply) override; + Status HandlePushMutableObject(const rpc::PushMutableObjectRequest &request, + rpc::PushMutableObjectReply *reply) override; Status WriteAcquire(const ObjectID &object_id, int64_t data_size, diff --git a/src/ray/core_worker/experimental_mutable_object_provider_interface.h b/src/ray/core_worker/experimental_mutable_object_provider_interface.h index 9abe2dd321f8..5e57bddfa820 100644 --- a/src/ray/core_worker/experimental_mutable_object_provider_interface.h +++ b/src/ray/core_worker/experimental_mutable_object_provider_interface.h @@ -59,8 +59,12 @@ class MutableObjectProviderInterface { /// RPC callback for when a writer pushes a mutable object over the network to a reader /// on this node. - virtual void HandlePushMutableObject(const rpc::PushMutableObjectRequest &request, - rpc::PushMutableObjectReply *reply) = 0; + /// \param[in] request The request containing the pushed object chunk. + /// \param[out] reply The reply indicating whether the complete object has been + /// received. \return The status of processing this chunk. A closed channel is reported + /// to the sender so it can stop polling the failed channel. + virtual Status HandlePushMutableObject(const rpc::PushMutableObjectRequest &request, + rpc::PushMutableObjectReply *reply) = 0; /// Acquires a write lock on the object that prevents readers from reading /// until we are done writing. This is safe for concurrent writers. diff --git a/src/ray/core_worker/fake_mutable_object_provider.h b/src/ray/core_worker/fake_mutable_object_provider.h index 9d330d62ffd3..d37b46194c7f 100644 --- a/src/ray/core_worker/fake_mutable_object_provider.h +++ b/src/ray/core_worker/fake_mutable_object_provider.h @@ -40,8 +40,10 @@ class FakeMutableObjectProvider : public MutableObjectProviderInterface { int64_t num_readers, const ObjectID &reader_object_id) override {} - void HandlePushMutableObject(const rpc::PushMutableObjectRequest &request, - rpc::PushMutableObjectReply *reply) override {} + Status HandlePushMutableObject(const rpc::PushMutableObjectRequest &request, + rpc::PushMutableObjectReply *reply) override { + return Status::OK(); + } Status WriteAcquire(const ObjectID &object_id, int64_t data_size, diff --git a/src/ray/core_worker/tests/BUILD.bazel b/src/ray/core_worker/tests/BUILD.bazel index 995a0eb8ad02..085eb6aa803c 100644 --- a/src/ray/core_worker/tests/BUILD.bazel +++ b/src/ray/core_worker/tests/BUILD.bazel @@ -221,6 +221,8 @@ ray_cc_test( "@com_google_absl//absl/functional:bind_front", "@com_google_absl//absl/random", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/time", ], ) diff --git a/src/ray/core_worker/tests/mutable_object_provider_test.cc b/src/ray/core_worker/tests/mutable_object_provider_test.cc index 9916b2f27560..98e45457935b 100644 --- a/src/ray/core_worker/tests/mutable_object_provider_test.cc +++ b/src/ray/core_worker/tests/mutable_object_provider_test.cc @@ -14,13 +14,17 @@ #include #include +#include #include #include +#include #include #include "absl/functional/bind_front.h" #include "absl/random/random.h" #include "absl/strings/str_format.h" +#include "absl/synchronization/notification.h" +#include "absl/time/time.h" #include "gtest/gtest.h" #include "ray/core_worker/experimental_mutable_object_provider.h" #include "ray/object_manager/common.h" @@ -84,8 +88,80 @@ class FakeRayletClientForTest : public rpc::FakeRayletClient { void *data, void *metadata, const rpc::ClientCallback &callback) override { + std::optional status; + { + absl::MutexLock guard(&lock_); + pushed_objects_.push_back(object_id); + status = push_status_; + push_callbacks_.push_back(status.has_value() ? nullptr : callback); + } + if (status.has_value()) { + rpc::PushMutableObjectReply reply; + callback(*status, std::move(reply)); + MarkPushCompleted(); + } + } + + void SetPushStatus(Status status) { absl::MutexLock guard(&lock_); - pushed_objects_.push_back(object_id); + push_status_ = std::move(status); + } + + bool WaitForPush() { return WaitForCompletedPushCount(1); } + + bool WaitForPushCount(size_t count) { + struct PushCountCondition { + const FakeRayletClientForTest *client; + size_t count; + }; + PushCountCondition condition{this, count}; + absl::MutexLock guard(&lock_); + return lock_.AwaitWithTimeout( + absl::Condition( + +[](PushCountCondition *condition) { + return condition->client->pushed_objects_.size() >= condition->count; + }, + &condition), + absl::Seconds(5)); + } + + bool WaitForCompletedPushCount(size_t count) { + struct PushCountCondition { + const FakeRayletClientForTest *client; + size_t count; + }; + PushCountCondition condition{this, count}; + absl::MutexLock guard(&lock_); + return lock_.AwaitWithTimeout(absl::Condition( + +[](PushCountCondition *condition) { + condition->client->lock_.AssertHeld(); + return condition->client->num_completed_pushes_ >= + condition->count; + }, + &condition), + absl::Seconds(5)); + } + + void CompletePush(size_t index, Status status) { + rpc::ClientCallback callback; + { + absl::MutexLock guard(&lock_); + RAY_CHECK_LT(index, push_callbacks_.size()); + RAY_CHECK(push_callbacks_[index]); + callback = std::move(push_callbacks_[index]); + push_callbacks_[index] = nullptr; + } + rpc::PushMutableObjectReply reply; + callback(status, std::move(reply)); + MarkPushCompleted(); + } + + void ClearPendingPushes() { + std::vector> callbacks; + { + absl::MutexLock guard(&lock_); + callbacks.swap(push_callbacks_); + } } std::vector pushed_objects() { @@ -94,8 +170,17 @@ class FakeRayletClientForTest : public rpc::FakeRayletClient { } private: + void MarkPushCompleted() { + absl::MutexLock guard(&lock_); + ++num_completed_pushes_; + } + absl::Mutex lock_; + std::optional push_status_ ABSL_GUARDED_BY(lock_); std::vector pushed_objects_; + size_t num_completed_pushes_ ABSL_GUARDED_BY(lock_) = 0; + std::vector> push_callbacks_ + ABSL_GUARDED_BY(lock_); }; std::shared_ptr GetFakeRayletClient( @@ -105,6 +190,53 @@ std::shared_ptr GetFakeRayletClient( } // namespace +TEST(FakeRayletClientForTest, ReleasesCompletedPushCallbacks) { + for (bool complete_automatically : {false, true}) { + std::weak_ptr weak_interface; + { + auto interface = std::make_shared(); + weak_interface = interface; + if (complete_automatically) { + interface->SetPushStatus(Status::OK()); + } + interface->PushMutableObject( + ObjectID::FromRandom(), + 0, + 0, + nullptr, + nullptr, + [interface](const Status &, const rpc::PushMutableObjectReply &) { + EXPECT_EQ(interface->pushed_objects().size(), 1); + }); + if (!complete_automatically) { + interface->CompletePush(0, Status::OK()); + } + ASSERT_TRUE(interface->WaitForCompletedPushCount(1)); + } + EXPECT_TRUE(weak_interface.expired()); + } +} + +TEST(FakeRayletClientForTest, ReleasesPendingPushCallbacks) { + std::weak_ptr weak_interface; + { + auto interface = std::make_shared(); + weak_interface = interface; + interface->PushMutableObject( + ObjectID::FromRandom(), + 0, + 0, + nullptr, + nullptr, + [interface](const Status &, const rpc::PushMutableObjectReply &) { + EXPECT_EQ(interface->pushed_objects().size(), 1); + ADD_FAILURE() << "A cleared push must not invoke its callback"; + }); + interface->ClearPendingPushes(); + } + EXPECT_TRUE(weak_interface.expired()); +} + TEST(MutableObjectProvider, RegisterWriterChannel) { ObjectID object_id = ObjectID::FromRandom(); NodeID node_id = NodeID::FromRandom(); @@ -134,6 +266,7 @@ TEST(MutableObjectProvider, RegisterWriterChannel) { EXPECT_EQ(interface->pushed_objects().size(), 1); EXPECT_EQ(interface->pushed_objects().front(), object_id); + interface->ClearPendingPushes(); } TEST(MutableObjectProvider, MutableObjectBufferReadRelease) { @@ -190,7 +323,7 @@ TEST(MutableObjectProvider, HandlePushMutableObject) { request.set_total_metadata_size(0); ray::rpc::PushMutableObjectReply reply; - provider.HandlePushMutableObject(request, &reply); + EXPECT_EQ(provider.HandlePushMutableObject(request, &reply).code(), StatusCode::OK); std::shared_ptr result; EXPECT_EQ(provider.ReadAcquire(local_object_id, result).code(), StatusCode::OK); @@ -198,6 +331,137 @@ TEST(MutableObjectProvider, HandlePushMutableObject) { EXPECT_EQ(provider.ReadRelease(local_object_id).code(), StatusCode::OK); } +TEST(MutableObjectProvider, HandlePushMutableObjectOnClosedChannel) { + ObjectID object_id = ObjectID::FromRandom(); + ObjectID local_object_id = ObjectID::FromRandom(); + auto plasma = std::make_shared(); + MutableObjectProvider provider(plasma, /*factory=*/nullptr, nullptr); + provider.HandleRegisterMutableObject(object_id, /*num_readers=*/1, local_object_id); + ASSERT_EQ(provider.SetError(local_object_id).code(), StatusCode::OK); + + rpc::PushMutableObjectRequest request; + request.set_writer_object_id(object_id.Binary()); + request.set_total_data_size(1); + request.set_total_metadata_size(0); + request.set_offset(0); + request.set_chunk_size(1); + request.set_data("x"); + + rpc::PushMutableObjectReply reply; + EXPECT_EQ(provider.HandlePushMutableObject(request, &reply).code(), + StatusCode::ChannelError); + EXPECT_FALSE(reply.done()); +} + +TEST(MutableObjectProvider, RemotePushFailureClosesWriterChannel) { + ObjectID object_id = ObjectID::FromRandom(); + NodeID node_id = NodeID::FromRandom(); + auto plasma = std::make_shared(); + auto interface = std::make_shared(); + interface->SetPushStatus(Status::ChannelError("reader channel is closed")); + + MutableObjectProvider provider( + plasma, + /*factory=*/absl::bind_front(GetFakeRayletClient, interface), + nullptr); + provider.RegisterWriterChannel(object_id, {node_id}); + + std::shared_ptr data; + ASSERT_EQ(provider + .WriteAcquire(object_id, + /*data_size=*/1, + /*metadata=*/nullptr, + /*metadata_size=*/0, + /*num_readers=*/1, + data) + .code(), + StatusCode::OK); + data->Data()[0] = 'x'; + ASSERT_EQ(provider.WriteRelease(object_id).code(), StatusCode::OK); + + ASSERT_TRUE(interface->WaitForPush()); + EXPECT_EQ(provider.GetChannelStatus(object_id, /*is_reader=*/false).code(), + StatusCode::ChannelError); + EXPECT_EQ(interface->pushed_objects().size(), 1); +} + +TEST(MutableObjectProvider, RemotePushFailureClosesWriterBeforeOtherReadersReply) { + ObjectID object_id = ObjectID::FromRandom(); + NodeID first_node_id = NodeID::FromRandom(); + NodeID second_node_id = NodeID::FromRandom(); + NodeID third_node_id = NodeID::FromRandom(); + auto plasma = std::make_shared(); + auto interface = std::make_shared(); + + MutableObjectProvider provider( + plasma, + /*factory=*/absl::bind_front(GetFakeRayletClient, interface), + nullptr); + provider.RegisterWriterChannel(object_id, + {first_node_id, second_node_id, third_node_id}); + + std::shared_ptr data; + ASSERT_EQ(provider + .WriteAcquire(object_id, + /*data_size=*/1, + /*metadata=*/nullptr, + /*metadata_size=*/0, + /*num_readers=*/1, + data) + .code(), + StatusCode::OK); + data->Data()[0] = 'x'; + ASSERT_EQ(provider.WriteRelease(object_id).code(), StatusCode::OK); + ASSERT_TRUE(interface->WaitForPushCount(3)); + + interface->CompletePush(0, Status::ChannelError("first reader channel is closed")); + EXPECT_EQ(provider.GetChannelStatus(object_id, /*is_reader=*/false).code(), + StatusCode::ChannelError); + + interface->CompletePush(1, Status::ChannelError("second reader channel is closed")); + EXPECT_EQ(provider.GetChannelStatus(object_id, /*is_reader=*/false).code(), + StatusCode::ChannelError); + + interface->CompletePush(2, Status::OK()); + EXPECT_EQ(provider.GetChannelStatus(object_id, /*is_reader=*/false).code(), + StatusCode::ChannelError); + EXPECT_EQ(interface->pushed_objects().size(), 3); +} + +TEST(MutableObjectProvider, SuccessfulRemotePushesContinuePolling) { + ObjectID object_id = ObjectID::FromRandom(); + NodeID first_node_id = NodeID::FromRandom(); + NodeID second_node_id = NodeID::FromRandom(); + auto plasma = std::make_shared(); + auto interface = std::make_shared(); + interface->SetPushStatus(Status::OK()); + + MutableObjectProvider provider( + plasma, + /*factory=*/absl::bind_front(GetFakeRayletClient, interface), + nullptr); + provider.RegisterWriterChannel(object_id, {first_node_id, second_node_id}); + + for (size_t value = 0; value < 2; ++value) { + std::shared_ptr data; + ASSERT_EQ(provider + .WriteAcquire(object_id, + /*data_size=*/1, + /*metadata=*/nullptr, + /*metadata_size=*/0, + /*num_readers=*/1, + data) + .code(), + StatusCode::OK); + data->Data()[0] = 'x'; + ASSERT_EQ(provider.WriteRelease(object_id).code(), StatusCode::OK); + ASSERT_TRUE(interface->WaitForCompletedPushCount((value + 1) * 2)); + } + + EXPECT_EQ(provider.GetChannelStatus(object_id, /*is_reader=*/false).code(), + StatusCode::OK); +} + TEST(MutableObjectProvider, MutableObjectBufferSetError) { ObjectID object_id = ObjectID::FromRandom(); auto plasma = std::make_shared(); diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 4e4bfdc81c46..93efbdbf2e52 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -803,8 +803,8 @@ void NodeManager::HandleRegisterMutableObject( void NodeManager::HandlePushMutableObject(rpc::PushMutableObjectRequest request, rpc::PushMutableObjectReply *reply, rpc::SendReplyCallback send_reply_callback) { - mutable_object_provider_->HandlePushMutableObject(request, reply); - send_reply_callback(Status::OK(), nullptr, nullptr); + Status status = mutable_object_provider_->HandlePushMutableObject(request, reply); + send_reply_callback(status, nullptr, nullptr); } void NodeManager::QueryAllWorkerStates( diff --git a/src/ray/raylet_rpc_client/raylet_client.cc b/src/ray/raylet_rpc_client/raylet_client.cc index fd4a3eafec34..be867ad0b8bc 100644 --- a/src/ray/raylet_rpc_client/raylet_client.cc +++ b/src/ray/raylet_rpc_client/raylet_client.cc @@ -14,6 +14,8 @@ #include "ray/raylet_rpc_client/raylet_client.h" +#include +#include #include #include #include @@ -164,6 +166,10 @@ void RayletClient::PushMutableObject( total_num_chunks++; } + // Send an empty-data chunk so metadata-only values can complete the push. + total_num_chunks = std::max(total_num_chunks, static_cast(1)); + + auto callback_called = std::make_shared>(false); for (uint64_t i = 0; i < total_num_chunks; i++) { rpc::PushMutableObjectRequest request; request.set_writer_object_id(writer_object_id.Binary()); @@ -181,20 +187,25 @@ void RayletClient::PushMutableObject( request.set_metadata(static_cast(metadata), metadata_size); // TODO(jackhumphries): Add failure recovery, retries, and timeout. - INVOKE_RPC_CALL( - NodeManagerService, - PushMutableObject, - request, - [callback](const Status &status, rpc::PushMutableObjectReply &&reply) { - RAY_LOG_IF_ERROR(ERROR, status) << "Error pushing mutable object: " << status; - if (reply.done()) { - // The callback is only executed once the receiver node receives all chunks - // for the mutable object write. - callback(status, std::move(reply)); - } - }, - grpc_client_, - /*method_timeout_ms*/ -1); + auto rpc_callback = [callback, callback_called](const Status &status, + rpc::PushMutableObjectReply &&reply) { + RAY_LOG_IF_ERROR(ERROR, status) << "Error pushing mutable object: " << status; + if (!status.ok() || reply.done()) { + bool expected = false; + if (!callback_called->compare_exchange_strong(expected, true)) { + return; + } + // The callback is only executed once the receiver node receives all chunks + // for the mutable object write, or when any chunk fails. + callback(status, std::move(reply)); + } + }; + INVOKE_RPC_CALL(NodeManagerService, + PushMutableObject, + request, + rpc_callback, + grpc_client_, + /*method_timeout_ms*/ -1); } }