Let worker get worker address and object store address from scheduler (#350)

src/ipc.cc

@@ -52,7 +52,7 @@ bool MessageQueue<>::connect(const std::string& name, bool create, size_t messag
     }
   }
   catch (bip::interprocess_exception &ex) {
-    RAY_CHECK(false, "boost::interprocess exception: " << ex.what());
+    RAY_CHECK(false, "name = " << name_ << ", create = " << create << ", boost::interprocess exception: " << ex.what());
   }
   return true;
 }

src/raylib.cc

@@ -642,19 +642,24 @@ static PyObject* deserialize_task(PyObject* worker_capsule, const Task& task) {
 // Ray Python API
 
 static PyObject* create_worker(PyObject* self, PyObject* args) {
-  const char* scheduler_addr;
-  const char* objstore_addr;
-  const char* worker_addr;
+  const char* node_ip_address;
+  const char* scheduler_address;
+  // The object store address can be the empty string, in which case the
+  // scheduler will choose the object store address.
+  const char* objstore_address;
   PyObject* is_driver_obj;
-  if (!PyArg_ParseTuple(args, "sssO", &scheduler_addr, &objstore_addr, &worker_addr, &is_driver_obj)) {
+  if (!PyArg_ParseTuple(args, "sssO", &node_ip_address, &scheduler_address, &objstore_address, &is_driver_obj)) {
     return NULL;
   }
   bool is_driver = PyObject_IsTrue(is_driver_obj);
-  auto scheduler_channel = grpc::CreateChannel(scheduler_addr, grpc::InsecureChannelCredentials());
-  auto objstore_channel = grpc::CreateChannel(objstore_addr, grpc::InsecureChannelCredentials());
-  Worker* worker = new Worker(std::string(worker_addr), scheduler_channel, objstore_channel);
-  worker->register_worker(std::string(worker_addr), std::string(objstore_addr), is_driver);
-  return PyCapsule_New(static_cast<void*>(worker), "worker", &WorkerCapsule_Destructor);
+  Worker* worker = new Worker(std::string(scheduler_address));
+  worker->register_worker(std::string(node_ip_address), std::string(objstore_address), is_driver);
+
+  PyObject* t = PyTuple_New(2);
+  PyObject* worker_capsule = PyCapsule_New(static_cast<void*>(worker), "worker", &WorkerCapsule_Destructor);
+  PyTuple_SetItem(t, 0, worker_capsule);
+  PyTuple_SetItem(t, 1, PyString_FromString(worker->get_worker_address()));
+  return t;
 }
 
 static PyObject* disconnect(PyObject* self, PyObject* args) {

src/scheduler.cc

@@ -215,12 +215,66 @@ Status SchedulerService::RegisterObjStore(ServerContext* context, const Register
 }
 
 Status SchedulerService::RegisterWorker(ServerContext* context, const RegisterWorkerRequest* request, RegisterWorkerReply* reply) {
-  std::pair<WorkerId, ObjStoreId> info = register_worker(request->worker_address(), request->objstore_address(), request->is_driver());
-  WorkerId workerid = info.first;
-  ObjStoreId objstoreid = info.second;
-  RAY_LOG(RAY_INFO, "registered worker with workerid " << workerid);
+  std::string objstore_address = request->objstore_address();
+  std::string node_ip_address = request->node_ip_address();
+  bool is_driver = request->is_driver();
+  RAY_LOG(RAY_INFO, "Registering a worker from node with IP address " << node_ip_address);
+  // Find the object store to connect to. We use the max size to indicate that
+  // the object store for this worker has not been found.
+  ObjStoreId objstoreid = std::numeric_limits<size_t>::max();
+  // TODO: HACK: num_attempts is a hack
+  for (int num_attempts = 0; num_attempts < 30; ++num_attempts) {
+    auto objstores = GET(objstores_);
+    for (size_t i = 0; i < objstores->size(); ++i) {
+      if (objstore_address != "" && (*objstores)[i].address == objstore_address) {
+        // This object store address is the same as the provided object store
+        // address.
+        objstoreid = i;
+      }
+      if ((*objstores)[i].address.compare(0, node_ip_address.size(), node_ip_address) == 0) {
+        // The object store address was not provided and this object store
+        // address has node_ip_address as a prefix, so it is on the same machine
+        // as the worker that is registering.
+        objstoreid = i;
+        objstore_address = (*objstores)[i].address;
+      }
+    }
+    if (objstoreid == std::numeric_limits<size_t>::max()) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(100));
+    } else {
+      break;
+    }
+  }
+  if (objstore_address.empty()) {
+    RAY_CHECK_NEQ(objstoreid, std::numeric_limits<size_t>::max(), "No object store with IP address " << node_ip_address << " has registered.");
+  } else {
+    RAY_CHECK_NEQ(objstoreid, std::numeric_limits<size_t>::max(), "Object store with address " << objstore_address << " not yet registered.");
+  }
+  // Populate the worker information and generate a worker address.
+  WorkerId workerid;
+  std::string worker_address;
+  {
+    auto workers = GET(workers_);
+    workerid = workers->size();
+    worker_address = node_ip_address + ":" + std::to_string(40000 + workerid);
+    workers->push_back(WorkerHandle());
+    auto channel = grpc::CreateChannel(worker_address, grpc::InsecureChannelCredentials());
+    (*workers)[workerid].channel = channel;
+    (*workers)[workerid].objstoreid = objstoreid;
+    (*workers)[workerid].worker_stub = WorkerService::NewStub(channel);
+    (*workers)[workerid].worker_address = worker_address;
+    (*workers)[workerid].initialized = false;
+    if (is_driver) {
+      (*workers)[workerid].current_task = ROOT_OPERATION; // We use this field to identify which workers are drivers.
+    } else {
+      (*workers)[workerid].current_task = NO_OPERATION;
+    }
+  }
+  RAY_LOG(RAY_INFO, "Finished registering worker with workerid " << workerid << ", worker address " << worker_address << " on node with IP address " << node_ip_address << ", is_driver = " << is_driver << ", assigned to object store with id " << objstoreid << " and address " << objstore_address);
   reply->set_workerid(workerid);
   reply->set_objstoreid(objstoreid);
+  reply->set_worker_address(worker_address);
+  reply->set_objstore_address(objstore_address);
   schedule();
   return Status::OK;
 }
@@ -540,42 +594,6 @@ bool SchedulerService::can_run(const Task& task) {
   return true;
 }
 
-std::pair<WorkerId, ObjStoreId> SchedulerService::register_worker(const std::string& worker_address, const std::string& objstore_address, bool is_driver) {
-  RAY_LOG(RAY_INFO, "registering worker " << worker_address << " connected to object store " << objstore_address);
-  ObjStoreId objstoreid = std::numeric_limits<size_t>::max();
-  // TODO: HACK: num_attempts is a hack
-  for (int num_attempts = 0; num_attempts < 30; ++num_attempts) {
-    auto objstores = GET(objstores_);
-    for (size_t i = 0; i < objstores->size(); ++i) {
-      if ((*objstores)[i].address == objstore_address) {
-        objstoreid = i;
-      }
-    }
-    if (objstoreid == std::numeric_limits<size_t>::max()) {
-      std::this_thread::sleep_for(std::chrono::milliseconds(100));
-    }
-  }
-  RAY_CHECK_NEQ(objstoreid, std::numeric_limits<size_t>::max(), "object store with address " << objstore_address << " not yet registered");
-  WorkerId workerid;
-  {
-    auto workers = GET(workers_);
-    workerid = workers->size();
-    workers->push_back(WorkerHandle());
-    auto channel = grpc::CreateChannel(worker_address, grpc::InsecureChannelCredentials());
-    (*workers)[workerid].channel = channel;
-    (*workers)[workerid].objstoreid = objstoreid;
-    (*workers)[workerid].worker_stub = WorkerService::NewStub(channel);
-    (*workers)[workerid].worker_address = worker_address;
-    (*workers)[workerid].initialized = false;
-    if (is_driver) {
-      (*workers)[workerid].current_task = ROOT_OPERATION; // We use this field to identify which workers are drivers.
-    } else {
-      (*workers)[workerid].current_task = NO_OPERATION;
-    }
-  }
-  return std::make_pair(workerid, objstoreid);
-}
-
 ObjectID SchedulerService::register_new_object() {
   // If we don't simultaneously lock objtable_ and target_objectids_, we will probably get errors.
   // TODO(rkn): increment/decrement_reference_count also acquire reference_counts_lock_ and target_objectids_lock_ (through has_canonical_objectid()), which caused deadlock in the past

src/scheduler.h

@@ -99,8 +99,6 @@ public:
   void assign_task(OperationId operationid, WorkerId workerid, const MySynchronizedPtr<ComputationGraph> &computation_graph);
   // 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)
-  std::pair<WorkerId, ObjStoreId> register_worker(const std::string& worker_address, const std::string& objstore_address, bool is_driver);
   // register a new object with the scheduler and return its object ID
   ObjectID register_new_object();
   // register the location of the object ID in the object table

src/worker.cc

@@ -56,11 +56,10 @@ Status WorkerServiceImpl::Die(ServerContext* context, const DieRequest* request,
   return Status::OK;
 }
 
-Worker::Worker(const std::string& worker_address, std::shared_ptr<Channel> scheduler_channel, std::shared_ptr<Channel> objstore_channel)
-    : worker_address_(worker_address),
-      scheduler_stub_(Scheduler::NewStub(scheduler_channel)) {
-  RAY_CHECK(receive_queue_.connect(worker_address_, true), "error connecting receive_queue_");
-  connected_ = true;
+Worker::Worker(const std::string& scheduler_address)
+    : scheduler_address_(scheduler_address) {
+  auto scheduler_channel = grpc::CreateChannel(scheduler_address, grpc::InsecureChannelCredentials());
+  scheduler_stub_ = Scheduler::NewStub(scheduler_channel);
 }
 
 SubmitTaskReply Worker::submit_task(SubmitTaskRequest* request, int max_retries, int retry_wait_milliseconds) {
@@ -87,10 +86,12 @@ bool Worker::kill_workers(ClientContext &context) {
   return reply.success();
 }
 
-void Worker::register_worker(const std::string& worker_address, const std::string& objstore_address, bool is_driver) {
+void Worker::register_worker(const std::string& node_ip_address, const std::string& objstore_address, bool is_driver) {
   unsigned int retry_wait_milliseconds = 20;
   RegisterWorkerRequest request;
-  request.set_worker_address(worker_address);
+  request.set_node_ip_address(node_ip_address);
+  // The object store address can be the empty string, in which case the
+  // scheduler will assign an object store address.
   request.set_objstore_address(objstore_address);
   request.set_is_driver(is_driver);
   RegisterWorkerReply reply;
@@ -108,9 +109,13 @@ void Worker::register_worker(const std::string& worker_address, const std::strin
   }
   workerid_ = reply.workerid();
   objstoreid_ = reply.objstoreid();
+  objstore_address_ = reply.objstore_address();
+  worker_address_ = reply.worker_address();
   segmentpool_ = std::make_shared<MemorySegmentPool>(objstoreid_, false);
-  RAY_CHECK(request_obj_queue_.connect(std::string("queue:") + objstore_address + std::string(":obj"), false), "error connecting request_obj_queue_");
-  RAY_CHECK(receive_obj_queue_.connect(std::string("queue:") + objstore_address + std::string(":worker:") + std::to_string(workerid_) + std::string(":obj"), true), "error connecting receive_obj_queue_");
+  RAY_CHECK(receive_queue_.connect(worker_address_, true), "error connecting receive_queue_");
+  RAY_CHECK(request_obj_queue_.connect(std::string("queue:") + objstore_address_ + std::string(":obj"), false), "error connecting request_obj_queue_");
+  RAY_CHECK(receive_obj_queue_.connect(std::string("queue:") + objstore_address_ + std::string(":worker:") + std::to_string(workerid_) + std::string(":obj"), true), "error connecting receive_obj_queue_");
+  connected_ = true;
   return;
 }
 

src/worker.h

@@ -37,7 +37,7 @@ private:
 
 class Worker {
  public:
-  Worker(const std::string& worker_address, std::shared_ptr<Channel> scheduler_channel, std::shared_ptr<Channel> objstore_channel);
+  Worker(const std::string& scheduler_address);
 
   // Submit a remote task to the scheduler. If the function in the task is not
   // registered with the scheduler, we will sleep for retry_wait_milliseconds
@@ -46,7 +46,7 @@ class Worker {
   // 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);
+  void register_worker(const std::string& ip_address, const std::string& objstore_address, bool is_driver);
   // get a new object ID that is registered with the scheduler
   ObjectID get_objectid();
   // request an object to be delivered to the local object store
@@ -94,6 +94,8 @@ class Worker {
   bool export_function(const std::string& function);
   // export reusable variable to workers
   void export_reusable_variable(const std::string& name, const std::string& initializer, const std::string& reinitializer);
+  // return the worker address
+  const char* get_worker_address() { return worker_address_.c_str(); }
 
  private:
   bool connected_;
@@ -104,6 +106,8 @@ class Worker {
   bip::managed_shared_memory segment_;
   WorkerId workerid_;
   ObjStoreId objstoreid_;
+  std::string scheduler_address_;
+  std::string objstore_address_;
   std::string worker_address_;
   MessageQueue<ObjRequest> request_obj_queue_;
   MessageQueue<ObjHandle> receive_obj_queue_;