From 1649f1370e51281db232d8416e64587a34ea95a7 Mon Sep 17 00:00:00 2001 From: Zhijun Fu <37800433+zhijunfu@users.noreply.github.com> Date: Fri, 12 Jul 2019 02:01:32 +0800 Subject: [PATCH] [direct call] changes raylet to push tasks to worker (#5140) * refactor grpc server * format * change GetTask() to PushTask() * change PushTask to AssignTask * format * add resource_ids * move done_callback to server call * remove SetTaskHandler and initialize it in task receiver's constructor * format * resolve comments * update * update * Update src/ray/core_worker/core_worker.cc Co-Authored-By: Stephanie Wang * resolve comments * format * Update src/ray/core_worker/transport/raylet_transport.cc Co-Authored-By: Hao Chen * resolve comments * resolve comments * fix build * format * fix * format * noop --- src/ray/common/task/scheduling_resources.cc | 7 + src/ray/common/task/scheduling_resources.h | 6 + src/ray/core_worker/core_worker.cc | 11 +- src/ray/core_worker/core_worker.h | 3 +- src/ray/core_worker/core_worker_test.cc | 12 +- src/ray/core_worker/mock_worker.cc | 55 +++--- src/ray/core_worker/task_execution.cc | 83 ++++---- src/ray/core_worker/task_execution.h | 19 +- .../core_worker/transport/raylet_transport.cc | 36 ++-- .../core_worker/transport/raylet_transport.h | 13 +- src/ray/core_worker/transport/transport.h | 3 - src/ray/protobuf/worker.proto | 4 + src/ray/raylet/node_manager.cc | 183 ++++++++++-------- src/ray/raylet/node_manager.h | 17 +- src/ray/raylet/raylet.h | 1 - src/ray/raylet/scheduling_policy.cc | 2 +- src/ray/raylet/worker.cc | 49 ++++- src/ray/raylet/worker.h | 14 +- src/ray/raylet/worker_pool_test.cc | 10 +- 19 files changed, 322 insertions(+), 206 deletions(-) diff --git a/src/ray/common/task/scheduling_resources.cc b/src/ray/common/task/scheduling_resources.cc index bbcdd6fec..5463b0933 100644 --- a/src/ray/common/task/scheduling_resources.cc +++ b/src/ray/common/task/scheduling_resources.cc @@ -700,6 +700,13 @@ std::vector> ResourceIdSet::ToF return return_message; } +const std::string ResourceIdSet::Serialize() const { + flatbuffers::FlatBufferBuilder fbb; + auto resource_id_set_flatbuf = ToFlatbuf(fbb); + fbb.Finish(fbb.CreateVector(resource_id_set_flatbuf)); + return std::string(fbb.GetBufferPointer(), fbb.GetBufferPointer() + fbb.GetSize()); +} + /// SchedulingResources class implementation SchedulingResources::SchedulingResources() diff --git a/src/ray/common/task/scheduling_resources.h b/src/ray/common/task/scheduling_resources.h index 355ad7573..045c0307b 100644 --- a/src/ray/common/task/scheduling_resources.h +++ b/src/ray/common/task/scheduling_resources.h @@ -429,6 +429,12 @@ class ResourceIdSet { std::vector> ToFlatbuf( flatbuffers::FlatBufferBuilder &fbb) const; + /// \brief Serialize this object as a string. + /// + /// \return A serialized string of this object. + /// TODO(zhijunfu): this can be removed after raylet client is migrated to grpc. + const std::string Serialize() const; + private: /// A mapping from reosurce name to a set of resource IDs for that resource. std::unordered_map available_resources_; diff --git a/src/ray/core_worker/core_worker.cc b/src/ray/core_worker/core_worker.cc index 5b14795fd..6fa560f27 100644 --- a/src/ray/core_worker/core_worker.cc +++ b/src/ray/core_worker/core_worker.cc @@ -3,9 +3,11 @@ namespace ray { -CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, - const std::string &store_socket, const std::string &raylet_socket, - const JobID &job_id) +CoreWorker::CoreWorker( + const WorkerType worker_type, const Language language, + const std::string &store_socket, const std::string &raylet_socket, + const JobID &job_id, + const CoreWorkerTaskExecutionInterface::TaskExecutor &execution_callback) : worker_type_(worker_type), language_(language), raylet_socket_(raylet_socket), @@ -14,9 +16,10 @@ CoreWorker::CoreWorker(const WorkerType worker_type, const Language language, object_interface_(worker_context_, raylet_client_, store_socket) { int rpc_server_port = 0; if (worker_type_ == ray::WorkerType::WORKER) { + RAY_CHECK(execution_callback != nullptr); task_execution_interface_ = std::unique_ptr( new CoreWorkerTaskExecutionInterface(worker_context_, raylet_client_, - object_interface_)); + object_interface_, execution_callback)); rpc_server_port = task_execution_interface_->worker_server_.GetPort(); } // TODO(zhijunfu): currently RayletClient would crash in its constructor if it cannot diff --git a/src/ray/core_worker/core_worker.h b/src/ray/core_worker/core_worker.h index 3798b8c4f..97816bdda 100644 --- a/src/ray/core_worker/core_worker.h +++ b/src/ray/core_worker/core_worker.h @@ -24,7 +24,8 @@ class CoreWorker { /// NOTE(zhijunfu): the constructor would throw if a failure happens. CoreWorker(const WorkerType worker_type, const Language language, const std::string &store_socket, const std::string &raylet_socket, - const JobID &job_id = JobID::Nil()); + const JobID &job_id, + const CoreWorkerTaskExecutionInterface::TaskExecutor &execution_callback); /// Type of this worker. WorkerType GetWorkerType() const { return worker_type_; } diff --git a/src/ray/core_worker/core_worker_test.cc b/src/ray/core_worker/core_worker_test.cc index 07aedad36..c01941d77 100644 --- a/src/ray/core_worker/core_worker_test.cc +++ b/src/ray/core_worker/core_worker_test.cc @@ -125,7 +125,7 @@ class CoreWorkerTest : public ::testing::Test { void TestNormalTask(const std::unordered_map &resources) { CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], - raylet_socket_names_[0], JobID::FromInt(1)); + raylet_socket_names_[0], JobID::FromInt(1), nullptr); // Test pass by value. { @@ -184,7 +184,7 @@ class CoreWorkerTest : public ::testing::Test { void TestActorTask(const std::unordered_map &resources) { CoreWorker driver(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], - raylet_socket_names_[0], JobID::FromInt(1)); + raylet_socket_names_[0], JobID::FromInt(1), nullptr); std::unique_ptr actor_handle; @@ -335,7 +335,7 @@ TEST_F(ZeroNodeTest, TestActorHandle) { TEST_F(SingleNodeTest, TestObjectInterface) { CoreWorker core_worker(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], raylet_socket_names_[0], - JobID::JobID::FromInt(1)); + JobID::FromInt(1), nullptr); uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8}; uint8_t array2[] = {10, 11, 12, 13, 14, 15}; @@ -398,10 +398,10 @@ TEST_F(SingleNodeTest, TestObjectInterface) { TEST_F(TwoNodeTest, TestObjectInterfaceCrossNodes) { CoreWorker worker1(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[0], - raylet_socket_names_[0], JobID::JobID::FromInt(1)); + raylet_socket_names_[0], JobID::FromInt(1), nullptr); CoreWorker worker2(WorkerType::DRIVER, Language::PYTHON, raylet_store_socket_names_[1], - raylet_socket_names_[1], JobID::JobID::FromInt(1)); + raylet_socket_names_[1], JobID::FromInt(1), nullptr); uint8_t array1[] = {1, 2, 3, 4, 5, 6, 7, 8}; uint8_t array2[] = {10, 11, 12, 13, 14, 15}; @@ -487,7 +487,7 @@ TEST_F(TwoNodeTest, TestActorTaskCrossNodes) { TEST_F(SingleNodeTest, TestCoreWorkerConstructorFailure) { try { CoreWorker core_worker(WorkerType::DRIVER, Language::PYTHON, "", - raylet_socket_names_[0], JobID::FromInt(1)); + raylet_socket_names_[0], JobID::FromInt(1), nullptr); } catch (const std::exception &e) { std::cout << "Caught exception when constructing core worker: " << e.what(); } diff --git a/src/ray/core_worker/mock_worker.cc b/src/ray/core_worker/mock_worker.cc index c171e5df8..d77931c87 100644 --- a/src/ray/core_worker/mock_worker.cc +++ b/src/ray/core_worker/mock_worker.cc @@ -3,6 +3,8 @@ #include "ray/core_worker/store_provider/store_provider.h" #include "ray/core_worker/task_execution.h" +using namespace std::placeholders; + namespace ray { /// A mock C++ worker used by core_worker_test.cc to verify the task submission/execution @@ -18,38 +20,39 @@ class MockWorker { public: MockWorker(const std::string &store_socket, const std::string &raylet_socket) : worker_(WorkerType::WORKER, Language::PYTHON, store_socket, raylet_socket, - JobID::JobID::FromInt(1)) {} + JobID::JobID::FromInt(1), + std::bind(&MockWorker::ExecuteTask, this, _1, _2, _3, _4)) {} void Run() { - auto executor_func = [this](const RayFunction &ray_function, - const std::vector> &args, - const TaskInfo &task_info, int num_returns) { - // Note that this doesn't include dummy object id. - RAY_CHECK(num_returns >= 0); - - // Merge all the content from input args. - std::vector buffer; - for (const auto &arg : args) { - auto &data = arg->GetData(); - buffer.insert(buffer.end(), data->Data(), data->Data() + data->Size()); - } - - auto return_value = RayObject( - std::make_shared(buffer.data(), buffer.size()), nullptr); - - // Write the merged content to each of return ids. - for (int i = 0; i < num_returns; i++) { - ObjectID id = ObjectID::ForTaskReturn(task_info.task_id, i + 1); - RAY_CHECK_OK(worker_.Objects().Put(return_value, id)); - } - return Status::OK(); - }; - // Start executing tasks. - worker_.Execution().Run(executor_func); + worker_.Execution().Run(); } private: + Status ExecuteTask(const RayFunction &ray_function, + const std::vector> &args, + const TaskInfo &task_info, int num_returns) { + // Note that this doesn't include dummy object id. + RAY_CHECK(num_returns >= 0); + + // Merge all the content from input args. + std::vector buffer; + for (const auto &arg : args) { + auto &data = arg->GetData(); + buffer.insert(buffer.end(), data->Data(), data->Data() + data->Size()); + } + + auto return_value = RayObject( + std::make_shared(buffer.data(), buffer.size()), nullptr); + + // Write the merged content to each of return ids. + for (int i = 0; i < num_returns; i++) { + ObjectID id = ObjectID::ForTaskReturn(task_info.task_id, i + 1); + RAY_CHECK_OK(worker_.Objects().Put(return_value, id)); + } + return Status::OK(); + } + CoreWorker worker_; }; diff --git a/src/ray/core_worker/task_execution.cc b/src/ray/core_worker/task_execution.cc index 19ed37f69..813e789eb 100644 --- a/src/ray/core_worker/task_execution.cc +++ b/src/ray/core_worker/task_execution.cc @@ -7,67 +7,64 @@ namespace ray { CoreWorkerTaskExecutionInterface::CoreWorkerTaskExecutionInterface( WorkerContext &worker_context, std::unique_ptr &raylet_client, - CoreWorkerObjectInterface &object_interface) + CoreWorkerObjectInterface &object_interface, const TaskExecutor &executor) : worker_context_(worker_context), object_interface_(object_interface), + execution_callback_(executor), worker_server_("Worker", 0 /* let grpc choose port */), main_work_(main_service_) { + RAY_CHECK(execution_callback_ != nullptr); + + auto func = std::bind(&CoreWorkerTaskExecutionInterface::ExecuteTask, this, + std::placeholders::_1); task_receivers_.emplace( static_cast(TaskTransportType::RAYLET), std::unique_ptr(new CoreWorkerRayletTaskReceiver( - raylet_client, main_service_, worker_server_))); + raylet_client, main_service_, worker_server_, func))); // Start RPC server after all the task receivers are properly initialized. worker_server_.Run(); } -Status CoreWorkerTaskExecutionInterface::Run(const TaskExecutor &executor) { - while (true) { - std::vector tasks; - auto status = - task_receivers_[static_cast(TaskTransportType::RAYLET)]->GetTasks(&tasks); - if (!status.ok()) { - RAY_LOG(ERROR) << "Getting task failed with error: " - << ray::Status::IOError(status.message()); - return status; - } +Status CoreWorkerTaskExecutionInterface::ExecuteTask(const TaskSpecification &spec) { + worker_context_.SetCurrentTask(spec); - for (const auto &task : tasks) { - const auto &spec = task.GetTaskSpecification(); - worker_context_.SetCurrentTask(spec); + RayFunction func{spec.GetLanguage(), spec.FunctionDescriptor()}; - RayFunction func{spec.GetLanguage(), spec.FunctionDescriptor()}; + std::vector> args; + RAY_CHECK_OK(BuildArgsForExecutor(spec, &args)); - std::vector> args; - RAY_CHECK_OK(BuildArgsForExecutor(spec, &args)); - - TaskType task_type; - if (spec.IsActorCreationTask()) { - task_type = TaskType::ACTOR_CREATION_TASK; - } else if (spec.IsActorTask()) { - task_type = TaskType::ACTOR_TASK; - } else { - task_type = TaskType::NORMAL_TASK; - } - - TaskInfo task_info{spec.TaskId(), spec.JobId(), task_type}; - - auto num_returns = spec.NumReturns(); - if (spec.IsActorCreationTask() || spec.IsActorTask()) { - RAY_CHECK(num_returns > 0); - // Decrease to account for the dummy object id. - num_returns--; - } - - status = executor(func, args, task_info, num_returns); - // TODO(zhijunfu): - // 1. Check and handle failure. - // 2. Save or load checkpoint. - } + TaskType task_type; + if (spec.IsActorCreationTask()) { + task_type = TaskType::ACTOR_CREATION_TASK; + } else if (spec.IsActorTask()) { + task_type = TaskType::ACTOR_TASK; + } else { + task_type = TaskType::NORMAL_TASK; } + TaskInfo task_info{spec.TaskId(), spec.JobId(), task_type}; + + auto num_returns = spec.NumReturns(); + if (spec.IsActorCreationTask() || spec.IsActorTask()) { + RAY_CHECK(num_returns > 0); + // Decrease to account for the dummy object id. + num_returns--; + } + + auto status = execution_callback_(func, args, task_info, num_returns); + // TODO(zhijunfu): + // 1. Check and handle failure. + // 2. Save or load checkpoint. + return status; +} + +void CoreWorkerTaskExecutionInterface::Run() { + // Run main IO service. + main_service_.run(); + // should never reach here. - return Status::OK(); + RAY_LOG(FATAL) << "should never reach here after running main io service"; } Status CoreWorkerTaskExecutionInterface::BuildArgsForExecutor( diff --git a/src/ray/core_worker/task_execution.h b/src/ray/core_worker/task_execution.h index b154b8f01..599bd71d7 100644 --- a/src/ray/core_worker/task_execution.h +++ b/src/ray/core_worker/task_execution.h @@ -23,10 +23,6 @@ class TaskSpecification; /// execution. class CoreWorkerTaskExecutionInterface { public: - CoreWorkerTaskExecutionInterface(WorkerContext &worker_context, - std::unique_ptr &raylet_client, - CoreWorkerObjectInterface &object_interface); - /// The callback provided app-language workers that executes tasks. /// /// \param ray_function[in] Information about the function to execute. @@ -37,9 +33,14 @@ class CoreWorkerTaskExecutionInterface { const std::vector> &args, const TaskInfo &task_info, int num_returns)>; + CoreWorkerTaskExecutionInterface(WorkerContext &worker_context, + std::unique_ptr &raylet_client, + CoreWorkerObjectInterface &object_interface, + const TaskExecutor &executor); + /// Start receving and executes tasks in a infinite loop. - /// \return Status. - Status Run(const TaskExecutor &executor); + /// \return void. + void Run(); private: /// Build arguments for task executor. This would loop through all the arguments @@ -53,11 +54,17 @@ class CoreWorkerTaskExecutionInterface { Status BuildArgsForExecutor(const TaskSpecification &spec, std::vector> *args); + /// Execute a task. + Status ExecuteTask(const TaskSpecification &spec); + /// Reference to the parent CoreWorker's context. WorkerContext &worker_context_; /// Reference to the parent CoreWorker's objects interface. CoreWorkerObjectInterface &object_interface_; + // Task execution callback. + TaskExecutor execution_callback_; + /// All the task task receivers supported. std::unordered_map> task_receivers_; diff --git a/src/ray/core_worker/transport/raylet_transport.cc b/src/ray/core_worker/transport/raylet_transport.cc index ece982701..c260c6d51 100644 --- a/src/ray/core_worker/transport/raylet_transport.cc +++ b/src/ray/core_worker/transport/raylet_transport.cc @@ -9,29 +9,16 @@ CoreWorkerRayletTaskSubmitter::CoreWorkerRayletTaskSubmitter( : raylet_client_(raylet_client) {} Status CoreWorkerRayletTaskSubmitter::SubmitTask(const TaskSpec &task) { + RAY_CHECK(raylet_client_ != nullptr); return raylet_client_->SubmitTask(task.GetDependencies(), task.GetTaskSpecification()); } -Status CoreWorkerRayletTaskReceiver::GetTasks(std::vector *tasks) { - std::unique_ptr task_spec; - auto status = raylet_client_->GetTask(&task_spec); - if (!status.ok()) { - RAY_LOG(ERROR) << "Get task from raylet failed with error: " - << ray::Status::IOError(status.message()); - return status; - } - - std::vector dependencies; - RAY_CHECK((*tasks).empty()); - (*tasks).emplace_back(*task_spec, dependencies); - - return Status::OK(); -} - CoreWorkerRayletTaskReceiver::CoreWorkerRayletTaskReceiver( std::unique_ptr &raylet_client, boost::asio::io_service &io_service, - rpc::GrpcServer &server) - : raylet_client_(raylet_client), task_service_(io_service, *this) { + rpc::GrpcServer &server, const TaskHandler &task_handler) + : raylet_client_(raylet_client), + task_service_(io_service, *this), + task_handler_(task_handler) { server.RegisterService(task_service_); } @@ -41,6 +28,19 @@ void CoreWorkerRayletTaskReceiver::HandleAssignTask( const Task task(request.task()); const auto &spec = task.GetTaskSpecification(); auto status = task_handler_(spec); + // Notify raylet that current task is done via a `TaskDone` message. This is to + // ensure that the task is marked as finished by raylet only after previous + // raylet client calls are completed. For example, if the worker sends a + // NotifyUnblocked message that it is no longer blocked in a `ray.get` + // on the normal raylet socket, then completes an assigned task, we + // need to guarantee that raylet gets the former message first before + // marking the task as completed. This is why a `TaskDone` message + // is required - without it, it's possible that raylet receives + // rpc reply first before the NotifyUnblocked message arrives, + // as they use different connections, the `TaskDone` message is sent + // to raylet via the same connection so the order is guaranteed. + raylet_client_->TaskDone(); + // send rpc reply. send_reply_callback(status, nullptr, nullptr); } diff --git a/src/ray/core_worker/transport/raylet_transport.h b/src/ray/core_worker/transport/raylet_transport.h index 7e8fea1ee..1b375305d 100644 --- a/src/ray/core_worker/transport/raylet_transport.h +++ b/src/ray/core_worker/transport/raylet_transport.h @@ -33,15 +33,8 @@ class CoreWorkerRayletTaskReceiver : public CoreWorkerTaskReceiver, public: CoreWorkerRayletTaskReceiver(std::unique_ptr &raylet_client, boost::asio::io_service &io_service, - rpc::GrpcServer &server); + rpc::GrpcServer &server, const TaskHandler &task_handler); - // Get tasks for execution from raylet. - virtual Status GetTasks(std::vector *tasks) override; - - /// TODO(zhijunfu): This is currently unused. Later when we migrate from worker "get - /// task" to raylet "assign task", this method will be used and the `GetTask` above will - /// be removed. - /// /// Handle a `AssignTask` request. /// The implementation can handle this request asynchronously. When hanling is done, the /// `send_reply_callback` should be called. @@ -56,10 +49,10 @@ class CoreWorkerRayletTaskReceiver : public CoreWorkerTaskReceiver, private: /// Raylet client. std::unique_ptr &raylet_client_; - /// The callback function to process a task. - TaskHandler task_handler_; /// The rpc service for `WorkerTaskService`. rpc::WorkerTaskGrpcService task_service_; + /// The callback function to process a task. + TaskHandler task_handler_; }; } // namespace ray diff --git a/src/ray/core_worker/transport/transport.h b/src/ray/core_worker/transport/transport.h index fc67a2d41..5376817cf 100644 --- a/src/ray/core_worker/transport/transport.h +++ b/src/ray/core_worker/transport/transport.h @@ -33,9 +33,6 @@ class CoreWorkerTaskSubmitter { class CoreWorkerTaskReceiver { public: using TaskHandler = std::function; - - // Get tasks for execution. - virtual Status GetTasks(std::vector *tasks) = 0; }; } // namespace ray diff --git a/src/ray/protobuf/worker.proto b/src/ray/protobuf/worker.proto index c2affa1ac..3c2d30ab6 100644 --- a/src/ray/protobuf/worker.proto +++ b/src/ray/protobuf/worker.proto @@ -7,6 +7,10 @@ import "src/ray/protobuf/common.proto"; message AssignTaskRequest { // The task to be pushed. Task task = 1; + // 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; } message AssignTaskReply { diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index d6d15437b..ca9006e95 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -777,7 +777,12 @@ void NodeManager::ProcessClientMessage( ProcessRegisterClientRequestMessage(client, message_data); } break; case protocol::MessageType::GetTask: { - ProcessGetTaskMessage(client); + RAY_CHECK(!registered_worker->UsePush()); + HandleWorkerAvailable(client); + } break; + case protocol::MessageType::TaskDone: { + RAY_CHECK(registered_worker->UsePush()); + HandleWorkerAvailable(client); } break; case protocol::MessageType::DisconnectClient: { ProcessDisconnectClientMessage(client); @@ -849,14 +854,20 @@ void NodeManager::ProcessClientMessage( void NodeManager::ProcessRegisterClientRequestMessage( const std::shared_ptr &client, const uint8_t *message_data) { auto message = flatbuffers::GetRoot(message_data); - client->SetClientID(from_flatbuf(*message->worker_id())); + auto client_id = from_flatbuf(*message->worker_id()); + client->SetClientID(client_id); Language language = static_cast(message->language()); - auto worker = - std::make_shared(message->worker_pid(), language, message->port(), client); + auto worker = std::make_shared(message->worker_pid(), language, message->port(), + client, client_call_manager_); if (message->is_worker()) { // Register the new worker. + bool use_push_task = worker->UsePush(); + auto connection = worker->Connection(); worker_pool_.RegisterWorker(std::move(worker)); - DispatchTasks(local_queues_.GetReadyTasksWithResources()); + if (use_push_task) { + // only call `HandleWorkerAvailable` when push mode is used. + HandleWorkerAvailable(connection); + } } else { // Register the new driver. const WorkerID driver_id = from_flatbuf(*message->worker_id()); @@ -917,7 +928,7 @@ void NodeManager::HandleDisconnectedActor(const ActorID &actor_id, bool was_loca PublishActorStateTransition(actor_id, new_actor_data, failure_callback); } -void NodeManager::ProcessGetTaskMessage( +void NodeManager::HandleWorkerAvailable( const std::shared_ptr &client) { std::shared_ptr worker = worker_pool_.GetRegisteredWorker(client); RAY_CHECK(worker); @@ -925,6 +936,7 @@ void NodeManager::ProcessGetTaskMessage( if (!worker->GetAssignedTaskId().IsNil()) { FinishAssignedTask(*worker); } + // Return the worker to the idle pool. worker_pool_.PushWorker(std::move(worker)); // Local resource availability changed: invoke scheduling policy for local node. @@ -1762,81 +1774,30 @@ bool NodeManager::AssignTask(const Task &task) { worker->SetTaskResourceIds(acquired_resources); } + auto task_id = spec.TaskId(); + auto finish_assign_task_callback = [this, worker, task_id](Status status) { + if (worker->UsePush()) { + // NOTE: we cannot directly call `FinishAssignTask` here because + // it assumes the task is in SWAP queue, thus we need to delay invoking this + // function after the assigned tasks are moved from READY queue to SWAP queue + // in `DispatchTasks`. + // Another option is to move the tasks to SWAP queue here just before calling + // `FinishAssignTask` so we can save an io_service post, at the + // expense of calling `MoveTask` for each of the assigned tasks. + // TODO(zhijunfu): after all workers are fully migrated to push mode, the + // `post` below and swap queue can be removed. + io_service_.post([this, status, worker, task_id]() { + FinishAssignTask(task_id, *worker, status.ok()); + }); + } else { + FinishAssignTask(task_id, *worker, status.ok()); + } + }; + ResourceIdSet resource_id_set = worker->GetTaskResourceIds().Plus(worker->GetLifetimeResourceIds()); - auto resource_id_set_flatbuf = resource_id_set.ToFlatbuf(fbb); - auto message = protocol::CreateGetTaskReply(fbb, fbb.CreateString(spec.Serialize()), - fbb.CreateVector(resource_id_set_flatbuf)); - fbb.Finish(message); - const auto &task_id = spec.TaskId(); - worker->Connection()->WriteMessageAsync( - static_cast(protocol::MessageType::ExecuteTask), fbb.GetSize(), - fbb.GetBufferPointer(), [this, worker, task_id](ray::Status status) { - // Remove the ASSIGNED task from the SWAP queue. - Task assigned_task; - TaskState state; - if (!local_queues_.RemoveTask(task_id, &assigned_task, &state)) { - return; - } - - RAY_CHECK(state == TaskState::SWAP); - - if (status.ok()) { - auto spec = assigned_task.GetTaskSpecification(); - // We successfully assigned the task to the worker. - worker->AssignTaskId(spec.TaskId()); - worker->AssignJobId(spec.JobId()); - // Actor tasks require extra accounting to track the actor's state. - if (spec.IsActorTask()) { - auto actor_entry = actor_registry_.find(spec.ActorId()); - RAY_CHECK(actor_entry != actor_registry_.end()); - // Process any new actor handles that were created since the - // previous task on this handle was executed. The first task - // submitted on a new actor handle will depend on the dummy object - // returned by the previous task, so the dependency will not be - // released until this first task is submitted. - for (auto &new_handle_id : spec.NewActorHandles()) { - // Get the execution dependency for the first task submitted on the new - // actor handle. Since the new actor handle was created after this task - // began and before this task finished, it must have the same execution - // dependency. - const auto &execution_dependencies = - assigned_task.GetTaskExecutionSpec().ExecutionDependencies(); - // TODO(swang): We expect this task to have exactly 1 execution dependency, - // the dummy object returned by the previous actor task. However, this - // leaks information about the TaskExecutionSpecification implementation. - RAY_CHECK(execution_dependencies.size() == 1); - const ObjectID &execution_dependency = execution_dependencies.front(); - // Add the new handle and give it a reference to the finished task's - // execution dependency. - actor_entry->second.AddHandle(new_handle_id, execution_dependency); - } - - // TODO(swang): For actors with multiple actor handles, to - // guarantee that tasks are replayed in the same order after a - // failure, we must update the task's execution dependency to be - // the actor's current execution dependency. - } - - // Mark the task as running. - // (See design_docs/task_states.rst for the state transition diagram.) - local_queues_.QueueTasks({assigned_task}, TaskState::RUNNING); - // Notify the task dependency manager that we no longer need this task's - // object dependencies. - RAY_CHECK(task_dependency_manager_.UnsubscribeDependencies(spec.TaskId())); - } else { - RAY_LOG(WARNING) << "Failed to send task to worker, disconnecting client"; - // We failed to send the task to the worker, so disconnect the worker. - ProcessDisconnectClientMessage(worker->Connection()); - // Queue this task for future assignment. We need to do this since - // DispatchTasks() removed it from the ready queue. The task will be - // assigned to a worker once one becomes available. - // (See design_docs/task_states.rst for the state transition diagram.) - local_queues_.QueueTasks({assigned_task}, TaskState::READY); - DispatchTasks(MakeTasksWithResources({assigned_task})); - } - }); + worker->AssignTask(task, resource_id_set, finish_assign_task_callback); // We assigned this task to a worker. // (Note this means that we sent the task to the worker. The assignment @@ -2294,6 +2255,72 @@ void NodeManager::ForwardTask( }); } +void NodeManager::FinishAssignTask(const TaskID &task_id, Worker &worker, bool success) { + // Remove the ASSIGNED task from the SWAP queue. + Task assigned_task; + TaskState state; + if (!local_queues_.RemoveTask(task_id, &assigned_task, &state)) { + return; + } + + RAY_CHECK(state == TaskState::SWAP); + + if (success) { + auto spec = assigned_task.GetTaskSpecification(); + // We successfully assigned the task to the worker. + worker.AssignTaskId(spec.TaskId()); + worker.AssignJobId(spec.JobId()); + // Actor tasks require extra accounting to track the actor's state. + if (spec.IsActorTask()) { + auto actor_entry = actor_registry_.find(spec.ActorId()); + RAY_CHECK(actor_entry != actor_registry_.end()); + // Process any new actor handles that were created since the + // previous task on this handle was executed. The first task + // submitted on a new actor handle will depend on the dummy object + // returned by the previous task, so the dependency will not be + // released until this first task is submitted. + for (auto &new_handle_id : spec.NewActorHandles()) { + // Get the execution dependency for the first task submitted on the new + // actor handle. Since the new actor handle was created after this task + // began and before this task finished, it must have the same execution + // dependency. + const auto &execution_dependencies = + assigned_task.GetTaskExecutionSpec().ExecutionDependencies(); + // TODO(swang): We expect this task to have exactly 1 execution dependency, + // the dummy object returned by the previous actor task. However, this + // leaks information about the TaskExecutionSpecification implementation. + RAY_CHECK(execution_dependencies.size() == 1); + const ObjectID &execution_dependency = execution_dependencies.front(); + // Add the new handle and give it a reference to the finished task's + // execution dependency. + actor_entry->second.AddHandle(new_handle_id, execution_dependency); + } + + // TODO(swang): For actors with multiple actor handles, to + // guarantee that tasks are replayed in the same order after a + // failure, we must update the task's execution dependency to be + // the actor's current execution dependency. + } + + // Mark the task as running. + // (See design_docs/task_states.rst for the state transition diagram.) + local_queues_.QueueTasks({assigned_task}, TaskState::RUNNING); + // Notify the task dependency manager that we no longer need this task's + // object dependencies. + RAY_CHECK(task_dependency_manager_.UnsubscribeDependencies(spec.TaskId())); + } else { + RAY_LOG(WARNING) << "Failed to send task to worker, disconnecting client"; + // We failed to send the task to the worker, so disconnect the worker. + ProcessDisconnectClientMessage(worker.Connection()); + // Queue this task for future assignment. We need to do this since + // DispatchTasks() removed it from the ready queue. The task will be + // assigned to a worker once one becomes available. + // (See design_docs/task_states.rst for the state transition diagram.) + local_queues_.QueueTasks({assigned_task}, TaskState::READY); + DispatchTasks(MakeTasksWithResources({assigned_task})); + } +} + void NodeManager::DumpDebugState() const { std::fstream fs; fs.open(initial_config_.session_dir + "/debug_state.txt", diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index a3a0243b6..fa92fa73e 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -375,11 +375,11 @@ class NodeManager : public rpc::NodeManagerServiceHandler { void ProcessRegisterClientRequestMessage( const std::shared_ptr &client, const uint8_t *message_data); - /// Process client message of GetTask + /// Handle the case that a worker is available. /// - /// \param client The client that sent the message. + /// \param client The connection for the worker. /// \return Void. - void ProcessGetTaskMessage(const std::shared_ptr &client); + void HandleWorkerAvailable(const std::shared_ptr &client); /// Handle a client that has disconnected. This can be called multiple times /// on the same client because this is triggered both when a client @@ -459,6 +459,14 @@ class NodeManager : public rpc::NodeManagerServiceHandler { void HandleDisconnectedActor(const ActorID &actor_id, bool was_local, bool intentional_disconnect); + /// Finish assigning a task to a worker. + /// + /// \param task_id Id of the task. + /// \param worker Worker which the task is assigned to. + /// \param success Whether the task is successfully assigned to the worker. + /// \return void. + void FinishAssignTask(const TaskID &task_id, Worker &worker, bool success); + /// Handle a `ForwardTask` request. void HandleForwardTask(const rpc::ForwardTaskRequest &request, rpc::ForwardTaskReply *reply, @@ -523,7 +531,8 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// The RPC service. rpc::NodeManagerGrpcService node_manager_service_; - /// The `ClientCallManager` object that is shared by all `NodeManagerClient`s. + /// The `ClientCallManager` object that is shared by all `NodeManagerClient`s + /// as well as all `WorkerTaskClient`s. rpc::ClientCallManager client_call_manager_; /// Map from node ids to clients of the remote node managers. diff --git a/src/ray/raylet/raylet.h b/src/ray/raylet/raylet.h index 01395141f..d0ae57c7f 100644 --- a/src/ray/raylet/raylet.h +++ b/src/ray/raylet/raylet.h @@ -18,7 +18,6 @@ namespace raylet { using rpc::ClientTableData; -class Task; class NodeManager; class Raylet { diff --git a/src/ray/raylet/scheduling_policy.cc b/src/ray/raylet/scheduling_policy.cc index 32d4e4f03..c87394140 100644 --- a/src/ray/raylet/scheduling_policy.cc +++ b/src/ray/raylet/scheduling_policy.cc @@ -51,7 +51,7 @@ std::unordered_map SchedulingPolicy::Schedule( ResourceSet(node_resources.GetAvailableResources()); // We have to subtract the current "load" because we set the current "load" // to be the resources used by tasks that are in the - // `SchedulingQueue::ready_queue_` in NodeManager::ProcessGetTaskMessage's + // `SchedulingQueue::ready_queue_` in NodeManager::HandleWorkerAvailable's // call to SchedulingQueue::GetResourceLoad. available_node_resources.SubtractResources(node_resources.GetLoadResources()); RAY_LOG(DEBUG) << "client_id " << node_client_id diff --git a/src/ray/raylet/worker.cc b/src/ray/raylet/worker.cc index d3d9833a3..8e1bd4092 100644 --- a/src/ray/raylet/worker.cc +++ b/src/ray/raylet/worker.cc @@ -11,13 +11,20 @@ namespace raylet { /// A constructor responsible for initializing the state of a worker. Worker::Worker(pid_t pid, const Language &language, int port, - std::shared_ptr connection) + std::shared_ptr connection, + rpc::ClientCallManager &client_call_manager) : pid_(pid), language_(language), port_(port), connection_(connection), dead_(false), - blocked_(false) {} + blocked_(false), + client_call_manager_(client_call_manager) { + if (port_ > 0) { + rpc_client_ = std::unique_ptr( + new rpc::WorkerTaskClient("127.0.0.1", port_, client_call_manager_)); + } +} void Worker::MarkDead() { dead_ = true; } @@ -103,6 +110,44 @@ void Worker::AcquireTaskCpuResources(const ResourceIdSet &cpu_resources) { task_resource_ids_.Release(cpu_resources); } +bool Worker::UsePush() const { return rpc_client_ != nullptr; } + +void Worker::AssignTask(const Task &task, const ResourceIdSet &resource_id_set, + const std::function finish_assign_callback) { + const TaskSpecification &spec = task.GetTaskSpecification(); + if (rpc_client_ != nullptr) { + // Use push mode. + RAY_CHECK(port_ > 0); + rpc::AssignTaskRequest request; + request.mutable_task()->mutable_task_spec()->CopyFrom( + task.GetTaskSpecification().GetMessage()); + request.mutable_task()->mutable_task_execution_spec()->CopyFrom( + task.GetTaskExecutionSpec().GetMessage()); + request.set_resource_ids(resource_id_set.Serialize()); + + auto status = rpc_client_->AssignTask( + request, [](Status status, const rpc::AssignTaskReply &reply) { + // Worker has finished this task. There's nothing to do here + // and assigning new task will be done when raylet receives + // `TaskDone` message. + }); + finish_assign_callback(status); + } else { + // Use pull mode. This corresponds to existing python/java workers that haven't been + // migrated to core worker architecture. + flatbuffers::FlatBufferBuilder fbb; + auto resource_id_set_flatbuf = resource_id_set.ToFlatbuf(fbb); + + auto message = + protocol::CreateGetTaskReply(fbb, fbb.CreateString(spec.Serialize()), + fbb.CreateVector(resource_id_set_flatbuf)); + fbb.Finish(message); + Connection()->WriteMessageAsync( + static_cast(protocol::MessageType::ExecuteTask), fbb.GetSize(), + fbb.GetBufferPointer(), finish_assign_callback); + } +} + } // namespace raylet } // end namespace ray diff --git a/src/ray/raylet/worker.h b/src/ray/raylet/worker.h index 81758f992..571249072 100644 --- a/src/ray/raylet/worker.h +++ b/src/ray/raylet/worker.h @@ -6,7 +6,9 @@ #include "ray/common/client_connection.h" #include "ray/common/id.h" #include "ray/common/task/scheduling_resources.h" +#include "ray/common/task/task.h" #include "ray/common/task/task_common.h" +#include "ray/rpc/worker/worker_client.h" namespace ray { @@ -19,7 +21,8 @@ class Worker { public: /// A constructor that initializes a worker object. Worker(pid_t pid, const Language &language, int port, - std::shared_ptr connection); + std::shared_ptr connection, + rpc::ClientCallManager &client_call_manager); /// A destructor responsible for freeing all worker state. ~Worker() {} void MarkDead(); @@ -53,6 +56,10 @@ class Worker { ResourceIdSet ReleaseTaskCpuResources(); void AcquireTaskCpuResources(const ResourceIdSet &cpu_resources); + bool UsePush() const; + void AssignTask(const Task &task, const ResourceIdSet &resource_id_set, + const std::function finish_assign_callback); + private: /// The worker's PID. pid_t pid_; @@ -81,6 +88,11 @@ class Worker { // of a task. ResourceIdSet task_resource_ids_; std::unordered_set blocked_task_ids_; + /// The `ClientCallManager` object that is shared by `WorkerTaskClient` from all + /// workers. + rpc::ClientCallManager &client_call_manager_; + /// The rpc client to send tasks to this worker. + std::unique_ptr rpc_client_; }; } // namespace raylet diff --git a/src/ray/raylet/worker_pool_test.cc b/src/ray/raylet/worker_pool_test.cc index 4b4579459..386bec31b 100644 --- a/src/ray/raylet/worker_pool_test.cc +++ b/src/ray/raylet/worker_pool_test.cc @@ -70,7 +70,11 @@ class WorkerPoolMock : public WorkerPool { class WorkerPoolTest : public ::testing::Test { public: - WorkerPoolTest() : worker_pool_(), io_service_(), error_message_type_(1) {} + WorkerPoolTest() + : worker_pool_(), + io_service_(), + error_message_type_(1), + client_call_manager_(io_service_) {} std::shared_ptr CreateWorker(pid_t pid, const Language &language = Language::PYTHON) { @@ -85,7 +89,8 @@ class WorkerPoolTest : public ::testing::Test { auto client = LocalClientConnection::Create(client_handler, message_handler, std::move(socket), "worker", {}, error_message_type_); - return std::shared_ptr(new Worker(pid, language, -1, client)); + return std::shared_ptr( + new Worker(pid, language, -1, client, client_call_manager_)); } void SetWorkerCommands(const WorkerCommandMap &worker_commands) { @@ -97,6 +102,7 @@ class WorkerPoolTest : public ::testing::Test { WorkerPoolMock worker_pool_; boost::asio::io_service io_service_; int64_t error_message_type_; + rpc::ClientCallManager client_call_manager_; private: void HandleNewClient(LocalClientConnection &){};