track more task info (#133)

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 {

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;

src/computation_graph.cc

@@ -13,6 +13,8 @@ OperationId ComputationGraph::add_operation(std::unique_ptr<Operation> 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();

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;
 }
 

src/scheduler.cc

@@ -38,8 +38,10 @@ Status SchedulerService::SubmitTask(ServerContext* context, const SubmitTaskRequ
 
     auto operation = std::unique_ptr<Operation>(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<std::mutex> 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<std::mutex> lock(avail_workers_lock_);
-    avail_workers_.push_back(request->workerid());
-  }
   if (request->has_previous_task_info()) {
+    OperationId operationid;
+    {
+      std::lock_guard<std::mutex> workers_lock(workers_lock_);
+      operationid = workers_[request->workerid()].current_task;
+    }
+    std::string task_name;
+    {
+      std::lock_guard<std::mutex> computation_graph_lock(computation_graph_lock_);
+      task_name = computation_graph_.get_task(operationid).name();
+    }
     TaskStatus info;
     {
       std::lock_guard<std::mutex> 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<std::mutex> failed_tasks_lock(failed_tasks_lock_);
       failed_tasks_.push_back(info);
+    } else {
+      std::lock_guard<std::mutex> successful_tasks_lock(successful_tasks_lock_);
+      successful_tasks_.push_back(info.operationid());
     }
     // TODO(rkn): Handle task failure
   }
+  {
+    std::lock_guard<std::mutex> 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<std::mutex> successful_tasks_lock(successful_tasks_lock_);
   std::lock_guard<std::mutex> failed_tasks_lock(failed_tasks_lock_);
-  for (size_t i = 0; i != failed_tasks_.size(); ++i) {
+  std::lock_guard<std::mutex> computation_graph_lock(computation_graph_lock_);
+  std::lock_guard<std::mutex> 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<ObjRef>
 // 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_,

src/scheduler.h

@@ -177,9 +177,12 @@ private:
   // List of pending pull calls.
   std::vector<std::pair<WorkerId, ObjRef> > pull_queue_;
   std::mutex pull_queue_lock_;
-  // List of failed workers
+  // List of failed tasks
   std::vector<TaskStatus> failed_tasks_;
   std::mutex failed_tasks_lock_;
+  // List of the IDs of successful tasks
+  std::vector<OperationId> 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<std::pair<ObjStoreId, std::pair<ObjRef, ObjRef> > > alias_notification_queue_;
   std::mutex alias_notification_queue_lock_;