From 335dade1e67975ecb42303e4f769dbe6be960415 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Fri, 13 Dec 2019 10:15:56 -0800 Subject: [PATCH] Check worker id for all core worker RPCs (#6472) * check worker id * fix test * owner * fix tests * comments --- src/ray/core_worker/core_worker.cc | 36 ++++++++--- src/ray/core_worker/core_worker.h | 21 ++++++ src/ray/core_worker/future_resolver.cc | 4 +- src/ray/core_worker/test/core_worker_test.cc | 4 +- .../test/direct_actor_transport_test.cc | 2 +- .../test/direct_task_transport_test.cc | 64 +++++++++---------- .../transport/direct_actor_transport.cc | 12 ++-- .../transport/direct_actor_transport.h | 4 ++ .../transport/direct_task_transport.cc | 12 ++-- .../transport/direct_task_transport.h | 3 +- .../core_worker/transport/raylet_transport.cc | 13 ---- src/ray/protobuf/core_worker.proto | 20 ++++-- src/ray/raylet/node_manager.cc | 3 +- src/ray/raylet/worker.cc | 3 +- src/ray/rpc/worker/core_worker_client.h | 3 +- 15 files changed, 123 insertions(+), 81 deletions(-) diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index 501e34832..d2bbd28a3 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -212,9 +212,9 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, SetCurrentTaskId(task_id); } - auto client_factory = [this](const rpc::WorkerAddress &addr) { + auto client_factory = [this](const std::string ip_address, int port) { return std::shared_ptr( - new rpc::CoreWorkerClient(addr.ip_address, addr.port, *client_call_manager_)); + new rpc::CoreWorkerClient(ip_address, port, *client_call_manager_)); }; direct_actor_submitter_ = std::unique_ptr( new CoreWorkerDirectActorTaskSubmitter(client_factory, memory_store_, @@ -223,9 +223,9 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, direct_task_submitter_ = std::unique_ptr(new CoreWorkerDirectTaskSubmitter( local_raylet_client_, client_factory, - [this](const rpc::Address &address) { - auto grpc_client = rpc::NodeManagerWorkerClient::make( - address.ip_address(), address.port(), *client_call_manager_); + [this](const std::string ip_address, int port) { + auto grpc_client = rpc::NodeManagerWorkerClient::make(ip_address, port, + *client_call_manager_); return std::shared_ptr( new raylet::RayletClient(std::move(grpc_client))); }, @@ -467,12 +467,10 @@ Status CoreWorker::Get(const std::vector &ids, const int64_t timeout_m // object. will_throw_exception = true; } - // If we got the result for this plasma ObjectID, the task that created it must + // If we got the result for this ObjectID, the task that created it must // have finished. Therefore, we can safely remove its reference counting // dependencies. - if (!ids[i].IsDirectCallType()) { - RemoveObjectIDDependencies(ids[i]); - } + RemoveObjectIDDependencies(ids[i]); } else { missing_result = true; } @@ -799,7 +797,10 @@ bool CoreWorker::AddActorHandle(std::unique_ptr actor_handle) { RAY_LOG(INFO) << "received notification on actor, state=" << static_cast(actor_data.state()) << ", actor_id: " << actor_id << ", ip address: " << actor_data.address().ip_address() - << ", port: " << actor_data.address().port(); + << ", port: " << actor_data.address().port() << ", worker_id: " + << WorkerID::FromBinary(actor_data.address().worker_id()) + << ", raylet_id: " + << ClientID::FromBinary(actor_data.address().raylet_id()); }; RAY_CHECK_OK(direct_actor_table_subscriber_->AsyncSubscribe( @@ -983,6 +984,11 @@ Status CoreWorker::BuildArgsForExecutor(const TaskSpecification &task, void CoreWorker::HandleAssignTask(const rpc::AssignTaskRequest &request, rpc::AssignTaskReply *reply, rpc::SendReplyCallback send_reply_callback) { + if (HandleWrongRecipient(WorkerID::FromBinary(request.intended_worker_id()), + send_reply_callback)) { + return; + } + if (worker_context_.CurrentActorIsDirectCall()) { send_reply_callback(Status::Invalid("This actor only accepts direct calls."), nullptr, nullptr); @@ -997,6 +1003,11 @@ void CoreWorker::HandleAssignTask(const rpc::AssignTaskRequest &request, void CoreWorker::HandlePushTask(const rpc::PushTaskRequest &request, rpc::PushTaskReply *reply, rpc::SendReplyCallback send_reply_callback) { + if (HandleWrongRecipient(WorkerID::FromBinary(request.intended_worker_id()), + send_reply_callback)) { + return; + } + task_execution_service_.post([=] { direct_task_receiver_->HandlePushTask(request, reply, send_reply_callback); }); @@ -1006,6 +1017,11 @@ void CoreWorker::HandleDirectActorCallArgWaitComplete( const rpc::DirectActorCallArgWaitCompleteRequest &request, rpc::DirectActorCallArgWaitCompleteReply *reply, rpc::SendReplyCallback send_reply_callback) { + if (HandleWrongRecipient(WorkerID::FromBinary(request.intended_worker_id()), + send_reply_callback)) { + return; + } + task_execution_service_.post([=] { direct_task_receiver_->HandleDirectActorCallArgWaitComplete(request, reply, send_reply_callback); diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index a8d03b4f4..edb4c078a 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -1,6 +1,7 @@ #ifndef RAY_CORE_WORKER_CORE_WORKER_H #define RAY_CORE_WORKER_CORE_WORKER_H +#include "absl/base/optimization.h" #include "absl/container/flat_hash_map.h" #include "ray/common/buffer.h" #include "ray/core_worker/actor_handle.h" @@ -484,6 +485,26 @@ class CoreWorker { } } + /// Returns whether the message was sent to the wrong worker. The right error reply + /// is sent automatically. Messages end up on the wrong worker when a worker dies + /// and a new one takes its place with the same place. In this situation, we want + /// the new worker to reject messages meant for the old one. + bool HandleWrongRecipient(const WorkerID &intended_worker_id, + rpc::SendReplyCallback send_reply_callback) { + if (intended_worker_id != worker_context_.GetWorkerID()) { + std::ostringstream stream; + stream << "Mismatched WorkerID: ignoring RPC for previous worker " + << intended_worker_id + << ", current worker ID: " << worker_context_.GetWorkerID(); + auto msg = stream.str(); + RAY_LOG(ERROR) << msg; + send_reply_callback(Status::Invalid(msg), nullptr, nullptr); + return true; + } else { + return false; + } + } + /// Type of this worker (i.e., DRIVER or WORKER). const WorkerType worker_type_; diff --git a/src/ray/core_worker/future_resolver.cc b/src/ray/core_worker/future_resolver.cc index 84629c698..c31d15993 100644 --- a/src/ray/core_worker/future_resolver.cc +++ b/src/ray/core_worker/future_resolver.cc @@ -8,8 +8,8 @@ void FutureResolver::ResolveFutureAsync(const ObjectID &object_id, const TaskID absl::MutexLock lock(&mu_); auto it = owner_clients_.find(owner_id); if (it == owner_clients_.end()) { - auto client = std::shared_ptr(client_factory_( - {owner_address.ip_address(), owner_address.port(), WorkerID::Nil()})); + auto client = std::shared_ptr( + client_factory_(owner_address.ip_address(), owner_address.port())); it = owner_clients_.emplace(owner_id, std::move(client)).first; } diff --git a/src/ray/core_worker/test/core_worker_test.cc b/src/ray/core_worker/test/core_worker_test.cc index 24ee7875a..8253492f2 100644 --- a/src/ray/core_worker/test/core_worker_test.cc +++ b/src/ray/core_worker/test/core_worker_test.cc @@ -993,13 +993,13 @@ TEST_F(TwoNodeTest, TestDirectActorTaskCrossNodes) { TestActorTask(resources, true); } -// TODO(ekl) support reconstruction for direct call actors +// TODO(ekl) re-enable once reconstruction is implemented // TEST_F(SingleNodeTest, TestDirectActorTaskLocalReconstruction) { // std::unordered_map resources; // TestActorReconstruction(resources, true); //} -// TODO(ekl) support reconstruction for direct call actors +// TODO(ekl) re-enable once reconstruction is implemented // TEST_F(TwoNodeTest, TestDirectActorTaskCrossNodesReconstruction) { // std::unordered_map resources; // resources.emplace("resource1", 1); 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 2ac27b521..6e3740b93 100644 --- a/src/ray/core_worker/test/direct_actor_transport_test.cc +++ b/src/ray/core_worker/test/direct_actor_transport_test.cc @@ -62,7 +62,7 @@ class DirectActorTransportTest : public ::testing::Test { : worker_client_(std::shared_ptr(new MockWorkerClient())), store_(std::shared_ptr(new CoreWorkerMemoryStore())), task_finisher_(std::make_shared()), - submitter_([&](const rpc::WorkerAddress &addr) { return worker_client_; }, store_, + submitter_([&](const std::string ip, int port) { return worker_client_; }, store_, task_finisher_) {} std::shared_ptr worker_client_; 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 a8fbc20b6..3582f51d7 100644 --- a/src/ray/core_worker/test/direct_task_transport_test.cc +++ b/src/ray/core_worker/test/direct_task_transport_test.cc @@ -235,7 +235,7 @@ TEST(DirectTaskTransportTest, TestSubmitOneTask) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; }; + auto factory = [&](const std::string &addr, int port) { return worker_client; }; auto task_finisher = std::make_shared(); CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store, task_finisher, ClientID::Nil(), kLongTimeout); @@ -265,7 +265,7 @@ TEST(DirectTaskTransportTest, TestHandleTaskFailure) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; }; + auto factory = [&](const std::string &addr, int port) { return worker_client; }; auto task_finisher = std::make_shared(); CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store, task_finisher, ClientID::Nil(), kLongTimeout); @@ -288,7 +288,7 @@ TEST(DirectTaskTransportTest, TestConcurrentWorkerLeases) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; }; + auto factory = [&](const std::string &addr, int port) { return worker_client; }; auto task_finisher = std::make_shared(); CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store, task_finisher, ClientID::Nil(), kLongTimeout); @@ -332,7 +332,7 @@ TEST(DirectTaskTransportTest, TestReuseWorkerLease) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; }; + auto factory = [&](const std::string &addr, int port) { return worker_client; }; auto task_finisher = std::make_shared(); CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store, task_finisher, ClientID::Nil(), kLongTimeout); @@ -379,7 +379,7 @@ TEST(DirectTaskTransportTest, TestWorkerNotReusedOnError) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; }; + auto factory = [&](const std::string &addr, int port) { return worker_client; }; auto task_finisher = std::make_shared(); CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store, task_finisher, ClientID::Nil(), kLongTimeout); @@ -416,15 +416,14 @@ TEST(DirectTaskTransportTest, TestSpillback) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; }; + auto factory = [&](const std::string &addr, int port) { return worker_client; }; - std::unordered_map> remote_lease_clients; - auto lease_client_factory = [&](const rpc::Address &addr) { - ClientID raylet_id = ClientID::FromBinary(addr.raylet_id()); + std::unordered_map> remote_lease_clients; + auto lease_client_factory = [&](const std::string &ip, int port) { // We should not create a connection to the same raylet more than once. - RAY_CHECK(remote_lease_clients.count(raylet_id) == 0); + RAY_CHECK(remote_lease_clients.count(port) == 0); auto client = std::make_shared(); - remote_lease_clients[raylet_id] = client; + remote_lease_clients[port] = client; return client; }; auto task_finisher = std::make_shared(); @@ -443,20 +442,20 @@ TEST(DirectTaskTransportTest, TestSpillback) { // Spillback to a remote node. auto remote_raylet_id = ClientID::FromRandom(); - ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1234, remote_raylet_id)); - ASSERT_EQ(remote_lease_clients.count(remote_raylet_id), 1); + ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 7777, remote_raylet_id)); + ASSERT_EQ(remote_lease_clients.count(7777), 1); // There should be no more callbacks on the local client. ASSERT_FALSE(raylet_client->GrantWorkerLease("remote", 1234, ClientID::Nil())); // Trigger retry at the remote node. - ASSERT_TRUE(remote_lease_clients[remote_raylet_id]->GrantWorkerLease("remote", 1234, - ClientID::Nil())); + ASSERT_TRUE( + remote_lease_clients[7777]->GrantWorkerLease("remote", 1234, ClientID::Nil())); // The worker is returned to the remote node, not the local one. ASSERT_TRUE(worker_client->ReplyPushTask()); ASSERT_EQ(raylet_client->num_workers_returned, 0); - ASSERT_EQ(remote_lease_clients[remote_raylet_id]->num_workers_returned, 1); + ASSERT_EQ(remote_lease_clients[7777]->num_workers_returned, 1); ASSERT_EQ(raylet_client->num_workers_disconnected, 0); - ASSERT_EQ(remote_lease_clients[remote_raylet_id]->num_workers_disconnected, 0); + ASSERT_EQ(remote_lease_clients[7777]->num_workers_disconnected, 0); ASSERT_EQ(task_finisher->num_tasks_complete, 1); ASSERT_EQ(task_finisher->num_tasks_failed, 0); } @@ -465,15 +464,14 @@ TEST(DirectTaskTransportTest, TestSpillbackRoundTrip) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; }; + auto factory = [&](const std::string &addr, int port) { return worker_client; }; - std::unordered_map> remote_lease_clients; - auto lease_client_factory = [&](const rpc::Address &addr) { - ClientID raylet_id = ClientID::FromBinary(addr.raylet_id()); + std::unordered_map> remote_lease_clients; + auto lease_client_factory = [&](const std::string &ip, int port) { // We should not create a connection to the same raylet more than once. - RAY_CHECK(remote_lease_clients.count(raylet_id) == 0); + RAY_CHECK(remote_lease_clients.count(port) == 0); auto client = std::make_shared(); - remote_lease_clients[raylet_id] = client; + remote_lease_clients[port] = client; return client; }; auto task_finisher = std::make_shared(); @@ -493,25 +491,25 @@ TEST(DirectTaskTransportTest, TestSpillbackRoundTrip) { // Spillback to a remote node. auto remote_raylet_id = ClientID::FromRandom(); - ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 1234, remote_raylet_id)); - ASSERT_EQ(remote_lease_clients.count(remote_raylet_id), 1); + ASSERT_TRUE(raylet_client->GrantWorkerLease("localhost", 7777, remote_raylet_id)); + ASSERT_EQ(remote_lease_clients.count(7777), 1); ASSERT_FALSE(raylet_client->GrantWorkerLease("remote", 1234, ClientID::Nil())); // Trigger a spillback back to the local node. - ASSERT_TRUE(remote_lease_clients[remote_raylet_id]->GrantWorkerLease("local", 1234, - local_raylet_id)); + ASSERT_TRUE( + remote_lease_clients[7777]->GrantWorkerLease("local", 1234, local_raylet_id)); // We should not have created another lease client to the local raylet. ASSERT_EQ(remote_lease_clients.size(), 1); // There should be no more callbacks on the remote node. - ASSERT_FALSE(remote_lease_clients[remote_raylet_id]->GrantWorkerLease("remote", 1234, - ClientID::Nil())); + ASSERT_FALSE( + remote_lease_clients[7777]->GrantWorkerLease("remote", 1234, ClientID::Nil())); // The worker is returned to the local node. ASSERT_TRUE(raylet_client->GrantWorkerLease("local", 1234, ClientID::Nil())); ASSERT_TRUE(worker_client->ReplyPushTask()); ASSERT_EQ(raylet_client->num_workers_returned, 1); - ASSERT_EQ(remote_lease_clients[remote_raylet_id]->num_workers_returned, 0); + ASSERT_EQ(remote_lease_clients[7777]->num_workers_returned, 0); ASSERT_EQ(raylet_client->num_workers_disconnected, 0); - ASSERT_EQ(remote_lease_clients[remote_raylet_id]->num_workers_disconnected, 0); + ASSERT_EQ(remote_lease_clients[7777]->num_workers_disconnected, 0); ASSERT_EQ(task_finisher->num_tasks_complete, 1); ASSERT_EQ(task_finisher->num_tasks_failed, 0); } @@ -523,7 +521,7 @@ void TestSchedulingKey(const std::shared_ptr store, const TaskSpecification &different) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); - auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; }; + auto factory = [&](const std::string &addr, int port) { return worker_client; }; auto task_finisher = std::make_shared(); CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store, task_finisher, ClientID::Nil(), kLongTimeout); @@ -622,7 +620,7 @@ TEST(DirectTaskTransportTest, TestWorkerLeaseTimeout) { auto raylet_client = std::make_shared(); auto worker_client = std::make_shared(); auto store = std::make_shared(); - auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; }; + auto factory = [&](const std::string &addr, int port) { return worker_client; }; auto task_finisher = std::make_shared(); CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store, task_finisher, ClientID::Nil(), diff --git a/src/ray/core_worker/transport/direct_actor_transport.cc b/src/ray/core_worker/transport/direct_actor_transport.cc index a8b3e7ef7..bc6ce5647 100644 --- a/src/ray/core_worker/transport/direct_actor_transport.cc +++ b/src/ray/core_worker/transport/direct_actor_transport.cc @@ -55,11 +55,12 @@ Status CoreWorkerDirectActorTaskSubmitter::SubmitTask(TaskSpecification task_spe void CoreWorkerDirectActorTaskSubmitter::ConnectActor(const ActorID &actor_id, const rpc::Address &address) { absl::MutexLock lock(&mu_); + // Update the mapping so new RPCs go out with the right intended worker id. + worker_ids_[actor_id] = address.worker_id(); // Create a new connection to the actor. if (rpc_clients_.count(actor_id) == 0) { - rpc::WorkerAddress addr = {address.ip_address(), address.port()}; - rpc_clients_[actor_id] = - std::shared_ptr(client_factory_(addr)); + rpc_clients_[actor_id] = std::shared_ptr( + client_factory_(address.ip_address(), address.port())); } if (pending_requests_.count(actor_id) > 0) { SendPendingTasks(actor_id); @@ -74,6 +75,7 @@ void CoreWorkerDirectActorTaskSubmitter::DisconnectActor(const ActorID &actor_id // will be inserted once actor reconstruction completes. We don't erase the // client when the actor is DEAD, so that all further tasks will be failed. rpc_clients_.erase(actor_id); + worker_ids_.erase(actor_id); } else { RAY_LOG(INFO) << "Failing pending tasks for actor " << actor_id; // If there are pending requests, treat the pending tasks as failed. @@ -117,7 +119,9 @@ void CoreWorkerDirectActorTaskSubmitter::PushActorTask( const ActorID &actor_id, const TaskID &task_id, int num_returns) { RAY_LOG(DEBUG) << "Pushing task " << task_id << " to actor " << actor_id; next_send_position_[actor_id]++; - + auto it = worker_ids_.find(actor_id); + RAY_CHECK(it != worker_ids_.end()) << "Actor worker id not found " << actor_id.Hex(); + request->set_intended_worker_id(it->second); RAY_CHECK_OK(client.PushActorTask( std::move(request), [this, task_id](Status status, const rpc::PushTaskReply &reply) { diff --git a/src/ray/core_worker/transport/direct_actor_transport.h b/src/ray/core_worker/transport/direct_actor_transport.h index c873a523e..9e8e0ae5d 100644 --- a/src/ray/core_worker/transport/direct_actor_transport.h +++ b/src/ray/core_worker/transport/direct_actor_transport.h @@ -103,6 +103,10 @@ class CoreWorkerDirectActorTaskSubmitter { absl::flat_hash_map> rpc_clients_ GUARDED_BY(mu_); + /// Map from actor ids to worker ids. TODO(ekl) consider unifying this with the + /// rpc_clients_ map. + absl::flat_hash_map worker_ids_ GUARDED_BY(mu_); + /// Map from actor id to the actor's pending requests. Each actor's requests /// are ordered by the task number in the request. absl::flat_hash_map>> diff --git a/src/ray/core_worker/transport/direct_task_transport.cc b/src/ray/core_worker/transport/direct_task_transport.cc index 6a58decd6..f2298a78d 100644 --- a/src/ray/core_worker/transport/direct_task_transport.cc +++ b/src/ray/core_worker/transport/direct_task_transport.cc @@ -27,8 +27,8 @@ void CoreWorkerDirectTaskSubmitter::AddWorkerLeaseClient( const rpc::WorkerAddress &addr, std::shared_ptr lease_client) { auto it = client_cache_.find(addr); if (it == client_cache_.end()) { - client_cache_[addr] = - std::shared_ptr(client_factory_(addr)); + client_cache_[addr] = std::shared_ptr( + client_factory_(addr.ip_address, addr.port)); RAY_LOG(INFO) << "Connected to " << addr.ip_address << ":" << addr.port; } int64_t expiration = current_time_ms() + lease_timeout_ms_; @@ -76,9 +76,10 @@ CoreWorkerDirectTaskSubmitter::GetOrConnectLeaseClient( auto it = remote_lease_clients_.find(raylet_id); if (it == remote_lease_clients_.end()) { RAY_LOG(DEBUG) << "Connecting to raylet " << raylet_id; - it = - remote_lease_clients_.emplace(raylet_id, lease_client_factory_(*raylet_address)) - .first; + it = remote_lease_clients_ + .emplace(raylet_id, lease_client_factory_(raylet_address->ip_address(), + raylet_address->port())) + .first; } lease_client = it->second; } else { @@ -169,6 +170,7 @@ void CoreWorkerDirectTaskSubmitter::PushNormalTask( // access the task. request->mutable_task_spec()->CopyFrom(task_spec.GetMessage()); request->mutable_resource_mapping()->CopyFrom(assigned_resources); + request->set_intended_worker_id(addr.worker_id.Binary()); auto status = client.PushNormalTask( std::move(request), [this, task_id, is_actor, is_actor_creation, scheduling_key, addr, diff --git a/src/ray/core_worker/transport/direct_task_transport.h b/src/ray/core_worker/transport/direct_task_transport.h index e0283d61f..2df661d8d 100644 --- a/src/ray/core_worker/transport/direct_task_transport.h +++ b/src/ray/core_worker/transport/direct_task_transport.h @@ -18,7 +18,8 @@ namespace ray { -typedef std::function(const rpc::Address &)> +typedef std::function(const std::string &ip_address, + int port)> LeaseClientFactoryFn; // The task queues are keyed on resource shape & function descriptor diff --git a/src/ray/core_worker/transport/raylet_transport.cc b/src/ray/core_worker/transport/raylet_transport.cc index 941531621..50c16a18f 100644 --- a/src/ray/core_worker/transport/raylet_transport.cc +++ b/src/ray/core_worker/transport/raylet_transport.cc @@ -16,19 +16,6 @@ CoreWorkerRayletTaskReceiver::CoreWorkerRayletTaskReceiver( void CoreWorkerRayletTaskReceiver::HandleAssignTask( const rpc::AssignTaskRequest &request, rpc::AssignTaskReply *reply, rpc::SendReplyCallback send_reply_callback) { - // Check that the message was intended for our WorkerID and drop it if not. - // This handles the case where a message is delayed so we get an AssignTask - // bound for a previous worker that is now dead because we bound to the same - // port. Note that returning the status here doesn't actually cause the raylet - // to fail the task, that happens due to the unintentional disconnect from the - // asio connection when the previous worker died. - WorkerID intended_worker_id = WorkerID::FromBinary(request.worker_id()); - if (intended_worker_id != worker_id_) { - RAY_LOG(WARNING) << "Received task for mismatched WorkerID " << intended_worker_id; - send_reply_callback(Status::Invalid("Mismatched WorkerID"), nullptr, nullptr); - return; - } - const Task task(request.task()); const auto &task_spec = task.GetTaskSpecification(); RAY_LOG(DEBUG) << "Received task " << task_spec.TaskId() << " is create " diff --git a/src/ray/protobuf/core_worker.proto b/src/ray/protobuf/core_worker.proto index fc75eb495..96e193989 100644 --- a/src/ray/protobuf/core_worker.proto +++ b/src/ray/protobuf/core_worker.proto @@ -36,7 +36,7 @@ message AssignTaskRequest { // The ID of the worker this message is intended for. This is used to // ensure that workers don't try to execute tasks assigned to workers // that used to be bound to the same port. - bytes worker_id = 1; + bytes intended_worker_id = 1; // The task to be pushed. Task task = 2; @@ -63,22 +63,24 @@ message ReturnObject { } message PushTaskRequest { + // The ID of the worker this message is intended for. + bytes intended_worker_id = 1; // The task to be pushed. - TaskSpec task_spec = 1; + TaskSpec task_spec = 2; // The sequence number of the task for this client. This must increase // sequentially starting from zero for each actor handle. The server // will guarantee tasks execute in this sequence, waiting for any // out-of-order request messages to arrive as necessary. // If set to -1, ordering is disabled and the task executes immediately. // This mode of behaviour is used for direct task submission only. - int64 sequence_number = 2; + int64 sequence_number = 3; // The max sequence number the client has processed responses for. This // is a performance optimization that allows the client to tell the server // to cancel any PushTaskRequests with seqno <= this value, rather than // waiting for the server to time out waiting for missing messages. - int64 client_processed_up_to = 3; + int64 client_processed_up_to = 4; // Resource mapping ids assigned to the worker executing the task. - repeated ResourceMapEntry resource_mapping = 4; + repeated ResourceMapEntry resource_mapping = 5; } message PushTaskReply { @@ -87,16 +89,20 @@ message PushTaskReply { } message DirectActorCallArgWaitCompleteRequest { + // The ID of the worker this message is intended for. + bytes intended_worker_id = 1; // Id used to uniquely identify this request. This is sent back to the core // worker to notify the wait has completed. - int64 tag = 1; + int64 tag = 2; } message DirectActorCallArgWaitCompleteReply { } message GetObjectStatusRequest { - // The owner of the object. + // The owner of the object. Note that we do not need to include + // intended_worker_id since the new worker can service this request too by + // inspecting the owner_id field. bytes owner_id = 1; // Wait for this object's status. bytes object_id = 2; diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index bc0ef2ef2..921752d78 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -1607,7 +1607,8 @@ void NodeManager::HandleWorkerLeaseRequest(const rpc::WorkerLeaseRequest &reques // TODO(swang): Kill worker if other end hangs up. // TODO(swang): Implement a lease term by which the owner needs to return the // worker. - RAY_CHECK(leased_workers_.find(worker_id) == leased_workers_.end()); + RAY_CHECK(leased_workers_.find(worker_id) == leased_workers_.end()) + << "Worker is already leased out " << worker_id; leased_workers_[worker_id] = std::static_pointer_cast(granted); }); task.OnSpillbackInstead( diff --git a/src/ray/raylet/worker.cc b/src/ray/raylet/worker.cc index b0bf67a12..a02e65a78 100644 --- a/src/ray/raylet/worker.cc +++ b/src/ray/raylet/worker.cc @@ -131,7 +131,7 @@ void Worker::SetActiveObjectIds(const std::unordered_set &&object_ids) Status Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set) { RAY_CHECK(port_ > 0); rpc::AssignTaskRequest request; - request.set_worker_id(worker_id_.Binary()); + request.set_intended_worker_id(worker_id_.Binary()); request.mutable_task()->mutable_task_spec()->CopyFrom( task.GetTaskSpecification().GetMessage()); request.mutable_task()->mutable_task_execution_spec()->CopyFrom( @@ -153,6 +153,7 @@ void Worker::DirectActorCallArgWaitComplete(int64_t tag) { RAY_CHECK(port_ > 0); rpc::DirectActorCallArgWaitCompleteRequest request; request.set_tag(tag); + request.set_intended_worker_id(worker_id_.Binary()); auto status = rpc_client_->DirectActorCallArgWaitComplete( request, [](Status status, const rpc::DirectActorCallArgWaitCompleteReply &reply) { if (!status.ok()) { diff --git a/src/ray/rpc/worker/core_worker_client.h b/src/ray/rpc/worker/core_worker_client.h index e3ca41439..b0a5a2472 100644 --- a/src/ray/rpc/worker/core_worker_client.h +++ b/src/ray/rpc/worker/core_worker_client.h @@ -69,7 +69,8 @@ class WorkerAddress { const ClientID raylet_id; }; -typedef std::function(const WorkerAddress &)> +typedef std::function(const std::string &, + int)> ClientFactoryFn; /// Abstract client interface for testing.