diff --git a/BUILD.bazel b/BUILD.bazel index f6da3a253..e87a25fe1 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -213,6 +213,7 @@ cc_library( ":ray_common", ":worker_cc_grpc", "@boost//:asio", + "@boost//:thread", "@com_github_grpc_grpc//:grpc++", ], ) diff --git a/python/ray/ray_perf.py b/python/ray/ray_perf.py index 96f22d838..f3ec0d85b 100644 --- a/python/ray/ray_perf.py +++ b/python/ray/ray_perf.py @@ -1,27 +1,46 @@ """This is the script for `ray microbenchmark`.""" +import os import time import numpy as np import multiprocessing import ray +# Only run tests matching this filter pattern. +filter_pattern = os.environ.get("TESTS_TO_RUN", "") + @ray.remote class Actor(object): def small_value(self): return 0 + def small_value_arg(self, x): + return 0 + def small_value_batch(self, n): ray.get([small_value.remote() for _ in range(n)]) @ray.remote class Client(object): - def __init__(self, server): - self.server = server + def __init__(self, servers): + if not isinstance(servers, list): + servers = [servers] + self.servers = servers def small_value_batch(self, n): - ray.get([self.server.small_value.remote() for _ in range(n)]) + results = [] + for s in self.servers: + results.extend([s.small_value.remote() for _ in range(n)]) + ray.get(results) + + def small_value_batch_arg(self, n): + x = ray.put(0) + results = [] + for s in self.servers: + results.extend([s.small_value_arg.remote(x) for _ in range(n)]) + ray.get(results) @ray.remote @@ -37,6 +56,8 @@ def small_value_batch(n): def timeit(name, fn, multiplier=1): + if filter_pattern not in name: + return # warmup start = time.time() while time.time() - start < 1: @@ -56,6 +77,7 @@ def timeit(name, fn, multiplier=1): def main(): + print("Tip: set TESTS_TO_RUN='pattern' to run a subset of benchmarks") ray.init() value = ray.put(0) arr = np.zeros(100 * 1024 * 1024, dtype=np.int64) @@ -141,23 +163,17 @@ def main(): timeit("multi client actor calls async", actor_multi2, m * n) - a = Actor._remote(is_direct_call=True) - - def actor_sync_direct(): - ray.get(a.small_value.remote()) - - timeit("single client direct actor calls sync", actor_sync_direct) - - a = Actor._remote(is_direct_call=True) - - def actor_async_direct(): - ray.get([a.small_value.remote() for _ in range(1000)]) - - timeit("single client direct actor calls async", actor_async_direct, 1000) - n = 5000 n_cpu = multiprocessing.cpu_count() // 2 actors = [Actor._remote(is_direct_call=True) for _ in range(n_cpu)] + client = Client.remote(actors) + + def actor_async_direct(): + ray.get(client.small_value_batch.remote(n)) + + timeit("single client direct actor calls async", actor_async_direct, + n * len(actors)) + clients = [Client.remote(a) for a in actors] def actor_multi2_direct(): @@ -166,6 +182,16 @@ def main(): timeit("multi client direct actor calls async", actor_multi2_direct, n * len(clients)) + n = 1000 + actors = [Actor._remote(is_direct_call=True) for _ in range(n_cpu)] + clients = [Client.remote(a) for a in actors] + + def actor_multi2_direct_arg(): + ray.get([c.small_value_batch_arg.remote(n) for c in clients]) + + timeit("multi client direct actor calls with arg async", + actor_multi2_direct_arg, n * len(clients)) + if __name__ == "__main__": main() diff --git a/python/ray/tests/test_basic.py b/python/ray/tests/test_basic.py index 6be4de0ec..ee18e5878 100644 --- a/python/ray/tests/test_basic.py +++ b/python/ray/tests/test_basic.py @@ -1223,9 +1223,70 @@ def test_direct_actor_errors(ray_start_regular): with pytest.raises(Exception): ray.get(f.remote([a.f.remote(2)])) - # by ref args not implemented - with pytest.raises(ray.exceptions.RayletError): - a.f.remote(f.remote(2)) + +def test_direct_actor_pass_by_ref(ray_start_regular): + @ray.remote + class Actor(object): + def __init__(self): + pass + + def f(self, x): + return x * 2 + + @ray.remote + def f(x): + return x + + @ray.remote + def error(): + sys.exit(0) + + a = Actor._remote(is_direct_call=True) + assert ray.get(a.f.remote(f.remote(1))) == 2 + + fut = [a.f.remote(f.remote(i)) for i in range(100)] + assert ray.get(fut) == [i * 2 for i in range(100)] + + # propagates errors for pass by ref + with pytest.raises(Exception): + ray.get(a.f.remote(error.remote())) + + +def test_direct_actor_pass_by_ref_order_optimization(shutdown_only): + ray.init(num_cpus=4) + + @ray.remote + class Actor(object): + def __init__(self): + pass + + def f(self, x): + pass + + a = Actor._remote(is_direct_call=True) + + @ray.remote + def fast_value(): + print("fast value") + pass + + @ray.remote + def slow_value(): + print("start sleep") + time.sleep(30) + + @ray.remote + def runner(f): + print("runner", a, f) + return ray.get(a.f.remote(f.remote())) + + runner.remote(slow_value) + time.sleep(1) + x2 = runner.remote(fast_value) + start = time.time() + ray.get(x2) + delta = time.time() - start + assert delta < 10, "did not skip slow value" def test_direct_actor_recursive(ray_start_regular): diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index 3ce84d331..753333d9b 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -131,6 +131,10 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, raylet_socket, WorkerID::FromBinary(worker_context_.GetWorkerID().Binary()), (worker_type_ == ray::WorkerType::WORKER), worker_context_.GetCurrentJobID(), language_, worker_server_.GetPort())); + // Unfortunately the raylet client has to be constructed after the receivers. + if (direct_actor_task_receiver_ != nullptr) { + direct_actor_task_receiver_->Init(*raylet_client_); + } // Set timer to periodically send heartbeats containing active object IDs to the raylet. // If the heartbeat timeout is < 0, the heartbeats are disabled. diff --git a/src/ray/core_worker/test/core_worker_test.cc b/src/ray/core_worker/test/core_worker_test.cc index cd78b30cb..d0a67e2bd 100644 --- a/src/ray/core_worker/test/core_worker_test.cc +++ b/src/ray/core_worker/test/core_worker_test.cc @@ -316,12 +316,6 @@ void CoreWorkerTest::TestActorTask(std::unordered_map &reso std::vector return_ids; RayFunction func(ray::Language::PYTHON, {}); auto status = driver.SubmitActorTask(actor_id, func, args, options, &return_ids); - if (is_direct_call) { - // For direct actor call, submitting a task with by-reference arguments - // would fail. - ASSERT_TRUE(!status.ok()); - return; - } ASSERT_TRUE(status.ok()); ASSERT_EQ(return_ids.size(), 1); diff --git a/src/ray/core_worker/test/scheduling_queue_test.cc b/src/ray/core_worker/test/scheduling_queue_test.cc index 4fa3f4dcd..46b78832d 100644 --- a/src/ray/core_worker/test/scheduling_queue_test.cc +++ b/src/ray/core_worker/test/scheduling_queue_test.cc @@ -5,9 +5,25 @@ namespace ray { +class MockWaiter : public DependencyWaiter { + public: + MockWaiter() {} + + void Wait(const std::vector &dependencies, + std::function on_dependencies_available) override { + callbacks_.push_back([on_dependencies_available]() { on_dependencies_available(); }); + } + + void Complete(int index) { callbacks_[index](); } + + private: + std::vector> callbacks_; +}; + TEST(SchedulingQueueTest, TestInOrder) { boost::asio::io_service io_service; - SchedulingQueue queue(io_service, 0); + MockWaiter waiter; + SchedulingQueue queue(io_service, waiter, 0); int n_ok = 0; int n_rej = 0; auto fn_ok = [&n_ok]() { n_ok++; }; @@ -21,9 +37,55 @@ TEST(SchedulingQueueTest, TestInOrder) { ASSERT_EQ(n_rej, 0); } +TEST(SchedulingQueueTest, TestWaitForObjects) { + ObjectID obj1 = ObjectID::FromRandom(); + ObjectID obj2 = ObjectID::FromRandom(); + ObjectID obj3 = ObjectID::FromRandom(); + boost::asio::io_service io_service; + MockWaiter waiter; + SchedulingQueue queue(io_service, waiter, 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, {obj1}); + queue.Add(2, -1, fn_ok, fn_rej, {obj2}); + queue.Add(3, -1, fn_ok, fn_rej, {obj3}); + ASSERT_EQ(n_ok, 1); + + waiter.Complete(0); + ASSERT_EQ(n_ok, 2); + + waiter.Complete(2); + ASSERT_EQ(n_ok, 2); + + waiter.Complete(1); + ASSERT_EQ(n_ok, 4); +} + +TEST(SchedulingQueueTest, TestWaitForObjectsNotSubjectToSeqTimeout) { + ObjectID obj1 = ObjectID::FromRandom(); + boost::asio::io_service io_service; + MockWaiter waiter; + SchedulingQueue queue(io_service, waiter, 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, {obj1}); + ASSERT_EQ(n_ok, 1); + io_service.run(); + ASSERT_EQ(n_rej, 0); + waiter.Complete(0); + ASSERT_EQ(n_ok, 2); +} + TEST(SchedulingQueueTest, TestOutOfOrder) { boost::asio::io_service io_service; - SchedulingQueue queue(io_service, 0); + MockWaiter waiter; + SchedulingQueue queue(io_service, waiter, 0); int n_ok = 0; int n_rej = 0; auto fn_ok = [&n_ok]() { n_ok++; }; @@ -37,9 +99,10 @@ TEST(SchedulingQueueTest, TestOutOfOrder) { ASSERT_EQ(n_rej, 0); } -TEST(SchedulingQueueTest, TestDepWaitTimeout) { +TEST(SchedulingQueueTest, TestSeqWaitTimeout) { boost::asio::io_service io_service; - SchedulingQueue queue(io_service, 0); + MockWaiter waiter; + SchedulingQueue queue(io_service, waiter, 0); int n_ok = 0; int n_rej = 0; auto fn_ok = [&n_ok]() { n_ok++; }; @@ -60,7 +123,8 @@ TEST(SchedulingQueueTest, TestDepWaitTimeout) { TEST(SchedulingQueueTest, TestSkipAlreadyProcessedByClient) { boost::asio::io_service io_service; - SchedulingQueue queue(io_service, 0); + MockWaiter waiter; + SchedulingQueue queue(io_service, waiter, 0); int n_ok = 0; int n_rej = 0; auto fn_ok = [&n_ok]() { n_ok++; }; diff --git a/src/ray/core_worker/transport/direct_actor_transport.cc b/src/ray/core_worker/transport/direct_actor_transport.cc index 75b81700c..016a2ee97 100644 --- a/src/ray/core_worker/transport/direct_actor_transport.cc +++ b/src/ray/core_worker/transport/direct_actor_transport.cc @@ -24,9 +24,6 @@ CoreWorkerDirectActorTaskSubmitter::CoreWorkerDirectActorTaskSubmitter( Status CoreWorkerDirectActorTaskSubmitter::SubmitTask( const TaskSpecification &task_spec) { RAY_LOG(DEBUG) << "Submitting task " << task_spec.TaskId(); - if (HasByReferenceArgs(task_spec)) { - return Status::Invalid("Direct actor call only supports by-value arguments"); - } RAY_CHECK(task_spec.IsActorTask()); const auto &actor_id = task_spec.ActorId(); @@ -84,7 +81,7 @@ void CoreWorkerDirectActorTaskSubmitter::HandleActorUpdate( // Check if this actor is the one that we're interested, if we already have // a connection to the actor, or have pending requests for it, we should // create a new connection. - if (pending_requests_.count(actor_id) > 0) { + if (pending_requests_.count(actor_id) > 0 && rpc_clients_.count(actor_id) == 0) { ConnectAndSendPendingTasks(actor_id, actor_data.ip_address(), actor_data.port()); } } else { @@ -147,6 +144,9 @@ void CoreWorkerDirectActorTaskSubmitter::PushTask( waiting_reply_tasks_[actor_id].erase(task_id); } if (!status.ok()) { + // Note that this might be the __ray_terminate__ task, so we don't log + // loudly with ERROR here. + RAY_LOG(DEBUG) << "Task failed with error: " << status; TreatTaskAsFailed(task_id, num_returns, rpc::ErrorType::ACTOR_DIED); return; } @@ -208,28 +208,38 @@ CoreWorkerDirectActorTaskReceiver::CoreWorkerDirectActorTaskReceiver( server.RegisterService(task_service_); } +void CoreWorkerDirectActorTaskReceiver::Init(RayletClient &raylet_client) { + waiter_.reset(new DependencyWaiterImpl(raylet_client)); +} + void CoreWorkerDirectActorTaskReceiver::HandlePushTask( const rpc::PushTaskRequest &request, rpc::PushTaskReply *reply, rpc::SendReplyCallback send_reply_callback) { + RAY_CHECK(waiter_ != nullptr) << "Must call init() prior to use"; const TaskSpecification task_spec(request.task_spec()); RAY_LOG(DEBUG) << "Received task " << task_spec.TaskId(); - if (HasByReferenceArgs(task_spec)) { - send_reply_callback( - Status::Invalid("Direct actor call only supports by value arguments"), nullptr, - nullptr); - return; - } if (task_spec.IsActorTask() && !worker_context_.CurrentActorUseDirectCall()) { send_reply_callback(Status::Invalid("This actor doesn't accept direct calls."), nullptr, nullptr); return; } + // TODO(ekl) resolving object dependencies is expensive and requires an IPC to + // the raylet, which is a central bottleneck. In the future, we should inline + // dependencies that are small and already known to be local to the client. + std::vector dependencies; + for (size_t i = 0; i < task_spec.NumArgs(); ++i) { + int count = task_spec.ArgIdCount(i); + for (int j = 0; j < count; j++) { + dependencies.push_back(task_spec.ArgId(i, j)); + } + } + auto it = scheduling_queue_.find(task_spec.CallerId()); if (it == scheduling_queue_.end()) { auto result = scheduling_queue_.emplace( - task_spec.CallerId(), - std::unique_ptr(new SchedulingQueue(task_main_io_service_))); + task_spec.CallerId(), std::unique_ptr( + new SchedulingQueue(task_main_io_service_, *waiter_))); it = result.first; } it->second->Add( @@ -268,7 +278,17 @@ void CoreWorkerDirectActorTaskReceiver::HandlePushTask( }, [send_reply_callback]() { send_reply_callback(Status::Invalid("client cancelled rpc"), nullptr, nullptr); - }); + }, + dependencies); +} + +void CoreWorkerDirectActorTaskReceiver::HandleDirectActorCallArgWaitComplete( + const rpc::DirectActorCallArgWaitCompleteRequest &request, + rpc::DirectActorCallArgWaitCompleteReply *reply, + rpc::SendReplyCallback send_reply_callback) { + RAY_LOG(DEBUG) << "Arg wait complete for tag " << request.tag(); + waiter_->OnWaitComplete(request.tag()); + send_reply_callback(Status::OK(), 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 d1407aac7..0fbf217f6 100644 --- a/src/ray/core_worker/transport/direct_actor_transport.h +++ b/src/ray/core_worker/transport/direct_actor_transport.h @@ -1,6 +1,7 @@ #ifndef RAY_CORE_WORKER_DIRECT_ACTOR_TRANSPORT_H #define RAY_CORE_WORKER_DIRECT_ACTOR_TRANSPORT_H +#include #include #include #include @@ -125,60 +126,139 @@ class CoreWorkerDirectActorTaskSubmitter { friend class CoreWorkerTest; }; +/// Object dependency and RPC state of an inbound request. +class InboundRequest { + public: + InboundRequest(){}; + InboundRequest(std::function accept_callback, + std::function reject_callback, bool has_dependencies) + : accept_callback_(accept_callback), + reject_callback_(reject_callback), + has_pending_dependencies_(has_dependencies) {} + + void Accept() { accept_callback_(); } + void Cancel() { reject_callback_(); } + bool CanExecute() const { return !has_pending_dependencies_; } + void MarkDependenciesSatisfied() { has_pending_dependencies_ = false; } + + private: + std::function accept_callback_; + std::function reject_callback_; + bool has_pending_dependencies_; +}; + +/// Waits for an object dependency to become available. Abstract for testing. +class DependencyWaiter { + public: + /// Calls `callback` once the specified objects become available. + virtual void Wait(const std::vector &dependencies, + std::function on_dependencies_available) = 0; +}; + +class DependencyWaiterImpl : public DependencyWaiter { + public: + DependencyWaiterImpl(RayletClient &raylet_client) : raylet_client_(raylet_client) {} + + void Wait(const std::vector &dependencies, + std::function on_dependencies_available) override { + auto tag = next_request_id_++; + requests_[tag] = on_dependencies_available; + raylet_client_.WaitForDirectActorCallArgs(dependencies, tag); + } + + /// Fulfills the callback stored by Wait(). + void OnWaitComplete(int64_t tag) { + auto it = requests_.find(tag); + RAY_CHECK(it != requests_.end()); + it->second(); + requests_.erase(it); + } + + private: + int64_t next_request_id_ = 0; + std::unordered_map> requests_; + RayletClient &raylet_client_; +}; + /// 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, + SchedulingQueue(boost::asio::io_service &main_io_service, DependencyWaiter &waiter, int64_t reorder_wait_seconds = kMaxReorderWaitSeconds) - : wait_timer_(io_service), reorder_wait_seconds_(reorder_wait_seconds) {} + : wait_timer_(main_io_service), + waiter_(waiter), + reorder_wait_seconds_(reorder_wait_seconds), + main_thread_id_(boost::this_thread::get_id()) {} void Add(int64_t seq_no, int64_t client_processed_up_to, - std::function accept_request, std::function reject_request) { + std::function accept_request, std::function reject_request, + const std::vector &dependencies = {}) { + RAY_CHECK(boost::this_thread::get_id() == main_thread_id_); if (client_processed_up_to >= next_seq_no_) { - RAY_LOG(DEBUG) << "client skipping requests " << next_seq_no_ << " to " + RAY_LOG(ERROR) << "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); + pending_tasks_[seq_no] = + InboundRequest(accept_request, reject_request, dependencies.size() > 0); + if (dependencies.size() > 0) { + waiter_.Wait(dependencies, [seq_no, this]() { + RAY_CHECK(boost::this_thread::get_id() == main_thread_id_); + auto it = pending_tasks_.find(seq_no); + if (it != pending_tasks_.end()) { + it->second.MarkDependenciesSatisfied(); + ScheduleRequests(); + } + }); + } + ScheduleRequests(); + } - // Reject any stale requests that the client doesn't need any longer. + private: + /// Schedules as many requests as possible in sequence. + void ScheduleRequests() { + // Cancel 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 + head->second.Cancel(); pending_tasks_.erase(head); } // Process as many in-order requests as we can. - while (!pending_tasks_.empty() && pending_tasks_.begin()->first == next_seq_no_) { + while (!pending_tasks_.empty() && pending_tasks_.begin()->first == next_seq_no_ && + pending_tasks_.begin()->second.CanExecute()) { auto head = pending_tasks_.begin(); - head->second.first(); // accept_request + head->second.Accept(); 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()) { + if (pending_tasks_.empty() || !pending_tasks_.begin()->second.CanExecute()) { + // No timeout for object dependency waits. + wait_timer_.cancel(); + } else { + // 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_)); 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(); + OnSequencingWaitTimeout(); }); } } - private: - /// Called when we time out waiting for a task dependency to show up. - void OnDependencyWaitTimeout() { + /// Called when we time out waiting for an earlier task to show up. + void OnSequencingWaitTimeout() { + RAY_CHECK(boost::this_thread::get_id() == main_thread_id_); 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 + head->second.Cancel(); pending_tasks_.erase(head); next_seq_no_ = std::max(next_seq_no_, head->first + 1); } @@ -187,12 +267,15 @@ class SchedulingQueue { /// 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_; + std::map 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_; + /// The id of the thread that constructed this scheduling queue. + boost::thread::id main_thread_id_; + /// Reference to the waiter owned by the task receiver. + DependencyWaiter &waiter_; friend class SchedulingQueueTest; }; @@ -208,16 +291,27 @@ class CoreWorkerDirectActorTaskReceiver : public rpc::DirectActorHandler { rpc::GrpcServer &server, const TaskHandler &task_handler); + /// Initialize this receiver. This must be called prior to use. + void Init(RayletClient &client); + /// Handle a `PushTask` request. - /// The implementation can handle this request asynchronously. When hanling is done, the - /// `done_callback` should be called. /// /// \param[in] request The request message. /// \param[out] reply The reply message. - /// \param[in] done_callback The callback to be called when the request is done. + /// \param[in] send_reply_callback The callback to be called when the request is done. void HandlePushTask(const rpc::PushTaskRequest &request, rpc::PushTaskReply *reply, rpc::SendReplyCallback send_reply_callback) override; + /// Handle a `DirectActorCallArgWaitComplete` request. + /// + /// \param[in] request The request message. + /// \param[out] reply The reply message. + /// \param[in] send_reply_callback The callback to be called when the request is done. + void HandleDirectActorCallArgWaitComplete( + const rpc::DirectActorCallArgWaitCompleteRequest &request, + rpc::DirectActorCallArgWaitCompleteReply *reply, + rpc::SendReplyCallback send_reply_callback) override; + private: // Worker context. WorkerContext &worker_context_; @@ -227,6 +321,8 @@ class CoreWorkerDirectActorTaskReceiver : public rpc::DirectActorHandler { TaskHandler task_handler_; /// The IO event loop for running tasks on. boost::asio::io_service &task_main_io_service_; + /// Shared waiter for dependencies required by incoming tasks. + std::unique_ptr waiter_; /// Queue of pending requests per actor handle. /// TODO(ekl) GC these queues once the handle is no longer active. std::unordered_map> scheduling_queue_; diff --git a/src/ray/core_worker/transport/raylet_transport.h b/src/ray/core_worker/transport/raylet_transport.h index 5597813ba..99735128c 100644 --- a/src/ray/core_worker/transport/raylet_transport.h +++ b/src/ray/core_worker/transport/raylet_transport.h @@ -22,8 +22,8 @@ class CoreWorkerRayletTaskReceiver : public rpc::WorkerTaskHandler { rpc::GrpcServer &server, const TaskHandler &task_handler); /// Handle a `AssignTask` request. - /// The implementation can handle this request asynchronously. When hanling is done, the - /// `send_reply_callback` should be called. + /// The implementation can handle this request asynchronously. When handling is done, + /// the `send_reply_callback` should be called. /// /// \param[in] request The request message. /// \param[out] reply The reply message. @@ -41,6 +41,8 @@ class CoreWorkerRayletTaskReceiver : public rpc::WorkerTaskHandler { rpc::WorkerTaskGrpcService task_service_; /// The callback function to process a task. TaskHandler task_handler_; + /// The callback to process arg wait complete. + std::function on_wait_complete_; }; } // namespace ray diff --git a/src/ray/protobuf/direct_actor.proto b/src/ray/protobuf/direct_actor.proto index ca691b3cd..b767f1db1 100644 --- a/src/ray/protobuf/direct_actor.proto +++ b/src/ray/protobuf/direct_actor.proto @@ -33,8 +33,20 @@ message PushTaskReply { repeated ReturnObject return_objects = 1; } +message DirectActorCallArgWaitCompleteRequest { + // Id used to uniquely identify this request. This is sent back to the core + // worker to notify the wait has completed. + int64 tag = 1; +} + +message DirectActorCallArgWaitCompleteReply { +} + // Service for direct actor. service DirectActorService { // Push a task to a worker. rpc PushTask(PushTaskRequest) returns (PushTaskReply); + // Notify wait for direct actor call args has completed + rpc DirectActorCallArgWaitComplete(DirectActorCallArgWaitCompleteRequest) + returns (DirectActorCallArgWaitCompleteReply); } diff --git a/src/ray/raylet/format/node_manager.fbs b/src/ray/raylet/format/node_manager.fbs index e0d07f79f..f32ac148f 100644 --- a/src/ray/raylet/format/node_manager.fbs +++ b/src/ray/raylet/format/node_manager.fbs @@ -55,6 +55,8 @@ enum MessageType:int { // The response message to WaitRequest; replies with the objects found and objects // remaining. WaitReply, + // Wait for objects asynchronously. The reply will be sent back via gRPC push. + WaitForDirectActorCallArgsRequest, // Push an error to the relevant driver. This is sent from a worker to the // node manager. PushErrorRequest, @@ -187,6 +189,14 @@ table WaitReply { remaining: [string]; } +table WaitForDirectActorCallArgsRequest { + // List of object ids we'll be waiting on. + object_ids: [string]; + // Id used to uniquely identify this request. This is sent back to the core + // worker to notify the wait has completed. + tag: int; +} + // This struct is the same as ErrorTableData. table PushErrorRequest { // The ID of the job that the error is for. diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 3942d3387..3743c3d0d 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -898,6 +898,9 @@ void NodeManager::ProcessClientMessage( case protocol::MessageType::WaitRequest: { ProcessWaitRequestMessage(client, message_data); } break; + case protocol::MessageType::WaitForDirectActorCallArgsRequest: { + ProcessWaitForDirectActorCallArgsRequestMessage(client, message_data); + } break; case protocol::MessageType::PushErrorRequest: { ProcessPushErrorRequestMessage(message_data); } break; @@ -1253,6 +1256,37 @@ void NodeManager::ProcessWaitRequestMessage( RAY_CHECK_OK(status); } +void NodeManager::ProcessWaitForDirectActorCallArgsRequestMessage( + const std::shared_ptr &client, const uint8_t *message_data) { + // Read the data. + auto message = + flatbuffers::GetRoot(message_data); + int64_t tag = message->tag(); + std::vector object_ids = from_flatbuf(*message->object_ids()); + std::vector required_object_ids; + for (auto const &object_id : object_ids) { + if (!task_dependency_manager_.CheckObjectLocal(object_id)) { + // Add any missing objects to the list to subscribe to in the task + // dependency manager. These objects will be pulled from remote node + // managers and reconstructed if necessary. + required_object_ids.push_back(object_id); + } + } + + ray::Status status = object_manager_.Wait( + object_ids, -1, object_ids.size(), false, + [this, client, tag](std::vector found, std::vector remaining) { + RAY_CHECK(remaining.empty()); + std::shared_ptr worker = worker_pool_.GetRegisteredWorker(client); + if (worker == nullptr) { + RAY_LOG(ERROR) << "Lost worker for wait request " << client; + } else { + worker->DirectActorCallArgWaitComplete(tag); + } + }); + RAY_CHECK_OK(status); +} + void NodeManager::ProcessPushErrorRequestMessage(const uint8_t *message_data) { auto message = flatbuffers::GetRoot(message_data); diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index 4185d240b..fda201d26 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -428,6 +428,14 @@ class NodeManager : public rpc::NodeManagerServiceHandler { void ProcessWaitRequestMessage(const std::shared_ptr &client, const uint8_t *message_data); + /// Process client message of WaitForDirectActorCallArgsRequest + /// + /// \param client The client that sent the message. + /// \param message_data A pointer to the message data. + /// \return Void. + void ProcessWaitForDirectActorCallArgsRequestMessage( + const std::shared_ptr &client, const uint8_t *message_data); + /// Process client message of PushErrorRequest /// /// \param message_data A pointer to the message data. diff --git a/src/ray/raylet/raylet_client.cc b/src/ray/raylet/raylet_client.cc index e37a7dab3..88cbb3909 100644 --- a/src/ray/raylet/raylet_client.cc +++ b/src/ray/raylet/raylet_client.cc @@ -279,6 +279,15 @@ ray::Status RayletClient::Wait(const std::vector &object_ids, int num_ return ray::Status::OK(); } +ray::Status RayletClient::WaitForDirectActorCallArgs( + const std::vector &object_ids, int64_t tag) { + flatbuffers::FlatBufferBuilder fbb; + auto message = ray::protocol::CreateWaitForDirectActorCallArgsRequest( + fbb, to_flatbuf(fbb, object_ids), tag); + fbb.Finish(message); + return conn_->WriteMessage(MessageType::WaitForDirectActorCallArgsRequest, &fbb); +} + ray::Status RayletClient::PushError(const ray::JobID &job_id, const std::string &type, const std::string &error_message, double timestamp) { flatbuffers::FlatBufferBuilder fbb; diff --git a/src/ray/raylet/raylet_client.h b/src/ray/raylet/raylet_client.h index 3e36c2078..0869c501a 100644 --- a/src/ray/raylet/raylet_client.h +++ b/src/ray/raylet/raylet_client.h @@ -118,6 +118,15 @@ class RayletClient { int64_t timeout_milliseconds, bool wait_local, const TaskID ¤t_task_id, WaitResultPair *result); + /// Wait for the given objects, asynchronously. The core worker is notified when + /// the wait completes. + /// + /// \param object_ids The objects to wait for. + /// \param tag Value that will be sent to the core worker via gRPC on completion. + /// \return ray::Status. + ray::Status WaitForDirectActorCallArgs(const std::vector &object_ids, + int64_t tag); + /// Push an error to the relevant driver. /// /// \param The ID of the job_id that the error is for. diff --git a/src/ray/raylet/worker.cc b/src/ray/raylet/worker.cc index c4d66971c..f9182f43a 100644 --- a/src/ray/raylet/worker.cc +++ b/src/ray/raylet/worker.cc @@ -4,6 +4,8 @@ #include "ray/raylet/format/node_manager_generated.h" #include "ray/raylet/raylet.h" +#include "src/ray/protobuf/direct_actor.grpc.pb.h" +#include "src/ray/protobuf/direct_actor.pb.h" namespace ray { @@ -24,6 +26,8 @@ Worker::Worker(const WorkerID &worker_id, pid_t pid, const Language &language, i if (port_ > 0) { rpc_client_ = std::unique_ptr( new rpc::WorkerTaskClient("127.0.0.1", port_, client_call_manager_)); + direct_rpc_client_ = std::unique_ptr( + new rpc::DirectActorClient("127.0.0.1", port_, client_call_manager_)); } } @@ -150,6 +154,21 @@ void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set, } } +void Worker::DirectActorCallArgWaitComplete(int64_t tag) { + RAY_CHECK(port_ > 0); + rpc::DirectActorCallArgWaitCompleteRequest request; + request.set_tag(tag); + auto status = direct_rpc_client_->DirectActorCallArgWaitComplete( + request, [](Status status, const rpc::DirectActorCallArgWaitCompleteReply &reply) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to send wait complete: " << status.ToString(); + } + }); + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to send wait complete: " << status.ToString(); + } +} + } // namespace raylet } // end namespace ray diff --git a/src/ray/raylet/worker.h b/src/ray/raylet/worker.h index 45e7b5e7b..878fc9f8e 100644 --- a/src/ray/raylet/worker.h +++ b/src/ray/raylet/worker.h @@ -8,6 +8,7 @@ #include "ray/common/task/scheduling_resources.h" #include "ray/common/task/task.h" #include "ray/common/task/task_common.h" +#include "ray/rpc/worker/direct_actor_client.h" #include "ray/rpc/worker/worker_client.h" namespace ray { @@ -62,6 +63,7 @@ class Worker { void AssignTask(const Task &task, const ResourceIdSet &resource_id_set, const std::function finish_assign_callback); + void DirectActorCallArgWaitComplete(int64_t tag); private: /// The worker's ID. @@ -100,6 +102,8 @@ class Worker { rpc::ClientCallManager &client_call_manager_; /// The rpc client to send tasks to this worker. std::unique_ptr rpc_client_; + /// The rpc client to send tasks to the direct actor service. + std::unique_ptr direct_rpc_client_; }; } // namespace raylet diff --git a/src/ray/rpc/worker/direct_actor_client.h b/src/ray/rpc/worker/direct_actor_client.h index 05eb49264..5ca280739 100644 --- a/src/ray/rpc/worker/direct_actor_client.h +++ b/src/ray/rpc/worker/direct_actor_client.h @@ -34,6 +34,19 @@ class DirectActorClient : public std::enable_shared_from_this return std::shared_ptr(instance); } + /// 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. + DirectActorClient(const std::string &address, const int port, + ClientCallManager &client_call_manager) + : client_call_manager_(client_call_manager) { + std::shared_ptr channel = grpc::CreateChannel( + address + ":" + std::to_string(port), grpc::InsecureChannelCredentials()); + stub_ = DirectActorService::NewStub(channel); + }; + /// Push a task. /// /// \param[in] request The request message. @@ -44,12 +57,33 @@ class DirectActorClient : public std::enable_shared_from_this request->set_sequence_number(request->task_spec().actor_task_spec().actor_counter()); { std::lock_guard lock(mutex_); + if (request->task_spec().caller_id() != cur_caller_id_) { + // We are running a new task, reset the seq no counter. + max_finished_seq_no_ = -1; + cur_caller_id_ = request->task_spec().caller_id(); + } send_queue_.push_back(std::make_pair(std::move(request), callback)); } SendRequests(); return ray::Status::OK(); } + /// Notify a wait has completed for direct actor call arguments. + /// + /// \param[in] request The request message. + /// \param[in] callback The callback function that handles reply. + /// \return if the rpc call succeeds + ray::Status DirectActorCallArgWaitComplete( + const DirectActorCallArgWaitCompleteRequest &request, + const ClientCallback &callback) { + auto call = client_call_manager_.CreateCall( + *stub_, &DirectActorService::Stub::PrepareAsyncDirectActorCallArgWaitComplete, + request, callback); + return call->GetStatus(); + } + /// 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 @@ -93,19 +127,6 @@ class DirectActorClient : public std::enable_shared_from_this } private: - /// 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. - DirectActorClient(const std::string &address, const int port, - ClientCallManager &client_call_manager) - : client_call_manager_(client_call_manager) { - std::shared_ptr channel = grpc::CreateChannel( - address + ":" + std::to_string(port), grpc::InsecureChannelCredentials()); - stub_ = DirectActorService::NewStub(channel); - }; - /// Protects against unsafe concurrent access from the callback thread. std::mutex mutex_; @@ -124,6 +145,10 @@ class DirectActorClient : public std::enable_shared_from_this /// The max sequence number we have processed responses for. int64_t max_finished_seq_no_ GUARDED_BY(mutex_) = -1; + + /// The task id we are currently sending requests for. When this changes, + /// the max finished seq no counter is reset. + std::string cur_caller_id_; }; } // namespace rpc diff --git a/src/ray/rpc/worker/direct_actor_server.h b/src/ray/rpc/worker/direct_actor_server.h index a0670aa2a..dc6c6b4fe 100644 --- a/src/ray/rpc/worker/direct_actor_server.h +++ b/src/ray/rpc/worker/direct_actor_server.h @@ -22,6 +22,16 @@ class DirectActorHandler { /// \param[in] done_callback The callback to be called when the request is done. virtual void HandlePushTask(const PushTaskRequest &request, PushTaskReply *reply, SendReplyCallback send_reply_callback) = 0; + + /// Handle a wait reply for direct actor call arg dependencies. + /// + /// \param[in] request The request message. + /// \param[out] reply The reply message. + /// \param[in] send_replay_callback The callback to be called when the request is done. + virtual void HandleDirectActorCallArgWaitComplete( + const rpc::DirectActorCallArgWaitCompleteRequest &request, + rpc::DirectActorCallArgWaitCompleteReply *reply, + rpc::SendReplyCallback send_reply_callback) = 0; }; /// The `GrpcServer` for `WorkerService`. @@ -48,10 +58,20 @@ class DirectActorGrpcService : public GrpcService { PushTaskReply>( service_, &DirectActorService::AsyncService::RequestPushTask, service_handler_, &DirectActorHandler::HandlePushTask, cq, main_service_)); - - // Set `PushTask`'s accept concurrency to 100. server_call_factories_and_concurrencies->emplace_back( std::move(push_task_call_Factory), 100); + + // Initialize the Factory for `DirectActorCallArgWaitComplete` requests. + std::unique_ptr wait_complete_call_Factory( + new ServerCallFactoryImpl( + service_, + &DirectActorService::AsyncService::RequestDirectActorCallArgWaitComplete, + service_handler_, &DirectActorHandler::HandleDirectActorCallArgWaitComplete, + cq, main_service_)); + server_call_factories_and_concurrencies->emplace_back( + std::move(wait_complete_call_Factory), 100); } private: