From f03caa4532131f020410225af6b4b13781152d5c Mon Sep 17 00:00:00 2001 From: Barak Michener Date: Fri, 21 Aug 2020 11:01:22 -0700 Subject: [PATCH] rpc: Follow-up by sharing the core worker client pool within the core worker. (#10206) * Share CoreWorkerClientPool * Format --- src/ray/core_worker/core_worker.cc | 20 +-- src/ray/core_worker/core_worker.h | 3 + src/ray/core_worker/future_resolver.cc | 2 +- src/ray/core_worker/future_resolver.h | 7 +- .../test/direct_actor_transport_test.cc | 8 +- .../test/direct_task_transport_test.cc | 144 ++++++++++-------- .../transport/direct_actor_transport.cc | 8 +- .../transport/direct_actor_transport.h | 19 +-- .../transport/direct_task_transport.cc | 8 +- .../transport/direct_task_transport.h | 7 +- 10 files changed, 127 insertions(+), 99 deletions(-) diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index 30d9cbf16..007822da5 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -446,10 +446,10 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_ } SetCurrentTaskId(task_id); } - auto client_factory = [this](const rpc::Address &addr) { - return std::shared_ptr( - new rpc::CoreWorkerClient(addr, *client_call_manager_)); - }; + + core_worker_client_pool_ = + std::make_shared(*client_call_manager_); + auto raylet_client_factory = [this](const std::string ip_address, int port) { auto grpc_client = rpc::NodeManagerWorkerClient::make(ip_address, port, *client_call_manager_); @@ -461,22 +461,24 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_ std::make_shared(gcs_client_); direct_actor_submitter_ = std::shared_ptr( - new CoreWorkerDirectActorTaskSubmitter(client_factory, memory_store_, + new CoreWorkerDirectActorTaskSubmitter(core_worker_client_pool_, memory_store_, task_manager_)); direct_task_submitter_ = std::unique_ptr(new CoreWorkerDirectTaskSubmitter( - rpc_address_, local_raylet_client_, client_factory, raylet_client_factory, - memory_store_, task_manager_, local_raylet_id, + rpc_address_, local_raylet_client_, core_worker_client_pool_, + raylet_client_factory, memory_store_, task_manager_, local_raylet_id, RayConfig::instance().worker_lease_timeout_milliseconds(), std::move(actor_creator), RayConfig::instance().max_tasks_in_flight_per_worker(), boost::asio::steady_timer(io_service_))); - future_resolver_.reset(new FutureResolver(memory_store_, client_factory, rpc_address_)); + future_resolver_.reset( + new FutureResolver(memory_store_, core_worker_client_pool_, rpc_address_)); // Unfortunately the raylet client has to be constructed after the receivers. if (direct_task_receiver_ != nullptr) { task_argument_waiter_.reset(new DependencyWaiterImpl(*local_raylet_client_)); - direct_task_receiver_->Init(client_factory, rpc_address_, task_argument_waiter_); + direct_task_receiver_->Init(core_worker_client_pool_, rpc_address_, + task_argument_waiter_); } actor_manager_ = std::unique_ptr( diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 387700933..af2e14cdc 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -1040,6 +1040,9 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler { /// Shared client call manager. std::unique_ptr client_call_manager_; + /// Shared core worker client pool. + std::shared_ptr core_worker_client_pool_; + /// Timer used to periodically check if the raylet has died. boost::asio::steady_timer death_check_timer_; diff --git a/src/ray/core_worker/future_resolver.cc b/src/ray/core_worker/future_resolver.cc index db9fdc9fa..8a1cc3f07 100644 --- a/src/ray/core_worker/future_resolver.cc +++ b/src/ray/core_worker/future_resolver.cc @@ -23,7 +23,7 @@ void FutureResolver::ResolveFutureAsync(const ObjectID &object_id, // with a borrowed reference executes on the object's owning worker. return; } - auto conn = owner_clients_.GetOrConnect(owner_address); + auto conn = owner_clients_->GetOrConnect(owner_address); rpc::GetObjectStatusRequest request; request.set_object_id(object_id.Binary()); diff --git a/src/ray/core_worker/future_resolver.h b/src/ray/core_worker/future_resolver.h index 142742fd1..be504a582 100644 --- a/src/ray/core_worker/future_resolver.h +++ b/src/ray/core_worker/future_resolver.h @@ -29,9 +29,10 @@ namespace ray { class FutureResolver { public: FutureResolver(std::shared_ptr store, - rpc::ClientFactoryFn client_factory, const rpc::Address &rpc_address) + std::shared_ptr core_worker_client_pool, + const rpc::Address &rpc_address) : in_memory_store_(store), - owner_clients_(client_factory), + owner_clients_(core_worker_client_pool), rpc_address_(rpc_address) {} /// Resolve the value for a future. This will periodically contact the given @@ -48,7 +49,7 @@ class FutureResolver { /// Used to store values of resolved futures. std::shared_ptr in_memory_store_; - rpc::CoreWorkerClientPool owner_clients_; + std::shared_ptr owner_clients_; /// Address of our RPC server. Used to notify borrowed objects' owners of our /// address, so the owner can contact us to ask when our reference to the diff --git a/src/ray/core_worker/test/direct_actor_transport_test.cc b/src/ray/core_worker/test/direct_actor_transport_test.cc index f6962b1a8..fb894565f 100644 --- a/src/ray/core_worker/test/direct_actor_transport_test.cc +++ b/src/ray/core_worker/test/direct_actor_transport_test.cc @@ -102,10 +102,10 @@ class DirectActorSubmitterTest : public ::testing::Test { store_(std::shared_ptr(new CoreWorkerMemoryStore())), task_finisher_(std::make_shared()), submitter_( - [&](const rpc::Address &addr) { + std::make_shared([&](const rpc::Address &addr) { num_clients_connected_++; return worker_client_; - }, + }), store_, task_finisher_) {} int num_clients_connected_ = 0; @@ -364,6 +364,7 @@ TEST_F(DirectActorSubmitterTest, TestActorRestartOutOfOrderGcs) { // We receive the RESTART message late. Nothing happens. submitter_.DisconnectActor(actor_id, 0, /*dead=*/false); + ASSERT_EQ(num_clients_connected_, 2); // Submit a task. task = CreateActorTaskHelper(actor_id, worker_id, 2); @@ -428,7 +429,8 @@ class DirectActorReceiverTest : public ::testing::Test { receiver_ = std::unique_ptr( new CoreWorkerDirectTaskReceiver(worker_context_, main_io_service_, execute_task, [] { return Status::OK(); })); - receiver_->Init([&](const rpc::Address &addr) { return worker_client_; }, + receiver_->Init(std::make_shared( + [&](const rpc::Address &addr) { return worker_client_; }), rpc_address_, dependency_waiter_); } diff --git a/src/ray/core_worker/test/direct_task_transport_test.cc b/src/ray/core_worker/test/direct_task_transport_test.cc index 6d2ad0444..388eeb36b 100644 --- a/src/ray/core_worker/test/direct_task_transport_test.cc +++ b/src/ray/core_worker/test/direct_task_transport_test.cc @@ -335,12 +335,13 @@ TEST(DirectTaskTransportTest, TestSubmitOneTask) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = @@ -371,12 +372,13 @@ TEST(DirectTaskTransportTest, TestHandleTaskFailure) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""); @@ -400,12 +402,13 @@ TEST(DirectTaskTransportTest, TestConcurrentWorkerLeases) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""); @@ -450,12 +453,13 @@ TEST(DirectTaskTransportTest, TestReuseWorkerLease) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""); @@ -507,12 +511,13 @@ TEST(DirectTaskTransportTest, TestRetryLeaseCancellation) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""); @@ -562,12 +567,13 @@ TEST(DirectTaskTransportTest, TestConcurrentCancellationAndSubmission) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""); @@ -614,12 +620,13 @@ TEST(DirectTaskTransportTest, TestWorkerNotReusedOnError) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""); @@ -657,12 +664,13 @@ TEST(DirectTaskTransportTest, TestWorkerNotReturnedOnExit) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""); @@ -690,7 +698,8 @@ TEST(DirectTaskTransportTest, TestSpillback) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); std::unordered_map> remote_lease_clients; auto lease_client_factory = [&](const std::string &ip, int port) { @@ -702,7 +711,7 @@ TEST(DirectTaskTransportTest, TestSpillback) { }; auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, lease_client_factory, store, task_finisher, ClientID::Nil(), kLongTimeout, actor_creator); std::unordered_map empty_resources; @@ -747,7 +756,8 @@ TEST(DirectTaskTransportTest, TestSpillbackRoundTrip) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); std::unordered_map> remote_lease_clients; auto lease_client_factory = [&](const std::string &ip, int port) { @@ -760,7 +770,7 @@ TEST(DirectTaskTransportTest, TestSpillbackRoundTrip) { auto task_finisher = std::make_shared(); auto local_raylet_id = ClientID::FromRandom(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, lease_client_factory, store, task_finisher, local_raylet_id, kLongTimeout, actor_creator); std::unordered_map empty_resources; @@ -813,12 +823,13 @@ void TestSchedulingKey(const std::shared_ptr store, rpc::Address address; auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); ASSERT_TRUE(submitter.SubmitTask(same1).ok()); ASSERT_TRUE(submitter.SubmitTask(same2).ok()); @@ -925,11 +936,12 @@ TEST(DirectTaskTransportTest, TestWorkerLeaseTimeout) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), /*lease_timeout_ms=*/5, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = @@ -978,12 +990,14 @@ TEST(DirectTaskTransportTest, TestKillExecutingTask) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); + auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""); @@ -1025,12 +1039,13 @@ TEST(DirectTaskTransportTest, TestKillPendingTask) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""); @@ -1053,12 +1068,13 @@ TEST(DirectTaskTransportTest, TestKillResolvingTask) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator); + CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, client_pool, nullptr, + store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator); std::unordered_map empty_resources; ray::FunctionDescriptor empty_descriptor = ray::FunctionDescriptorBuilder::BuildPython("", "", "", ""); @@ -1083,7 +1099,8 @@ TEST(DirectTaskTransportTest, TestPipeliningConcurrentWorkerLeases) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); @@ -1091,9 +1108,9 @@ TEST(DirectTaskTransportTest, TestPipeliningConcurrentWorkerLeases) { // of task submissions. This is done by passing a max_tasks_in_flight_per_worker // parameter to the CoreWorkerDirectTaskSubmitter. uint32_t max_tasks_in_flight_per_worker = 10; - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator, max_tasks_in_flight_per_worker); + CoreWorkerDirectTaskSubmitter submitter( + address, raylet_client, client_pool, nullptr, store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator, max_tasks_in_flight_per_worker); // Prepare 20 tasks and save them in a vector. std::unordered_map empty_resources; @@ -1152,7 +1169,8 @@ TEST(DirectTaskTransportTest, TestPipeliningReuseWorkerLease) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::Address &addr) { return worker_client; }; + auto client_pool = std::make_shared( + [&](const rpc::Address &addr) { return worker_client; }); auto task_finisher = std::make_shared(); auto actor_creator = std::make_shared(); @@ -1160,9 +1178,9 @@ TEST(DirectTaskTransportTest, TestPipeliningReuseWorkerLease) { // of task submissions. This is done by passing a max_tasks_in_flight_per_worker // parameter to the CoreWorkerDirectTaskSubmitter. uint32_t max_tasks_in_flight_per_worker = 10; - CoreWorkerDirectTaskSubmitter submitter(address, raylet_client, factory, nullptr, store, - task_finisher, ClientID::Nil(), kLongTimeout, - actor_creator, max_tasks_in_flight_per_worker); + CoreWorkerDirectTaskSubmitter submitter( + address, raylet_client, client_pool, nullptr, store, task_finisher, ClientID::Nil(), + kLongTimeout, actor_creator, max_tasks_in_flight_per_worker); // prepare 30 tasks and save them in a vector std::unordered_map empty_resources; diff --git a/src/ray/core_worker/transport/direct_actor_transport.cc b/src/ray/core_worker/transport/direct_actor_transport.cc index fe3bceef2..eb7569b5c 100644 --- a/src/ray/core_worker/transport/direct_actor_transport.cc +++ b/src/ray/core_worker/transport/direct_actor_transport.cc @@ -114,6 +114,7 @@ Status CoreWorkerDirectActorTaskSubmitter::SubmitTask(TaskSpecification task_spe void CoreWorkerDirectActorTaskSubmitter::DisconnectRpcClient(ClientQueue &queue) { queue.rpc_client = nullptr; + core_worker_client_pool_->Disconnect(ray::WorkerID::FromBinary(queue.worker_id)); queue.worker_id.clear(); queue.pending_force_kill.reset(); } @@ -149,8 +150,7 @@ void CoreWorkerDirectActorTaskSubmitter::ConnectActor(const ActorID &actor_id, // Update the mapping so new RPCs go out with the right intended worker id. queue->second.worker_id = address.worker_id(); // Create a new connection to the actor. - queue->second.rpc_client = - std::shared_ptr(client_factory_(address)); + queue->second.rpc_client = core_worker_client_pool_->GetOrConnect(address); // TODO(swang): This assumes that all replies from the previous incarnation // of the actor have been received. Fix this by setting an epoch for each // actor task, so we can ignore completed tasks from old epochs. @@ -293,11 +293,11 @@ bool CoreWorkerDirectActorTaskSubmitter::IsActorAlive(const ActorID &actor_id) c } void CoreWorkerDirectTaskReceiver::Init( - rpc::ClientFactoryFn client_factory, rpc::Address rpc_address, + std::shared_ptr client_pool, rpc::Address rpc_address, std::shared_ptr dependency_waiter) { waiter_ = std::move(dependency_waiter); rpc_address_ = rpc_address; - client_factory_ = client_factory; + client_pool_ = client_pool; } void CoreWorkerDirectTaskReceiver::HandlePushTask( diff --git a/src/ray/core_worker/transport/direct_actor_transport.h b/src/ray/core_worker/transport/direct_actor_transport.h index ada51d242..afaaa93f9 100644 --- a/src/ray/core_worker/transport/direct_actor_transport.h +++ b/src/ray/core_worker/transport/direct_actor_transport.h @@ -63,10 +63,11 @@ class CoreWorkerDirectActorTaskSubmitterInterface { class CoreWorkerDirectActorTaskSubmitter : public CoreWorkerDirectActorTaskSubmitterInterface { public: - CoreWorkerDirectActorTaskSubmitter(rpc::ClientFactoryFn client_factory, - std::shared_ptr store, - std::shared_ptr task_finisher) - : client_factory_(client_factory), + CoreWorkerDirectActorTaskSubmitter( + std::shared_ptr core_worker_client_pool, + std::shared_ptr store, + std::shared_ptr task_finisher) + : core_worker_client_pool_(core_worker_client_pool), resolver_(store, task_finisher), task_finisher_(task_finisher) {} @@ -217,8 +218,8 @@ class CoreWorkerDirectActorTaskSubmitter /// \return Whether this actor is alive. bool IsActorAlive(const ActorID &actor_id) const; - /// Factory for producing new core worker clients. - rpc::ClientFactoryFn client_factory_; + /// Pool for producing new core worker clients. + std::shared_ptr core_worker_client_pool_; /// Mutex to protect the various maps below. mutable absl::Mutex mu_; @@ -493,7 +494,7 @@ class CoreWorkerDirectTaskReceiver { task_done_(task_done) {} /// Initialize this receiver. This must be called prior to use. - void Init(rpc::ClientFactoryFn client_factory, rpc::Address rpc_address, + void Init(std::shared_ptr, rpc::Address rpc_address, std::shared_ptr dependency_waiter); /// Handle a `PushTask` request. @@ -513,8 +514,8 @@ class CoreWorkerDirectTaskReceiver { boost::asio::io_service &task_main_io_service_; /// The callback function to be invoked when finishing a task. OnTaskDone task_done_; - /// Factory for producing new core worker clients. - rpc::ClientFactoryFn client_factory_; + /// Shared pool for producing new core worker clients. + std::shared_ptr client_pool_; /// Address of our RPC server. rpc::Address rpc_address_; /// Shared waiter for dependencies required by incoming tasks. diff --git a/src/ray/core_worker/transport/direct_task_transport.cc b/src/ray/core_worker/transport/direct_task_transport.cc index 4564f48f5..79f9cd733 100644 --- a/src/ray/core_worker/transport/direct_task_transport.cc +++ b/src/ray/core_worker/transport/direct_task_transport.cc @@ -94,7 +94,7 @@ Status CoreWorkerDirectTaskSubmitter::SubmitTask(TaskSpecification task_spec) { void CoreWorkerDirectTaskSubmitter::AddWorkerLeaseClient( const rpc::WorkerAddress &addr, std::shared_ptr lease_client) { - client_cache_.GetOrConnect(addr.ToProto()); + client_cache_->GetOrConnect(addr.ToProto()); int64_t expiration = current_time_ms() + lease_timeout_ms_; LeaseEntry new_lease_entry = LeaseEntry(std::move(lease_client), expiration, 0); worker_to_lease_entry_.emplace(addr, new_lease_entry); @@ -126,7 +126,7 @@ void CoreWorkerDirectTaskSubmitter::OnWorkerIdle( } } else { - auto &client = *client_cache_.GetOrConnect(addr.ToProto()); + auto &client = *client_cache_->GetOrConnect(addr.ToProto()); while (!queue_entry->second.empty() && lease_entry.tasks_in_flight_ < max_tasks_in_flight_per_worker_) { @@ -370,7 +370,7 @@ Status CoreWorkerDirectTaskSubmitter::CancelTask(TaskSpecification task_spec, return Status::OK(); } // Looks for an RPC handle for the worker executing the task. - auto maybe_client = client_cache_.GetByID(rpc_client->second.worker_id); + auto maybe_client = client_cache_->GetByID(rpc_client->second.worker_id); if (!maybe_client.has_value()) { // If we don't have a connection to that worker, we can't cancel it. // This case is reached for tasks that have unresolved dependencies. @@ -409,7 +409,7 @@ Status CoreWorkerDirectTaskSubmitter::CancelTask(TaskSpecification task_spec, Status CoreWorkerDirectTaskSubmitter::CancelRemoteTask(const ObjectID &object_id, const rpc::Address &worker_addr, bool force_kill) { - auto maybe_client = client_cache_.GetByID(rpc::WorkerAddress(worker_addr).worker_id); + auto maybe_client = client_cache_->GetByID(rpc::WorkerAddress(worker_addr).worker_id); if (!maybe_client.has_value()) { return Status::Invalid("No remote worker found"); diff --git a/src/ray/core_worker/transport/direct_task_transport.h b/src/ray/core_worker/transport/direct_task_transport.h index 6db472c0b..c038bb75c 100644 --- a/src/ray/core_worker/transport/direct_task_transport.h +++ b/src/ray/core_worker/transport/direct_task_transport.h @@ -52,7 +52,8 @@ class CoreWorkerDirectTaskSubmitter { public: explicit CoreWorkerDirectTaskSubmitter( rpc::Address rpc_address, std::shared_ptr lease_client, - rpc::ClientFactoryFn client_factory, LeaseClientFactoryFn lease_client_factory, + std::shared_ptr core_worker_client_pool, + LeaseClientFactoryFn lease_client_factory, std::shared_ptr store, std::shared_ptr task_finisher, ClientID local_raylet_id, int64_t lease_timeout_ms, std::shared_ptr actor_creator, @@ -67,7 +68,7 @@ class CoreWorkerDirectTaskSubmitter { lease_timeout_ms_(lease_timeout_ms), local_raylet_id_(local_raylet_id), actor_creator_(std::move(actor_creator)), - client_cache_(client_factory), + client_cache_(core_worker_client_pool), max_tasks_in_flight_per_worker_(max_tasks_in_flight_per_worker), cancel_retry_timer_(std::move(cancel_timer)) {} @@ -167,7 +168,7 @@ class CoreWorkerDirectTaskSubmitter { absl::Mutex mu_; /// Cache of gRPC clients to other workers. - rpc::CoreWorkerClientPool client_cache_; + std::shared_ptr client_cache_; // max_tasks_in_flight_per_worker_ limits the number of tasks that can be pipelined to a // worker using a single lease.