Automatically add relevant directories to Python paths of workers (#380)

* Make ray.init set python paths of workers.

* Decouple starting cluster from copying user source code

* also add current directory to path

* Add comments about deallocation.

* Add test for new code path.

src/raylib.cc

@@ -718,7 +718,8 @@ static PyObject* wait_for_next_message(PyObject* self, PyObject* args) {
     bool task_present = !message->task().name().empty();
     bool function_present = !message->function().implementation().empty();
     bool reusable_variable_present = !message->reusable_variable().name().empty();
-    RAY_CHECK(task_present + function_present + reusable_variable_present <= 1, "The worker message should contain at most one item.");
+    bool function_to_run_present = !message->function_to_run().implementation().empty();
+    RAY_CHECK(task_present + function_present + reusable_variable_present + function_to_run_present <= 1, "The worker message should contain at most one item.");
     PyObject* t = PyTuple_New(2);
     if (task_present) {
       PyTuple_SetItem(t, 0, PyString_FromString("task"));
@@ -736,6 +737,9 @@ static PyObject* wait_for_next_message(PyObject* self, PyObject* args) {
       PyTuple_SetItem(reusable_variable, 1, PyString_FromStringAndSize(message->reusable_variable().initializer().implementation().data(), static_cast<ssize_t>(message->reusable_variable().initializer().implementation().size())));
       PyTuple_SetItem(reusable_variable, 2, PyString_FromStringAndSize(message->reusable_variable().reinitializer().implementation().data(), static_cast<ssize_t>(message->reusable_variable().reinitializer().implementation().size())));
       PyTuple_SetItem(t, 1, reusable_variable);
+    } else if (function_to_run_present) {
+      PyTuple_SetItem(t, 0, PyString_FromString("function_to_run"));
+      PyTuple_SetItem(t, 1, PyString_FromStringAndSize(message->function_to_run().implementation().data(), static_cast<ssize_t>(message->function_to_run().implementation().size())));
     } else {
       PyTuple_SetItem(t, 0, PyString_FromString("die"));
       Py_INCREF(Py_None);
@@ -747,6 +751,17 @@ static PyObject* wait_for_next_message(PyObject* self, PyObject* args) {
   Py_RETURN_NONE;
 }
 
+static PyObject* run_function_on_all_workers(PyObject* self, PyObject* args) {
+  Worker* worker;
+  const char* function;
+  int function_size;
+  if (!PyArg_ParseTuple(args, "O&s#", &PyObjectToWorker, &worker, &function, &function_size)) {
+    return NULL;
+  }
+  worker->run_function_on_all_workers(std::string(function, static_cast<size_t>(function_size)));
+  Py_RETURN_NONE;
+}
+
 static PyObject* export_remote_function(PyObject* self, PyObject* args) {
   Worker* worker;
   const char* function_name;
@@ -1088,6 +1103,7 @@ static PyMethodDef RayLibMethods[] = {
  { "ready_for_new_task", ready_for_new_task, METH_VARARGS, "notify the scheduler that the worker is ready for a new task" },
  { "scheduler_info", scheduler_info, METH_VARARGS, "get info about scheduler state" },
  { "task_info", task_info, METH_VARARGS, "get information about task statuses and failures" },
+ { "run_function_on_all_workers", run_function_on_all_workers, METH_VARARGS, "run an arbitrary function on all workers" },
  { "export_remote_function", export_remote_function, METH_VARARGS, "export a remote function to workers" },
  { "export_reusable_variable", export_reusable_variable, METH_VARARGS, "export a reusable variable to the workers" },
  { "dump_computation_graph", dump_computation_graph, METH_VARARGS, "dump the current computation graph to a file" },

src/scheduler.cc

@@ -365,6 +365,8 @@ Status SchedulerService::ReadyForNewTask(ServerContext* context, const ReadyForN
       // all of the exported functions and all of the exported reusable variables.
       if (!(*workers)[workerid].initialized) {
         // This should only happen once.
+        // Queue up all functions to run on the worker.
+        add_all_functions_to_run_to_worker_queue(workerid);
         // Queue up all remote functions to be imported on the worker.
         add_all_remote_functions_to_worker_export_queue(workerid);
         // Queue up all reusable variables to be imported on the worker.
@@ -513,6 +515,22 @@ Status SchedulerService::KillWorkers(ServerContext* context, const KillWorkersRe
   return Status::OK;
 }
 
+Status SchedulerService::RunFunctionOnAllWorkers(ServerContext* context, const RunFunctionOnAllWorkersRequest* request, AckReply* reply) {
+  {
+    auto workers = GET(workers_);
+    auto function_to_run_queue = GET(function_to_run_queue_);
+    auto exported_functions_to_run = GET(exported_functions_to_run_);
+    exported_functions_to_run->push_back(std::unique_ptr<Function>(new Function(request->function())));
+    for (WorkerId workerid = 0; workerid < workers->size(); ++workerid) {
+      if ((*workers)[workerid].current_task != ROOT_OPERATION) {
+        function_to_run_queue->push(std::make_pair(workerid, exported_functions_to_run->size() - 1));
+      }
+    }
+  }
+  schedule();
+  return Status::OK;
+}
+
 Status SchedulerService::ExportRemoteFunction(ServerContext* context, const ExportRemoteFunctionRequest* request, AckReply* reply) {
   {
     auto workers = GET(workers_);
@@ -609,6 +627,10 @@ void SchedulerService::deliver_object_async(ObjectID canonical_objectid, ObjStor
 }
 
 void SchedulerService::schedule() {
+  // Run functions on workers. This must happen before we schedule tasks in
+  // order to guarantee that remote function calls use the most up to date
+  // environment.
+  perform_functions_to_run();
   // Export remote functions to the workers. This must happen before we schedule
   // tasks in order to guarantee that remote function calls use the most up to
   // date definitions.
@@ -801,6 +823,17 @@ bool SchedulerService::is_canonical(ObjectID objectid) {
   return objectid == (*target_objectids)[objectid];
 }
 
+void SchedulerService::perform_functions_to_run() {
+  auto workers = GET(workers_);
+  auto function_to_run_queue = GET(function_to_run_queue_);
+  auto exported_functions_to_run = GET(exported_functions_to_run_);
+  while (!function_to_run_queue->empty()) {
+    std::pair<WorkerId, int> workerid_functionid_pair = function_to_run_queue->front();
+    export_function_to_run_to_worker(workerid_functionid_pair.first, workerid_functionid_pair.second, workers, exported_functions_to_run);
+    function_to_run_queue->pop();
+  }
+}
+
 void SchedulerService::perform_remote_function_exports() {
   auto workers = GET(workers_);
   auto remote_function_export_queue = GET(remote_function_export_queue_);
@@ -1084,22 +1117,39 @@ void SchedulerService::get_equivalent_objectids(ObjectID objectid, std::vector<O
 }
 
 
+void SchedulerService::export_function_to_run_to_worker(WorkerId workerid, int function_index, MySynchronizedPtr<std::vector<WorkerHandle> > &workers, const MySynchronizedPtr<std::vector<std::unique_ptr<Function> > > &exported_functions_to_run) {
+  RAY_LOG(RAY_INFO, "exporting function to run with index " << function_index << " to worker " << workerid);
+  ClientContext context;
+  RunFunctionOnWorkerRequest request;
+  request.mutable_function()->CopyFrom(*(*exported_functions_to_run)[function_index].get());
+  AckReply reply;
+  RAY_CHECK_GRPC((*workers)[workerid].worker_stub->RunFunctionOnWorker(&context, request, &reply));
+}
+
 void SchedulerService::export_function_to_worker(WorkerId workerid, int function_index, MySynchronizedPtr<std::vector<WorkerHandle> > &workers, const MySynchronizedPtr<std::vector<std::unique_ptr<Function> > > &exported_functions) {
-  RAY_LOG(RAY_INFO, "exporting function with index " << function_index << " to worker " << workerid);
-  ClientContext import_context;
-  ImportRemoteFunctionRequest import_request;
-  import_request.mutable_function()->CopyFrom(*(*exported_functions)[function_index].get());
-  AckReply import_reply;
-  RAY_CHECK_GRPC((*workers)[workerid].worker_stub->ImportRemoteFunction(&import_context, import_request, &import_reply));
+  RAY_LOG(RAY_INFO, "exporting remote function with index " << function_index << " to worker " << workerid);
+  ClientContext context;
+  ImportRemoteFunctionRequest request;
+  request.mutable_function()->CopyFrom(*(*exported_functions)[function_index].get());
+  AckReply reply;
+  RAY_CHECK_GRPC((*workers)[workerid].worker_stub->ImportRemoteFunction(&context, request, &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) {
   RAY_LOG(RAY_INFO, "exporting reusable variable with index " << reusable_variable_index << " to worker " << workerid);
-  ClientContext import_context;
-  ImportReusableVariableRequest import_request;
-  import_request.mutable_reusable_variable()->CopyFrom(*(*exported_reusable_variables)[reusable_variable_index].get());
-  AckReply import_reply;
-  RAY_CHECK_GRPC((*workers)[workerid].worker_stub->ImportReusableVariable(&import_context, import_request, &import_reply));
+  ClientContext context;
+  ImportReusableVariableRequest request;
+  request.mutable_reusable_variable()->CopyFrom(*(*exported_reusable_variables)[reusable_variable_index].get());
+  AckReply reply;
+  RAY_CHECK_GRPC((*workers)[workerid].worker_stub->ImportReusableVariable(&context, request, &reply));
+}
+
+void SchedulerService::add_all_functions_to_run_to_worker_queue(WorkerId workerid) {
+  auto function_to_run_queue = GET(function_to_run_queue_);
+  auto exported_functions_to_run = GET(exported_functions_to_run_);
+  for (int i = 0; i < exported_functions_to_run->size(); ++i) {
+    function_to_run_queue->push(std::make_pair(workerid, i));
+  }
 }
 
 void SchedulerService::add_all_remote_functions_to_worker_export_queue(WorkerId workerid) {

src/scheduler.h

@@ -75,6 +75,7 @@ public:
   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;
+  Status RunFunctionOnAllWorkers(ServerContext* context, const RunFunctionOnAllWorkersRequest* request, AckReply* reply) override;
   Status ExportRemoteFunction(ServerContext* context, const ExportRemoteFunctionRequest* request, AckReply* reply) override;
   Status ExportReusableVariable(ServerContext* context, const ExportReusableVariableRequest* request, AckReply* reply) override;
   Status NotifyFailure(ServerContext*, const NotifyFailureRequest* request, AckReply* reply) override;
@@ -119,10 +120,13 @@ private:
   ObjStoreId pick_objstore(ObjectID objectid);
   // checks if objectid is a canonical objectid
   bool is_canonical(ObjectID objectid);
+  // Export all queued up functions to run.
+  void perform_functions_to_run();
   // Export all queued up remote functions.
   void perform_remote_function_exports();
   // Export all queued up reusable variables.
   void perform_reusable_variable_exports();
+  // Perform all queued up gets that can be performed.
   void perform_gets();
   // schedule tasks using the naive algorithm
   void schedule_tasks_naively();
@@ -148,10 +152,15 @@ private:
   void upstream_objectids(ObjectID objectid, std::vector<ObjectID> &objectids, const MySynchronizedPtr<std::vector<std::vector<ObjectID> > > &reverse_target_objectids);
   // Find all of the object IDs that refer to the same object as objectid (as best as we can determine at the moment). The information may be incomplete because not all of the aliases may be known.
   void get_equivalent_objectids(ObjectID objectid, std::vector<ObjectID> &equivalent_objectids);
+  // Export a function to run to a worker.
+  void export_function_to_run_to_worker(WorkerId workerid, int function_index, MySynchronizedPtr<std::vector<WorkerHandle> > &workers, const MySynchronizedPtr<std::vector<std::unique_ptr<Function> > > &exported_functions_to_run);
   // Export a remote function to a worker.
   void export_function_to_worker(WorkerId workerid, int function_index, MySynchronizedPtr<std::vector<WorkerHandle> > &workers, const MySynchronizedPtr<std::vector<std::unique_ptr<Function> > > &exported_functions);
   // Export a reusable variable to a worker
   void 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);
+  // Add to the function to run export queue the job of exporting all functions
+  // to run to the given worker. This is used when a new worker registers.
+  void add_all_functions_to_run_to_worker_queue(WorkerId workerid);
   // Add to the remote function export queue the job of exporting all remote
   // functions to the given worker. This is used when a new worker registers.
   void add_all_remote_functions_to_worker_export_queue(WorkerId workerid);
@@ -226,6 +235,11 @@ private:
   // lock (objects_lock_). // TODO(rkn): Consider making this part of the
   // objtable data structure.
   std::vector<std::vector<ObjectID> > objects_in_transit_;
+  // List of pending functions to run on workers. These should be processed in a
+  // first in first out manner. The first element of each pair is the ID of the
+  // worker to run the function on, and the second element of each pair is the
+  // index of the function to run.
+  Synchronized<std::queue<std::pair<WorkerId, int> > > function_to_run_queue_;
   // List of pending remote function exports. These should be processed in a
   // first in first out manner. The first element of each pair is the ID of the
   // worker to export the remote function to, and the second element of each
@@ -236,6 +250,8 @@ private:
   // worker to export the reusable variable to, and the second element of each
   // pair is the index of the reusable variable to export.
   Synchronized<std::queue<std::pair<WorkerId, int> > > reusable_variable_export_queue_;
+  // All of the functions that have been exported to the workers to run.
+  Synchronized<std::vector<std::unique_ptr<Function> > > exported_functions_to_run_;
   // All of the remote functions that have been exported to the workers.
   Synchronized<std::vector<std::unique_ptr<Function> > > exported_functions_;
   // All of the reusable variables that have been exported to the workers.

src/worker.cc

@@ -26,6 +26,21 @@ Status WorkerServiceImpl::ExecuteTask(ServerContext* context, const ExecuteTaskR
     WorkerMessage* message_ptr = message.get();
     RAY_CHECK(send_queue_.send(&message_ptr), "Failed to send message from the worker service to the worker because the message queue was full.");
   }
+  // The message will get deleted in receive_next_message().
+  message.release();
+  return Status::OK;
+}
+
+Status WorkerServiceImpl::RunFunctionOnWorker(ServerContext* context, const RunFunctionOnWorkerRequest* request, AckReply* reply) {
+  RAY_CHECK(mode_ == Mode::WORKER_MODE, "RunFunctionOnWorker can only be called on workers.");
+  std::unique_ptr<WorkerMessage> message(new WorkerMessage());
+  message->mutable_function_to_run()->CopyFrom(request->function());
+  RAY_LOG(RAY_INFO, "Running function on worker.");
+  {
+    WorkerMessage* message_ptr = message.get();
+    RAY_CHECK(send_queue_.send(&message_ptr), "Failed to send message from the worker service to the worker because the message queue was full.");
+  }
+  // The message will get deleted in receive_next_message().
   message.release();
   return Status::OK;
 }
@@ -39,6 +54,7 @@ Status WorkerServiceImpl::ImportRemoteFunction(ServerContext* context, const Imp
     WorkerMessage* message_ptr = message.get();
     RAY_CHECK(send_queue_.send(&message_ptr), "Failed to send message from the worker service to the worker because the message queue was full.");
   }
+  // The message will get deleted in receive_next_message().
   message.release();
   return Status::OK;
 }
@@ -52,6 +68,7 @@ Status WorkerServiceImpl::ImportReusableVariable(ServerContext* context, const I
     WorkerMessage* message_ptr = message.get();
     RAY_CHECK(send_queue_.send(&message_ptr), "Failed to send message from the worker service to the worker because the message queue was full.");
   }
+  // The message will get deleted in receive_next_message().
   message.release();
   return Status::OK;
 }
@@ -440,6 +457,15 @@ std::vector<int> Worker::select(std::vector<ObjectID>& objectids) {
   return result;
 }
 
+void Worker::run_function_on_all_workers(const std::string& function) {
+  RAY_CHECK(connected_, "Attempted to run function on all workers but failed.");
+  ClientContext context;
+  RunFunctionOnAllWorkersRequest request;
+  request.mutable_function()->set_implementation(function);
+  AckReply reply;
+  RAY_CHECK_GRPC(scheduler_stub_->RunFunctionOnAllWorkers(&context, request, &reply));
+}
+
 bool Worker::export_remote_function(const std::string& function_name, const std::string& function) {
   RAY_CHECK(connected_, "Attempted to export function but failed.");
   ClientContext context;

src/worker.h

@@ -32,6 +32,7 @@ class WorkerServiceImpl final : public WorkerService::Service {
 public:
   WorkerServiceImpl(const std::string& worker_address, Mode mode);
   Status ExecuteTask(ServerContext* context, const ExecuteTaskRequest* request, AckReply* reply) override;
+  Status RunFunctionOnWorker(ServerContext* context, const RunFunctionOnWorkerRequest* request, AckReply* reply) override;
   Status ImportRemoteFunction(ServerContext* context, const ImportRemoteFunctionRequest* request, AckReply* reply) override;
   Status Die(ServerContext* context, const DieRequest* request, AckReply* reply) override;
   Status ImportReusableVariable(ServerContext* context, const ImportReusableVariableRequest* request, AckReply* reply) override;
@@ -104,6 +105,8 @@ class Worker {
   void task_info(ClientContext &context, TaskInfoRequest &request, TaskInfoReply &reply);
   // gets indices of available objects
   std::vector<int> select(std::vector<ObjectID>& objectids);
+  // Export a function to be run on all workers.
+  void run_function_on_all_workers(const std::string& function);
   // export function to workers
   bool export_remote_function(const std::string& function_name, const std::string& function);
   // export reusable variable to workers