[xray] Implement timeline and profiling API. (#2306)

* Add profile table and store profiling information there.

* Code for dumping timeline.

* Improve color scheme.

* Push timeline events on driver only for raylet.

* Improvements to profiling and timeline visualization

* Some linting

* Small fix.

* Linting

* Propagate node IP address through profiling events.

* Fix test.

* object_id.hex() should return byte string in python 2.

* Include gcs.fbs in node_manager.fbs.

* Remove flatbuffer definition duplication.

* Decode to unicode in Python 3 and bytes in Python 2.

* Minor

* Submit profile events in a batch. Revert some CMake changes.

* Fix

* Workaround test failure.

* Fix linting

* Linting

* Don't return anything from chrome_tracing_dump when filename is provided.

* Remove some redundancy from profile table.

* Linting

* Move TODOs out of docstring.

* Minor

src/common/lib/python/common_extension.cc

@@ -165,7 +165,11 @@ static PyObject *PyObjectID_id(PyObject *self) {
 static PyObject *PyObjectID_hex(PyObject *self) {
   PyObjectID *s = (PyObjectID *) self;
   std::string hex_id = s->object_id.hex();
-  PyObject *result = PyUnicode_FromString(hex_id.c_str());
+#if PY_MAJOR_VERSION >= 3
+  PyObject *result = PyUnicode_FromStringAndSize(hex_id.data(), hex_id.size());
+#else
+  PyObject *result = PyBytes_FromStringAndSize(hex_id.data(), hex_id.size());
+#endif
   return result;
 }
 

src/common/redis_module/ray_redis_module.cc

@@ -695,6 +695,8 @@ int TableAppend_DoWrite(RedisModuleCtx *ctx,
     // Check that we actually add a new entry during the append. This is only
     // necessary since we implement the log with a sorted set, so all entries
     // must be unique, or else we will have gaps in the log.
+    // TODO(rkn): We need to get rid of this uniqueness requirement. We can
+    // easily have multiple log events with the same message.
     RAY_CHECK(flags == REDISMODULE_ZADD_ADDED) << "Appended a duplicate entry";
     return REDISMODULE_OK;
   } else {

src/local_scheduler/lib/python/local_scheduler_extension.cc

@@ -309,6 +309,109 @@ static PyObject *PyLocalSchedulerClient_push_error(PyObject *self,
   Py_RETURN_NONE;
 }
 
+int PyBytes_or_PyUnicode_to_string(PyObject *py_string, std::string &out) {
+  // Handle the case where the key is a bytes object and the case where it
+  // is a unicode object.
+  if (PyUnicode_Check(py_string)) {
+    PyObject *ascii_string = PyUnicode_AsASCIIString(py_string);
+    out =
+        std::string(PyBytes_AsString(ascii_string), PyBytes_Size(ascii_string));
+    Py_DECREF(ascii_string);
+  } else if (PyBytes_Check(py_string)) {
+    out = std::string(PyBytes_AsString(py_string), PyBytes_Size(py_string));
+  } else {
+    return -1;
+  }
+
+  return 0;
+}
+
+static PyObject *PyLocalSchedulerClient_push_profile_events(PyObject *self,
+                                                            PyObject *args) {
+  const char *component_type;
+  int component_type_length;
+  UniqueID component_id;
+  PyObject *profile_data;
+  const char *node_ip_address;
+  int node_ip_address_length;
+
+  if (!PyArg_ParseTuple(args, "s#O&s#O", &component_type,
+                        &component_type_length, &PyObjectToUniqueID,
+                        &component_id, &node_ip_address,
+                        &node_ip_address_length, &profile_data)) {
+    return NULL;
+  }
+
+  ProfileTableDataT profile_info;
+  profile_info.component_type =
+      std::string(component_type, component_type_length);
+  profile_info.component_id = component_id.binary();
+  profile_info.node_ip_address =
+      std::string(node_ip_address, node_ip_address_length);
+
+  if (PyList_Size(profile_data) == 0) {
+    // Short circuit if there are no profile events.
+    Py_RETURN_NONE;
+  }
+
+  for (int64_t i = 0; i < PyList_Size(profile_data); ++i) {
+    ProfileEventT profile_event;
+    PyObject *py_profile_event = PyList_GetItem(profile_data, i);
+
+    if (!PyDict_CheckExact(py_profile_event)) {
+      return NULL;
+    }
+
+    PyObject *key, *val;
+    Py_ssize_t pos = 0;
+    while (PyDict_Next(py_profile_event, &pos, &key, &val)) {
+      std::string key_string;
+      if (PyBytes_or_PyUnicode_to_string(key, key_string) == -1) {
+        return NULL;
+      }
+
+      // TODO(rkn): If the dictionary is formatted incorrectly, that could lead
+      // to errors. E.g., if any of the strings are empty, that will cause
+      // segfaults in the node manager.
+
+      if (key_string == std::string("event_type")) {
+        if (PyBytes_or_PyUnicode_to_string(val, profile_event.event_type) ==
+            -1) {
+          return NULL;
+        }
+        if (profile_event.event_type.size() == 0) {
+          return NULL;
+        }
+      } else if (key_string == std::string("start_time")) {
+        profile_event.start_time = PyFloat_AsDouble(val);
+      } else if (key_string == std::string("end_time")) {
+        profile_event.end_time = PyFloat_AsDouble(val);
+      } else if (key_string == std::string("extra_data")) {
+        if (PyBytes_or_PyUnicode_to_string(val, profile_event.extra_data) ==
+            -1) {
+          return NULL;
+        }
+        if (profile_event.extra_data.size() == 0) {
+          return NULL;
+        }
+      } else {
+        return NULL;
+      }
+    }
+
+    // Note that profile_info.profile_events is a vector of unique pointers, so
+    // profile_event will be deallocated when profile_info goes out of scope.
+    profile_info.profile_events.emplace_back(new ProfileEventT(profile_event));
+  }
+
+  local_scheduler_push_profile_events(
+      reinterpret_cast<PyLocalSchedulerClient *>(self)
+          ->local_scheduler_connection,
+      profile_info);
+
+  Py_RETURN_NONE;
+}
+
 static PyMethodDef PyLocalSchedulerClient_methods[] = {
     {"disconnect", (PyCFunction) PyLocalSchedulerClient_disconnect, METH_NOARGS,
      "Notify the local scheduler that this client is exiting gracefully."},
@@ -338,6 +441,9 @@ static PyMethodDef PyLocalSchedulerClient_methods[] = {
      "Wait for a list of objects to be created."},
     {"push_error", (PyCFunction) PyLocalSchedulerClient_push_error,
      METH_VARARGS, "Push an error message to the relevant driver."},
+    {"push_profile_events",
+     (PyCFunction) PyLocalSchedulerClient_push_profile_events, METH_VARARGS,
+     "Store some profiling events in the GCS."},
     {NULL} /* Sentinel */
 };
 

src/local_scheduler/local_scheduler_client.cc

@@ -322,3 +322,17 @@ void local_scheduler_push_error(LocalSchedulerConnection *conn,
                                 ray::protocol::MessageType::PushErrorRequest),
                 fbb.GetSize(), fbb.GetBufferPointer());
 }
+
+void local_scheduler_push_profile_events(
+    LocalSchedulerConnection *conn,
+    const ProfileTableDataT &profile_events) {
+  flatbuffers::FlatBufferBuilder fbb;
+
+  auto message = CreateProfileTableData(fbb, &profile_events);
+  fbb.Finish(message);
+
+  write_message(conn->conn,
+                static_cast<int64_t>(
+                    ray::protocol::MessageType::PushProfileEventsRequest),
+                fbb.GetSize(), fbb.GetBufferPointer());
+}

src/local_scheduler/local_scheduler_client.h

@@ -225,4 +225,13 @@ void local_scheduler_push_error(LocalSchedulerConnection *conn,
                                 const std::string &error_message,
                                 double timestamp);
 
+/// Store some profile events in the GCS.
+///
+/// \param conn The connection information.
+/// \param profile_events A batch of profiling event information.
+/// \return Void.
+void local_scheduler_push_profile_events(
+    LocalSchedulerConnection *conn,
+    const ProfileTableDataT &profile_events);
+
 #endif

src/ray/gcs/client.cc

@@ -17,6 +17,7 @@ AsyncGcsClient::AsyncGcsClient(const ClientID &client_id, CommandType command_ty
   task_reconstruction_log_.reset(new TaskReconstructionLog(context_, this));
   heartbeat_table_.reset(new HeartbeatTable(context_, this));
   error_table_.reset(new ErrorTable(primary_context_, this));
+  profile_table_.reset(new ProfileTable(context_, this));
   command_type_ = command_type;
 }
 
@@ -84,6 +85,8 @@ HeartbeatTable &AsyncGcsClient::heartbeat_table() { return *heartbeat_table_; }
 
 ErrorTable &AsyncGcsClient::error_table() { return *error_table_; }
 
+ProfileTable &AsyncGcsClient::profile_table() { return *profile_table_; }
+
 }  // namespace gcs
 
 }  // namespace ray

src/ray/gcs/client.h

@@ -59,6 +59,7 @@ class RAY_EXPORT AsyncGcsClient {
   ClientTable &client_table();
   HeartbeatTable &heartbeat_table();
   ErrorTable &error_table();
+  ProfileTable &profile_table();
 
   // We also need something to export generic code to run on workers from the
   // driver (to set the PYTHONPATH)
@@ -81,6 +82,7 @@ class RAY_EXPORT AsyncGcsClient {
   std::unique_ptr<TaskReconstructionLog> task_reconstruction_log_;
   std::unique_ptr<HeartbeatTable> heartbeat_table_;
   std::unique_ptr<ErrorTable> error_table_;
+  std::unique_ptr<ProfileTable> profile_table_;
   std::unique_ptr<ClientTable> client_table_;
   // The following contexts write to the data shard
   std::shared_ptr<RedisContext> context_;

src/ray/gcs/format/gcs.fbs

@@ -15,6 +15,7 @@ enum TablePrefix:int {
   TASK_RECONSTRUCTION,
   HEARTBEAT,
   ERROR_INFO,
+  PROFILE,
 }
 
 // The channel that Add operations to the Table should be published on, if any.
@@ -121,6 +122,33 @@ table CustomSerializerData {
 table ConfigTableData {
 }
 
+table ProfileEvent {
+  // The type of the event.
+  event_type: string;
+  // The start time of the event.
+  start_time: double;
+  // The end time of the event. If the event is a point event, then this should
+  // be the same as the start time.
+  end_time: double;
+  // Additional data associated with the event. This data must be serialized
+  // using JSON.
+  extra_data: string;
+}
+
+table ProfileTableData {
+  // The type of the component that generated the event, e.g., worker or
+  // object_manager, or node_manager.
+  component_type: string;
+  // An identifier for the component that generated the event.
+  component_id: string;
+  // An identifier for the node that generated the event.
+  node_ip_address: string;
+  // This is a batch of profiling events. We batch these together for
+  // performance reasons because a single task may generate many events, and
+  // we don't want each event to require a GCS command.
+  profile_events: [ProfileEvent];
+}
+
 table RayResource {
   // The type of the resource.
   resource_name: string;

src/ray/gcs/tables.cc

@@ -219,6 +219,45 @@ Status ErrorTable::PushErrorToDriver(const JobID &job_id, const std::string &typ
   });
 }
 
+Status ProfileTable::AddProfileEvent(const std::string &event_type,
+                                     const std::string &component_type,
+                                     const UniqueID &component_id,
+                                     const std::string &node_ip_address,
+                                     double start_time, double end_time,
+                                     const std::string &extra_data) {
+  auto data = std::make_shared<ProfileTableDataT>();
+
+  ProfileEventT profile_event;
+  profile_event.event_type = event_type;
+  profile_event.start_time = start_time;
+  profile_event.end_time = end_time;
+  profile_event.extra_data = extra_data;
+
+  data->component_type = component_type;
+  data->component_id = component_id.binary();
+  data->node_ip_address = node_ip_address;
+  data->profile_events.emplace_back(new ProfileEventT(profile_event));
+
+  return Append(JobID::nil(), component_id, data,
+                [](ray::gcs::AsyncGcsClient *client, const JobID &id,
+                   const ProfileTableDataT &data) {
+                  RAY_LOG(DEBUG) << "Profile message pushed callback";
+                });
+}
+
+Status ProfileTable::AddProfileEventBatch(const ProfileTableData &profile_events) {
+  auto data = std::make_shared<ProfileTableDataT>();
+  // There is some room for optimization here because the Append function will just
+  // call "Pack" and undo the "UnPack".
+  profile_events.UnPackTo(data.get());
+
+  return Append(JobID::nil(), from_flatbuf(*profile_events.component_id()), data,
+                [](ray::gcs::AsyncGcsClient *client, const JobID &id,
+                   const ProfileTableDataT &data) {
+                  RAY_LOG(DEBUG) << "Profile message pushed callback";
+                });
+}
+
 void ClientTable::RegisterClientAddedCallback(const ClientTableCallback &callback) {
   client_added_callback_ = callback;
   // Call the callback for any added clients that are cached.
@@ -371,6 +410,7 @@ template class Log<TaskID, TaskReconstructionData>;
 template class Table<ClientID, HeartbeatTableData>;
 template class Log<JobID, ErrorTableData>;
 template class Log<UniqueID, ClientTableData>;
+template class Log<UniqueID, ProfileTableData>;
 
 }  // namespace gcs
 

src/ray/gcs/tables.h

@@ -12,6 +12,7 @@
 
 #include "ray/gcs/format/gcs_generated.h"
 #include "ray/gcs/redis_context.h"
+// TODO(rkn): Remove this include.
 #include "ray/raylet/format/node_manager_generated.h"
 
 // TODO(pcm): Remove this
@@ -95,7 +96,8 @@ class Log : virtual public PubsubInterface<ID> {
   ///
   /// \param job_id The ID of the job (= driver).
   /// \param id The ID of the data that is added to the GCS.
-  /// \param data Data to append to the log.
+  /// \param data Data to append to the log. TODO(rkn): This can be made const,
+  /// right?
   /// \param done Callback that is called once the data has been written to the
   /// GCS.
   /// \return Status
@@ -438,7 +440,8 @@ class ErrorTable : private Log<JobID, ErrorTableData> {
   /// Push an error message for a specific job.
   ///
   /// TODO(rkn): We need to make sure that the errors are unique because
-  /// duplicate messages currently cause failures (the GCS doesn't allow it).
+  /// duplicate messages currently cause failures (the GCS doesn't allow it). A
+  /// natural way to do this is to have finer-grained time stamps.
   ///
   /// \param job_id The ID of the job that generated the error. If the error
   /// should be pushed to all jobs, then this should be nil.
@@ -450,6 +453,37 @@ class ErrorTable : private Log<JobID, ErrorTableData> {
                            const std::string &error_message, double timestamp);
 };
 
+class ProfileTable : private Log<UniqueID, ProfileTableData> {
+ public:
+  ProfileTable(const std::shared_ptr<RedisContext> &context, AsyncGcsClient *client)
+      : Log(context, client) {
+    prefix_ = TablePrefix::PROFILE;
+  };
+
+  /// Add a single profile event to the profile table.
+  ///
+  /// \param event_type The type of the event.
+  /// \param component_type The type of the component that the event came from.
+  /// \param component_id An identifier for the component that generated the event.
+  /// \param node_ip_address The IP address of the node that generated the event.
+  /// \param start_time The timestamp of the event start, this should be in seconds since
+  /// the Unix epoch.
+  /// \param end_time The timestamp of the event end, this should be in seconds since
+  /// the Unix epoch. If the event is a point event, this should be equal to start_time.
+  /// \param extra_data Additional data to associate with the event.
+  /// \return Status.
+  Status AddProfileEvent(const std::string &event_type, const std::string &component_type,
+                         const UniqueID &component_id, const std::string &node_ip_address,
+                         double start_time, double end_time,
+                         const std::string &extra_data);
+
+  /// Add a batch of profiling events to the profile table.
+  ///
+  /// \param profile_events The profile events to record.
+  /// \return Status.
+  Status AddProfileEventBatch(const ProfileTableData &profile_events);
+};
+
 using CustomSerializerTable = Table<ClassID, CustomSerializerData>;
 
 using ConfigTable = Table<ConfigID, ConfigTableData>;

src/ray/raylet/CMakeLists.txt

@@ -12,7 +12,7 @@ add_custom_command(
         # flatbuffers message Message, which can be used to store deserialized
         # messages in data structures. This is currently used for ObjectInfo for
         # example.
-        COMMAND ${FLATBUFFERS_COMPILER} -c -o ${OUTPUT_DIR} ${NODE_MANAGER_FBS_SRC} --cpp --gen-object-api --gen-mutable --scoped-enums
+        COMMAND ${FLATBUFFERS_COMPILER} -c -o ${OUTPUT_DIR} -I ${GCS_FBS_OUTPUT_DIRECTORY} ${NODE_MANAGER_FBS_SRC} --cpp --gen-object-api --gen-mutable --scoped-enums
         DEPENDS ${FBS_DEPENDS}
         COMMENT "Running flatc compiler on ${NODE_MANAGER_FBS_SRC}"
         VERBATIM)
@@ -23,7 +23,7 @@ add_custom_target(gen_node_manager_fbs DEPENDS ${NODE_MANAGER_FBS_OUTPUT_FILES})
 set(PYTHON_OUTPUT_DIR ${CMAKE_CURRENT_LIST_DIR}/../../../python/ray/core/generated/)
 add_custom_command(
     TARGET gen_node_manager_fbs
-    COMMAND ${FLATBUFFERS_COMPILER} -p -o ${PYTHON_OUTPUT_DIR} ${NODE_MANAGER_FBS_SRC}
+    COMMAND ${FLATBUFFERS_COMPILER} -p -o ${PYTHON_OUTPUT_DIR} -I ${GCS_FBS_OUTPUT_DIRECTORY} ${NODE_MANAGER_FBS_SRC}
     DEPENDS ${FBS_DEPENDS}
     COMMENT "Running flatc compiler on ${NODE_MANAGER_FBS_SRC}"
     VERBATIM)
@@ -38,6 +38,7 @@ ADD_RAY_TEST(task_test STATIC_LINK_LIBS ray_static gtest gtest_main gmock_main p
 ADD_RAY_TEST(lineage_cache_test STATIC_LINK_LIBS ray_static gtest gtest_main gmock_main pthread ${Boost_SYSTEM_LIBRARY})
 ADD_RAY_TEST(task_dependency_manager_test STATIC_LINK_LIBS ray_static gtest gtest_main gmock_main pthread ${Boost_SYSTEM_LIBRARY})
 
+include_directories(${GCS_FBS_OUTPUT_DIRECTORY})
 add_library(rayletlib raylet.cc ${NODE_MANAGER_FBS_OUTPUT_FILES})
 target_link_libraries(rayletlib ray_static ${Boost_SYSTEM_LIBRARY})
 

src/ray/raylet/format/node_manager.fbs

@@ -1,5 +1,8 @@
 // Local scheduler protocol specification
 
+include "gcs.fbs";
+
+
 // TODO(swang): We put the flatbuffer types in a separate namespace for now to
 // avoid conflicts with legacy Ray types.
 namespace ray.protocol;
@@ -62,6 +65,9 @@ enum MessageType:int {
   // Push an error to the relevant driver. This is sent from a worker to the
   // node manager.
   PushErrorRequest,
+  // Push some profiling events to the GCS. When sending this message to the
+  // node manager, the message itself is serialized as a ProfileTableData object.
+  PushProfileEventsRequest,
 }
 
 table TaskExecutionSpecification {

src/ray/raylet/node_manager.cc

@@ -552,6 +552,11 @@ void NodeManager::ProcessClientMessage(
     RAY_CHECK_OK(gcs_client_->error_table().PushErrorToDriver(job_id, type, error_message,
                                                               timestamp));
   } break;
+  case protocol::MessageType::PushProfileEventsRequest: {
+    auto message = flatbuffers::GetRoot<ProfileTableData>(message_data);
+
+    RAY_CHECK_OK(gcs_client_->profile_table().AddProfileEventBatch(*message));
+  } break;
 
   default:
     RAY_LOG(FATAL) << "Received unexpected message type " << message_type;