clean up, mainly the scheduler

src/orchlib.cc

@@ -12,8 +12,12 @@ size_t orch_remote_call(Worker* worker, RemoteCallRequest* request) {
   return worker->remote_call(request);
 }
 
-Call* orch_main_loop(Worker* worker) {
-  return worker->main_loop();
+void orch_start_worker_service(Worker* worker) {
+  worker->start_worker_service();
+}
+
+Call* orch_wait_for_next_task(Worker* worker) {
+  return worker->receive_next_task();
 }
 
 size_t orch_push(Worker* worker, Obj* obj) {

src/orchlib.h

@@ -1,3 +1,6 @@
+// A minimal C API that is used for implementing Orchestra workers in C based
+// languages (Python at the moment, in the future potentially Julia, R, MATLAB)
+
 extern "C" {
 
 struct slice {
@@ -9,10 +12,14 @@ struct Worker;
 struct RemoteCallRequest;
 struct Value;
 
+// connect to the scheduler and the object store
 Worker* orch_create_context(const char* server_addr, const char* worker_addr, const char* objstore_addr);
-size_t orch_remote_call(Worker* context, RemoteCallRequest* request);
-size_t orch_push(Worker* context, Obj* value);
-Call* orch_main_loop(Worker* worker);
+// start the worker service for this worker
+void orch_start_worker_service(Worker* worker);
+// Submit a function call to the scheduler
+size_t orch_remote_call(Worker* worker, RemoteCallRequest* request);
+size_t orch_push(Worker* worker, Obj* value);
+Call* orch_wait_for_next_task(Worker* worker);
 slice orch_get_serialized_obj(Worker* worker, size_t objref);
 void orch_register_function(Worker* worker, const char* name, size_t num_return_vals);
 

src/scheduler.cc

@@ -1,17 +1,54 @@
 #include "scheduler.h"
 
-size_t Scheduler::add_task(const Call& task) {
+Status SchedulerService::RemoteCall(ServerContext* context, const RemoteCallRequest* request, RemoteCallReply* reply) {
+  std::unique_ptr<Call> task(new Call(request->call())); // need to copy, because request is const
   fntable_lock_.lock();
-  size_t num_return_vals = fntable_[task.name()].num_return_vals();
+  size_t num_return_vals = fntable_[task->name()].num_return_vals();
   fntable_lock_.unlock();
-  std::unique_ptr<Call> task_ptr(new Call(task));
+
+  for (size_t i = 0; i < num_return_vals; ++i) {
+    ObjRef result = register_new_object();
+    reply->add_result(result);
+    task->add_result(result);
+  }
+
   tasks_lock_.lock();
-  tasks_.emplace_back(std::move(task_ptr));
+  tasks_.emplace_back(std::move(task));
   tasks_lock_.unlock();
-  return num_return_vals;
+  return Status::OK;
 }
 
-void Scheduler::schedule() {
+Status SchedulerService::PushObj(ServerContext* context, const PushObjRequest* request, PushObjReply* reply) {
+  ObjRef objref = register_new_object();
+  ObjStoreId objstoreid = get_store(request->workerid());
+  add_location(objref, objstoreid);
+  reply->set_objref(objref);
+  return Status::OK;
+}
+
+Status SchedulerService::PullObj(ServerContext* context, const PullObjRequest* request, AckReply* reply) {
+  return Status::OK;
+}
+
+Status SchedulerService::RegisterWorker(ServerContext* context, const RegisterWorkerRequest* request, RegisterWorkerReply* reply) {
+  WorkerId workerid = register_worker(request->worker_address(), request->objstore_address());
+  std::cout << "registered worker with workerid" << workerid << std::endl;
+  reply->set_workerid(workerid);
+  return Status::OK;
+}
+
+Status SchedulerService::RegisterFunction(ServerContext* context, const RegisterFunctionRequest* request, AckReply* reply) {
+  std::cout << "RegisterFunction: workerid is" << request->workerid() << std::endl;
+  register_function(request->fnname(), request->workerid(), request->num_return_vals());
+  return Status::OK;
+}
+
+Status SchedulerService::GetDebugInfo(ServerContext* context, const GetDebugInfoRequest* request, GetDebugInfoReply* reply) {
+  debug_info(*request, reply);
+  return Status::OK;
+}
+
+void SchedulerService::schedule() {
   // TODO: work out a better strategy here
   WorkerId workerid = 0;
   {
@@ -34,7 +71,7 @@ void Scheduler::schedule() {
   }
 }
 
-void Scheduler::submit_task(std::unique_ptr<Call> call, WorkerId workerid) {
+void SchedulerService::submit_task(std::unique_ptr<Call> call, WorkerId workerid) {
   ClientContext context;
   InvokeCallRequest request;
   InvokeCallReply reply;
@@ -59,7 +96,7 @@ void Scheduler::submit_task(std::unique_ptr<Call> call, WorkerId workerid) {
   Status status = workers_[workerid].worker_stub->InvokeCall(&context, request, &reply);
 }
 
-bool Scheduler::can_run(const Call& task) {
+bool SchedulerService::can_run(const Call& task) {
   std::lock_guard<std::mutex> lock(objtable_lock_);
   for (int i = 0; i < task.arg_size(); ++i) {
     if (!task.arg(i).has_obj()) {
@@ -71,7 +108,7 @@ bool Scheduler::can_run(const Call& task) {
   return true;
 }
 
-WorkerId Scheduler::register_worker(const std::string& worker_address, const std::string& objstore_address) {
+WorkerId SchedulerService::register_worker(const std::string& worker_address, const std::string& objstore_address) {
   ObjStoreId objstoreid = std::numeric_limits<size_t>::max();
   objstores_lock_.lock();
   for (size_t i = 0; i < objstores_.size(); ++i) {
@@ -97,7 +134,7 @@ WorkerId Scheduler::register_worker(const std::string& worker_address, const std
   auto channel = grpc::CreateChannel(worker_address, grpc::InsecureChannelCredentials());
   workers_[workerid].channel = channel;
   workers_[workerid].objstoreid = objstoreid;
-  workers_[workerid].worker_stub = WorkerServer::NewStub(channel);
+  workers_[workerid].worker_stub = WorkerService::NewStub(channel);
   workers_lock_.unlock();
   avail_workers_lock_.lock();
   avail_workers_.push_back(workerid);
@@ -105,7 +142,7 @@ WorkerId Scheduler::register_worker(const std::string& worker_address, const std
   return workerid;
 }
 
-ObjRef Scheduler::register_new_object() {
+ObjRef SchedulerService::register_new_object() {
   objtable_lock_.lock();
   ObjRef result = objtable_.size();
   objtable_.push_back(std::vector<ObjStoreId>());
@@ -113,7 +150,7 @@ ObjRef Scheduler::register_new_object() {
   return result;
 }
 
-void Scheduler::add_location(ObjRef objref, ObjStoreId objstoreid) {
+void SchedulerService::add_location(ObjRef objref, ObjStoreId objstoreid) {
   objtable_lock_.lock();
   // do a binary search
   auto pos = std::lower_bound(objtable_[objref].begin(), objtable_[objref].end(), objstoreid);
@@ -123,14 +160,14 @@ void Scheduler::add_location(ObjRef objref, ObjStoreId objstoreid) {
   objtable_lock_.unlock();
 }
 
-ObjStoreId Scheduler::get_store(WorkerId workerid) {
+ObjStoreId SchedulerService::get_store(WorkerId workerid) {
   workers_lock_.lock();
   ObjStoreId result = workers_[workerid].objstoreid;
   workers_lock_.unlock();
   return result;
 }
 
-void Scheduler::register_function(const std::string& name, WorkerId workerid, size_t num_return_vals) {
+void SchedulerService::register_function(const std::string& name, WorkerId workerid, size_t num_return_vals) {
   fntable_lock_.lock();
   FnInfo& info = fntable_[name];
   info.set_num_return_vals(num_return_vals);
@@ -138,7 +175,7 @@ void Scheduler::register_function(const std::string& name, WorkerId workerid, si
   fntable_lock_.unlock();
 }
 
-void Scheduler::debug_info(const GetDebugInfoRequest& request, GetDebugInfoReply* reply) {
+void SchedulerService::debug_info(const GetDebugInfoRequest& request, GetDebugInfoReply* reply) {
   if (request.do_scheduling()) {
     schedule();
   }
@@ -163,3 +200,19 @@ void Scheduler::debug_info(const GetDebugInfoRequest& request, GetDebugInfoReply
   }
   avail_workers_lock_.unlock();
 }
+
+void start_scheduler_service(const char* server_address) {
+  SchedulerService service;
+  ServerBuilder builder;
+  builder.AddListeningPort(std::string(server_address), grpc::InsecureServerCredentials());
+  builder.RegisterService(&service);
+  std::unique_ptr<Server> server(builder.BuildAndStart());
+  server->Wait();
+}
+
+int main(int argc, char** argv) {
+  if (argc != 2)
+    return 1;
+  start_scheduler_service(argv[1]);
+  return 0;
+}

src/scheduler.h

@@ -1,6 +1,7 @@
 #ifndef ORCHESTRA_SCHEDULER_H
 #define ORCHESTRA_SCHEDULER_H
 
+
 #include <deque>
 #include <memory>
 #include <algorithm>
@@ -24,7 +25,7 @@ using grpc::Channel;
 
 struct WorkerHandle {
   std::shared_ptr<Channel> channel;
-  std::unique_ptr<WorkerServer::Stub> worker_stub;
+  std::unique_ptr<WorkerService::Stub> worker_stub;
   ObjStoreId objstoreid;
 };
 
@@ -34,10 +35,15 @@ struct ObjStoreHandle {
   std::string address;
 };
 
-class Scheduler {
+class SchedulerService : public Scheduler::Service {
 public:
-  // returns number of return values of task
-  size_t add_task(const Call& task);
+  Status RemoteCall(ServerContext* context, const RemoteCallRequest* request, RemoteCallReply* reply) override;
+  Status PushObj(ServerContext* context, const PushObjRequest* request, PushObjReply* reply) override;
+  Status PullObj(ServerContext* context, const PullObjRequest* request, AckReply* reply) override;
+  Status RegisterWorker(ServerContext* context, const RegisterWorkerRequest* request, RegisterWorkerReply* reply) override;
+  Status RegisterFunction(ServerContext* context, const RegisterFunctionRequest* request, AckReply* reply) override;
+  Status GetDebugInfo(ServerContext* context, const GetDebugInfoRequest* request, GetDebugInfoReply* reply) override;
+
   // assign a task to a worker
   void schedule();
   // execute a task on a worker and ship required object references

src/worker.cc

@@ -92,29 +92,26 @@ void Worker::register_function(const std::string& name, size_t num_return_vals)
   scheduler_stub_->RegisterFunction(&context, request, &reply);
 }
 
-void start_worker_server(const char* server_address) {
-  WorkerServiceImpl service(server_address);
-  ServerBuilder builder;
-  builder.AddListeningPort(server_address, grpc::InsecureServerCredentials());
-  builder.RegisterService(&service);
-  std::unique_ptr<Server> server(builder.BuildAndStart());
-  std::cout << "Server listening on " << server_address << std::endl;
-  server->Wait();
-}
-
 // Communication between the WorkerServer and the Worker happens via a message
 // queue. This is because the Python interpreter needs to be single threaded
 // (in our case running in the main thread), whereas the WorkerService will
 // run in a separate thread and potentially utilize multiple threads.
-Call* Worker::main_loop() {
-  // start the worker server
-  worker_server_thread_ = std::thread(start_worker_server, worker_address_.c_str());
-  // process the next call
-  return receive(worker_address_.c_str());
+void Worker::start_worker_service() {
+  const char* server_address = worker_address_.c_str();
+  worker_server_thread_ = std::thread([server_address]() {
+    WorkerServiceImpl service(server_address);
+    ServerBuilder builder;
+    builder.AddListeningPort(server_address, grpc::InsecureServerCredentials());
+    builder.RegisterService(&service);
+    std::unique_ptr<Server> server(builder.BuildAndStart());
+    std::cout << "WorkerServer listening on " << server_address << std::endl;
+    server->Wait();
+  });
 }
 
-Call* receive(const char* message_queue_name) {
-	try {
+Call* Worker::receive_next_task() {
+  const char* message_queue_name = worker_address_.c_str();
+  try {
     message_queue::remove(message_queue_name);
     message_queue mq(create_only, message_queue_name, 1, sizeof(Call*));
     unsigned int priority;

src/worker.h

@@ -27,7 +27,7 @@ using grpc::Channel;
 using grpc::ClientContext;
 using grpc::ClientWriter;
 
-class WorkerServiceImpl final : public WorkerServer::Service {
+class WorkerServiceImpl final : public WorkerService::Service {
 public:
   WorkerServiceImpl(const std::string& worker_address)
     : worker_address_(worker_address) {}
@@ -37,15 +37,11 @@ private:
   Call call_; // copy of the current call
 };
 
-void start_worker_server(const char* worker_addr);
-
-Call* receive(const char* worker_addr);
-
 class Worker {
  public:
   Worker(const std::string& worker_address, std::shared_ptr<Channel> scheduler_channel, std::shared_ptr<Channel> objstore_channel)
       : worker_address_(worker_address),
-        scheduler_stub_(SchedulerServer::NewStub(scheduler_channel)),
+        scheduler_stub_(Scheduler::NewStub(scheduler_channel)),
         objstore_stub_(ObjStore::NewStub(objstore_channel))
     {}
 
@@ -59,12 +55,15 @@ class Worker {
   slice get_serialized_obj(ObjRef objref);
   // register function with scheduler
   void register_function(const std::string& name, size_t num_return_vals);
-  // start the main loop
-  Call* main_loop();
+  // start the worker server which accepts tasks from the scheduler and stores
+  // it in the message queue, which is read by the Python interpreter
+  void start_worker_service();
+  // wait for next task from the RPC system
+  Call* receive_next_task();
 
  private:
   const size_t CHUNK_SIZE = 8 * 1024;
-  std::unique_ptr<SchedulerServer::Stub> scheduler_stub_;
+  std::unique_ptr<Scheduler::Stub> scheduler_stub_;
   std::unique_ptr<ObjStore::Stub> objstore_stub_;
   std::thread worker_server_thread_;
   std::thread other_thread_;