Check for GRPC status codes (#366)

* check for GRPC status codes

* more checkss

include/ray/logging.h

@@ -2,6 +2,8 @@
 #include <iostream>
 #include <fstream>
 
+#include <grpc++/grpc++.h>
+
 struct RayConfig {
   bool log_to_file = false;
   std::ofstream logfile;
@@ -54,3 +56,9 @@ extern "C" __declspec(dllimport) int __stdcall IsDebuggerPresent();
 #define RAY_CHECK_LT(var1, var2, message) RAY_CHECK((var1) < (var2), message)
 #define RAY_CHECK_GE(var1, var2, message) RAY_CHECK((var1) >= (var2), message)
 #define RAY_CHECK_GT(var1, var2, message) RAY_CHECK((var1) > (var2), message)
+
+#define RAY_CHECK_GRPC(expr) \
+  do { \
+    grpc::Status _s = (expr); \
+    RAY_CHECK(_s.ok(), "grpc call failed with message " << _s.error_message()); \
+  } while (0);

src/objstore.cc

@@ -31,7 +31,7 @@ void ObjStoreService::get_data_from(ObjectID objectid, ObjStore::Stub& stub) {
     data += chunk.data().size();
     num_bytes += chunk.data().size();
   } while (reader->Read(&chunk));
-  Status status = reader->Finish(); // Right now we don't use the status.
+  RAY_CHECK_GRPC(reader->Finish());
 
   // finalize object
   RAY_CHECK_EQ(num_bytes, total_size, "Streamed objectid " << objectid << ", but num_bytes != total_size");
@@ -58,7 +58,7 @@ void ObjStoreService::register_objstore() {
   RegisterObjStoreRequest request;
   request.set_objstore_address(objstore_address_);
   RegisterObjStoreReply reply;
-  scheduler_stub_->RegisterObjStore(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->RegisterObjStore(&context, request, &reply));
   objstoreid_ = reply.objstoreid();
   segmentpool_ = std::make_shared<MemorySegmentPool>(objstoreid_, objstore_address_, true);
 }
@@ -321,7 +321,7 @@ void ObjStoreService::object_ready(ObjectID objectid, size_t metadata_offset) {
   objready_request.set_objectid(objectid);
   objready_request.set_objstoreid(objstoreid_);
   AckReply objready_reply;
-  scheduler_stub_->ObjReady(&objready_context, objready_request, &objready_reply);
+  RAY_CHECK_GRPC(scheduler_stub_->ObjReady(&objready_context, objready_request, &objready_reply));
 }
 
 void ObjStoreService::start_objstore_service() {

src/scheduler.cc

@@ -328,7 +328,7 @@ Status SchedulerService::NotifyFailure(ServerContext* context, const NotifyFailu
         PrintErrorMessageRequest print_request;
         print_request.mutable_failure()->CopyFrom(request->failure());
         AckReply print_reply;
-        Status status = worker->worker_stub->PrintErrorMessage(&client_context, print_request, &print_reply);
+        RAY_CHECK_GRPC(worker->worker_stub->PrintErrorMessage(&client_context, print_request, &print_reply));
       }
     }
   }
@@ -500,7 +500,7 @@ Status SchedulerService::KillWorkers(ServerContext* context, const KillWorkersRe
       DieRequest die_request;
       AckReply 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);
+      RAY_CHECK_GRPC(idle_worker->worker_stub->Die(&client_context, die_request, &die_reply));
       idle_worker->worker_stub.reset();
     }
     avail_workers->clear();
@@ -580,7 +580,7 @@ void SchedulerService::deliver_object_async(ObjectID canonical_objectid, ObjStor
   request.set_objectid(canonical_objectid);
   auto objstores = GET(objstores_);
   request.set_objstore_address((*objstores)[from].address);
-  (*objstores)[to].objstore_stub->StartDelivery(&context, request, &reply);
+  RAY_CHECK_GRPC((*objstores)[to].objstore_stub->StartDelivery(&context, request, &reply));
 }
 
 void SchedulerService::schedule() {
@@ -622,7 +622,7 @@ void SchedulerService::assign_task(OperationId operationid, WorkerId workerid, c
     auto workers = GET(workers_);
     (*workers)[workerid].current_task = operationid;
     request.mutable_task()->CopyFrom(task); // TODO(rkn): Is ownership handled properly here?
-    Status status = (*workers)[workerid].worker_stub->ExecuteTask(&context, request, &reply);
+    RAY_CHECK_GRPC((*workers)[workerid].worker_stub->ExecuteTask(&context, request, &reply));
   }
 }
 
@@ -935,7 +935,7 @@ bool SchedulerService::attempt_notify_alias(ObjStoreId objstoreid, ObjectID alia
   NotifyAliasRequest request;
   request.set_alias_objectid(alias_objectid);
   request.set_canonical_objectid(canonical_objectid);
-  (*GET(objstores_))[objstoreid].objstore_stub->NotifyAlias(&context, request, &reply);
+  RAY_CHECK_GRPC((*GET(objstores_))[objstoreid].objstore_stub->NotifyAlias(&context, request, &reply));
   return true;
 }
 
@@ -957,7 +957,7 @@ void SchedulerService::deallocate_object(ObjectID canonical_objectid, const MySy
       request.set_canonical_objectid(canonical_objectid);
       ObjStoreId objstoreid = locations[i];
       RAY_LOG(RAY_REFCOUNT, "Attempting to deallocate canonical_objectid " << canonical_objectid << " from objstore " << objstoreid);
-      (*objstores)[objstoreid].objstore_stub->DeallocateObject(&context, request, &reply);
+      RAY_CHECK_GRPC((*objstores)[objstoreid].objstore_stub->DeallocateObject(&context, request, &reply));
     }
     locations.clear();
   }
@@ -1035,7 +1035,7 @@ void SchedulerService::export_function_to_worker(WorkerId workerid, int function
   ImportRemoteFunctionRequest import_request;
   import_request.mutable_function()->CopyFrom(*(*exported_functions)[function_index].get());
   AckReply import_reply;
-  (*workers)[workerid].worker_stub->ImportRemoteFunction(&import_context, import_request, &import_reply);
+  RAY_CHECK_GRPC((*workers)[workerid].worker_stub->ImportRemoteFunction(&import_context, import_request, &import_reply));
 }
 
 void SchedulerService::export_reusable_variable_to_worker(WorkerId workerid, int reusable_variable_index, MySynchronizedPtr<std::vector<WorkerHandle> > &workers, const MySynchronizedPtr<std::vector<std::unique_ptr<ReusableVar> > > &exported_reusable_variables) {
@@ -1044,7 +1044,7 @@ void SchedulerService::export_reusable_variable_to_worker(WorkerId workerid, int
   ImportReusableVariableRequest import_request;
   import_request.mutable_reusable_variable()->CopyFrom(*(*exported_reusable_variables)[reusable_variable_index].get());
   AckReply import_reply;
-  (*workers)[workerid].worker_stub->ImportReusableVariable(&import_context, import_request, &import_reply);
+  RAY_CHECK_GRPC((*workers)[workerid].worker_stub->ImportReusableVariable(&import_context, import_request, &import_reply));
 }
 
 void SchedulerService::export_all_functions_to_worker(WorkerId workerid, MySynchronizedPtr<std::vector<WorkerHandle> > &workers, const MySynchronizedPtr<std::vector<std::unique_ptr<Function> > > &exported_functions) {

src/worker.cc

@@ -107,11 +107,10 @@ Worker::Worker(const std::string& node_ip_address, const std::string& scheduler_
 SubmitTaskReply Worker::submit_task(SubmitTaskRequest* request, int max_retries, int retry_wait_milliseconds) {
   RAY_CHECK(connected_, "Attempted to perform submit_task but failed.");
   SubmitTaskReply reply;
-  Status status;
   request->set_workerid(workerid_);
   for (int i = 0; i < 1 + max_retries; ++i) {
     ClientContext context;
-    status = scheduler_stub_->SubmitTask(&context, *request, &reply);
+    RAY_CHECK_GRPC(scheduler_stub_->SubmitTask(&context, *request, &reply));
     if (reply.function_registered()) {
       break;
     }
@@ -124,7 +123,7 @@ SubmitTaskReply Worker::submit_task(SubmitTaskRequest* request, int max_retries,
 bool Worker::kill_workers(ClientContext &context) {
   KillWorkersRequest request;
   KillWorkersReply reply;
-  Status status = scheduler_stub_->KillWorkers(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->KillWorkers(&context, request, &reply));
   return reply.success();
 }
 
@@ -138,18 +137,19 @@ void Worker::register_worker(const std::string& node_ip_address, const std::stri
   request.set_objstore_address(objstore_address);
   request.set_is_driver(is_driver);
   RegisterWorkerReply reply;
-  grpc::StatusCode status_code = grpc::UNAVAILABLE;
+  Status status;
   // TODO: HACK: retrying is a hack
   for (int i = 0; i < 5; ++i) {
     ClientContext context;
-    status_code = scheduler_stub_->RegisterWorker(&context, request, &reply).error_code();
-    if (status_code != grpc::UNAVAILABLE) {
+    status = scheduler_stub_->RegisterWorker(&context, request, &reply);
+    if (status.error_code() != grpc::UNAVAILABLE) {
       break;
     }
     // Note that each pass through the loop may take substantially longer than
     // retry_wait_milliseconds because grpc may do its own retrying.
     std::this_thread::sleep_for(std::chrono::milliseconds(retry_wait_milliseconds));
   }
+  RAY_CHECK_GRPC(status);
   workerid_ = reply.workerid();
   objstoreid_ = reply.objstoreid();
   objstore_address_ = reply.objstore_address();
@@ -173,7 +173,7 @@ void Worker::request_object(ObjectID objectid) {
   request.set_objectid(objectid);
   AckReply reply;
   ClientContext context;
-  Status status = scheduler_stub_->RequestObj(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->RequestObj(&context, request, &reply));
   return;
 }
 
@@ -184,7 +184,7 @@ ObjectID Worker::get_objectid() {
   request.set_workerid(workerid_);
   PutObjReply reply;
   ClientContext context;
-  Status status = scheduler_stub_->PutObj(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->PutObj(&context, request, &reply));
   return reply.objectid();
 }
 
@@ -240,7 +240,7 @@ void Worker::put_object(ObjectID objectid, const Obj* obj, std::vector<ObjectID>
   }
   AckReply reply;
   ClientContext context;
-  scheduler_stub_->AddContainedObjectIDs(&context, contained_objectids_request, &reply);
+   RAY_CHECK_GRPC(scheduler_stub_->AddContainedObjectIDs(&context, contained_objectids_request, &reply));
 }
 
 #define CHECK_ARROW_STATUS(s, msg)                              \
@@ -322,7 +322,7 @@ void Worker::alias_objectids(ObjectID alias_objectid, ObjectID target_objectid)
   request.set_alias_objectid(alias_objectid);
   request.set_target_objectid(target_objectid);
   AckReply reply;
-  scheduler_stub_->AliasObjectIDs(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->AliasObjectIDs(&context, request, &reply));
 }
 
 void Worker::increment_reference_count(std::vector<ObjectID> &objectids) {
@@ -338,7 +338,7 @@ void Worker::increment_reference_count(std::vector<ObjectID> &objectids) {
       request.add_objectid(objectids[i]);
     }
     AckReply reply;
-    scheduler_stub_->IncrementRefCount(&context, request, &reply);
+    RAY_CHECK_GRPC(scheduler_stub_->IncrementRefCount(&context, request, &reply));
   }
 }
 
@@ -355,7 +355,7 @@ void Worker::decrement_reference_count(std::vector<ObjectID> &objectids) {
       request.add_objectid(objectids[i]);
     }
     AckReply reply;
-    scheduler_stub_->DecrementRefCount(&context, request, &reply);
+    RAY_CHECK_GRPC(scheduler_stub_->DecrementRefCount(&context, request, &reply));
   }
 }
 
@@ -367,7 +367,7 @@ void Worker::register_remote_function(const std::string& name, size_t num_return
   request.set_function_name(name);
   request.set_num_return_vals(num_return_vals);
   AckReply reply;
-  scheduler_stub_->RegisterRemoteFunction(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->RegisterRemoteFunction(&context, request, &reply));
 }
 
 void Worker::notify_failure(FailedType type, const std::string& name, const std::string& error_message) {
@@ -380,7 +380,7 @@ void Worker::notify_failure(FailedType type, const std::string& name, const std:
   request.mutable_failure()->set_name(name);
   request.mutable_failure()->set_error_message(error_message);
   AckReply reply;
-  scheduler_stub_->NotifyFailure(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->NotifyFailure(&context, request, &reply));
 }
 
 std::unique_ptr<WorkerMessage> Worker::receive_next_message() {
@@ -395,7 +395,7 @@ void Worker::ready_for_new_task() {
   ReadyForNewTaskRequest request;
   request.set_workerid(workerid_);
   AckReply reply;
-  scheduler_stub_->ReadyForNewTask(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->ReadyForNewTask(&context, request, &reply));
 }
 
 void Worker::disconnect() {
@@ -410,12 +410,12 @@ void Worker::disconnect() {
 // TODO(rkn): Should we be using pointers or references? And should they be const?
 void Worker::scheduler_info(ClientContext &context, SchedulerInfoRequest &request, SchedulerInfoReply &reply) {
   RAY_CHECK(connected_, "Attempted to get scheduler info but failed.");
-  scheduler_stub_->SchedulerInfo(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->SchedulerInfo(&context, request, &reply));
 }
 
 void Worker::task_info(ClientContext &context, TaskInfoRequest &request, TaskInfoReply &reply) {
   RAY_CHECK(connected_, "Attempted to get worker info but failed.");
-  scheduler_stub_->TaskInfo(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->TaskInfo(&context, request, &reply));
 }
 
 bool Worker::export_remote_function(const std::string& function_name, const std::string& function) {
@@ -425,7 +425,7 @@ bool Worker::export_remote_function(const std::string& function_name, const std:
   request.mutable_function()->set_name(function_name);
   request.mutable_function()->set_implementation(function);
   AckReply reply;
-  Status status = scheduler_stub_->ExportRemoteFunction(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->ExportRemoteFunction(&context, request, &reply));
   return true;
 }
 
@@ -437,7 +437,7 @@ void Worker::export_reusable_variable(const std::string& name, const std::string
   request.mutable_reusable_variable()->mutable_initializer()->set_implementation(initializer);
   request.mutable_reusable_variable()->mutable_reinitializer()->set_implementation(reinitializer);
   AckReply reply;
-  Status status = scheduler_stub_->ExportReusableVariable(&context, request, &reply);
+  RAY_CHECK_GRPC(scheduler_stub_->ExportReusableVariable(&context, request, &reply));
 }
 
 // Communication between the WorkerServer and the Worker happens via a message