add computation graph (#75)

CMakeLists.txt

@@ -80,7 +80,7 @@ endif()
 
 add_executable(objstore src/objstore.cc src/ipc.cc ${GENERATED_PROTOBUF_FILES})
 target_link_libraries(objstore arrow numbuf pynumbuf)
-add_executable(scheduler src/scheduler.cc ${GENERATED_PROTOBUF_FILES})
+add_executable(scheduler src/scheduler.cc src/computation_graph.cc ${GENERATED_PROTOBUF_FILES})
 add_library(orchpylib SHARED src/orchpylib.cc src/worker.cc src/ipc.cc ${GENERATED_PROTOBUF_FILES})
 target_link_libraries(orchpylib arrow numbuf pynumbuf)
 

include/orchestra/orchestra.h

@@ -7,6 +7,7 @@
 typedef size_t ObjRef;
 typedef size_t WorkerId;
 typedef size_t ObjStoreId;
+typedef size_t OperationId;
 
 class FnInfo {
   size_t num_return_vals_;

protos/orchestra.proto

@@ -139,7 +139,7 @@ message FnTableEntry {
 }
 
 message SchedulerInfoReply {
-  repeated Task task = 1; // Tasks on the task queue
+  repeated uint64 operationid = 1; // OperationIds of the tasks on the task queue
   repeated uint64 avail_worker = 3; // List of workers waiting to get a task assigned
   map<string, FnTableEntry> function_table = 2; // Table of all available remote function
   repeated uint64 target_objref = 4; // The target_objrefs_ data structure

protos/types.proto

@@ -56,6 +56,20 @@ message Task {
   repeated uint64 result = 3; // Object references for result
 }
 
+message Push {
+  uint64 objref = 1; // The objref for the pushed object
+}
+
+// This is used internally by the scheduler. From the scheduler's perspective,
+// the submission of tasks (via SubmitTask) and the submission of pushes (via
+// PushObj) look very similar, and so it is useful to be able to handle them
+// together (for example in the computation graph).
+message Operation {
+  Task task = 1;
+  Push push = 2;
+  uint64 creator_operationid = 3; // The id of the task that called this task or push.
+}
+
 message Array {
   repeated uint64 shape = 1;
   sint64 dtype = 2;

src/computation_graph.cc

@@ -0,0 +1,25 @@
+#include "computation_graph.h"
+
+OperationId ComputationGraph::add_operation(std::unique_ptr<Operation> operation) {
+  OperationId operationid = operations_.size();
+  OperationId creator_operationid = operation->creator_operationid();
+  if (spawned_operations_.size() != operationid) {
+    ORCH_LOG(ORCH_FATAL, "ComputationGraph is attempting to call add_operation, but spawned_operations_.size() != operationid.");
+  }
+  operations_.emplace_back(std::move(operation));
+  if (creator_operationid != NO_OPERATION && creator_operationid != ROOT_OPERATION) {
+    spawned_operations_[creator_operationid].push_back(operationid);
+  }
+  spawned_operations_.push_back(std::vector<OperationId>());
+  return operationid;
+}
+
+const Task& ComputationGraph::get_task(OperationId operationid) {
+  if (operationid >= operations_.size()) {
+    ORCH_LOG(ORCH_FATAL, "ComputationGraph attempting to get_task with operationid " << operationid << ", but operationid >= operations_.size().");
+  }
+  if (!operations_[operationid]->has_task()) {
+    ORCH_LOG(ORCH_FATAL, "Calling get_task with operationid " << operationid << ", but this corresponds to a push not a task.");
+  }
+  return operations_[operationid]->task();
+}

src/computation_graph.h

@@ -0,0 +1,32 @@
+#ifndef ORCHESTRA_COMPUTATIONGRAPH_H
+#define ORCHESTRA_COMPUTATIONGRAPH_H
+
+#include <iostream>
+#include <limits>
+
+#include "orchestra/orchestra.h"
+#include "orchestra.grpc.pb.h"
+#include "types.pb.h"
+
+// used to represent the root operation (that is, the driver code)
+const OperationId ROOT_OPERATION = std::numeric_limits<OperationId>::max();
+// used to represent the absence of an operation
+const OperationId NO_OPERATION = std::numeric_limits<OperationId>::max() - 1;
+
+class ComputationGraph {
+public:
+  // Add an operation to the computation graph, this returns the OperationId for
+  // the new operation. This method takes ownership over operation.
+  OperationId add_operation(std::unique_ptr<Operation> operation);
+  // Return the task corresponding to a particular OperationId. If operationid
+  // corresponds to a push, then fail.
+  const Task& get_task(OperationId operationid);
+private:
+  // maps an OperationId to the corresponding task or push
+  std::vector<std::unique_ptr<Operation> > operations_;
+  // spawned_operations_[operationid] is a vector of the OperationIds of the
+  // operations spawned by the task with OperationId operationid
+  std::vector<std::vector<OperationId> > spawned_operations_;
+};
+
+#endif

src/scheduler.cc

@@ -33,8 +33,16 @@ Status SchedulerService::SubmitTask(ServerContext* context, const SubmitTaskRequ
     increment_ref_count(result_objrefs); // We increment once so the objrefs don't go out of scope before the task is scheduled on the worker. The corresponding decrement will happen in deserialize_task in orchpylib.
   }
 
+  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);
+  computation_graph_lock_.lock();
+  OperationId operationid = computation_graph_.add_operation(std::move(operation));
+  computation_graph_lock_.unlock();
+
   task_queue_lock_.lock();
-  task_queue_.emplace_back(std::move(task));
+  task_queue_.push_back(operationid);
   task_queue_lock_.unlock();
 
   schedule();
@@ -232,15 +240,17 @@ void SchedulerService::schedule() {
   perform_notify_aliases(); // See what we can do in alias_notification_queue_
 }
 
-void SchedulerService::assign_task(std::unique_ptr<Task> task, WorkerId workerid) {
-  // submit task assumes that the canonical objrefs for its arguments are all ready, that is has_canonical_objref() is true for all of the call's arguments
+// assign_task assumes that computation_graph_lock_ has been acquired.
+// assign_task assumes that the canonical objrefs for its arguments are all ready, that is has_canonical_objref() is true for all of the call's arguments
+void SchedulerService::assign_task(OperationId operationid, WorkerId workerid) {
+  const Task& task = computation_graph_.get_task(operationid);
   ClientContext context;
   ExecuteTaskRequest request;
   ExecuteTaskReply reply;
   ORCH_LOG(ORCH_INFO, "starting to send arguments");
-  for (size_t i = 0; i < task->arg_size(); ++i) {
-    if (!task->arg(i).has_obj()) {
-      ObjRef objref = task->arg(i).ref();
+  for (size_t i = 0; i < task.arg_size(); ++i) {
+    if (!task.arg(i).has_obj()) {
+      ObjRef objref = task.arg(i).ref();
       ObjRef canonical_objref = get_canonical_objref(objref);
       {
         // Notify the relevant objstore about potential aliasing when it's ready
@@ -258,7 +268,7 @@ void SchedulerService::assign_task(std::unique_ptr<Task> task, WorkerId workerid
       }
     }
   }
-  request.set_allocated_task(task.release()); // protobuf object takes ownership
+  request.mutable_task()->CopyFrom(task); // TODO(rkn): Is ownership handled properly here?
   Status status = workers_[workerid].worker_stub->ExecuteTask(&context, request, &reply);
 }
 
@@ -410,8 +420,7 @@ void SchedulerService::get_info(const SchedulerInfoRequest& request, SchedulerIn
     }
   }
   for (const auto& entry : task_queue_) {
-    Task* task = reply->add_task();
-    task->CopyFrom(*entry);
+    reply->add_operationid(entry);
   }
   for (const WorkerId& entry : avail_workers_) {
     reply->add_avail_worker(entry);
@@ -481,6 +490,7 @@ void SchedulerService::perform_pulls() {
 }
 
 void SchedulerService::schedule_tasks_naively() {
+  std::lock_guard<std::mutex> computation_graph_lock(computation_graph_lock_);
   std::lock_guard<std::mutex> fntable_lock(fntable_lock_);
   std::lock_guard<std::mutex> avail_workers_lock(avail_workers_lock_);
   std::lock_guard<std::mutex> task_queue_lock(task_queue_lock_);
@@ -491,10 +501,11 @@ void SchedulerService::schedule_tasks_naively() {
       // The use of erase(it) below invalidates the iterator, but we
       // immediately break out of the inner loop, so the iterator is not used
       // after the erase
-      const Task& task = *(*it);
+      const OperationId operationid = *it;
+      const Task& task = computation_graph_.get_task(operationid);
       auto& workers = fntable_[task.name()].workers();
       if (std::binary_search(workers.begin(), workers.end(), workerid) && can_run(task)) {
-        assign_task(std::move(*it), workerid);
+        assign_task(operationid, workerid);
         task_queue_.erase(it);
         std::swap(avail_workers_[i], avail_workers_[avail_workers_.size() - 1]);
         avail_workers_.pop_back();
@@ -506,6 +517,7 @@ void SchedulerService::schedule_tasks_naively() {
 }
 
 void SchedulerService::schedule_tasks_location_aware() {
+  std::lock_guard<std::mutex> computation_graph_lock(computation_graph_lock_);
   std::lock_guard<std::mutex> fntable_lock(fntable_lock_);
   std::lock_guard<std::mutex> avail_workers_lock(avail_workers_lock_);
   std::lock_guard<std::mutex> task_queue_lock(task_queue_lock_);
@@ -516,7 +528,8 @@ void SchedulerService::schedule_tasks_location_aware() {
     auto bestit = task_queue_.end(); // keep track of the task that fits the worker best so far
     size_t min_num_shipped_objects = std::numeric_limits<size_t>::max(); // number of objects that need to be transfered for this worker
     for (auto it = task_queue_.begin(); it != task_queue_.end(); ++it) {
-      const Task& task = *(*it);
+      OperationId operationid = *it;
+      const Task& task = computation_graph_.get_task(operationid);
       auto& workers = fntable_[task.name()].workers();
       if (std::binary_search(workers.begin(), workers.end(), workerid) && can_run(task)) {
         // determine how many objects would need to be shipped
@@ -542,7 +555,7 @@ void SchedulerService::schedule_tasks_location_aware() {
     }
     // if we found a suitable task
     if (bestit != task_queue_.end()) {
-      assign_task(std::move(*bestit), workerid);
+      assign_task(*bestit, workerid);
       task_queue_.erase(bestit);
       std::swap(avail_workers_[i], avail_workers_[avail_workers_.size() - 1]);
       avail_workers_.pop_back();

src/scheduler.h

@@ -14,6 +14,8 @@
 #include "orchestra.grpc.pb.h"
 #include "types.pb.h"
 
+#include "computation_graph.h"
+
 using grpc::Server;
 using grpc::ServerBuilder;
 using grpc::ServerReader;
@@ -69,7 +71,7 @@ public:
   // assign a task to a worker
   void schedule();
   // execute a task on a worker and ship required object references
-  void assign_task(std::unique_ptr<Task> task, WorkerId workerid);
+  void assign_task(OperationId operationid, WorkerId workerid);
   // checks if the dependencies of the task are met
   bool can_run(const Task& task);
   // register a worker and its object store (if it has not been registered yet)
@@ -116,6 +118,10 @@ private:
   // Find all of the object references that refer to the same object as objref (as best as we can determine at the moment). The information may be incomplete because not all of the aliases may be known.
   void get_equivalent_objrefs(ObjRef objref, std::vector<ObjRef> &equivalent_objrefs);
 
+  // The computation graph tracks the operations that have been submitted to the
+  // scheduler and is mostly used for fault tolerance.
+  ComputationGraph computation_graph_;
+  std::mutex computation_graph_lock_;
   // Vector of all workers registered in the system. Their index in this vector
   // is the workerid.
   std::vector<WorkerHandle> workers_;
@@ -144,7 +150,7 @@ private:
   FnTable fntable_;
   std::mutex fntable_lock_;
   // List of pending tasks.
-  std::deque<std::unique_ptr<Task> > task_queue_;
+  std::deque<OperationId> task_queue_;
   std::mutex task_queue_lock_;
   // List of pending pull calls.
   std::vector<std::pair<WorkerId, ObjRef> > pull_queue_;