diff --git a/src/ray/protobuf/node_manager.proto b/src/ray/protobuf/node_manager.proto index 5661b5296..c6dd3b23e 100644 --- a/src/ray/protobuf/node_manager.proto +++ b/src/ray/protobuf/node_manager.proto @@ -4,14 +4,6 @@ package ray.rpc; import "src/ray/protobuf/common.proto"; -// Submit a task for execution. -message SubmitTaskRequest { - TaskSpec task_spec = 1; -} - -message SubmitTaskReply { -} - // Request a worker from the raylet with the specified resources. message WorkerLeaseRequest { // TaskSpec containing the requested resources. @@ -90,8 +82,6 @@ message NodeStatsReply { // Service for inter-node-manager communication. service NodeManagerService { - // Submit a task (from a local or remote worker) to the node manager. - rpc SubmitTask(SubmitTaskRequest) returns (SubmitTaskReply); // Request a worker from the raylet. rpc RequestWorkerLease(WorkerLeaseRequest) returns (WorkerLeaseReply); // Release a worker back to its raylet. diff --git a/src/ray/raylet/format/node_manager.fbs b/src/ray/raylet/format/node_manager.fbs index 23b7406e5..cccbba91a 100644 --- a/src/ray/raylet/format/node_manager.fbs +++ b/src/ray/raylet/format/node_manager.fbs @@ -5,9 +5,12 @@ namespace ray.protocol; enum MessageType:int { + // Task is submitted to the raylet. This is sent from a worker to a + // raylet. + SubmitTask = 1, // Notify the raylet that a task has finished. This is sent from a // worker to a raylet. - TaskDone = 1, + TaskDone, // Log a message to the event table. This is sent from a worker to a raylet. EventLogMessage, // Send an initial connection message to the raylet. This is sent @@ -96,6 +99,10 @@ table Task { task_execution_spec: TaskExecutionSpecification; } +table SubmitTaskRequest { + task_spec: string; +} + // This message describes a given resource that is reserved for a worker. table ResourceIdSetInfo { // The name of the resource. diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index d41918e1e..f833924c4 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -900,6 +900,12 @@ void NodeManager::ProcessClientMessage( // because it's already disconnected. return; } break; + case protocol::MessageType::SubmitTask: { + // For tasks submitted via the raylet path, we must make sure to order the + // task submission so that tasks are always submitted after the tasks that + // they depend on. + ProcessSubmitTaskMessage(message_data); + } break; case protocol::MessageType::SetResourceRequest: { ProcessSetResourceRequest(client, message_data); } break; @@ -1408,21 +1414,16 @@ void NodeManager::ProcessReportActiveObjectIDs( unordered_set_from_flatbuf(*message->object_ids())); } -void NodeManager::HandleSubmitTask(const rpc::SubmitTaskRequest &request, - rpc::SubmitTaskReply *reply, - rpc::SendReplyCallback send_reply_callback) { - rpc::Task task; - task.mutable_task_spec()->CopyFrom(request.task_spec()); - // Set the caller's node ID. - if (task.task_spec().caller_address().raylet_id() == "") { - task.mutable_task_spec()->mutable_caller_address()->set_raylet_id( - gcs_client_->client_table().GetLocalClientId().Binary()); - } +void NodeManager::ProcessSubmitTaskMessage(const uint8_t *message_data) { + // Read the task submitted by the client. + auto fbs_message = flatbuffers::GetRoot(message_data); + rpc::Task task_message; + RAY_CHECK(task_message.mutable_task_spec()->ParseFromArray( + fbs_message->task_spec()->data(), fbs_message->task_spec()->size())); // Submit the task to the raylet. Since the task was submitted // locally, there is no uncommitted lineage. - SubmitTask(Task(task), Lineage()); - send_reply_callback(Status::OK(), nullptr, nullptr); + SubmitTask(Task(task_message), Lineage()); } void NodeManager::HandleWorkerLeaseRequest(const rpc::WorkerLeaseRequest &request, diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index f49ea2f20..ad4416c1f 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -395,6 +395,12 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// \return True if the invariants are satisfied and false otherwise. bool CheckDependencyManagerInvariant() const; + /// Process client message of SubmitTask + /// + /// \param message_data A pointer to the message data. + /// \return Void. + void ProcessSubmitTaskMessage(const uint8_t *message_data); + /// Process client message of RegisterClientRequest /// /// \param client The client that sent the message. @@ -510,11 +516,6 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// \return void. void FinishAssignTask(const TaskID &task_id, Worker &worker, bool success); - /// Handle a `SubmitTask` request. - void HandleSubmitTask(const rpc::SubmitTaskRequest &request, - rpc::SubmitTaskReply *reply, - rpc::SendReplyCallback send_reply_callback) override; - /// Handle a `WorkerLease` request. void HandleWorkerLeaseRequest(const rpc::WorkerLeaseRequest &request, rpc::WorkerLeaseReply *reply, diff --git a/src/ray/raylet/raylet_client.cc b/src/ray/raylet/raylet_client.cc index 75e4f970d..1c12d2ad8 100644 --- a/src/ray/raylet/raylet_client.cc +++ b/src/ray/raylet/raylet_client.cc @@ -229,7 +229,6 @@ RayletClient::RayletClient(std::shared_ptr gr } ray::Status RayletClient::SubmitTask(const ray::TaskSpecification &task_spec) { - ray::rpc::SubmitTaskRequest request; for (size_t i = 0; i < task_spec.NumArgs(); i++) { if (task_spec.ArgByRef(i)) { for (size_t j = 0; j < task_spec.ArgIdCount(i); j++) { @@ -238,8 +237,11 @@ ray::Status RayletClient::SubmitTask(const ray::TaskSpecification &task_spec) { } } } - request.mutable_task_spec()->CopyFrom(task_spec.GetMessage()); - return grpc_client_->SubmitTask(request, /*callback=*/nullptr); + flatbuffers::FlatBufferBuilder fbb; + auto message = ray::protocol::CreateSubmitTaskRequest( + fbb, fbb.CreateString(task_spec.Serialize())); + fbb.Finish(message); + return conn_->WriteMessage(MessageType::SubmitTask, &fbb); } ray::Status RayletClient::TaskDone() { diff --git a/src/ray/rpc/node_manager/node_manager_client.h b/src/ray/rpc/node_manager/node_manager_client.h index 686548ace..06a6723e9 100644 --- a/src/ray/rpc/node_manager/node_manager_client.h +++ b/src/ray/rpc/node_manager/node_manager_client.h @@ -73,16 +73,6 @@ class NodeManagerWorkerClient return std::shared_ptr(instance); } - /// Submit a task. - ray::Status SubmitTask(const SubmitTaskRequest &request, - const ClientCallback &callback) { - auto call = client_call_manager_ - .CreateCall( - *stub_, &NodeManagerService::Stub::PrepareAsyncSubmitTask, - request, callback); - return call->GetStatus(); - } - /// Request a worker lease. ray::Status RequestWorkerLease(const WorkerLeaseRequest &request, const ClientCallback &callback) { diff --git a/src/ray/rpc/node_manager/node_manager_server.h b/src/ray/rpc/node_manager/node_manager_server.h index 7d89b320f..a8e916e61 100644 --- a/src/ray/rpc/node_manager/node_manager_server.h +++ b/src/ray/rpc/node_manager/node_manager_server.h @@ -24,9 +24,6 @@ class NodeManagerServiceHandler { /// \param[out] reply The reply message. /// \param[in] send_reply_callback The callback to be called when the request is done. - virtual void HandleSubmitTask(const SubmitTaskRequest &request, SubmitTaskReply *reply, - SendReplyCallback send_reply_callback) = 0; - virtual void HandleWorkerLeaseRequest(const WorkerLeaseRequest &request, WorkerLeaseReply *reply, SendReplyCallback send_reply_callback) = 0; @@ -63,13 +60,6 @@ class NodeManagerGrpcService : public GrpcService { std::vector, int>> *server_call_factories_and_concurrencies) override { // Initialize the factory for requests. - std::unique_ptr submit_task_call_factory( - new ServerCallFactoryImpl( - service_, &NodeManagerService::AsyncService::RequestSubmitTask, - service_handler_, &NodeManagerServiceHandler::HandleSubmitTask, cq, - main_service_)); - std::unique_ptr request_worker_lease_call_factory( new ServerCallFactoryImpl( @@ -99,8 +89,6 @@ class NodeManagerGrpcService : public GrpcService { main_service_)); // Set accept concurrency. - server_call_factories_and_concurrencies->emplace_back( - std::move(submit_task_call_factory), 100); server_call_factories_and_concurrencies->emplace_back( std::move(request_worker_lease_call_factory), 100); server_call_factories_and_concurrencies->emplace_back(