Kill workers (#148)

lib/python/ray/__init__.py

@@ -8,6 +8,6 @@ PYTHON_MODE = 3
 
 import libraylib as lib
 import serialization
-from worker import scheduler_info, dump_computation_graph, task_info, register_module, connect, disconnect, get, put, remote
+from worker import scheduler_info, dump_computation_graph, task_info, register_module, connect, disconnect, get, put, remote, kill_workers
 from libraylib import ObjRef
 import internal

lib/python/ray/worker.py

@@ -193,6 +193,12 @@ def put(value, worker=global_worker):
     print_task_info(ray.lib.task_info(worker.handle), worker.mode)
   return objref
 
+def kill_workers(worker=global_worker):
+  success = ray.lib.kill_workers(worker.handle)
+  if not success:
+    print "Could not kill all workers; check that there are no tasks currently running."
+  return success
+
 def main_loop(worker=global_worker):
   if not ray.lib.connected(worker.handle):
     raise Exception("Worker is attempting to enter main_loop but has not been connected yet.")
@@ -216,6 +222,11 @@ def main_loop(worker=global_worker):
       ray.lib.notify_task_completed(worker.handle, True, "") # notify the scheduler that the task completed successfully
   while True:
     task = ray.lib.wait_for_next_task(worker.handle)
+    if task is None:
+      # We use this as a mechanism to allow the scheduler to kill workers. When
+      # the scheduler wants to kill a worker, it gives the worker a null task,
+      # causing the worker program to exit the main loop here.
+      break
     process_task(task)
 
 def remote(arg_types, return_types, worker=global_worker):

protos/ray.proto

@@ -50,6 +50,8 @@ service Scheduler {
   rpc SchedulerInfo(SchedulerInfoRequest) returns (SchedulerInfoReply);
   // Get information about tasks
   rpc TaskInfo(TaskInfoRequest) returns (TaskInfoReply);
+  // Kills the workers
+  rpc KillWorkers(KillWorkersRequest) returns (KillWorkersReply);
 }
 
 message AckReply {
@@ -219,6 +221,13 @@ message TaskInfoReply {
   // TODO(mehrdadn): We'll want to return information from computation_graph since it's important for visualizing tasks that have been completed etc.
 }
 
+message KillWorkersRequest {
+}
+
+message KillWorkersReply {
+  bool success = 1;  // Currently, the only reason to fail is if there are workers still executing tasks
+}
+
 // These messages are for getting information about the object store state
 
 message ObjStoreInfoRequest {
@@ -234,6 +243,7 @@ message ObjStoreInfoReply {
 
 service WorkerService {
   rpc ExecuteTask(ExecuteTaskRequest) returns (ExecuteTaskReply); // Scheduler calls a function from the worker
+  rpc Die(DieRequest) returns (DieReply); // Kills this worker
 }
 
 message ExecuteTaskRequest {
@@ -241,5 +251,10 @@ message ExecuteTaskRequest {
 }
 
 message ExecuteTaskReply {
+}
 
+message DieRequest {
+}
+
+message DieReply {
 }

src/ipc.cc

@@ -48,11 +48,11 @@ void MemorySegmentPool::open_segment(SegmentId segmentid, size_t size) {
   std::string segment_name = get_segment_name(segmentid);
   if (create_mode_) {
     assert(size > 0);
-    shared_memory_object::remove(segment_name.c_str()); // remove segment if it has not been properly removed from last run
+    bip::shared_memory_object::remove(segment_name.c_str()); // remove segment if it has not been properly removed from last run
     size_t new_size = (size / page_size_ + 2) * page_size_; // additional room for boost's bookkeeping
-    segments_[segmentid] = std::make_pair(std::unique_ptr<managed_shared_memory>(new managed_shared_memory(create_only, segment_name.c_str(), new_size)), SegmentStatusType::OPENED);
+    segments_[segmentid] = std::make_pair(std::unique_ptr<bip::managed_shared_memory>(new bip::managed_shared_memory(bip::create_only, segment_name.c_str(), new_size)), SegmentStatusType::OPENED);
   } else {
-    segments_[segmentid] = std::make_pair(std::unique_ptr<managed_shared_memory>(new managed_shared_memory(open_only, segment_name.c_str())), SegmentStatusType::OPENED);
+    segments_[segmentid] = std::make_pair(std::unique_ptr<bip::managed_shared_memory>(new bip::managed_shared_memory(bip::open_only, segment_name.c_str())), SegmentStatusType::OPENED);
   }
 }
 
@@ -64,7 +64,7 @@ void MemorySegmentPool::unmap_segment(SegmentId segmentid) {
 void MemorySegmentPool::close_segment(SegmentId segmentid) {
   RAY_LOG(RAY_DEBUG, "closing segmentid " << segmentid);
   std::string segment_name = get_segment_name(segmentid);
-  shared_memory_object::remove(segment_name.c_str());
+  bip::shared_memory_object::remove(segment_name.c_str());
   segments_[segmentid].first.reset();
   segments_[segmentid].second = SegmentStatusType::CLOSED;
 }
@@ -93,7 +93,7 @@ uint8_t* MemorySegmentPool::get_address(ObjHandle pointer) {
   if (!create_mode_) {
     open_segment(pointer.segmentid());
   }
-  managed_shared_memory* segment = segments_[pointer.segmentid()].first.get();
+  bip::managed_shared_memory* segment = segments_[pointer.segmentid()].first.get();
   return static_cast<uint8_t*>(segment->get_address_from_handle(pointer.ipcpointer()));
 }
 
@@ -106,6 +106,6 @@ MemorySegmentPool::~MemorySegmentPool() {
   for (size_t segmentid = 0; segmentid < segments_.size(); ++segmentid) {
     std::string segment_name = get_segment_name(segmentid);
     segments_[segmentid].first.reset();
-    shared_memory_object::remove(segment_name.c_str());
+    bip::shared_memory_object::remove(segment_name.c_str());
   }
 }

src/ipc.h

@@ -12,7 +12,7 @@
 
 #include "ray/ray.h"
 
-using namespace boost::interprocess;
+namespace bip = boost::interprocess;
 
 // Methods for inter process communication (abstracts from the shared memory implementation)
 
@@ -24,7 +24,7 @@ public:
   MessageQueue() {};
 
   ~MessageQueue() {
-    message_queue::remove(name_.c_str());
+    bip::message_queue::remove(name_.c_str());
   }
 
   MessageQueue(MessageQueue<T>&& other) noexcept
@@ -36,12 +36,12 @@ public:
     name_ = name;
     try {
       if (create) {
-        message_queue::remove(name.c_str()); // remove queue if it has not been properly removed from last run
-        queue_ = std::unique_ptr<message_queue>(new message_queue(create_only, name.c_str(), 100, sizeof(T)));
+        bip::message_queue::remove(name.c_str()); // remove queue if it has not been properly removed from last run
+        queue_ = std::unique_ptr<bip::message_queue>(new bip::message_queue(bip::create_only, name.c_str(), 100, sizeof(T)));
       } else {
-        queue_ = std::unique_ptr<message_queue>(new message_queue(open_only, name.c_str()));
+        queue_ = std::unique_ptr<bip::message_queue>(new bip::message_queue(bip::open_only, name.c_str()));
       }
-    } catch(interprocess_exception &ex) {
+    } catch(bip::interprocess_exception &ex) {
       RAY_CHECK(false, "boost::interprocess exception: " << ex.what());
     }
     return true;
@@ -54,7 +54,7 @@ public:
   bool send(const T* object) {
     try {
       queue_->send(object, sizeof(T), 0);
-    } catch(interprocess_exception &ex) {
+    } catch(bip::interprocess_exception &ex) {
       RAY_CHECK(false, "boost::interprocess exception: " << ex.what());
     }
     return true;
@@ -62,10 +62,10 @@ public:
 
   bool receive(T* object) {
     unsigned int priority;
-    message_queue::size_type recvd_size;
+    bip::message_queue::size_type recvd_size;
     try {
       queue_->receive(object, sizeof(T), recvd_size, priority);
-    } catch(interprocess_exception &ex) {
+    } catch(bip::interprocess_exception &ex) {
       RAY_CHECK(false, "boost::interprocess exception: " << ex.what());
     }
     return true;
@@ -73,7 +73,7 @@ public:
 
 private:
   std::string name_;
-  std::unique_ptr<message_queue> queue_;
+  std::unique_ptr<bip::message_queue> queue_;
 };
 
 // Object Queues
@@ -101,7 +101,7 @@ struct ObjRequest {
 };
 
 typedef size_t SegmentId; // index into a memory segment table
-typedef managed_shared_memory::handle_t IpcPointer;
+typedef bip::managed_shared_memory::handle_t IpcPointer;
 
 // Object handle: Handle to object that can be passed around between processes
 // that are connected to the same object store
@@ -158,8 +158,8 @@ private:
   void close_segment(SegmentId segmentid); // close a segment
   bool create_mode_; // true in the object stores, false on the workers
   ObjStoreId objstoreid_; // the identity of the associated object store
-  size_t page_size_ = mapped_region::get_page_size();
-  std::vector<std::pair<std::unique_ptr<managed_shared_memory>, SegmentStatusType> > segments_;
+  size_t page_size_ = bip::mapped_region::get_page_size();
+  std::vector<std::pair<std::unique_ptr<bip::managed_shared_memory>, SegmentStatusType> > segments_;
 };
 
 #endif

src/raylib.cc

@@ -631,8 +631,12 @@ PyObject* wait_for_next_task(PyObject* self, PyObject* args) {
   if (!PyArg_ParseTuple(args, "O&", &PyObjectToWorker, &worker)) {
     return NULL;
   }
-  Task* task = worker->receive_next_task();
-  return PyCapsule_New(static_cast<void*>(task), "task", NULL); // This task is owned by the C++ worker class, so we do not deallocate it.
+  if (std::unique_ptr<Task> task = worker->receive_next_task()) {
+    PyObject* pyobj = PyCapsule_New(task.get(), "task", TaskCapsule_Destructor);
+    task.release(); // Now that the wrapper object was constructed successfully, release ownership
+    return pyobj;
+  }
+  Py_RETURN_NONE;
 }
 
 PyObject* submit_task(PyObject* self, PyObject* args) {
@@ -852,6 +856,19 @@ PyObject* set_log_config(PyObject* self, PyObject* args) {
   Py_RETURN_NONE;
 }
 
+PyObject* kill_workers(PyObject* self, PyObject* args) {
+  Worker* worker;
+  if (!PyArg_ParseTuple(args, "O&", &PyObjectToWorker, &worker)) {
+    return NULL;
+  }
+  ClientContext context;
+  if (worker->kill_workers(context)) {
+    Py_RETURN_TRUE;
+  } else {
+    Py_RETURN_FALSE;
+  }
+}
+
 static PyMethodDef RayLibMethods[] = {
  { "serialize_object", serialize_object, METH_VARARGS, "serialize an object to protocol buffers" },
  { "deserialize_object", deserialize_object, METH_VARARGS, "deserialize an object from protocol buffers" },
@@ -878,6 +895,7 @@ static PyMethodDef RayLibMethods[] = {
  { "task_info", task_info, METH_VARARGS, "get task statuses" },
  { "dump_computation_graph", dump_computation_graph, METH_VARARGS, "dump the current computation graph to a file" },
  { "set_log_config", set_log_config, METH_VARARGS, "set filename for raylib logging" },
+ { "kill_workers", kill_workers, METH_VARARGS, "kills all of the workers" },
  { NULL, NULL, 0, NULL }
 };
 

src/scheduler.cc

@@ -129,7 +129,7 @@ Status SchedulerService::ObjReady(ServerContext* context, const ObjReadyRequest*
     // If this is the first time that ObjReady has been called for this objref,
     // the corresponding increment was done in register_new_object in the
     // scheduler. For all subsequent calls to ObjReady, the corresponding
-    // increment was done in deliver_object_if_necessary in the scheduler.
+    // increment was done in deliver_object_async_if_necessary in the scheduler.
     auto reference_counts = reference_counts_.get(); // we grab this lock because decrement_ref_count assumes it has been acquired
     auto contained_objrefs = contained_objrefs_.get(); // we grab this lock because decrement_ref_count assumes it has been acquired
     decrement_ref_count(std::vector<ObjRef>({objref}));
@@ -222,7 +222,7 @@ Status SchedulerService::TaskInfo(ServerContext* context, const TaskInfoRequest*
     TaskStatus* info = reply->add_failed_task();
     *info = (*failed_tasks)[i];
   }
-  for (int i = 0; i < workers->size(); ++i) {
+  for (size_t i = 0; i < workers->size(); ++i) {
     OperationId operationid = (*workers)[i].current_task;
     if (operationid != NO_OPERATION && operationid != ROOT_OPERATION) {
       const Task& task = computation_graph->get_task(operationid);
@@ -236,7 +236,55 @@ Status SchedulerService::TaskInfo(ServerContext* context, const TaskInfoRequest*
   return Status::OK;
 }
 
-void SchedulerService::deliver_object_if_necessary(ObjRef canonical_objref, ObjStoreId from, ObjStoreId to) {
+Status SchedulerService::KillWorkers(ServerContext* context, const KillWorkersRequest* request, KillWorkersReply* reply) {
+  // TODO: Update reference counts
+  auto failed_tasks = failed_tasks_.get();
+  auto get_queue = get_queue_.get();
+  auto computation_graph = computation_graph_.get();
+  auto fntable = fntable_.get();
+  auto avail_workers = avail_workers_.get();
+  auto task_queue = task_queue_.get();
+  auto workers = workers_.get();
+  size_t busy_workers = 0;
+  std::vector<WorkerHandle*> idle_workers;
+  RAY_LOG(RAY_INFO, "Attempting to kill workers.");
+  for (size_t i = 0; i < workers->size(); ++i) {
+    WorkerHandle* worker = &(*workers)[i];
+    if (worker->worker_stub) {
+      if (worker->current_task == NO_OPERATION) {
+        idle_workers.push_back(worker);
+        RAY_CHECK(std::find(avail_workers->begin(), avail_workers->end(), i) != avail_workers->end(), "Worker with workerid " << i << " is idle, but is not in avail_workers_");
+        RAY_LOG(RAY_INFO, "Worker with workerid " << i << " is idle.");
+      } else if (worker->current_task == ROOT_OPERATION) {
+        // Skip the driver
+        RAY_LOG(RAY_INFO, "Worker with workerid " << i << " is a driver.");
+      } else {
+        ++busy_workers;
+        RAY_LOG(RAY_INFO, "Worker with workerid " << i << " is running a task.");
+      }
+    }
+  }
+  if (task_queue->empty() && busy_workers == 0) {
+    RAY_LOG(RAY_INFO, "Killing " << idle_workers.size() << " idle workers.");
+    for (WorkerHandle* idle_worker : idle_workers) {
+      ClientContext client_context;
+      DieRequest die_request;
+      DieReply die_reply;
+      // TODO: Fault handling... what if a worker refuses to die? We just assume it dies here.
+      idle_worker->worker_stub->Die(&client_context, die_request, &die_reply);
+      idle_worker->worker_stub.reset();
+    }
+    avail_workers->clear();
+    fntable->clear();
+    reply->set_success(true);
+  } else {
+    RAY_LOG(RAY_INFO, "Either the task queue is not empty or there are still busy workers, so we are not killing any workers.");
+    reply->set_success(false);
+  }
+  return Status::OK;
+}
+
+void SchedulerService::deliver_object_async_if_necessary(ObjRef canonical_objref, ObjStoreId from, ObjStoreId to) {
   bool object_present_or_in_transit;
   {
     auto objtable = objtable_.get();
@@ -250,7 +298,7 @@ void SchedulerService::deliver_object_if_necessary(ObjRef canonical_objref, ObjS
     }
   }
   if (!object_present_or_in_transit) {
-    deliver_object(canonical_objref, from, to);
+    deliver_object_async(canonical_objref, from, to);
   }
 }
 
@@ -260,8 +308,8 @@ void SchedulerService::deliver_object_if_necessary(ObjRef canonical_objref, ObjS
 // delivery once. However, we may want to handle it in the scheduler in the
 // future.
 //
-// deliver_object assumes that the aliasing for objref has already been completed. That is, has_canonical_objref(objref) == true
-void SchedulerService::deliver_object(ObjRef canonical_objref, ObjStoreId from, ObjStoreId to) {
+// deliver_object_async assumes that the aliasing for objref has already been completed. That is, has_canonical_objref(objref) == true
+void SchedulerService::deliver_object_async(ObjRef canonical_objref, ObjStoreId from, ObjStoreId to) {
   RAY_CHECK_NEQ(from, to, "attempting to deliver canonical_objref " << canonical_objref << " from objstore " << from << " to itself.");
   RAY_CHECK(is_canonical(canonical_objref), "attempting to deliver objref " << canonical_objref << ", but this objref is not a canonical objref.");
   {
@@ -310,7 +358,7 @@ void SchedulerService::assign_task(OperationId operationid, WorkerId workerid) {
       alias_notification_queue_.get()->push_back(std::make_pair(objstoreid, std::make_pair(objref, canonical_objref)));
       attempt_notify_alias(objstoreid, objref, canonical_objref);
       RAY_LOG(RAY_DEBUG, "task contains object ref " << canonical_objref);
-      deliver_object_if_necessary(canonical_objref, pick_objstore(canonical_objref), objstoreid);
+      deliver_object_async_if_necessary(canonical_objref, pick_objstore(canonical_objref), objstoreid);
     }
   }
   {
@@ -496,12 +544,12 @@ void SchedulerService::perform_gets() {
       continue;
     }
     ObjRef canonical_objref = get_canonical_objref(objref);
-    RAY_LOG(RAY_DEBUG, "attempting to get objref " << get.second << " with canonical objref " << canonical_objref << " to objstore " << get_store(workerid));
+    RAY_LOG(RAY_DEBUG, "attempting to get objref " << get.second << " with canonical objref " << canonical_objref << " to objstore " << objstoreid);
     int num_stores = (*objtable_.get())[canonical_objref].size();
     if (num_stores > 0) {
-      deliver_object_if_necessary(canonical_objref, pick_objstore(canonical_objref), objstoreid);
+      deliver_object_async_if_necessary(canonical_objref, pick_objstore(canonical_objref), objstoreid);
       // Notify the relevant objstore about potential aliasing when it's ready
-      alias_notification_queue_.get()->push_back(std::make_pair(get_store(workerid), std::make_pair(objref, canonical_objref)));
+      alias_notification_queue_.get()->push_back(std::make_pair(objstoreid, std::make_pair(objref, canonical_objref)));
       // Remove the get task from the queue
       std::swap((*get_queue)[i], (*get_queue)[get_queue->size() - 1]);
       get_queue->pop_back();

src/scheduler.h

@@ -34,7 +34,7 @@ const RefCount DEALLOCATED = std::numeric_limits<RefCount>::max();
 
 struct WorkerHandle {
   std::shared_ptr<Channel> channel;
-  std::unique_ptr<WorkerService::Stub> worker_stub;
+  std::unique_ptr<WorkerService::Stub> worker_stub; // If null, the worker has died
   ObjStoreId objstoreid;
   std::string worker_address;
   OperationId current_task;
@@ -69,13 +69,14 @@ public:
   Status AddContainedObjRefs(ServerContext* context, const AddContainedObjRefsRequest* request, AckReply* reply) override;
   Status SchedulerInfo(ServerContext* context, const SchedulerInfoRequest* request, SchedulerInfoReply* reply) override;
   Status TaskInfo(ServerContext* context, const TaskInfoRequest* request, TaskInfoReply* reply) override;
+  Status KillWorkers(ServerContext* context, const KillWorkersRequest* request, KillWorkersReply* reply) override;
 
   // This will ask an object store to send an object to another object store if
   // the object is not already present in that object store and is not already
   // being transmitted.
-  void deliver_object_if_necessary(ObjRef objref, ObjStoreId from, ObjStoreId to);
+  void deliver_object_async_if_necessary(ObjRef objref, ObjStoreId from, ObjStoreId to);
   // ask an object store to send object to another object store
-  void deliver_object(ObjRef objref, ObjStoreId from, ObjStoreId to);
+  void deliver_object_async(ObjRef objref, ObjStoreId from, ObjStoreId to);
   // assign a task to a worker
   void schedule();
   // execute a task on a worker and ship required object references
@@ -157,7 +158,7 @@ private:
   // For each object store objstoreid, objects_in_transit_[objstoreid] is a
   // vector of the canonical object references that are being streamed to that
   // object store but are not yet present. Object references are added to this
-  // in deliver_object_if_necessary (to ensure that we do not attempt to deliver
+  // in deliver_object_async_if_necessary (to ensure that we do not attempt to deliver
   // the same object to a given object store twice), and object references are
   // removed when add_location is called (from ObjReady), and they are moved to
   // the objtable_. Note that objects_in_transit_ and objtable_ share the same

src/worker.cc

@@ -11,11 +11,22 @@ extern "C" {
   static PyObject *RayError;
 }
 
+inline WorkerServiceImpl::WorkerServiceImpl(const std::string& worker_address)
+  : worker_address_(worker_address) {
+  send_queue_.connect(worker_address_, false);
+}
+
 Status WorkerServiceImpl::ExecuteTask(ServerContext* context, const ExecuteTaskRequest* request, ExecuteTaskReply* reply) {
-  task_ = request->task(); // Copy task
+  task_ = std::unique_ptr<Task>(new Task(request->task())); // Copy task
   RAY_LOG(RAY_INFO, "invoked task " << request->task().name());
-  Task* taskptr = &task_;
-  send_queue_.send(&taskptr);
+  WorkerMessage message = { &task_ };
+  send_queue_.send(&message);
+  return Status::OK;
+}
+
+Status WorkerServiceImpl::Die(ServerContext* context, const DieRequest* request, DieReply* reply) {
+  WorkerMessage message = { NULL };
+  send_queue_.send(&message);
   return Status::OK;
 }
 
@@ -43,6 +54,13 @@ SubmitTaskReply Worker::submit_task(SubmitTaskRequest* request, int max_retries,
   return reply;
 }
 
+bool Worker::kill_workers(ClientContext &context) {
+  KillWorkersRequest request;
+  KillWorkersReply reply;
+  Status status = scheduler_stub_->KillWorkers(&context, request, &reply);
+  return reply.success();
+}
+
 void Worker::register_worker(const std::string& worker_address, const std::string& objstore_address, bool is_driver) {
   RegisterWorkerRequest request;
   request.set_worker_address(worker_address);
@@ -268,10 +286,10 @@ void Worker::register_function(const std::string& name, size_t num_return_vals)
   scheduler_stub_->RegisterFunction(&context, request, &reply);
 }
 
-Task* Worker::receive_next_task() {
-  Task* task;
-  receive_queue_.receive(&task);
-  return task;
+std::unique_ptr<Task> Worker::receive_next_task() {
+  WorkerMessage message;
+  receive_queue_.receive(&message);
+  return message.task ? std::move(*message.task) : std::unique_ptr<Task>();
 }
 
 void Worker::notify_task_completed(bool task_succeeded, std::string error_message) {

src/worker.h

@@ -23,17 +23,20 @@ using grpc::Channel;
 using grpc::ClientContext;
 using grpc::ClientWriter;
 
+struct WorkerMessage {
+  std::unique_ptr<Task>* task;
+};
+static_assert(std::is_pod<WorkerMessage>::value, "WorkerMessage must be memcpy-able");
+
 class WorkerServiceImpl final : public WorkerService::Service {
 public:
-  WorkerServiceImpl(const std::string& worker_address)
-    : worker_address_(worker_address) {
-    send_queue_.connect(worker_address_, false);
-  }
+  WorkerServiceImpl(const std::string& worker_address);
   Status ExecuteTask(ServerContext* context, const ExecuteTaskRequest* request, ExecuteTaskReply* reply) override;
+  Status Die(ServerContext* context, const DieRequest* request, DieReply* reply) override;
 private:
   std::string worker_address_;
-  Task task_; // copy of the current task
-  MessageQueue<Task*> send_queue_;
+  std::unique_ptr<Task> task_; // copy of the current task
+  MessageQueue<WorkerMessage> send_queue_;
 };
 
 class Worker {
@@ -44,6 +47,8 @@ class Worker {
   // registered with the scheduler, we will sleep for retry_wait_milliseconds
   // and try to resubmit the task to the scheduler up to max_retries more times.
   SubmitTaskReply submit_task(SubmitTaskRequest* request, int max_retries = 120, int retry_wait_milliseconds = 500);
+  // Requests the scheduler to kill workers
+  bool kill_workers(ClientContext &context);
   // send request to the scheduler to register this worker
   void register_worker(const std::string& worker_address, const std::string& objstore_address, bool is_driver);
   // get a new object reference that is registered with the scheduler
@@ -73,8 +78,8 @@ class Worker {
   // 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
-  Task* receive_next_task();
+  // wait for next task from the RPC system. If null, it means there are no more tasks and the worker should shut down.
+  std::unique_ptr<Task> receive_next_task();
   // tell the scheduler that we are done with the current task and request the
   // next one, if task_succeeded is false, this tells the scheduler that the
   // task threw an exception
@@ -93,8 +98,8 @@ class Worker {
   const size_t CHUNK_SIZE = 8 * 1024;
   std::unique_ptr<Scheduler::Stub> scheduler_stub_;
   std::thread worker_server_thread_;
-  MessageQueue<Task*> receive_queue_;
-  managed_shared_memory segment_;
+  MessageQueue<WorkerMessage> receive_queue_;
+  bip::managed_shared_memory segment_;
   WorkerId workerid_;
   ObjStoreId objstoreid_;
   std::string worker_address_;