implement key value store for sharing reusable variables

src/raylib.cc

@@ -714,13 +714,21 @@ static PyObject* wait_for_next_message(PyObject* self, PyObject* args) {
   Worker* worker;
   PyObjectToWorker(worker_capsule, &worker);
   if (std::unique_ptr<WorkerMessage> message = worker->receive_next_message()) {
+    PyObject* variable_info;
+    if (!message->reusable_variable.variable_name.empty()) {
+      variable_info = PyTuple_New(3);
+      PyTuple_SetItem(variable_info, 0, PyString_FromStringAndSize(message->reusable_variable.variable_name.data(), static_cast<ssize_t>(message->reusable_variable.variable_name.size())));
+      PyTuple_SetItem(variable_info, 1, PyString_FromStringAndSize(message->reusable_variable.initializer.data(), static_cast<ssize_t>(message->reusable_variable.initializer.size())));
+      PyTuple_SetItem(variable_info, 2, PyString_FromStringAndSize(message->reusable_variable.reinitializer.data(), static_cast<ssize_t>(message->reusable_variable.reinitializer.size())));
+    }
     // The tuple constructed below will take ownership of some None objects.
     // When the tuple goes out of scope, the reference count for None will be
     // decremented. Therefore, we need to increment the reference count for None
     // every time we put a None in the tuple.
-    PyObject* t = PyTuple_New(2); // We set the items of the tuple using PyTuple_SetItem, because that transfers ownership to the tuple.
+    PyObject* t = PyTuple_New(3); // We set the items of the tuple using PyTuple_SetItem, because that transfers ownership to the tuple.
     PyTuple_SetItem(t, 0, message->task.name().empty() ? Py_INCREF(Py_None), Py_None : deserialize_task(worker_capsule, &message->task));
     PyTuple_SetItem(t, 1, message->function.empty() ? Py_INCREF(Py_None), Py_None : PyString_FromStringAndSize(message->function.data(), static_cast<ssize_t>(message->function.size())));
+    PyTuple_SetItem(t, 2, message->reusable_variable.variable_name.empty() ? Py_INCREF(Py_None), Py_None : variable_info);
     return t;
   }
   Py_RETURN_NONE;
@@ -740,6 +748,24 @@ static PyObject* export_function(PyObject* self, PyObject* args) {
   }
 }
 
+static PyObject* export_reusable_variable(PyObject* self, PyObject* args) {
+  Worker* worker;
+  const char* name;
+  int name_size;
+  const char* initializer;
+  int initializer_size;
+  const char* reinitializer;
+  int reinitializer_size;
+  if (!PyArg_ParseTuple(args, "O&s#s#s#", &PyObjectToWorker, &worker, &name, &name_size, &initializer, &initializer_size, &reinitializer, &reinitializer_size)) {
+    return NULL;
+  }
+  std::string name_str(name, static_cast<size_t>(name_size));
+  std::string initializer_str(initializer, static_cast<size_t>(initializer_size));
+  std::string reinitializer_str(reinitializer, static_cast<size_t>(reinitializer_size));
+  worker->export_reusable_variable(name_str, initializer_str, reinitializer_str);
+  Py_RETURN_NONE;
+}
+
 static PyObject* submit_task(PyObject* self, PyObject* args) {
   PyObject* worker_capsule;
   Task* task;
@@ -997,6 +1023,7 @@ static PyMethodDef RayLibMethods[] = {
  { "scheduler_info", scheduler_info, METH_VARARGS, "get info about scheduler state" },
  { "task_info", task_info, METH_VARARGS, "get task statuses" },
  { "export_function", export_function, METH_VARARGS, "export 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" },
  { "set_log_config", set_log_config, METH_VARARGS, "set filename for raylib logging" },
  { "kill_workers", kill_workers, METH_VARARGS, "kills all of the workers" },

src/scheduler.cc

@@ -305,6 +305,22 @@ Status SchedulerService::ExportFunction(ServerContext* context, const ExportFunc
   return Status::OK;
 }
 
+Status SchedulerService::ExportReusableVariable(ServerContext* context, const ExportReusableVariableRequest* request, AckReply* reply) {
+  auto workers = workers_.get();
+  for (size_t i = 0; i < workers->size(); ++i) {
+    ClientContext import_context;
+    ImportReusableVariableRequest import_request;
+    import_request.set_name(request->name());
+    import_request.mutable_initializer()->set_implementation(request->initializer().implementation());
+    import_request.mutable_reinitializer()->set_implementation(request->reinitializer().implementation());
+    if ((*workers)[i].current_task != ROOT_OPERATION) {
+      AckReply import_reply;
+      (*workers)[i].worker_stub->ImportReusableVariable(&import_context, import_request, &import_reply);
+    }
+  }
+  return Status::OK;
+}
+
 void SchedulerService::deliver_object_async_if_necessary(ObjRef canonical_objref, ObjStoreId from, ObjStoreId to) {
   bool object_present_or_in_transit;
   {

src/scheduler.h

@@ -71,6 +71,7 @@ public:
   Status TaskInfo(ServerContext* context, const TaskInfoRequest* request, TaskInfoReply* reply) override;
   Status KillWorkers(ServerContext* context, const KillWorkersRequest* request, KillWorkersReply* reply) override;
   Status ExportFunction(ServerContext* context, const ExportFunctionRequest* request, ExportFunctionReply* reply) override;
+  Status ExportReusableVariable(ServerContext* context, const ExportReusableVariableRequest* request, AckReply* 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

src/worker.cc

@@ -40,6 +40,20 @@ Status WorkerServiceImpl::ImportFunction(ServerContext* context, const ImportFun
   return Status::OK;
 }
 
+Status WorkerServiceImpl::ImportReusableVariable(ServerContext* context, const ImportReusableVariableRequest* request, AckReply* reply) {
+  std::unique_ptr<WorkerMessage> message(new WorkerMessage());
+  message->reusable_variable.variable_name = request->name();
+  message->reusable_variable.initializer = request->initializer().implementation();
+  message->reusable_variable.reinitializer = request->reinitializer().implementation();
+  RAY_LOG(RAY_INFO, "importing reusable variable");
+  {
+    WorkerMessage* message_ptr = message.get();
+    RAY_CHECK(send_queue_.send(&message_ptr), "error sending over IPC");
+  }
+  message.release();
+  return Status::OK;
+}
+
 Status WorkerServiceImpl::Die(ServerContext* context, const DieRequest* request, DieReply* reply) {
   WorkerMessage* message_ptr = NULL;
   RAY_CHECK(send_queue_.send(&message_ptr), "error sending over IPC");
@@ -389,6 +403,17 @@ bool Worker::export_function(const std::string& function) {
   return true;
 }
 
+void Worker::export_reusable_variable(const std::string& name, const std::string& initializer, const std::string& reinitializer) {
+  RAY_CHECK(connected_, "Attempted to export reusable variable but failed.");
+  ClientContext context;
+  ExportReusableVariableRequest request;
+  request.set_name(name);
+  request.mutable_initializer()->set_implementation(initializer);
+  request.mutable_reinitializer()->set_implementation(reinitializer);
+  AckReply reply;
+  Status status = scheduler_stub_->ExportReusableVariable(&context, request, &reply);
+}
+
 // Communication between the WorkerServer and the Worker happens via a message
 // queue. This is because the Python interpreter needs to be single threaded
 // (in our case running in the main thread), whereas the WorkerService will

src/worker.h

@@ -23,9 +23,16 @@ using grpc::Channel;
 using grpc::ClientContext;
 using grpc::ClientWriter;
 
+struct ReusableVariable {
+  std::string variable_name;
+  std::string initializer;
+  std::string reinitializer;
+};
+
 struct WorkerMessage {
   Task task;
-  std::string function;
+  std::string function; // Used for importing remote functions.
+  ReusableVariable reusable_variable; // Used for importing reusable variables.
 };
 
 class WorkerServiceImpl final : public WorkerService::Service {
@@ -34,6 +41,7 @@ public:
   Status ExecuteTask(ServerContext* context, const ExecuteTaskRequest* request, ExecuteTaskReply* reply) override;
   Status ImportFunction(ServerContext* context, const ImportFunctionRequest* request, ImportFunctionReply* reply) override;
   Status Die(ServerContext* context, const DieRequest* request, DieReply* reply) override;
+  Status ImportReusableVariable(ServerContext* context, const ImportReusableVariableRequest* request, AckReply* reply) override;
 private:
   std::string worker_address_;
   MessageQueue<WorkerMessage*> send_queue_;
@@ -100,6 +108,8 @@ class Worker {
   void task_info(ClientContext &context, TaskInfoRequest &request, TaskInfoReply &reply);
   // export function to workers
   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);
 
  private:
   bool connected_;