diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index db9d7d3e9..287dcabde 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -106,8 +106,9 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, RAY_CHECK(task_execution_callback_ != nullptr); auto execute_task = std::bind(&CoreWorker::ExecuteTask, this, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3); - raylet_task_receiver_ = std::unique_ptr( - new CoreWorkerRayletTaskReceiver(raylet_client_, execute_task, exit_handler)); + raylet_task_receiver_ = + std::unique_ptr(new CoreWorkerRayletTaskReceiver( + worker_context_.GetWorkerID(), raylet_client_, execute_task, exit_handler)); direct_task_receiver_ = std::unique_ptr(new CoreWorkerDirectTaskReceiver( worker_context_, task_execution_service_, execute_task, exit_handler)); diff --git a/src/ray/core_worker/transport/raylet_transport.cc b/src/ray/core_worker/transport/raylet_transport.cc index ac17eca80..5ac3e0a68 100644 --- a/src/ray/core_worker/transport/raylet_transport.cc +++ b/src/ray/core_worker/transport/raylet_transport.cc @@ -6,15 +6,29 @@ namespace ray { CoreWorkerRayletTaskReceiver::CoreWorkerRayletTaskReceiver( - std::shared_ptr &raylet_client, const TaskHandler &task_handler, - const std::function &exit_handler) - : raylet_client_(raylet_client), + const WorkerID &worker_id, std::shared_ptr &raylet_client, + const TaskHandler &task_handler, const std::function &exit_handler) + : worker_id_(worker_id), + raylet_client_(raylet_client), task_handler_(task_handler), exit_handler_(exit_handler) {} 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/core_worker/transport/raylet_transport.h b/src/ray/core_worker/transport/raylet_transport.h index c2137f754..da36951b3 100644 --- a/src/ray/core_worker/transport/raylet_transport.h +++ b/src/ray/core_worker/transport/raylet_transport.h @@ -15,7 +15,8 @@ class CoreWorkerRayletTaskReceiver { const TaskSpecification &task_spec, const ResourceMappingType &resource_ids, std::vector> *return_objects)>; - CoreWorkerRayletTaskReceiver(std::shared_ptr &raylet_client, + CoreWorkerRayletTaskReceiver(const WorkerID &worker_id, + std::shared_ptr &raylet_client, const TaskHandler &task_handler, const std::function &exit_handler); @@ -31,6 +32,8 @@ class CoreWorkerRayletTaskReceiver { rpc::SendReplyCallback send_reply_callback); private: + // WorkerID of this worker. + WorkerID worker_id_; /// Reference to the core worker's raylet client. This is a pointer ref so that it /// can be initialized by core worker after this class is constructed. std::shared_ptr &raylet_client_; diff --git a/src/ray/protobuf/core_worker.proto b/src/ray/protobuf/core_worker.proto index 323624feb..106f4998e 100644 --- a/src/ray/protobuf/core_worker.proto +++ b/src/ray/protobuf/core_worker.proto @@ -33,12 +33,18 @@ message ActorHandle { } 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; + // The task to be pushed. - Task task = 1; + Task task = 2; + // A list of the resources reserved for this worker. // TODO(zhijunfu): `resource_ids` is represented as // flatbutters-serialized bytes, will be moved to protobuf later. - bytes resource_ids = 2; + bytes resource_ids = 3; } message AssignTaskReply { diff --git a/src/ray/raylet/worker.cc b/src/ray/raylet/worker.cc index 89288366a..b0bf67a12 100644 --- a/src/ray/raylet/worker.cc +++ b/src/ray/raylet/worker.cc @@ -131,6 +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.mutable_task()->mutable_task_spec()->CopyFrom( task.GetTaskSpecification().GetMessage()); request.mutable_task()->mutable_task_execution_spec()->CopyFrom(