wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-15 12:45:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

rpc: Follow-up by sharing the core worker client pool within the core worker. (#10206)

Browse Source 
* Share CoreWorkerClientPool

* Format
This commit is contained in:
Barak Michener
2020-08-21 11:01:22 -07:00
committed by GitHub
parent e968b52cb7
commit f03caa4532
10 changed files with 127 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ray/core_worker/core_worker.cc
+11 -9
View File
@@ -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<rpc::CoreWorkerClient>(
new rpc::CoreWorkerClient(addr, *client_call_manager_));
};
core_worker_client_pool_ =
std::make_shared<rpc::CoreWorkerClientPool>(*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<DefaultActorCreator>(gcs_client_);
direct_actor_submitter_ = std::shared_ptr<CoreWorkerDirectActorTaskSubmitter>(
new CoreWorkerDirectActorTaskSubmitter(client_factory, memory_store_,
new CoreWorkerDirectActorTaskSubmitter(core_worker_client_pool_, memory_store_,
task_manager_));
direct_task_submitter_ =
std::unique_ptr<CoreWorkerDirectTaskSubmitter>(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<ActorManager>(
src/ray/core_worker/core_worker.h
+3
View File
@@ -1040,6 +1040,9 @@ class CoreWorker : public rpc::CoreWorkerServiceHandler {
/// Shared client call manager.
std::unique_ptr<rpc::ClientCallManager> client_call_manager_;
/// Shared core worker client pool.
std::shared_ptr<rpc::CoreWorkerClientPool> core_worker_client_pool_;
/// Timer used to periodically check if the raylet has died.
boost::asio::steady_timer death_check_timer_;
src/ray/core_worker/future_resolver.cc
+1 -1
View File
@@ -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());
src/ray/core_worker/future_resolver.h
+4 -3
View File
@@ -29,9 +29,10 @@ namespace ray {
class FutureResolver {
public:
FutureResolver(std::shared_ptr<CoreWorkerMemoryStore> store,
rpc::ClientFactoryFn client_factory, const rpc::Address &rpc_address)
std::shared_ptr<rpc::CoreWorkerClientPool> 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<CoreWorkerMemoryStore> in_memory_store_;
rpc::CoreWorkerClientPool owner_clients_;
std::shared_ptr<rpc::CoreWorkerClientPool> 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
src/ray/core_worker/test/direct_actor_transport_test.cc
+5 -3
View File
@@ -102,10 +102,10 @@ class DirectActorSubmitterTest : public ::testing::Test {
store_(std::shared_ptr<CoreWorkerMemoryStore>(new CoreWorkerMemoryStore())),
task_finisher_(std::make_shared<MockTaskFinisher>()),
submitter_(
[&](const rpc::Address &addr) {
std::make_shared<rpc::CoreWorkerClientPool>([&](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<CoreWorkerDirectTaskReceiver>(
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<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client_; }),
rpc_address_, dependency_waiter_);
}
src/ray/core_worker/test/direct_task_transport_test.cc
+81 -63
View File
@@ -335,12 +335,13 @@ TEST(DirectTaskTransportTest, TestSubmitOneTask) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
@@ -371,12 +372,13 @@ TEST(DirectTaskTransportTest, TestHandleTaskFailure) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
@@ -400,12 +402,13 @@ TEST(DirectTaskTransportTest, TestConcurrentWorkerLeases) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
@@ -450,12 +453,13 @@ TEST(DirectTaskTransportTest, TestReuseWorkerLease) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
@@ -507,12 +511,13 @@ TEST(DirectTaskTransportTest, TestRetryLeaseCancellation) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
@@ -562,12 +567,13 @@ TEST(DirectTaskTransportTest, TestConcurrentCancellationAndSubmission) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
@@ -614,12 +620,13 @@ TEST(DirectTaskTransportTest, TestWorkerNotReusedOnError) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
@@ -657,12 +664,13 @@ TEST(DirectTaskTransportTest, TestWorkerNotReturnedOnExit) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
@@ -690,7 +698,8 @@ TEST(DirectTaskTransportTest, TestSpillback) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
std::unordered_map<int, std::shared_ptr<MockRayletClient>> 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<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
@@ -747,7 +756,8 @@ TEST(DirectTaskTransportTest, TestSpillbackRoundTrip) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
std::unordered_map<int, std::shared_ptr<MockRayletClient>> 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<MockTaskFinisher>();
auto local_raylet_id = ClientID::FromRandom();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
@@ -813,12 +823,13 @@ void TestSchedulingKey(const std::shared_ptr<CoreWorkerMemoryStore> store,
rpc::Address address;
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
@@ -978,12 +990,14 @@ TEST(DirectTaskTransportTest, TestKillExecutingTask) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
@@ -1025,12 +1039,13 @@ TEST(DirectTaskTransportTest, TestKillPendingTask) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
@@ -1053,12 +1068,13 @@ TEST(DirectTaskTransportTest, TestKillResolvingTask) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
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<std::string, double> empty_resources;
ray::FunctionDescriptor empty_descriptor =
ray::FunctionDescriptorBuilder::BuildPython("", "", "", "");
@@ -1083,7 +1099,8 @@ TEST(DirectTaskTransportTest, TestPipeliningConcurrentWorkerLeases) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
@@ -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<std::string, double> empty_resources;
@@ -1152,7 +1169,8 @@ TEST(DirectTaskTransportTest, TestPipeliningReuseWorkerLease) {
auto raylet_client = std::make_shared<MockRayletClient>();
auto worker_client = std::make_shared<MockWorkerClient>();
auto store = std::make_shared<CoreWorkerMemoryStore>();
auto factory = [&](const rpc::Address &addr) { return worker_client; };
auto client_pool = std::make_shared<rpc::CoreWorkerClientPool>(
[&](const rpc::Address &addr) { return worker_client; });
auto task_finisher = std::make_shared<MockTaskFinisher>();
auto actor_creator = std::make_shared<MockActorCreator>();
@@ -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<std::string, double> empty_resources;
src/ray/core_worker/transport/direct_actor_transport.cc
+4 -4
View File
@@ -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<rpc::CoreWorkerClientInterface>(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<rpc::CoreWorkerClientPool> client_pool, rpc::Address rpc_address,
std::shared_ptr<DependencyWaiter> dependency_waiter) {
waiter_ = std::move(dependency_waiter);
rpc_address_ = rpc_address;
client_factory_ = client_factory;
client_pool_ = client_pool;
}
void CoreWorkerDirectTaskReceiver::HandlePushTask(
src/ray/core_worker/transport/direct_actor_transport.h
+10 -9
View File
@@ -63,10 +63,11 @@ class CoreWorkerDirectActorTaskSubmitterInterface {
class CoreWorkerDirectActorTaskSubmitter
: public CoreWorkerDirectActorTaskSubmitterInterface {
public:
CoreWorkerDirectActorTaskSubmitter(rpc::ClientFactoryFn client_factory,
std::shared_ptr<CoreWorkerMemoryStore> store,
std::shared_ptr<TaskFinisherInterface> task_finisher)
: client_factory_(client_factory),
CoreWorkerDirectActorTaskSubmitter(
std::shared_ptr<rpc::CoreWorkerClientPool> core_worker_client_pool,
std::shared_ptr<CoreWorkerMemoryStore> store,
std::shared_ptr<TaskFinisherInterface> 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<rpc::CoreWorkerClientPool> 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::CoreWorkerClientPool>, rpc::Address rpc_address,
std::shared_ptr<DependencyWaiter> 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<rpc::CoreWorkerClientPool> client_pool_;
/// Address of our RPC server.
rpc::Address rpc_address_;
/// Shared waiter for dependencies required by incoming tasks.
src/ray/core_worker/transport/direct_task_transport.cc
+4 -4
View File
@@ -94,7 +94,7 @@ Status CoreWorkerDirectTaskSubmitter::SubmitTask(TaskSpecification task_spec) {
void CoreWorkerDirectTaskSubmitter::AddWorkerLeaseClient(
const rpc::WorkerAddress &addr, std::shared_ptr<WorkerLeaseInterface> 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");
src/ray/core_worker/transport/direct_task_transport.h
+4 -3
View File
@@ -52,7 +52,8 @@ class CoreWorkerDirectTaskSubmitter {
public:
explicit CoreWorkerDirectTaskSubmitter(
rpc::Address rpc_address, std::shared_ptr<WorkerLeaseInterface> lease_client,
rpc::ClientFactoryFn client_factory, LeaseClientFactoryFn lease_client_factory,
std::shared_ptr<rpc::CoreWorkerClientPool> core_worker_client_pool,
LeaseClientFactoryFn lease_client_factory,
std::shared_ptr<CoreWorkerMemoryStore> store,
std::shared_ptr<TaskFinisherInterface> task_finisher, ClientID local_raylet_id,
int64_t lease_timeout_ms, std::shared_ptr<ActorCreatorInterface> 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<rpc::CoreWorkerClientPool> client_cache_;
// max_tasks_in_flight_per_worker_ limits the number of tasks that can be pipelined to a
// worker using a single lease.