Use flatbuffers for some messages from Redis. (ray-project#341)

* Compile the Ray redis module with C++.

* Redo parsing of object table notifications with flatbuffers.

* Update redis module python tests.

* Redo parsing of task table notifications with flatbuffers.

* Fix linting.

* Redo parsing of db client notifications with flatbuffers.

* Redo publishing of local scheduler heartbeats with flatbuffers.

* Fix linting.

* Remove usage of fixed-width formatting of scheduling state in channel name.

* Reply with flatbuffer object to task table queries, also simplify redis string to flatbuffer string conversion.

* Fix linting and tests.

* fix

* cleanup

* simplify logic in ReplyWithTask

.travis/install-dependencies.sh

@@ -20,15 +20,15 @@ fi
 if [[ "$PYTHON" == "2.7" ]] && [[ "$platform" == "linux" ]]; then
   sudo apt-get update
   sudo apt-get install -y cmake build-essential autoconf curl libtool python-dev python-numpy python-pip libboost-all-dev unzip
-  sudo pip install cloudpickle funcsigs colorama psutil redis tensorflow
+  sudo pip install cloudpickle funcsigs colorama psutil redis tensorflow flatbuffers
 elif [[ "$PYTHON" == "3.5" ]] && [[ "$platform" == "linux" ]]; then
   sudo apt-get update
   sudo apt-get install -y cmake python-dev python-numpy build-essential autoconf curl libtool libboost-all-dev unzip
   # Install miniconda.
   wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh
   bash miniconda.sh -b -p $HOME/miniconda
   export PATH="$HOME/miniconda/bin:$PATH"
-  pip install numpy cloudpickle funcsigs colorama psutil redis tensorflow
+  pip install numpy cloudpickle funcsigs colorama psutil redis tensorflow flatbuffers
 elif [[ "$PYTHON" == "2.7" ]] && [[ "$platform" == "macosx" ]]; then
   # check that brew is installed
   which -s brew
@@ -41,7 +41,7 @@ elif [[ "$PYTHON" == "2.7" ]] && [[ "$platform" == "macosx" ]]; then
   fi
   brew install cmake automake autoconf libtool boost
   sudo easy_install pip
-  sudo pip install numpy cloudpickle funcsigs colorama psutil redis tensorflow --ignore-installed six
+  sudo pip install numpy cloudpickle funcsigs colorama psutil redis tensorflow flatbuffers --ignore-installed six
 elif [[ "$PYTHON" == "3.5" ]] && [[ "$platform" == "macosx" ]]; then
   # check that brew is installed
   which -s brew
@@ -57,7 +57,7 @@ elif [[ "$PYTHON" == "3.5" ]] && [[ "$platform" == "macosx" ]]; then
   wget https://repo.continuum.io/miniconda/Miniconda3-latest-MacOSX-x86_64.sh -O miniconda.sh
   bash miniconda.sh -b -p $HOME/miniconda
   export PATH="$HOME/miniconda/bin:$PATH"
-  pip install numpy cloudpickle funcsigs colorama psutil redis tensorflow
+  pip install numpy cloudpickle funcsigs colorama psutil redis tensorflow flatbuffers
 elif [[ "$LINT" == "1" ]]; then
   sudo apt-get update
   sudo apt-get install -y cmake build-essential autoconf curl libtool libboost-all-dev unzip

doc/source/install-on-macosx.md

@@ -11,7 +11,7 @@ To install Ray, first install the following dependencies. We recommend using
 brew update
 brew install cmake automake autoconf libtool boost wget
 
-pip install numpy cloudpickle funcsigs colorama psutil redis --ignore-installed six
+pip install numpy cloudpickle funcsigs colorama psutil redis flatbuffers --ignore-installed six
 ```
 
 If you are using Anaconda, you may also need to run the following.

doc/source/install-on-ubuntu.md

@@ -12,7 +12,7 @@ To install Ray, first install the following dependencies. We recommend using
 sudo apt-get update
 sudo apt-get install -y cmake build-essential autoconf curl libtool libboost-all-dev unzip python-dev python-pip  # If you're using Anaconda, then python-dev and python-pip are unnecessary.
 
-pip install numpy cloudpickle funcsigs colorama psutil redis
+pip install numpy cloudpickle funcsigs colorama psutil redis flatbuffers
 ```
 
 If you are using Anaconda, you may also need to run the following.

docker/base-deps/Dockerfile

@@ -11,4 +11,5 @@ RUN echo 'export PATH=/opt/conda/bin:$PATH' > /etc/profile.d/conda.sh \
     && rm /tmp/anaconda.sh
 ENV PATH "/opt/conda/bin:$PATH"
 RUN conda install -y libgcc
+RUN pip install flatbuffers
 RUN pip install --upgrade pip cloudpickle

python/ray/common/redis_module/runtest.py

@@ -11,6 +11,10 @@
 import redis
 import ray.services
 
+# Import flatbuffer bindings.
+from ray.core.generated.SubscribeToNotificationsReply import SubscribeToNotificationsReply
+from ray.core.generated.TaskReply import TaskReply
+
 OBJECT_INFO_PREFIX = "OI:"
 OBJECT_LOCATION_PREFIX = "OL:"
 OBJECT_SUBSCRIBE_PREFIX = "OS:"
@@ -142,6 +146,15 @@ def testObjectTableAddAndRemove(self):
     self.assertEqual(set(response), set())
 
   def testObjectTableSubscribeToNotifications(self):
+    # Define a helper method for checking the contents of object notifications.
+    def check_object_notification(notification_message, object_id, object_size, manager_ids):
+      notification_object = SubscribeToNotificationsReply.GetRootAsSubscribeToNotificationsReply(notification_message, 0)
+      self.assertEqual(notification_object.ObjectId(), object_id)
+      self.assertEqual(notification_object.ObjectSize(), object_size)
+      self.assertEqual(notification_object.ManagerIdsLength(), len(manager_ids))
+      for i in range(len(manager_ids)):
+        self.assertEqual(notification_object.ManagerIds(i), manager_ids[i])
+
     data_size = 0xf1f0
     p = self.redis.pubsub()
     # Subscribe to an object ID.
@@ -151,24 +164,37 @@ def testObjectTableSubscribeToNotifications(self):
     self.assertEqual(get_next_message(p)["data"], 1)
     # Request a notification and receive the data.
     self.redis.execute_command("RAY.OBJECT_TABLE_REQUEST_NOTIFICATIONS", "manager_id1", "object_id1")
-    self.assertEqual(get_next_message(p)["data"], b"object_id1 %s MANAGERS manager_id2"\
-                     %integerToAsciiHex(data_size, 8))
+    # Verify that the notification is correct.
+    check_object_notification(get_next_message(p)["data"],
+                              b"object_id1",
+                              data_size,
+                              [b"manager_id2"])
+
     # Request a notification for an object that isn't there. Then add the object
     # and receive the data. Only the first call to RAY.OBJECT_TABLE_ADD should
     # trigger notifications.
     self.redis.execute_command("RAY.OBJECT_TABLE_REQUEST_NOTIFICATIONS", "manager_id1", "object_id2", "object_id3")
     self.redis.execute_command("RAY.OBJECT_TABLE_ADD", "object_id3", data_size, "hash1", "manager_id1")
     self.redis.execute_command("RAY.OBJECT_TABLE_ADD", "object_id3", data_size, "hash1", "manager_id2")
     self.redis.execute_command("RAY.OBJECT_TABLE_ADD", "object_id3", data_size, "hash1", "manager_id3")
-    self.assertEqual(get_next_message(p)["data"], b"object_id3 %s MANAGERS manager_id1"\
-                     %integerToAsciiHex(data_size, 8))
+    # Verify that the notification is correct.
+    check_object_notification(get_next_message(p)["data"],
+                              b"object_id3",
+                              data_size,
+                              [b"manager_id1"])
     self.redis.execute_command("RAY.OBJECT_TABLE_ADD", "object_id2", data_size, "hash1", "manager_id3")
-    self.assertEqual(get_next_message(p)["data"], b"object_id2 %s MANAGERS manager_id3"\
-                     %integerToAsciiHex(data_size, 8))
+    # Verify that the notification is correct.
+    check_object_notification(get_next_message(p)["data"],
+                              b"object_id2",
+                              data_size,
+                              [b"manager_id3"])
     # Request notifications for object_id3 again.
     self.redis.execute_command("RAY.OBJECT_TABLE_REQUEST_NOTIFICATIONS", "manager_id1", "object_id3")
-    self.assertEqual(get_next_message(p)["data"], b"object_id3 %s MANAGERS manager_id1 manager_id2 manager_id3"\
-                     %integerToAsciiHex(data_size, 8))
+    # Verify that the notification is correct.
+    check_object_notification(get_next_message(p)["data"],
+                              b"object_id3",
+                              data_size,
+                              [b"manager_id1", b"manager_id2", b"manager_id3"])
 
   def testResultTableAddAndLookup(self):
     # Try looking up something in the result table before anything is added.
@@ -205,10 +231,6 @@ def testInvalidTaskTableAdd(self):
       # Non-integer scheduling states should not be added.
       self.redis.execute_command("RAY.TASK_TABLE_ADD", "task_id",
                                  "invalid_state", "node_id", "task_spec")
-    with self.assertRaises(redis.ResponseError):
-      # Scheduling states with invalid width should not be added.
-      self.redis.execute_command("RAY.TASK_TABLE_ADD", "task_id", 101,
-                                 "node_id", "task_spec")
     with self.assertRaises(redis.ResponseError):
       # Should not be able to update a non-existent task.
       self.redis.execute_command("RAY.TASK_TABLE_UPDATE", "task_id", 10,
@@ -219,39 +241,46 @@ def testTaskTableAddAndLookup(self):
     TASK_STATUS_SCHEDULED = 2
     TASK_STATUS_QUEUED = 4
 
+    def check_task_reply(message, task_args):
+      task_status, local_scheduler_id, task_spec = task_args
+      task_reply_object = TaskReply.GetRootAsTaskReply(message, 0)
+      self.assertEqual(task_reply_object.State(), task_status)
+      self.assertEqual(task_reply_object.LocalSchedulerId(), local_scheduler_id)
+      self.assertEqual(task_reply_object.TaskSpec(), task_spec)
+
     # Check that task table adds, updates, and lookups work correctly.
     task_args = [TASK_STATUS_WAITING, b"node_id", b"task_spec"]
     response = self.redis.execute_command("RAY.TASK_TABLE_ADD", "task_id",
                                           *task_args)
     response = self.redis.execute_command("RAY.TASK_TABLE_GET", "task_id")
-    self.assertEqual(response, task_args)
+    check_task_reply(response, task_args)
 
     task_args[0] = TASK_STATUS_SCHEDULED
     self.redis.execute_command("RAY.TASK_TABLE_UPDATE", "task_id", *task_args[:2])
     response = self.redis.execute_command("RAY.TASK_TABLE_GET", "task_id")
-    self.assertEqual(response, task_args)
+    check_task_reply(response, task_args)
 
     # If the current value, test value, and set value are all the same, the
     # update happens, and the response is still the same task.
     task_args = [task_args[0]] + task_args
     response = self.redis.execute_command("RAY.TASK_TABLE_TEST_AND_UPDATE",
                                           "task_id",
                                           *task_args[:3])
-    self.assertEqual(response, task_args[1:])
+    check_task_reply(response, task_args[1:])
     # Check that the task entry is still the same.
     get_response = self.redis.execute_command("RAY.TASK_TABLE_GET", "task_id")
-    self.assertEqual(get_response, task_args[1:])
+    check_task_reply(get_response, task_args[1:])
 
     # If the current value is the same as the test value, and the set value is
     # different, the update happens, and the response is the entire task.
     task_args[1] = TASK_STATUS_QUEUED
     response = self.redis.execute_command("RAY.TASK_TABLE_TEST_AND_UPDATE",
                                           "task_id",
                                           *task_args[:3])
-    self.assertEqual(response, task_args[1:])
+    check_task_reply(response, task_args[1:])
     # Check that the update happened.
     get_response = self.redis.execute_command("RAY.TASK_TABLE_GET", "task_id")
-    self.assertEqual(get_response, task_args[1:])
+    check_task_reply(get_response, task_args[1:])
 
     # If the current value is no longer the same as the test value, the
     # response is nil.
@@ -271,7 +300,7 @@ def testTaskTableAddAndLookup(self):
     response = self.redis.execute_command("RAY.TASK_TABLE_TEST_AND_UPDATE",
                                           "task_id",
                                           *task_args[:3])
-    self.assertEqual(response, task_args[1:])
+    check_task_reply(response, task_args[1:])
 
     # If the test value is a bitmask that does not match the current value, the
     # update does not happen.
@@ -288,24 +317,27 @@ def testTaskTableAddAndLookup(self):
 
   def testTaskTableSubscribe(self):
     scheduling_state = 1
-    node_id = "node_id"
+    local_scheduler_id = "local_scheduler_id"
     # Subscribe to the task table.
     p = self.redis.pubsub()
     p.psubscribe("{prefix}*:*".format(prefix=TASK_PREFIX))
-    p.psubscribe("{prefix}*:{state: >2}".format(prefix=TASK_PREFIX, state=scheduling_state))
-    p.psubscribe("{prefix}{node}:*".format(prefix=TASK_PREFIX, node=node_id))
-    task_args = [b"task_id", scheduling_state, node_id.encode("ascii"), b"task_spec"]
+    p.psubscribe("{prefix}*:{state}".format(prefix=TASK_PREFIX, state=scheduling_state))
+    p.psubscribe("{prefix}{local_scheduler_id}:*".format(prefix=TASK_PREFIX, local_scheduler_id=local_scheduler_id))
+    task_args = [b"task_id", scheduling_state, local_scheduler_id.encode("ascii"), b"task_spec"]
     self.redis.execute_command("RAY.TASK_TABLE_ADD", *task_args)
     # Receive the acknowledgement message.
     self.assertEqual(get_next_message(p)["data"], 1)
     self.assertEqual(get_next_message(p)["data"], 2)
     self.assertEqual(get_next_message(p)["data"], 3)
     # Receive the actual data.
     for i in range(3):
-        message = get_next_message(p)["data"]
-        message = message.split()
-        message[1] = int(message[1])
-        self.assertEqual(message, task_args)
+      message = get_next_message(p)["data"]
+      # Check that the notification object is correct.
+      notification_object = TaskReply.GetRootAsTaskReply(message, 0)
+      self.assertEqual(notification_object.TaskId(), b"task_id")
+      self.assertEqual(notification_object.State(), scheduling_state)
+      self.assertEqual(notification_object.LocalSchedulerId(), local_scheduler_id.encode("ascii"))
+      self.assertEqual(notification_object.TaskSpec(), b"task_spec")
 
 if __name__ == "__main__":
   unittest.main(verbosity=2)

python/ray/global_scheduler/test/test.py

@@ -34,7 +34,7 @@
 TASK_STATUS_RUNNING = 8
 TASK_STATUS_DONE = 16
 
-# These constants are an implementation detail of ray_redis_module.c, so this
+# These constants are an implementation detail of ray_redis_module.cc, so this
 # must be kept in sync with that file.
 DB_CLIENT_PREFIX = "CL:"
 TASK_PREFIX = "TT:"

python/ray/monitor.py

@@ -12,13 +12,16 @@
 from ray.services import get_ip_address
 from ray.services import get_port
 
+# Import flatbuffer bindings.
+from ray.core.generated.SubscribeToDBClientTableReply import SubscribeToDBClientTableReply
+
 # These variables must be kept in sync with the C codebase.
 # common/common.h
 DB_CLIENT_ID_SIZE = 20
 NIL_ID = b"\xff" * DB_CLIENT_ID_SIZE
 # common/task.h
 TASK_STATUS_LOST = 32
-# common/redis_module/ray_redis_module.c
+# common/redis_module/ray_redis_module.cc
 TASK_PREFIX = "TT:"
 DB_CLIENT_PREFIX = "CL:"
 DB_CLIENT_TABLE_NAME = b"db_clients"
@@ -89,14 +92,12 @@ def read_message(self):
 
     # Parse the message.
     data = message["data"]
-    db_client_id = data[:DB_CLIENT_ID_SIZE]
-    data = data[DB_CLIENT_ID_SIZE + 1:]
-    data = data.split(b" ")
-    client_type, auxiliary_address, is_insertion = data
-    is_insertion = int(is_insertion)
-    if is_insertion != 1 and is_insertion != 0:
-      raise Exception("Expected 0 or 1 for insertion field, got {} instead".format(is_insertion))
-    is_insertion = bool(is_insertion)
+
+    notification_object = SubscribeToDBClientTableReply.GetRootAsSubscribeToDBClientTableReply(data, 0)
+    db_client_id = notification_object.DbClientId()
+    client_type = notification_object.ClientType()
+    auxiliary_address = notification_object.AuxAddress()
+    is_insertion = notification_object.IsInsertion()
 
     return db_client_id, client_type, auxiliary_address, is_insertion
 

python/ray/worker.py

@@ -723,7 +723,7 @@ def get_address_info_from_redis_helper(redis_address, node_ip_address):
   # must have run "CONFIG SET protected-mode no".
   redis_client = redis.StrictRedis(host=redis_ip_address, port=int(redis_port))
   # The client table prefix must be kept in sync with the file
-  # "src/common/redis_module/ray_redis_module.c" where it is defined.
+  # "src/common/redis_module/ray_redis_module.cc" where it is defined.
   REDIS_CLIENT_TABLE_PREFIX = "CL:"
   client_keys = redis_client.keys("{}*".format(REDIS_CLIENT_TABLE_PREFIX))
   # Filter to clients on the same node and do some basic checking.

scripts/start_ray.py

@@ -21,7 +21,7 @@ def check_no_existing_redis_clients(node_ip_address, redis_address):
   redis_ip_address, redis_port = redis_address.split(":")
   redis_client = redis.StrictRedis(host=redis_ip_address, port=int(redis_port))
   # The client table prefix must be kept in sync with the file
-  # "src/common/redis_module/ray_redis_module.c" where it is defined.
+  # "src/common/redis_module/ray_redis_module.cc" where it is defined.
   REDIS_CLIENT_TABLE_PREFIX = "CL:"
   client_keys = redis_client.keys("{}*".format(REDIS_CLIENT_TABLE_PREFIX))
   # Filter to clients on the same node and do some basic checking.

src/common/CMakeLists.txt

@@ -24,6 +24,15 @@ add_custom_command(
   COMMENT "Running flatc compiler on ${COMMON_FBS_SRC}"
   VERBATIM)
 
+# Generate Python bindings for the flatbuffers objects.
+set(PYTHON_OUTPUT_DIR ${CMAKE_BINARY_DIR}/generated/)
+add_custom_command(
+  TARGET gen_common_fbs
+  COMMAND ${FLATBUFFERS_COMPILER} -p -o ${PYTHON_OUTPUT_DIR} ${COMMON_FBS_SRC}
+  DEPENDS ${FBS_DEPENDS}
+  COMMENT "Running flatc compiler on ${COMMON_FBS_SRC}"
+  VERBATIM)
+
 add_dependencies(gen_common_fbs flatbuffers_ep)
 
 add_custom_target(

src/common/format/common.fbs

@@ -50,3 +50,62 @@ table TaskInfo {
 }
 
 root_type TaskInfo;
+
+table SubscribeToNotificationsReply {
+  // The object ID of the object that the notification is about.
+  object_id: string;
+  // The size of the object.
+  object_size: long;
+  // The IDs of the managers that contain this object.
+  manager_ids: [string];
+}
+
+root_type SubscribeToNotificationsReply;
+
+table TaskReply {
+  // The task ID of the task that the message is about.
+  task_id: string;
+  // The state of the task. This is encoded as a bit mask of scheduling_state
+  // enum values in task.h.
+  state: long;
+  // A local scheduler ID.
+  local_scheduler_id: string;
+  // A string of bytes representing the task specification.
+  task_spec: string;
+}
+
+root_type TaskReply;
+
+table SubscribeToDBClientTableReply {
+  // The db client ID of the client that the message is about.
+  db_client_id: string;
+  // The type of the client.
+  client_type: string;
+  // If the client is a local scheduler, this is the address of the plasma
+  // manager that the local scheduler is connected to. Otherwise, it is empty.
+  aux_address: string;
+  // True if the message is about the addition of a client and false if it is
+  // about the deletion of a client.
+  is_insertion: bool;
+}
+
+root_type SubscribeToDBClientTableReply;
+
+table LocalSchedulerInfoMessage {
+  // The db client ID of the client that the message is about.
+  db_client_id: string;
+  // The total number of workers that are connected to this local scheduler.
+  total_num_workers: long;
+  // The number of tasks queued in this local scheduler.
+  task_queue_length: long;
+  // The number of workers that are available and waiting for tasks.
+  available_workers: long;
+  // The resource vector of resources generally available to this local
+  // scheduler.
+  static_resources: [double];
+  // The resource vector of resources currently available to this local
+  // scheduler.
+  dynamic_resources: [double];
+}
+
+root_type LocalSchedulerInfoMessage;