From 4979b8c4d9132dd2b176ef48b6c5652713530474 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Sat, 14 Sep 2019 22:31:33 -0700 Subject: [PATCH] Ordered execution of tasks per actor handle (#5664) --- BUILD.bazel | 10 ++ src/ray/core_worker/test/core_worker_test.cc | 15 ++- src/ray/core_worker/test/mock_worker.cc | 11 ++ .../core_worker/test/scheduling_queue_test.cc | 81 ++++++++++++++ .../transport/direct_actor_transport.cc | 95 +++++++++------- .../transport/direct_actor_transport.h | 93 +++++++++++++++- src/ray/protobuf/direct_actor.proto | 10 ++ src/ray/rpc/worker/direct_actor_client.h | 103 +++++++++++++++--- src/ray/rpc/worker/direct_actor_common.h | 28 +++++ src/ray/util/test_util.h | 3 + 10 files changed, 389 insertions(+), 60 deletions(-) create mode 100644 src/ray/core_worker/test/scheduling_queue_test.cc create mode 100644 src/ray/rpc/worker/direct_actor_common.h diff --git a/BUILD.bazel b/BUILD.bazel index 58b621290..5fcd1742d 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -416,6 +416,16 @@ cc_binary( ], ) +cc_test( + name = "scheduling_queue_test", + srcs = ["src/ray/core_worker/test/scheduling_queue_test.cc"], + copts = COPTS, + deps = [ + ":core_worker_lib", + "@com_google_googletest//:gtest_main", + ], +) + cc_test( name = "lineage_cache_test", srcs = ["src/ray/raylet/lineage_cache_test.cc"], diff --git a/src/ray/core_worker/test/core_worker_test.cc b/src/ray/core_worker/test/core_worker_test.cc index b5853463d..8d075f7b8 100644 --- a/src/ray/core_worker/test/core_worker_test.cc +++ b/src/ray/core_worker/test/core_worker_test.cc @@ -698,6 +698,7 @@ TEST_F(SingleNodeTest, TestDirectActorTaskSubmissionPerf) { raylet_socket_names_[0], JobID::FromInt(1), gcs_options_, "", nullptr); std::unique_ptr actor_handle; + std::vector object_ids; // Test creating actor. uint8_t array[] = {1, 2, 3}; @@ -715,11 +716,14 @@ TEST_F(SingleNodeTest, TestDirectActorTaskSubmissionPerf) { 30 * 1000 /* 30s */)); // Test submitting some tasks with by-value args for that actor. int64_t start_ms = current_time_ms(); - const int num_tasks = 10000; + const int num_tasks = 100000; RAY_LOG(INFO) << "start submitting " << num_tasks << " tasks"; for (int i = 0; i < num_tasks; i++) { // Create arguments with PassByValue. std::vector args; + int64_t array[] = {SHOULD_CHECK_MESSAGE_ORDER, i}; + auto buffer = std::make_shared(reinterpret_cast(array), + sizeof(array)); args.emplace_back(TaskArg::PassByValue(std::make_shared(buffer, nullptr))); TaskOptions options{1, resources}; @@ -729,9 +733,18 @@ TEST_F(SingleNodeTest, TestDirectActorTaskSubmissionPerf) { RAY_CHECK_OK( driver.Tasks().SubmitActorTask(*actor_handle, func, args, options, &return_ids)); ASSERT_EQ(return_ids.size(), 1); + object_ids.emplace_back(return_ids[0]); } RAY_LOG(INFO) << "finish submitting " << num_tasks << " tasks" << ", which takes " << current_time_ms() - start_ms << " ms"; + + for (const auto &object_id : object_ids) { + std::vector> results; + RAY_CHECK_OK(driver.Objects().Get({object_id}, -1, &results)); + ASSERT_EQ(results.size(), 1); + } + RAY_LOG(INFO) << "finish executing " << num_tasks << " tasks" + << ", which takes " << current_time_ms() - start_ms << " ms"; } TEST_F(ZeroNodeTest, TestWorkerContext) { diff --git a/src/ray/core_worker/test/mock_worker.cc b/src/ray/core_worker/test/mock_worker.cc index e93107bdc..789586591 100644 --- a/src/ray/core_worker/test/mock_worker.cc +++ b/src/ray/core_worker/test/mock_worker.cc @@ -44,6 +44,16 @@ class MockWorker { auto &data = arg->GetData(); buffer.insert(buffer.end(), data->Data(), data->Data() + data->Size()); } + if (buffer.size() >= 8) { + auto int_arr = reinterpret_cast(buffer.data()); + if (int_arr[0] == SHOULD_CHECK_MESSAGE_ORDER) { + auto seq_no = int_arr[1]; + if (seq_no > 0) { + RAY_CHECK(seq_no == prev_seq_no_ + 1) << seq_no << " vs " << prev_seq_no_; + } + prev_seq_no_ = seq_no; + } + } auto memory_buffer = std::make_shared(buffer.data(), buffer.size(), true); @@ -56,6 +66,7 @@ class MockWorker { } CoreWorker worker_; + int64_t prev_seq_no_ = 0; }; } // namespace ray diff --git a/src/ray/core_worker/test/scheduling_queue_test.cc b/src/ray/core_worker/test/scheduling_queue_test.cc new file mode 100644 index 000000000..4fa3f4dcd --- /dev/null +++ b/src/ray/core_worker/test/scheduling_queue_test.cc @@ -0,0 +1,81 @@ +#include +#include "gtest/gtest.h" + +#include "ray/core_worker/transport/direct_actor_transport.h" + +namespace ray { + +TEST(SchedulingQueueTest, TestInOrder) { + boost::asio::io_service io_service; + SchedulingQueue queue(io_service, 0); + int n_ok = 0; + int n_rej = 0; + auto fn_ok = [&n_ok]() { n_ok++; }; + auto fn_rej = [&n_rej]() { n_rej++; }; + queue.Add(0, -1, fn_ok, fn_rej); + queue.Add(1, -1, fn_ok, fn_rej); + queue.Add(2, -1, fn_ok, fn_rej); + queue.Add(3, -1, fn_ok, fn_rej); + io_service.run(); + ASSERT_EQ(n_ok, 4); + ASSERT_EQ(n_rej, 0); +} + +TEST(SchedulingQueueTest, TestOutOfOrder) { + boost::asio::io_service io_service; + SchedulingQueue queue(io_service, 0); + int n_ok = 0; + int n_rej = 0; + auto fn_ok = [&n_ok]() { n_ok++; }; + auto fn_rej = [&n_rej]() { n_rej++; }; + queue.Add(2, -1, fn_ok, fn_rej); + queue.Add(0, -1, fn_ok, fn_rej); + queue.Add(3, -1, fn_ok, fn_rej); + queue.Add(1, -1, fn_ok, fn_rej); + io_service.run(); + ASSERT_EQ(n_ok, 4); + ASSERT_EQ(n_rej, 0); +} + +TEST(SchedulingQueueTest, TestDepWaitTimeout) { + boost::asio::io_service io_service; + SchedulingQueue queue(io_service, 0); + int n_ok = 0; + int n_rej = 0; + auto fn_ok = [&n_ok]() { n_ok++; }; + auto fn_rej = [&n_rej]() { n_rej++; }; + queue.Add(2, -1, fn_ok, fn_rej); + queue.Add(0, -1, fn_ok, fn_rej); + queue.Add(3, -1, fn_ok, fn_rej); + ASSERT_EQ(n_ok, 1); + ASSERT_EQ(n_rej, 0); + io_service.run(); // immediately triggers timeout + ASSERT_EQ(n_ok, 1); + ASSERT_EQ(n_rej, 2); + queue.Add(4, -1, fn_ok, fn_rej); + queue.Add(5, -1, fn_ok, fn_rej); + ASSERT_EQ(n_ok, 3); + ASSERT_EQ(n_rej, 2); +} + +TEST(SchedulingQueueTest, TestSkipAlreadyProcessedByClient) { + boost::asio::io_service io_service; + SchedulingQueue queue(io_service, 0); + int n_ok = 0; + int n_rej = 0; + auto fn_ok = [&n_ok]() { n_ok++; }; + auto fn_rej = [&n_rej]() { n_rej++; }; + queue.Add(2, 2, fn_ok, fn_rej); + queue.Add(3, 2, fn_ok, fn_rej); + queue.Add(1, 2, fn_ok, fn_rej); + io_service.run(); + ASSERT_EQ(n_ok, 1); + ASSERT_EQ(n_rej, 2); +} + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/ray/core_worker/transport/direct_actor_transport.cc b/src/ray/core_worker/transport/direct_actor_transport.cc index 0a6a7d2db..b038d1d0a 100644 --- a/src/ray/core_worker/transport/direct_actor_transport.cc +++ b/src/ray/core_worker/transport/direct_actor_transport.cc @@ -1,4 +1,3 @@ - #include "ray/core_worker/transport/direct_actor_transport.h" #include "ray/common/task/task.h" @@ -68,7 +67,7 @@ Status CoreWorkerDirectActorTaskSubmitter::SubmitTask( // Submit request. auto &client = rpc_clients_[actor_id]; - PushTask(*client, *request, actor_id, task_id, num_returns); + PushTask(*client, std::move(request), actor_id, task_id, num_returns); return Status::OK(); } else { // Actor is dead, treat the task as failure. @@ -138,32 +137,31 @@ Status CoreWorkerDirectActorTaskSubmitter::SubscribeActorUpdates( void CoreWorkerDirectActorTaskSubmitter::ConnectAndSendPendingTasks( const ActorID &actor_id, std::string ip_address, int port) { - std::unique_ptr grpc_client( - new rpc::DirectActorClient(ip_address, port, client_call_manager_)); + std::shared_ptr grpc_client = + rpc::DirectActorClient::make(ip_address, port, client_call_manager_); RAY_CHECK(rpc_clients_.emplace(actor_id, std::move(grpc_client)).second); // Submit all pending requests. auto &client = rpc_clients_[actor_id]; auto &requests = pending_requests_[actor_id]; while (!requests.empty()) { - const auto &request = *requests.front(); - PushTask(*client, request, actor_id, - TaskID::FromBinary(request.task_spec().task_id()), - request.task_spec().num_returns()); + auto request = std::move(requests.front()); + auto num_returns = request->task_spec().num_returns(); + auto task_id = TaskID::FromBinary(request->task_spec().task_id()); + PushTask(*client, std::move(request), actor_id, task_id, num_returns); requests.pop_front(); } } -void CoreWorkerDirectActorTaskSubmitter::PushTask(rpc::DirectActorClient &client, - const rpc::PushTaskRequest &request, - const ActorID &actor_id, - const TaskID &task_id, - int num_returns) { +void CoreWorkerDirectActorTaskSubmitter::PushTask( + rpc::DirectActorClient &client, std::unique_ptr request, + const ActorID &actor_id, const TaskID &task_id, int num_returns) { RAY_LOG(DEBUG) << "Pushing task " << task_id << " to actor " << actor_id; waiting_reply_tasks_[actor_id].insert(std::make_pair(task_id, num_returns)); - auto status = - client.PushTask(request, [this, actor_id, task_id, num_returns]( - Status status, const rpc::PushTaskReply &reply) { + + auto status = client.PushTask( + std::move(request), [this, actor_id, task_id, num_returns]( + Status status, const rpc::PushTaskReply &reply) { { std::unique_lock guard(mutex_); waiting_reply_tasks_[actor_id].erase(task_id); @@ -230,6 +228,7 @@ CoreWorkerDirectActorTaskReceiver::CoreWorkerDirectActorTaskReceiver( boost::asio::io_service &io_service, rpc::GrpcServer &server, const TaskHandler &task_handler) : worker_context_(worker_context), + io_service_(io_service), object_interface_(object_interface), task_service_(io_service, *this), task_handler_(task_handler) { @@ -253,33 +252,47 @@ void CoreWorkerDirectActorTaskReceiver::HandlePushTask( return; } - auto num_returns = task_spec.NumReturns(); - RAY_CHECK(task_spec.IsActorCreationTask() || task_spec.IsActorTask()); - RAY_CHECK(num_returns > 0); - // Decrease to account for the dummy object id. - num_returns--; - - std::vector> results; - auto status = task_handler_(task_spec, &results); - RAY_CHECK(results.size() == num_returns) << results.size() << " " << num_returns; - - for (size_t i = 0; i < results.size(); i++) { - auto return_object = (*reply).add_return_objects(); - ObjectID id = ObjectID::ForTaskReturn( - task_spec.TaskId(), /*index=*/i + 1, - /*transport_type=*/static_cast(TaskTransportType::DIRECT_ACTOR)); - return_object->set_object_id(id.Binary()); - const auto &result = results[i]; - if (result->HasData()) { - return_object->set_data(result->GetData()->Data(), result->GetData()->Size()); - } - if (result->HasMetadata()) { - return_object->set_metadata(result->GetMetadata()->Data(), - result->GetMetadata()->Size()); - } + auto it = scheduling_queue_.find(task_spec.ActorHandleId()); + if (it == scheduling_queue_.end()) { + auto result = scheduling_queue_.emplace( + task_spec.ActorHandleId(), + std::unique_ptr(new SchedulingQueue(io_service_))); + it = result.first; } + it->second->Add( + request.sequence_number(), request.client_processed_up_to(), + [this, reply, send_reply_callback, task_spec]() { + auto num_returns = task_spec.NumReturns(); + RAY_CHECK(task_spec.IsActorCreationTask() || task_spec.IsActorTask()); + RAY_CHECK(num_returns > 0); + // Decrease to account for the dummy object id. + num_returns--; - send_reply_callback(status, nullptr, nullptr); + std::vector> results; + auto status = task_handler_(task_spec, &results); + RAY_CHECK(results.size() == num_returns) << results.size() << " " << num_returns; + + for (size_t i = 0; i < results.size(); i++) { + auto return_object = (*reply).add_return_objects(); + ObjectID id = ObjectID::ForTaskReturn( + task_spec.TaskId(), /*index=*/i + 1, + /*transport_type=*/static_cast(TaskTransportType::DIRECT_ACTOR)); + return_object->set_object_id(id.Binary()); + const auto &result = results[i]; + if (result->GetData() != nullptr) { + return_object->set_data(result->GetData()->Data(), result->GetData()->Size()); + } + if (result->GetMetadata() != nullptr) { + return_object->set_metadata(result->GetMetadata()->Data(), + result->GetMetadata()->Size()); + } + } + + send_reply_callback(status, nullptr, nullptr); + }, + [this, send_reply_callback]() { + send_reply_callback(Status::Invalid("client cancelled rpc"), nullptr, nullptr); + }); } } // namespace ray diff --git a/src/ray/core_worker/transport/direct_actor_transport.h b/src/ray/core_worker/transport/direct_actor_transport.h index effa7b01d..de8054fb1 100644 --- a/src/ray/core_worker/transport/direct_actor_transport.h +++ b/src/ray/core_worker/transport/direct_actor_transport.h @@ -3,7 +3,9 @@ #include #include +#include +#include "ray/common/id.h" #include "ray/core_worker/object_interface.h" #include "ray/core_worker/transport/transport.h" #include "ray/gcs/redis_gcs_client.h" @@ -12,6 +14,9 @@ namespace ray { +/// The max time to wait for out-of-order tasks. +const int kMaxReorderWaitSeconds = 30; + /// In direct actor call task submitter and receiver, a task is directly submitted /// to the actor that will execute it. @@ -53,8 +58,9 @@ class CoreWorkerDirectActorTaskSubmitter : public CoreWorkerTaskSubmitter { /// \param[in] task_id The ID of a task. /// \param[in] num_returns Number of return objects. /// \return Void. - void PushTask(rpc::DirectActorClient &client, const rpc::PushTaskRequest &request, - const ActorID &actor_id, const TaskID &task_id, int num_returns); + void PushTask(rpc::DirectActorClient &client, + std::unique_ptr request, const ActorID &actor_id, + const TaskID &task_id, int num_returns); /// Treat a task as failed. /// @@ -98,7 +104,11 @@ class CoreWorkerDirectActorTaskSubmitter : public CoreWorkerTaskSubmitter { std::unordered_map actor_states_; /// Map from actor id to rpc client. This only includes actors that we send tasks to. - std::unordered_map> rpc_clients_; + /// We use shared_ptr to enable shared_from_this for pending client callbacks. + /// + /// TODO(zhijunfu): this will be moved into `actor_states_` later when we can + /// subscribe updates for a specific actor. + std::unordered_map> rpc_clients_; /// Map from actor id to the actor's pending requests. std::unordered_map>> @@ -116,6 +126,78 @@ class CoreWorkerDirectActorTaskSubmitter : public CoreWorkerTaskSubmitter { friend class CoreWorkerTest; }; +/// Used to ensure serial order of task execution per actor handle. +/// See direct_actor.proto for a description of the ordering protocol. +class SchedulingQueue { + public: + SchedulingQueue(boost::asio::io_service &io_service, + int64_t reorder_wait_seconds = kMaxReorderWaitSeconds) + : wait_timer_(io_service), reorder_wait_seconds_(reorder_wait_seconds) {} + + void Add(int64_t seq_no, int64_t client_processed_up_to, + std::function accept_request, std::function reject_request) { + if (client_processed_up_to >= next_seq_no_) { + RAY_LOG(DEBUG) << "client skipping requests " << next_seq_no_ << " to " + << client_processed_up_to; + next_seq_no_ = client_processed_up_to + 1; + } + pending_tasks_[seq_no] = make_pair(accept_request, reject_request); + + // Reject any stale requests that the client doesn't need any longer. + while (!pending_tasks_.empty() && pending_tasks_.begin()->first < next_seq_no_) { + auto head = pending_tasks_.begin(); + head->second.second(); // reject_request + pending_tasks_.erase(head); + } + + // Process as many in-order requests as we can. + while (!pending_tasks_.empty() && pending_tasks_.begin()->first == next_seq_no_) { + auto head = pending_tasks_.begin(); + head->second.first(); // accept_request + pending_tasks_.erase(head); + next_seq_no_++; + } + + // Set a timeout on the queued tasks to avoid an infinite wait on failure. + wait_timer_.expires_from_now(boost::posix_time::seconds(reorder_wait_seconds_)); + if (!pending_tasks_.empty()) { + RAY_LOG(DEBUG) << "waiting for " << next_seq_no_ << " queue size " + << pending_tasks_.size(); + wait_timer_.async_wait([this](const boost::system::error_code &error) { + if (error == boost::asio::error::operation_aborted) { + return; // time deadline was adjusted + } + OnDependencyWaitTimeout(); + }); + } + } + + private: + /// Called when we time out waiting for a task dependency to show up. + void OnDependencyWaitTimeout() { + RAY_LOG(ERROR) << "timed out waiting for " << next_seq_no_ + << ", cancelling all queued tasks"; + while (!pending_tasks_.empty()) { + auto head = pending_tasks_.begin(); + head->second.second(); // reject_request + pending_tasks_.erase(head); + next_seq_no_ = std::max(next_seq_no_, head->first + 1); + } + } + + /// Max time in seconds to wait for dependencies to show up. + const int64_t reorder_wait_seconds_ = 0; + /// Sorted map of (accept, rej) task callbacks keyed by their sequence number. + std::map, std::function>> + pending_tasks_; + /// The next sequence number we are waiting for to arrive. + int64_t next_seq_no_ = 0; + /// Timer for waiting on dependencies. + boost::asio::deadline_timer wait_timer_; + + friend class SchedulingQueueTest; +}; + class CoreWorkerDirectActorTaskReceiver : public CoreWorkerTaskReceiver, public rpc::DirectActorHandler { public: @@ -138,12 +220,17 @@ class CoreWorkerDirectActorTaskReceiver : public CoreWorkerTaskReceiver, private: // Worker context. WorkerContext &worker_context_; + /// The IO event loop. + boost::asio::io_service &io_service_; // Object interface. CoreWorkerObjectInterface &object_interface_; /// The rpc service for `DirectActorService`. rpc::DirectActorGrpcService task_service_; /// The callback function to process a task. TaskHandler task_handler_; + /// Queue of pending requests per actor handle. + /// TODO(ekl) GC these queues once the handle is no longer active. + std::unordered_map> scheduling_queue_; }; } // namespace ray diff --git a/src/ray/protobuf/direct_actor.proto b/src/ray/protobuf/direct_actor.proto index 373af60f1..ca691b3cd 100644 --- a/src/ray/protobuf/direct_actor.proto +++ b/src/ray/protobuf/direct_actor.proto @@ -16,6 +16,16 @@ message ReturnObject { message PushTaskRequest { // The task to be pushed. TaskSpec task_spec = 1; + // The sequence number of the task for this client. This must increase + // sequentially starting from zero for each actor handle. The server + // will guarantee tasks execute in this sequence, waiting for any + // out-of-order request messages to arrive as necessary. + int64 sequence_number = 2; + // The max sequence number the client has processed responses for. This + // is a performance optimization that allows the client to tell the server + // to cancel any PushTaskRequests with seqno <= this value, rather than + // waiting for the server to time out waiting for missing messages. + int64 client_processed_up_to = 3; } message PushTaskReply { diff --git a/src/ray/rpc/worker/direct_actor_client.h b/src/ray/rpc/worker/direct_actor_client.h index 51e74ffd6..238100a90 100644 --- a/src/ray/rpc/worker/direct_actor_client.h +++ b/src/ray/rpc/worker/direct_actor_client.h @@ -1,12 +1,16 @@ #ifndef RAY_RPC_DIRECT_ACTOR_CLIENT_H #define RAY_RPC_DIRECT_ACTOR_CLIENT_H +#include +#include +#include #include #include #include "ray/common/status.h" #include "ray/rpc/client_call.h" +#include "ray/rpc/worker/direct_actor_common.h" #include "ray/util/logging.h" #include "src/ray/protobuf/direct_actor.grpc.pb.h" #include "src/ray/protobuf/direct_actor.pb.h" @@ -15,8 +19,76 @@ namespace ray { namespace rpc { /// Client used for communicating with a direct actor server. -class DirectActorClient { +class DirectActorClient : public std::enable_shared_from_this { public: + /// Constructor. + /// + /// \param[in] address Address of the direct actor server. + /// \param[in] port Port of the direct actor server. + /// \param[in] client_call_manager The `ClientCallManager` used for managing requests. + static std::shared_ptr make(const std::string &address, + const int port, + ClientCallManager &client_call_manager) { + auto instance = new DirectActorClient(address, port, client_call_manager); + return std::shared_ptr(instance); + } + + /// Push a task. + /// + /// \param[in] request The request message. + /// \param[in] callback The callback function that handles reply. + /// \return if the rpc call succeeds + ray::Status PushTask(std::unique_ptr request, + const ClientCallback &callback) { + request->set_sequence_number(next_seq_no_++); + send_queue_.push_back(std::make_pair(std::move(request), callback)); + SendRequests(); + return ray::Status::OK(); + } + + /// Send as many pending tasks as possible. This method is thread-safe. + /// + /// The client will guarantee no more than kMaxBytesInFlight bytes of RPCs are being + /// sent at once. This prevents the server scheduling queue from being overwhelmed. + /// See direct_actor.proto for a description of the ordering protocol. + void SendRequests() { + std::lock_guard lock(mutex_); + auto this_ptr = this->shared_from_this(); + + while (!send_queue_.empty() && rpc_bytes_in_flight_ < kMaxBytesInFlight) { + auto pair = std::move(*send_queue_.begin()); + send_queue_.pop_front(); + + auto request = std::move(pair.first); + auto callback = pair.second; + int64_t task_size = RequestSizeInBytes(*request); + int64_t seq_no = request->sequence_number(); + request->set_client_processed_up_to(max_finished_seq_no_); + rpc_bytes_in_flight_ += task_size; + + client_call_manager_.CreateCall( + *stub_, &DirectActorService::Stub::PrepareAsyncPushTask, *request, + [this, this_ptr, seq_no, task_size, callback](Status status, + const rpc::PushTaskReply &reply) { + { + std::lock_guard lock(mutex_); + if (seq_no > max_finished_seq_no_) { + max_finished_seq_no_ = seq_no; + } + rpc_bytes_in_flight_ -= task_size; + RAY_CHECK(rpc_bytes_in_flight_ >= 0); + } + SendRequests(); + callback(status, reply); + }); + } + + if (!send_queue_.empty()) { + RAY_LOG(DEBUG) << "client send queue size " << send_queue_.size(); + } + } + + private: /// Constructor. /// /// \param[in] address Address of the direct actor server. @@ -30,26 +102,27 @@ class DirectActorClient { stub_ = DirectActorService::NewStub(channel); }; - /// Push a task. - /// - /// \param[in] request The request message. - /// \param[in] callback The callback function that handles reply. - /// \return if the rpc call succeeds - ray::Status PushTask(const PushTaskRequest &request, - const ClientCallback &callback) { - auto call = client_call_manager_ - .CreateCall( - *stub_, &DirectActorService::Stub::PrepareAsyncPushTask, request, - callback); - return call->GetStatus(); - } + /// Protects against unsafe concurrent access from the callback thread. + std::mutex mutex_; - private: /// The gRPC-generated stub. std::unique_ptr stub_; /// The `ClientCallManager` used for managing requests. ClientCallManager &client_call_manager_; + + /// Queue of requests to send. + std::deque, ClientCallback>> + send_queue_; + + /// The next sequence number to assign to a task for this server. + int64_t next_seq_no_ = 0; + + /// The number of bytes currently in flight. + int64_t rpc_bytes_in_flight_ = 0; + + /// The max sequence number we have processed responses for. + int64_t max_finished_seq_no_ = -1; }; } // namespace rpc diff --git a/src/ray/rpc/worker/direct_actor_common.h b/src/ray/rpc/worker/direct_actor_common.h new file mode 100644 index 000000000..f63cd27fb --- /dev/null +++ b/src/ray/rpc/worker/direct_actor_common.h @@ -0,0 +1,28 @@ +#ifndef RAY_RPC_DIRECT_ACTOR_COMMON_H +#define RAY_RPC_DIRECT_ACTOR_COMMON_H + +#include "src/ray/protobuf/direct_actor.grpc.pb.h" +#include "src/ray/protobuf/direct_actor.pb.h" + +namespace ray { +namespace rpc { + +/// The maximum number of requests in flight per client. +const int64_t kMaxBytesInFlight = 16 * 1024 * 1024; + +/// The base size in bytes per request. +const int64_t kBaseRequestSize = 1024; + +/// Get the estimated size in bytes of the given task. +const static int64_t RequestSizeInBytes(const PushTaskRequest &request) { + int64_t size = kBaseRequestSize; + for (auto &arg : request.task_spec().args()) { + size += arg.data().size(); + } + return size; +} + +} // namespace rpc +} // namespace ray + +#endif // RAY_RPC_DIRECT_ACTOR_COMMON_H diff --git a/src/ray/util/test_util.h b/src/ray/util/test_util.h index 937dc0f06..a89acc7f1 100644 --- a/src/ray/util/test_util.h +++ b/src/ray/util/test_util.h @@ -7,6 +7,9 @@ namespace ray { +// Magic argument to signal to mock_worker we should check message order. +int64_t SHOULD_CHECK_MESSAGE_ORDER = 123450000; + /// Wait until the condition is met, or timeout is reached. /// /// \param[in] condition The condition to wait for.