diff --git a/protos/ray.proto b/protos/ray.proto index 827eb5f0d..a9ab18230 100644 --- a/protos/ray.proto +++ b/protos/ray.proto @@ -79,7 +79,8 @@ message RegisterFunctionRequest { } message SubmitTaskRequest { - Task task = 1; // Contains name of the function to be executed and arguments + uint64 workerid = 1; // The ID of the worker submitting the task + Task task = 2; // Contains name of the function to be executed and arguments } message SubmitTaskReply { @@ -210,15 +211,11 @@ message TaskInfoRequest { message TaskInfoReply { repeated TaskStatus failed_task = 1; + repeated TaskStatus running_task = 2; + uint64 num_succeeded = 3; // TODO(mehrdadn): We'll want to return information from computation_graph since it's important for visualizing tasks that have been completed etc. } -message TaskStatus { - uint64 operationid = 1; - string worker_address = 2; - string error_message = 3; -} - // These messages are for getting information about the object store state message ObjStoreInfoRequest { diff --git a/protos/types.proto b/protos/types.proto index da383377a..85d40f859 100644 --- a/protos/types.proto +++ b/protos/types.proto @@ -85,6 +85,13 @@ message Operation { uint64 creator_operationid = 3; // The id of the task that called this task or push. } +message TaskStatus { + uint64 operationid = 1; + string function_name = 2; + string worker_address = 3; + string error_message = 4; +} + message Array { repeated uint64 shape = 1; sint64 dtype = 2; diff --git a/src/computation_graph.cc b/src/computation_graph.cc index a0d553063..ce13bd6bc 100644 --- a/src/computation_graph.cc +++ b/src/computation_graph.cc @@ -13,6 +13,8 @@ OperationId ComputationGraph::add_operation(std::unique_ptr operation } const Task& ComputationGraph::get_task(OperationId operationid) { + RAY_CHECK_NEQ(operationid, ROOT_OPERATION, "ComputationGraph attempting to get_task with operationid == ROOT_OPERATION"); + RAY_CHECK_NEQ(operationid, NO_OPERATION, "ComputationGraph attempting to get_task with operationid == NO_OPERATION"); RAY_CHECK_LT(operationid, operations_.size(), "ComputationGraph attempting to get_task with operationid " << operationid << ", but operationid >= operations_.size()."); RAY_CHECK(operations_[operationid]->has_task(), "Calling get_task with operationid " << operationid << ", but this corresponds to a push not a task."); return operations_[operationid]->task(); diff --git a/src/raylib.cc b/src/raylib.cc index 87ca0d787..5a9e76800 100644 --- a/src/raylib.cc +++ b/src/raylib.cc @@ -821,13 +821,26 @@ PyObject* task_info(PyObject* self, PyObject* args) { const TaskStatus& info = reply.failed_task(i); PyObject* info_dict = PyDict_New(); set_dict_item_and_transfer_ownership(info_dict, PyString_FromString("worker_address"), PyString_FromStringAndSize(info.worker_address().data(), info.worker_address().size())); + set_dict_item_and_transfer_ownership(info_dict, PyString_FromString("function_name"), PyString_FromStringAndSize(info.function_name().data(), info.function_name().size())); set_dict_item_and_transfer_ownership(info_dict, PyString_FromString("operationid"), PyInt_FromLong(info.operationid())); set_dict_item_and_transfer_ownership(info_dict, PyString_FromString("error_message"), PyString_FromStringAndSize(info.error_message().data(), info.error_message().size())); PyList_SetItem(failed_tasks_list, i, info_dict); } + PyObject* running_tasks_list = PyList_New(reply.running_task_size()); + for (size_t i = 0; i < reply.running_task_size(); ++i) { + const TaskStatus& info = reply.running_task(i); + PyObject* info_dict = PyDict_New(); + set_dict_item_and_transfer_ownership(info_dict, PyString_FromString("worker_address"), PyString_FromStringAndSize(info.worker_address().data(), info.worker_address().size())); + set_dict_item_and_transfer_ownership(info_dict, PyString_FromString("function_name"), PyString_FromStringAndSize(info.function_name().data(), info.function_name().size())); + set_dict_item_and_transfer_ownership(info_dict, PyString_FromString("operationid"), PyInt_FromLong(info.operationid())); + PyList_SetItem(running_tasks_list, i, info_dict); + } + PyObject* dict = PyDict_New(); set_dict_item_and_transfer_ownership(dict, PyString_FromString("failed_tasks"), failed_tasks_list); + set_dict_item_and_transfer_ownership(dict, PyString_FromString("running_tasks"), running_tasks_list); + set_dict_item_and_transfer_ownership(dict, PyString_FromString("number_succeeded"), PyInt_FromLong(reply.num_succeeded())); return dict; } diff --git a/src/scheduler.cc b/src/scheduler.cc index 5f58dc16f..9f93f7ce1 100644 --- a/src/scheduler.cc +++ b/src/scheduler.cc @@ -38,8 +38,10 @@ Status SchedulerService::SubmitTask(ServerContext* context, const SubmitTaskRequ auto operation = std::unique_ptr(new Operation()); operation->set_allocated_task(task.release()); - OperationId creator_operationid = ROOT_OPERATION; // TODO(rkn): Later, this should be the ID of the task that spawned this current task. - operation->set_creator_operationid(creator_operationid); + { + std::lock_guard workers_lock(workers_lock_); + operation->set_creator_operationid(workers_[request->workerid()].current_task); + } OperationId operationid; { @@ -160,25 +162,41 @@ Status SchedulerService::ObjReady(ServerContext* context, const ObjReadyRequest* Status SchedulerService::ReadyForNewTask(ServerContext* context, const ReadyForNewTaskRequest* request, AckReply* reply) { RAY_LOG(RAY_INFO, "worker " << request->workerid() << " is ready for a new task"); - { - std::lock_guard lock(avail_workers_lock_); - avail_workers_.push_back(request->workerid()); - } if (request->has_previous_task_info()) { + OperationId operationid; + { + std::lock_guard workers_lock(workers_lock_); + operationid = workers_[request->workerid()].current_task; + } + std::string task_name; + { + std::lock_guard computation_graph_lock(computation_graph_lock_); + task_name = computation_graph_.get_task(operationid).name(); + } TaskStatus info; { std::lock_guard workers_lock(workers_lock_); - info.set_operationid(workers_[request->workerid()].current_task); + operationid = workers_[request->workerid()].current_task; + info.set_operationid(operationid); + info.set_function_name(task_name); info.set_worker_address(workers_[request->workerid()].worker_address); info.set_error_message(request->previous_task_info().error_message()); workers_[request->workerid()].current_task = NO_OPERATION; // clear operation ID } if (!request->previous_task_info().task_succeeded()) { + RAY_LOG(RAY_INFO, "Error: Task " << info.operationid() << " executing function " << info.function_name() << " on worker " << request->workerid() << " failed with error message: " << info.error_message()); std::lock_guard failed_tasks_lock(failed_tasks_lock_); failed_tasks_.push_back(info); + } else { + std::lock_guard successful_tasks_lock(successful_tasks_lock_); + successful_tasks_.push_back(info.operationid()); } // TODO(rkn): Handle task failure } + { + std::lock_guard lock(avail_workers_lock_); + avail_workers_.push_back(request->workerid()); + } schedule(); return Status::OK; } @@ -227,11 +245,25 @@ Status SchedulerService::SchedulerInfo(ServerContext* context, const SchedulerIn } Status SchedulerService::TaskInfo(ServerContext* context, const TaskInfoRequest* request, TaskInfoReply* reply) { + std::lock_guard successful_tasks_lock(successful_tasks_lock_); std::lock_guard failed_tasks_lock(failed_tasks_lock_); - for (size_t i = 0; i != failed_tasks_.size(); ++i) { + std::lock_guard computation_graph_lock(computation_graph_lock_); + std::lock_guard workers_lock(workers_lock_); + for (int i = 0; i < failed_tasks_.size(); ++i) { TaskStatus* info = reply->add_failed_task(); *info = failed_tasks_[i]; } + for (int i = 0; i < workers_.size(); ++i) { + OperationId operationid = workers_[i].current_task; + if (operationid != NO_OPERATION) { + const Task& task = computation_graph_.get_task(operationid); + TaskStatus* info = reply->add_running_task(); + info->set_operationid(operationid); + info->set_function_name(task.name()); + info->set_worker_address(workers_[i].worker_address); + } + } + reply->set_num_succeeded(successful_tasks_.size()); return Status::OK; } @@ -748,6 +780,7 @@ void SchedulerService::get_equivalent_objrefs(ObjRef objref, std::vector // This method defines the order in which locks should be acquired. void SchedulerService::do_on_locks(bool lock) { std::mutex *mutexes[] = { + &successful_tasks_lock_, &failed_tasks_lock_, &pull_queue_lock_, &computation_graph_lock_, diff --git a/src/scheduler.h b/src/scheduler.h index 02513b383..f04f86ee7 100644 --- a/src/scheduler.h +++ b/src/scheduler.h @@ -177,9 +177,12 @@ private: // List of pending pull calls. std::vector > pull_queue_; std::mutex pull_queue_lock_; - // List of failed workers + // List of failed tasks std::vector failed_tasks_; std::mutex failed_tasks_lock_; + // List of the IDs of successful tasks + std::vector successful_tasks_; // Right now, we only use this information in the TaskInfo call. + std::mutex successful_tasks_lock_; // List of pending alias notifications. Each element consists of (objstoreid, (alias_objref, canonical_objref)). std::vector > > alias_notification_queue_; std::mutex alias_notification_queue_lock_;