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

@@ -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<rpc::CoreWorkerClient>(
-        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<CoreWorkerDirectActorTaskSubmitter>(
       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<CoreWorkerDirectTaskSubmitter>(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<raylet::RayletClient>(
                 new raylet::RayletClient(std::move(grpc_client)));
           },
@@ -467,12 +467,10 @@ Status CoreWorker::Get(const std::vector<ObjectID> &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<ActorHandle> actor_handle) {
       RAY_LOG(INFO) << "received notification on actor, state="
                     << static_cast<int>(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);

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_;
 

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<rpc::CoreWorkerClientInterface>(client_factory_(
-        {owner_address.ip_address(), owner_address.port(), WorkerID::Nil()}));
+    auto client = std::shared_ptr<rpc::CoreWorkerClientInterface>(
+        client_factory_(owner_address.ip_address(), owner_address.port()));
     it = owner_clients_.emplace(owner_id, std::move(client)).first;
   }
 

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<std::string, double> 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<std::string, double> resources;
 //  resources.emplace("resource1", 1);

src/ray/core_worker/test/direct_actor_transport_test.cc

@@ -62,7 +62,7 @@ class DirectActorTransportTest : public ::testing::Test {
       : worker_client_(std::shared_ptr<MockWorkerClient>(new MockWorkerClient())),
         store_(std::shared_ptr<CoreWorkerMemoryStore>(new CoreWorkerMemoryStore())),
         task_finisher_(std::make_shared<MockTaskFinisher>()),
-        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<MockWorkerClient> worker_client_;

src/ray/core_worker/test/direct_task_transport_test.cc

@@ -235,7 +235,7 @@ 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::WorkerAddress &addr) { return worker_client; };
+  auto factory = [&](const std::string &addr, int port) { return worker_client; };
   auto task_finisher = std::make_shared<MockTaskFinisher>();
   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<MockRayletClient>();
   auto worker_client = std::make_shared<MockWorkerClient>();
   auto store = std::make_shared<CoreWorkerMemoryStore>();
-  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<MockTaskFinisher>();
   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<MockRayletClient>();
   auto worker_client = std::make_shared<MockWorkerClient>();
   auto store = std::make_shared<CoreWorkerMemoryStore>();
-  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<MockTaskFinisher>();
   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<MockRayletClient>();
   auto worker_client = std::make_shared<MockWorkerClient>();
   auto store = std::make_shared<CoreWorkerMemoryStore>();
-  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<MockTaskFinisher>();
   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<MockRayletClient>();
   auto worker_client = std::make_shared<MockWorkerClient>();
   auto store = std::make_shared<CoreWorkerMemoryStore>();
-  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<MockTaskFinisher>();
   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<MockRayletClient>();
   auto worker_client = std::make_shared<MockWorkerClient>();
   auto store = std::make_shared<CoreWorkerMemoryStore>();
-  auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; };
+  auto factory = [&](const std::string &addr, int port) { return worker_client; };
 
-  std::unordered_map<ClientID, std::shared_ptr<MockRayletClient>> remote_lease_clients;
-  auto lease_client_factory = [&](const rpc::Address &addr) {
-    ClientID raylet_id = ClientID::FromBinary(addr.raylet_id());
+  std::unordered_map<int, std::shared_ptr<MockRayletClient>> 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<MockRayletClient>();
-    remote_lease_clients[raylet_id] = client;
+    remote_lease_clients[port] = client;
     return client;
   };
   auto task_finisher = std::make_shared<MockTaskFinisher>();
@@ -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<MockRayletClient>();
   auto worker_client = std::make_shared<MockWorkerClient>();
   auto store = std::make_shared<CoreWorkerMemoryStore>();
-  auto factory = [&](const rpc::WorkerAddress &addr) { return worker_client; };
+  auto factory = [&](const std::string &addr, int port) { return worker_client; };
 
-  std::unordered_map<ClientID, std::shared_ptr<MockRayletClient>> remote_lease_clients;
-  auto lease_client_factory = [&](const rpc::Address &addr) {
-    ClientID raylet_id = ClientID::FromBinary(addr.raylet_id());
+  std::unordered_map<int, std::shared_ptr<MockRayletClient>> 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<MockRayletClient>();
-    remote_lease_clients[raylet_id] = client;
+    remote_lease_clients[port] = client;
     return client;
   };
   auto task_finisher = std::make_shared<MockTaskFinisher>();
@@ -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<CoreWorkerMemoryStore> store,
                        const TaskSpecification &different) {
   auto raylet_client = std::make_shared<MockRayletClient>();
   auto worker_client = std::make_shared<MockWorkerClient>();
-  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<MockTaskFinisher>();
   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<MockRayletClient>();
   auto worker_client = std::make_shared<MockWorkerClient>();
   auto store = std::make_shared<CoreWorkerMemoryStore>();
-  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<MockTaskFinisher>();
   CoreWorkerDirectTaskSubmitter submitter(raylet_client, factory, nullptr, store,
                                           task_finisher, ClientID::Nil(),

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<rpc::CoreWorkerClientInterface>(client_factory_(addr));
+    rpc_clients_[actor_id] = std::shared_ptr<rpc::CoreWorkerClientInterface>(
+        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) {

src/ray/core_worker/transport/direct_actor_transport.h

@@ -103,6 +103,10 @@ class CoreWorkerDirectActorTaskSubmitter {
   absl::flat_hash_map<ActorID, std::shared_ptr<rpc::CoreWorkerClientInterface>>
       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<ActorID, std::string> 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<ActorID, std::map<int64_t, std::unique_ptr<rpc::PushTaskRequest>>>

src/ray/core_worker/transport/direct_task_transport.cc

@@ -27,8 +27,8 @@ void CoreWorkerDirectTaskSubmitter::AddWorkerLeaseClient(
     const rpc::WorkerAddress &addr, std::shared_ptr<WorkerLeaseInterface> lease_client) {
   auto it = client_cache_.find(addr);
   if (it == client_cache_.end()) {
-    client_cache_[addr] =
-        std::shared_ptr<rpc::CoreWorkerClientInterface>(client_factory_(addr));
+    client_cache_[addr] = std::shared_ptr<rpc::CoreWorkerClientInterface>(
+        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,

src/ray/core_worker/transport/direct_task_transport.h

@@ -18,7 +18,8 @@
 
 namespace ray {
 
-typedef std::function<std::shared_ptr<WorkerLeaseInterface>(const rpc::Address &)>
+typedef std::function<std::shared_ptr<WorkerLeaseInterface>(const std::string &ip_address,
+                                                            int port)>
     LeaseClientFactoryFn;
 
 // The task queues are keyed on resource shape & function descriptor

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 "

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;

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<Worker>(granted);
       });
   task.OnSpillbackInstead(

src/ray/raylet/worker.cc

@@ -131,7 +131,7 @@ void Worker::SetActiveObjectIds(const std::unordered_set<ObjectID> &&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()) {

src/ray/rpc/worker/core_worker_client.h

@@ -69,7 +69,8 @@ class WorkerAddress {
   const ClientID raylet_id;
 };
 
-typedef std::function<std::shared_ptr<CoreWorkerClientInterface>(const WorkerAddress &)>
+typedef std::function<std::shared_ptr<CoreWorkerClientInterface>(const std::string &,
+                                                                 int)>
     ClientFactoryFn;
 
 /// Abstract client interface for testing.