Write computation graph to file

src/computation_graph.cc

@@ -19,3 +19,9 @@ const Task& ComputationGraph::get_task(OperationId operationid) {
   RAY_CHECK(operations_[operationid]->has_task(), "Calling get_task with operationid " << operationid << ", but this corresponds to a put not a task.");
   return operations_[operationid]->task();
 }
+
+void ComputationGraph::to_protobuf(CompGraph* computation_graph) {
+  for (OperationId id = 0; id < operations_.size(); ++id) {
+    computation_graph->add_operation()->CopyFrom(*operations_[id]);
+  }
+}

src/computation_graph.h

@@ -5,7 +5,8 @@
 #include <limits>
 
 #include "ray/ray.h"
-#include "ray.grpc.pb.h"
+
+#include "graph.pb.h"
 #include "types.pb.h"
 
 // used to represent the root operation (that is, the driver code)
@@ -21,6 +22,8 @@ public:
   // Return the task corresponding to a particular OperationId. If operationid
   // corresponds to a put, then fail.
   const Task& get_task(OperationId operationid);
+  // Serialize the computation graph to ProtoBuf and store it in computation_graph
+  void to_protobuf(CompGraph* computation_graph);
 private:
   // maps an OperationId to the corresponding task or put
   std::vector<std::unique_ptr<Operation> > operations_;

src/raylib.cc

@@ -824,6 +824,21 @@ PyObject* task_info(PyObject* self, PyObject* args) {
   return dict;
 }
 
+PyObject* dump_computation_graph(PyObject* self, PyObject* args) {
+  Worker* worker;
+  const char* output_file_name;
+  if (!PyArg_ParseTuple(args, "O&s", &PyObjectToWorker, &worker, &output_file_name)) {
+    return NULL;
+  }
+  ClientContext context;
+  SchedulerInfoRequest request;
+  SchedulerInfoReply reply;
+  worker->scheduler_info(context, request, reply);
+  std::fstream output(output_file_name, std::ios::out | std::ios::trunc | std::ios::binary);
+  RAY_CHECK(reply.computation_graph().SerializeToOstream(&output), "Cannot dump computation graph to file " << output_file_name);
+  Py_RETURN_NONE;
+}
+
 PyObject* set_log_config(PyObject* self, PyObject* args) {
   const char* log_file_name;
   if (!PyArg_ParseTuple(args, "s", &log_file_name)) {
@@ -859,6 +874,7 @@ static PyMethodDef RayLibMethods[] = {
  { "start_worker_service", start_worker_service, METH_VARARGS, "start the worker service" },
  { "scheduler_info", scheduler_info, METH_VARARGS, "get info about scheduler state" },
  { "task_info", task_info, METH_VARARGS, "get task statuses" },
+ { "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" },
  { NULL, NULL, 0, NULL }
 };

src/scheduler.cc

@@ -494,6 +494,7 @@ void SchedulerService::get_info(const SchedulerInfoRequest& request, SchedulerIn
   for (const WorkerId& entry : avail_workers_) {
     reply->add_avail_worker(entry);
   }
+  computation_graph_.to_protobuf(reply->mutable_computation_graph());
   release_all_locks();
 }