diff --git a/CMakeLists.txt b/CMakeLists.txt index 215f58047..d70d5a070 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/include/orchestra/orchestra.h b/include/orchestra/orchestra.h index 9abd728ef..f0fbc8b98 100644 --- a/include/orchestra/orchestra.h +++ b/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_; diff --git a/protos/orchestra.proto b/protos/orchestra.proto index 9c1a40adc..71fcfbdff 100644 --- a/protos/orchestra.proto +++ b/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 function_table = 2; // Table of all available remote function repeated uint64 target_objref = 4; // The target_objrefs_ data structure diff --git a/protos/types.proto b/protos/types.proto index af19503f3..6b08478e6 100644 --- a/protos/types.proto +++ b/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; diff --git a/src/computation_graph.cc b/src/computation_graph.cc new file mode 100644 index 000000000..bd309dec5 --- /dev/null +++ b/src/computation_graph.cc @@ -0,0 +1,25 @@ +#include "computation_graph.h" + +OperationId ComputationGraph::add_operation(std::unique_ptr 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()); + 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(); +} diff --git a/src/computation_graph.h b/src/computation_graph.h new file mode 100644 index 000000000..55a4923a7 --- /dev/null +++ b/src/computation_graph.h @@ -0,0 +1,32 @@ +#ifndef ORCHESTRA_COMPUTATIONGRAPH_H +#define ORCHESTRA_COMPUTATIONGRAPH_H + +#include +#include + +#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::max(); +// used to represent the absence of an operation +const OperationId NO_OPERATION = std::numeric_limits::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); + // 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 > operations_; + // spawned_operations_[operationid] is a vector of the OperationIds of the + // operations spawned by the task with OperationId operationid + std::vector > spawned_operations_; +}; + +#endif diff --git a/src/scheduler.cc b/src/scheduler.cc index d76c8b337..3aeaea056 100644 --- a/src/scheduler.cc +++ b/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(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, 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, 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 computation_graph_lock(computation_graph_lock_); std::lock_guard fntable_lock(fntable_lock_); std::lock_guard avail_workers_lock(avail_workers_lock_); std::lock_guard 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 computation_graph_lock(computation_graph_lock_); std::lock_guard fntable_lock(fntable_lock_); std::lock_guard avail_workers_lock(avail_workers_lock_); std::lock_guard 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::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(); diff --git a/src/scheduler.h b/src/scheduler.h index db299ec9a..d0bdf8182 100644 --- a/src/scheduler.h +++ b/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, 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 &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 workers_; @@ -144,7 +150,7 @@ private: FnTable fntable_; std::mutex fntable_lock_; // List of pending tasks. - std::deque > task_queue_; + std::deque task_queue_; std::mutex task_queue_lock_; // List of pending pull calls. std::vector > pull_queue_;