catch exceptions on workers and pass them to the scheduler (#93)

lib/python/ray/worker.py

@@ -90,9 +90,13 @@ def main_loop(worker=global_worker):
   def process_task(task): # wrapping these lines in a function should cause the local variables to go out of scope more quickly, which is useful for inspecting reference counts
     func_name, args, return_objrefs = serialization.deserialize_task(worker.handle, task)
     arguments = get_arguments_for_execution(worker.functions[func_name], args, worker) # get args from objstore
-    outputs = worker.functions[func_name].executor(arguments) # execute the function
-    store_outputs_in_objstore(return_objrefs, outputs, worker) # store output in local object store
-    ray.lib.notify_task_completed(worker.handle) # notify the scheduler that the task has completed
+    try:
+      outputs = worker.functions[func_name].executor(arguments) # execute the function
+    except Exception as e:
+      ray.lib.notify_task_completed(worker.handle, False, str(e)) # notify the scheduler that the task threw an exception
+    else:
+      store_outputs_in_objstore(return_objrefs, outputs, worker) # store output in local object store
+      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)
     process_task(task)

protos/ray.proto

@@ -44,7 +44,7 @@ service Scheduler {
   // Used by the worker to notify the scheduler about which objrefs a particular object contains
   rpc AddContainedObjRefs(AddContainedObjRefsRequest) returns (AckReply);
   // Used by the worker to report back and ask for more work
-  rpc WorkerReady(WorkerReadyRequest) returns (AckReply);
+  rpc NotifyTaskCompleted(NotifyTaskCompletedRequest) returns (AckReply);
   // Get information about the scheduler state
   rpc SchedulerInfo(SchedulerInfoRequest) returns (SchedulerInfoReply);
 }
@@ -120,8 +120,10 @@ message DecrementRefCountRequest {
   repeated uint64 objref = 1; // Object references whose reference count should be decremented. Duplicates will be decremented multiple times.
 }
 
-message WorkerReadyRequest {
-  uint64 workerid = 1; // ID of the worker which is ready
+message NotifyTaskCompletedRequest {
+  uint64 workerid = 1; // ID of the worker which executed the task
+  bool task_succeeded = 2; // True if the task succeeded, false if it threw an exception
+  string error_message = 3; // The contents of the exception, if the task threw an exception
 }
 
 message ChangeCountRequest {

src/raylib.cc

@@ -662,10 +662,14 @@ PyObject* submit_task(PyObject* self, PyObject* args) {
 
 PyObject* notify_task_completed(PyObject* self, PyObject* args) {
   Worker* worker;
-  if (!PyArg_ParseTuple(args, "O&", &PyObjectToWorker, &worker)) {
+  PyObject* task_succeeded_obj;
+  const char* error_message_ptr;
+  if (!PyArg_ParseTuple(args, "O&Os", &PyObjectToWorker, &worker, &task_succeeded_obj, &error_message_ptr)) {
     return NULL;
   }
-  worker->notify_task_completed();
+  std::string error_message(error_message_ptr);
+  bool task_succeeded = PyObject_IsTrue(task_succeeded_obj);
+  worker->notify_task_completed(task_succeeded, error_message);
   Py_RETURN_NONE;
 }
 

src/scheduler.cc

@@ -144,12 +144,15 @@ Status SchedulerService::ObjReady(ServerContext* context, const ObjReadyRequest*
   return Status::OK;
 }
 
-Status SchedulerService::WorkerReady(ServerContext* context, const WorkerReadyRequest* request, AckReply* reply) {
+Status SchedulerService::NotifyTaskCompleted(ServerContext* context, const NotifyTaskCompletedRequest* request, AckReply* reply) {
   RAY_LOG(RAY_INFO, "worker " << request->workerid() << " reported back");
   {
     std::lock_guard<std::mutex> lock(avail_workers_lock_);
     avail_workers_.push_back(request->workerid());
   }
+  if (!request->task_succeeded()) {
+    RAY_LOG(RAY_FATAL, "The task on worker " << request->workerid() << " threw an exception with the following error message: " << request->error_message());
+  }
   schedule();
   return Status::OK;
 }

src/scheduler.h

@@ -60,7 +60,7 @@ public:
   Status RegisterWorker(ServerContext* context, const RegisterWorkerRequest* request, RegisterWorkerReply* reply) override;
   Status RegisterFunction(ServerContext* context, const RegisterFunctionRequest* request, AckReply* reply) override;
   Status ObjReady(ServerContext* context, const ObjReadyRequest* request, AckReply* reply) override;
-  Status WorkerReady(ServerContext* context, const WorkerReadyRequest* request, AckReply* reply) override;
+  Status NotifyTaskCompleted(ServerContext* context, const NotifyTaskCompletedRequest* request, AckReply* reply) override;
   Status IncrementRefCount(ServerContext* context, const IncrementRefCountRequest* request, AckReply* reply) override;
   Status DecrementRefCount(ServerContext* context, const DecrementRefCountRequest* request, AckReply* reply) override;
   Status AddContainedObjRefs(ServerContext* context, const AddContainedObjRefsRequest* request, AckReply* reply) override;

src/worker.cc

@@ -263,15 +263,17 @@ Task* Worker::receive_next_task() {
   return task;
 }
 
-void Worker::notify_task_completed() {
+void Worker::notify_task_completed(bool task_succeeded, std::string error_message) {
   if (!connected_) {
     RAY_LOG(RAY_FATAL, "Attempting to perform notify_task_completed, but connected_ = " << connected_ << ".");
   }
   ClientContext context;
-  WorkerReadyRequest request;
+  NotifyTaskCompletedRequest request;
   request.set_workerid(workerid_);
+  request.set_task_succeeded(task_succeeded);
+  request.set_error_message(error_message);
   AckReply reply;
-  scheduler_stub_->WorkerReady(&context, request, &reply);
+  scheduler_stub_->NotifyTaskCompleted(&context, request, &reply);
 }
 
 void Worker::disconnect() {

src/worker.h

@@ -71,8 +71,10 @@ class Worker {
   void start_worker_service();
   // wait for next task from the RPC system
   Task* receive_next_task();
-  // tell the scheduler that we are done with the current task and request the next one
-  void notify_task_completed();
+  // 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
+  void notify_task_completed(bool task_succeeded, std::string error_message);
   // disconnect the worker
   void disconnect();
   // return connected_

test/test_functions.py

@@ -78,3 +78,9 @@ try:
   varargs_and_kwargs_exception_thrown = False
 except:
   varargs_and_kwargs_exception_thrown = True
+
+# test throwing an exception
+
+@ray.remote([], [])
+def throw_exception_fct():
+    raise Exception("Test function intentionally failed.")