diff --git a/src/common/redis_module/ray_redis_module.cc b/src/common/redis_module/ray_redis_module.cc index e38a69dcc..3db89c433 100644 --- a/src/common/redis_module/ray_redis_module.cc +++ b/src/common/redis_module/ray_redis_module.cc @@ -50,7 +50,7 @@ } static const char *table_prefixes[] = { - NULL, "TASK:", "CLIENT:", "OBJECT:", "FUNCTION:", + NULL, "TASK:", "TASK:", "CLIENT:", "OBJECT:", "FUNCTION:", }; /// Parse a Redis string into a TablePubsub channel. diff --git a/src/ray/common/client_connection.cc b/src/ray/common/client_connection.cc index b1e033664..e2a184654 100644 --- a/src/ray/common/client_connection.cc +++ b/src/ray/common/client_connection.cc @@ -40,7 +40,7 @@ template void ClientConnection::ProcessMessageHeader(const boost::system::error_code &error) { if (error) { // If there was an error, disconnect the client. - read_type_ = MessageType_DisconnectClient; + read_type_ = protocol::MessageType_DisconnectClient; read_length_ = 0; ProcessMessage(error); return; @@ -81,7 +81,7 @@ template void ClientConnection::ProcessMessage(const boost::system::error_code &error) { if (error) { // TODO(hme): Disconnect differently & remove dependency on node_manager_generated.h - read_type_ = MessageType_DisconnectClient; + read_type_ = protocol::MessageType_DisconnectClient; } manager_.ProcessClientMessage(this->shared_from_this(), read_type_, read_message_.data()); diff --git a/src/ray/gcs/client.cc b/src/ray/gcs/client.cc index 01bfc4566..485684daf 100644 --- a/src/ray/gcs/client.cc +++ b/src/ray/gcs/client.cc @@ -16,6 +16,7 @@ Status AsyncGcsClient::Connect(const std::string &address, int port, RAY_RETURN_NOT_OK(context_->Connect(address, port)); object_table_.reset(new ObjectTable(context_, this)); task_table_.reset(new TaskTable(context_, this)); + raylet_task_table_.reset(new raylet::TaskTable(context_, this)); client_table_.reset(new ClientTable(context_, this, client_info)); // TODO(swang): Call the client table's Connect() method here. To do this, // we need to make sure that we are attached to an event loop first. This @@ -41,6 +42,8 @@ ObjectTable &AsyncGcsClient::object_table() { return *object_table_; } TaskTable &AsyncGcsClient::task_table() { return *task_table_; } +raylet::TaskTable &AsyncGcsClient::raylet_task_table() { return *raylet_task_table_; } + ClientTable &AsyncGcsClient::client_table() { return *client_table_; } FunctionTable &AsyncGcsClient::function_table() { return *function_table_; } diff --git a/src/ray/gcs/client.h b/src/ray/gcs/client.h index 4b9b27719..b92c48be1 100644 --- a/src/ray/gcs/client.h +++ b/src/ray/gcs/client.h @@ -45,6 +45,7 @@ class RAY_EXPORT AsyncGcsClient { inline ConfigTable &config_table(); ObjectTable &object_table(); TaskTable &task_table(); + raylet::TaskTable &raylet_task_table(); ClientTable &client_table(); inline ErrorTable &error_table(); @@ -63,6 +64,7 @@ class RAY_EXPORT AsyncGcsClient { std::unique_ptr class_table_; std::unique_ptr object_table_; std::unique_ptr task_table_; + std::unique_ptr raylet_task_table_; std::unique_ptr client_table_; std::shared_ptr context_; std::unique_ptr asio_async_client_; diff --git a/src/ray/gcs/client_test.cc b/src/ray/gcs/client_test.cc index eeb6b1437..659817f60 100644 --- a/src/ray/gcs/client_test.cc +++ b/src/ray/gcs/client_test.cc @@ -109,26 +109,42 @@ void LookupFailed(gcs::AsyncGcsClient *client, const UniqueID &id) { test->Stop(); } -void TestObjectTable(const JobID &job_id, std::shared_ptr client) { - auto data = std::make_shared(); - data->managers.push_back("A"); - data->managers.push_back("B"); - ObjectID object_id = ObjectID::from_random(); - RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, &ObjectAdded)); - RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, &Lookup, &LookupFailed)); +void TestTableLookup(const JobID &job_id, std::shared_ptr client) { + TaskID task_id = TaskID::from_random(); + auto data = std::make_shared(); + data->task_specification = "123"; + + auto add_callback = [data](gcs::AsyncGcsClient *client, const UniqueID &id, + const protocol::TaskT &d) { + ASSERT_EQ(data->task_specification, d.task_specification); + }; + + auto lookup_callback = [data](gcs::AsyncGcsClient *client, const UniqueID &id, + const protocol::TaskT &d) { + ASSERT_EQ(data->task_specification, d.task_specification); + test->Stop(); + }; + + auto failure_callback = [](gcs::AsyncGcsClient *client, const UniqueID &id) { + RAY_CHECK(false); + }; + + RAY_CHECK_OK(client->raylet_task_table().Add(job_id, task_id, data, add_callback)); + RAY_CHECK_OK(client->raylet_task_table().Lookup(job_id, task_id, lookup_callback, + failure_callback)); // Run the event loop. The loop will only stop if the Lookup callback is // called (or an assertion failure). test->Start(); } -TEST_F(TestGcsWithAe, TestObjectTable) { +TEST_F(TestGcsWithAe, TestTableLookup) { test = this; - TestObjectTable(job_id_, client_); + TestTableLookup(job_id_, client_); } -TEST_F(TestGcsWithAsio, TestObjectTable) { +TEST_F(TestGcsWithAsio, TestTableLookup) { test = this; - TestObjectTable(job_id_, client_); + TestTableLookup(job_id_, client_); } void TestLookupFailure(const JobID &job_id, std::shared_ptr client) { diff --git a/src/ray/gcs/format/gcs.fbs b/src/ray/gcs/format/gcs.fbs index e9aa5169a..e0ef9ade2 100644 --- a/src/ray/gcs/format/gcs.fbs +++ b/src/ray/gcs/format/gcs.fbs @@ -7,6 +7,7 @@ enum Language:int { enum TablePrefix:int { UNUSED = 0, TASK, + RAYLET_TASK, CLIENT, OBJECT, FUNCTION @@ -16,6 +17,7 @@ enum TablePrefix:int { enum TablePubsub:int { NO_PUBLISH = 0, TASK, + RAYLET_TASK, CLIENT, OBJECT, ACTOR diff --git a/src/ray/gcs/tables.cc b/src/ray/gcs/tables.cc index e84e8d1d8..8233efefc 100644 --- a/src/ray/gcs/tables.cc +++ b/src/ray/gcs/tables.cc @@ -239,6 +239,7 @@ const ClientTableDataT &ClientTable::GetClient(const ClientID &client_id) { } } +template class Table; template class Table; template class Table; diff --git a/src/ray/gcs/tables.h b/src/ray/gcs/tables.h index 7f38055fe..54cf65627 100644 --- a/src/ray/gcs/tables.h +++ b/src/ray/gcs/tables.h @@ -12,6 +12,7 @@ #include "ray/gcs/format/gcs_generated.h" #include "ray/gcs/redis_context.h" +#include "ray/raylet/format/node_manager_generated.h" // TODO(pcm): Remove this #include "task.h" @@ -165,6 +166,18 @@ using ClassTable = Table; // TODO(swang): Set the pubsub channel for the actor table. using ActorTable = Table; +namespace raylet { + +class TaskTable : public Table { + public: + TaskTable(const std::shared_ptr &context, AsyncGcsClient *client) + : Table(context, client) { + pubsub_channel_ = TablePubsub_RAYLET_TASK; + prefix_ = TablePrefix_RAYLET_TASK; + } +}; +} + class TaskTable : public Table { public: TaskTable(const std::shared_ptr &context, AsyncGcsClient *client) diff --git a/src/ray/raylet/format/node_manager.fbs b/src/ray/raylet/format/node_manager.fbs index 7a5629f49..7de47dd31 100644 --- a/src/ray/raylet/format/node_manager.fbs +++ b/src/ray/raylet/format/node_manager.fbs @@ -1,5 +1,9 @@ // Local scheduler protocol specification +// TODO(swang): We put the flatbuffer types in a separate namespace for now to +// avoid conflicts with legacy Ray types. +namespace ray.protocol; + enum MessageType:int { // Task is submitted to the local scheduler. This is sent from a worker to a // local scheduler. @@ -50,6 +54,21 @@ enum MessageType:int { SetActorFrontier } +table TaskExecutionSpecification { + // A list of object IDs representing the dependencies of this task that may + // change at execution time. + dependencies: [string]; + // The last time this task was received for scheduling. + last_timestamp: double; + // The number of times this task was spilled back by local schedulers. + num_forwards: int; +} + +table Task { + task_specification: string; + task_execution_spec: TaskExecutionSpecification; +} + table SubmitTaskRequest { execution_dependencies: [string]; task_spec: string; diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index 0be0cc62b..d0e144b74 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -32,8 +32,8 @@ void NodeManager::ProcessClientMessage(std::shared_ptr cl RAY_LOG(DEBUG) << "Message of type " << message_type; switch (message_type) { - case MessageType_RegisterClientRequest: { - auto message = flatbuffers::GetRoot(message_data); + case protocol::MessageType_RegisterClientRequest: { + auto message = flatbuffers::GetRoot(message_data); if (message->is_worker()) { // Create a new worker from the registration request. std::shared_ptr worker(new Worker(message->worker_pid(), client)); @@ -45,14 +45,15 @@ void NodeManager::ProcessClientMessage(std::shared_ptr cl // is legacy code and should be removed once actor creation tasks are // implemented. flatbuffers::FlatBufferBuilder fbb; - auto reply = CreateRegisterClientReply(fbb, fbb.CreateVector(std::vector())); + auto reply = + protocol::CreateRegisterClientReply(fbb, fbb.CreateVector(std::vector())); fbb.Finish(reply); // Reply to the worker's registration request, then listen for more // messages. - client->WriteMessage(MessageType_RegisterClientReply, fbb.GetSize(), + client->WriteMessage(protocol::MessageType_RegisterClientReply, fbb.GetSize(), fbb.GetBufferPointer()); } break; - case MessageType_GetTask: { + case protocol::MessageType_GetTask: { const std::shared_ptr worker = worker_pool_.GetRegisteredWorker(client); RAY_CHECK(worker); // If the worker was assigned a task, mark it as finished. @@ -69,16 +70,16 @@ void NodeManager::ProcessClientMessage(std::shared_ptr cl AssignTask(scheduled_tasks.front()); } } break; - case MessageType_DisconnectClient: { + case protocol::MessageType_DisconnectClient: { // Remove the dead worker from the pool and stop listening for messages. const std::shared_ptr worker = worker_pool_.GetRegisteredWorker(client); if (worker) { worker_pool_.DisconnectWorker(worker); } } break; - case MessageType_SubmitTask: { + case protocol::MessageType_SubmitTask: { // Read the task submitted by the client. - auto message = flatbuffers::GetRoot(message_data); + auto message = flatbuffers::GetRoot(message_data); TaskExecutionSpecification task_execution_spec( from_flatbuf(*message->execution_dependencies())); TaskSpecification task_spec(*message->task_spec()); @@ -152,10 +153,10 @@ void NodeManager::AssignTask(const Task &task) { flatbuffers::FlatBufferBuilder fbb; const TaskSpecification &spec = task.GetTaskSpecification(); - auto message = CreateGetTaskReply(fbb, spec.ToFlatbuffer(fbb), - fbb.CreateVector(std::vector())); + auto message = protocol::CreateGetTaskReply(fbb, spec.ToFlatbuffer(fbb), + fbb.CreateVector(std::vector())); fbb.Finish(message); - worker->Connection()->WriteMessage(MessageType_ExecuteTask, fbb.GetSize(), + worker->Connection()->WriteMessage(protocol::MessageType_ExecuteTask, fbb.GetSize(), fbb.GetBufferPointer()); worker->AssignTaskId(spec.TaskId()); local_queues_.QueueRunningTasks(std::vector({task}));