Define common data structures with protobuf. (#5121)

python/ray/includes/common.pxd

@@ -87,17 +87,14 @@ cdef extern from "ray/common/id.h" namespace "ray" nogil:
                                  int parent_task_counter)
 
 
-cdef extern from "ray/gcs/format/gcs_generated.h" nogil:
-    cdef cppclass GCSArg "Arg":
-        pass
-
+cdef extern from "ray/protobuf/common.pb.h" nogil:
     cdef cppclass CLanguage "Language":
         pass
 
 
 # This is a workaround for C++ enum class since Cython has no corresponding
 # representation.
-cdef extern from "ray/gcs/format/gcs_generated.h" namespace "Language" nogil:
+cdef extern from "ray/protobuf/common.pb.h" namespace "Language" nogil:
     cdef CLanguage LANGUAGE_PYTHON "Language::PYTHON"
     cdef CLanguage LANGUAGE_CPP "Language::CPP"
     cdef CLanguage LANGUAGE_JAVA "Language::JAVA"

python/ray/includes/libraylet.pxd

@@ -19,23 +19,23 @@ from ray.includes.unique_ids cimport (
     CObjectID,
     CTaskID,
 )
-from ray.includes.task cimport CTaskSpecification
+from ray.includes.task cimport CTaskSpec
 
 
-cdef extern from "ray/gcs/format/gcs_generated.h" nogil:
-    cdef cppclass GCSProfileEventT "ProfileEventT":
-        c_string event_type
-        double start_time
-        double end_time
-        c_string extra_data
-        GCSProfileEventT()
+cdef extern from "ray/protobuf/gcs.pb.h" nogil:
+    cdef cppclass GCSProfileEvent "ProfileTableData::ProfileEvent":
+        void set_event_type(const c_string &value)
+        void set_start_time(double value)
+        void set_end_time(double value)
+        c_string set_extra_data(const c_string &value)
+        GCSProfileEvent()
 
-    cdef cppclass GCSProfileTableDataT "ProfileTableDataT":
-        c_string component_type
-        c_string component_id
-        c_string node_ip_address
-        c_vector[unique_ptr[GCSProfileEventT]] profile_events
-        GCSProfileTableDataT()
+    cdef cppclass GCSProfileTableData "ProfileTableData":
+        void set_component_type(const c_string &value)
+        void set_component_id(const c_string &value)
+        void set_node_ip_address(const c_string &value)
+        GCSProfileEvent *add_profile_events()
+        GCSProfileTableData()
 
 
 ctypedef unordered_map[c_string, c_vector[pair[int64_t, double]]] \
@@ -52,8 +52,8 @@ cdef extern from "ray/raylet/raylet_client.h" nogil:
         CRayStatus Disconnect()
         CRayStatus SubmitTask(
             const c_vector[CObjectID] &execution_dependencies,
-            const CTaskSpecification &task_spec)
-        CRayStatus GetTask(unique_ptr[CTaskSpecification] *task_spec)
+            const CTaskSpec &task_spec)
+        CRayStatus GetTask(unique_ptr[CTaskSpec] *task_spec)
         CRayStatus TaskDone()
         CRayStatus FetchOrReconstruct(c_vector[CObjectID] &object_ids,
                                       c_bool fetch_only,
@@ -66,7 +66,7 @@ cdef extern from "ray/raylet/raylet_client.h" nogil:
         CRayStatus PushError(const CJobID &job_id, const c_string &type,
                              const c_string &error_message, double timestamp)
         CRayStatus PushProfileEvents(
-            const GCSProfileTableDataT &profile_events)
+            const GCSProfileTableData &profile_events)
         CRayStatus FreeObjects(const c_vector[CObjectID] &object_ids,
                                c_bool local_only, c_bool delete_creating_tasks)
         CRayStatus PrepareActorCheckpoint(const CActorID &actor_id,

python/ray/includes/task.pxd

@@ -1,4 +1,4 @@
-from libc.stdint cimport int64_t, uint8_t
+from libc.stdint cimport uint8_t, uint64_t
 from libcpp cimport bool as c_bool
 from libcpp.memory cimport unique_ptr, shared_ptr
 from libcpp.string cimport string as c_string
@@ -17,72 +17,45 @@ from ray.includes.unique_ids cimport (
     CTaskID,
 )
 
+cdef extern from "ray/protobuf/common.pb.h" namespace "ray::rpc" nogil:
+    cdef cppclass RpcTaskSpec "ray::rpc::TaskSpec":
+        void CopyFrom(const RpcTaskSpec &value)
 
-cdef extern from "ray/raylet/task_execution_spec.h" \
-        namespace "ray::raylet" nogil:
-    cdef cppclass CTaskExecutionSpecification \
-            "ray::raylet::TaskExecutionSpecification":
-        CTaskExecutionSpecification(const c_vector[CObjectID] &&dependencies)
-        CTaskExecutionSpecification(
-            const c_vector[CObjectID] &&dependencies, int num_forwards)
-        c_vector[CObjectID] ExecutionDependencies() const
-        void SetExecutionDependencies(const c_vector[CObjectID] &dependencies)
-        int NumForwards() const
-        void IncrementNumForwards()
-        int64_t LastTimestamp() const
-        void SetLastTimestamp(int64_t new_timestamp)
+    cdef cppclass RpcTaskExecutionSpec "ray::rpc::TaskExecutionSpec":
+        void CopyFrom(const RpcTaskExecutionSpec &value)
+        void add_dependencies(const c_string &value)
+
+    cdef cppclass RpcTask "ray::rpc::Task":
+        RpcTaskSpec *mutable_task_spec()
+
+
+cdef extern from "ray/protobuf/gcs.pb.h" namespace "ray::rpc" nogil:
+    cdef cppclass TaskTableData "ray::rpc::TaskTableData":
+        RpcTask *mutable_task()
+        const c_string &SerializeAsString()
 
 
 cdef extern from "ray/raylet/task_spec.h" namespace "ray::raylet" nogil:
-    cdef cppclass CTaskArgument "ray::raylet::TaskArgument":
-        pass
-
-    cdef cppclass CTaskArgumentByReference \
-            "ray::raylet::TaskArgumentByReference":
-        CTaskArgumentByReference(const c_vector[CObjectID] &references)
-
-    cdef cppclass CTaskArgumentByValue "ray::raylet::TaskArgumentByValue":
-        CTaskArgumentByValue(const uint8_t *value, size_t length)
-
-    cdef cppclass CTaskSpecification "ray::raylet::TaskSpecification":
-        CTaskSpecification(
-            const CJobID &job_id, const CTaskID &parent_task_id,
-            int64_t parent_counter,
-            const c_vector[shared_ptr[CTaskArgument]] &task_arguments,
-            int64_t num_returns,
-            const unordered_map[c_string, double] &required_resources,
-            const CLanguage &language,
-            const c_vector[c_string] &function_descriptor)
-        CTaskSpecification(
-            const CJobID &job_id, const CTaskID &parent_task_id,
-            int64_t parent_counter, const CActorID &actor_creation_id,
-            const CObjectID &actor_creation_dummy_object_id,
-            int64_t max_actor_reconstructions, const CActorID &actor_id,
-            const CActorHandleID &actor_handle_id, int64_t actor_counter,
-            const c_vector[CActorHandleID] &new_actor_handles,
-            const c_vector[shared_ptr[CTaskArgument]] &task_arguments,
-            int64_t num_returns,
-            const unordered_map[c_string, double] &required_resources,
-            const unordered_map[c_string, double] &required_placement_res,
-            const CLanguage &language,
-            const c_vector[c_string] &function_descriptor)
-        CTaskSpecification(const c_string &string)
-        c_string SerializeAsString() const
+    cdef cppclass CTaskSpec "ray::raylet::TaskSpecification":
+        CTaskSpec(const RpcTaskSpec message)
+        CTaskSpec(const c_string &serialized_binary)
+        const RpcTaskSpec &GetMessage()
+        c_string Serialize() const
 
         CTaskID TaskId() const
         CJobID JobId() const
         CTaskID ParentTaskId() const
-        int64_t ParentCounter() const
+        uint64_t ParentCounter() const
         c_vector[c_string] FunctionDescriptor() const
         c_string FunctionDescriptorString() const
-        int64_t NumArgs() const
-        int64_t NumReturns() const
-        c_bool ArgByRef(int64_t arg_index) const
-        int ArgIdCount(int64_t arg_index) const
-        CObjectID ArgId(int64_t arg_index, int64_t id_index) const
-        CObjectID ReturnId(int64_t return_index) const
-        const uint8_t *ArgVal(int64_t arg_index) const
-        size_t ArgValLength(int64_t arg_index) const
+        uint64_t NumArgs() const
+        uint64_t NumReturns() const
+        c_bool ArgByRef(uint64_t arg_index) const
+        int ArgIdCount(uint64_t arg_index) const
+        CObjectID ArgId(uint64_t arg_index, uint64_t id_index) const
+        CObjectID ReturnId(uint64_t return_index) const
+        const uint8_t *ArgVal(uint64_t arg_index) const
+        size_t ArgValLength(uint64_t arg_index) const
         double GetRequiredResource(const c_string &resource_name) const
         const ResourceSet GetRequiredResources() const
         const ResourceSet GetRequiredPlacementResources() const
@@ -93,25 +66,46 @@ cdef extern from "ray/raylet/task_spec.h" namespace "ray::raylet" nogil:
         c_bool IsActorTask() const
         CActorID ActorCreationId() const
         CObjectID ActorCreationDummyObjectId() const
-        int64_t MaxActorReconstructions() const
+        uint64_t MaxActorReconstructions() const
         CActorID ActorId() const
         CActorHandleID ActorHandleId() const
-        int64_t ActorCounter() const
+        uint64_t ActorCounter() const
         CObjectID ActorDummyObject() const
         c_vector[CActorHandleID] NewActorHandles() const
 
 
+cdef extern from "ray/raylet/task_util.h" namespace "ray::raylet" nogil:
+    cdef cppclass TaskSpecBuilder "ray::raylet::TaskSpecBuilder":
+        TaskSpecBuilder &SetCommonTaskSpec(
+            const CLanguage &language, const c_vector[c_string] &function_descriptor,
+            const CJobID &job_id, const CTaskID &parent_task_id, uint64_t parent_counter,
+            uint64_t num_returns, const unordered_map[c_string, double] &required_resources,
+            const unordered_map[c_string, double] &required_placement_resources);
+
+        TaskSpecBuilder &AddByRefArg(const CObjectID &arg_id);
+
+        TaskSpecBuilder &AddByValueArg(const c_string &data);
+
+        TaskSpecBuilder &SetActorCreationTaskSpec(
+            const CActorID &actor_id, uint64_t max_reconstructions,
+            const c_vector[c_string] &dynamic_worker_options);
+
+        TaskSpecBuilder &SetActorTaskSpec(
+            const CActorID &actor_id, const CActorHandleID &actor_handle_id,
+            const CObjectID &actor_creation_dummy_object_id, uint64_t actor_counter,
+            const c_vector[CActorHandleID] &new_handle_ids);
+
+        RpcTaskSpec GetMessage();
+
+
+cdef extern from "ray/raylet/task_execution_spec.h" namespace "ray::raylet" nogil:
+    cdef cppclass CTaskExecutionSpec "ray::raylet::TaskExecutionSpecification":
+        CTaskExecutionSpec(RpcTaskExecutionSpec message)
+        CTaskExecutionSpec(const c_string &serialized_binary)
+        const RpcTaskExecutionSpec &GetMessage()
+        c_vector[CObjectID] ExecutionDependencies()
+        uint64_t NumForwards()
+
 cdef extern from "ray/raylet/task.h" namespace "ray::raylet" nogil:
     cdef cppclass CTask "ray::raylet::Task":
-        CTask(const CTaskExecutionSpecification &execution_spec,
-              const CTaskSpecification &task_spec)
-        const CTaskExecutionSpecification &GetTaskExecutionSpec() const
-        const CTaskSpecification &GetTaskSpecification() const
-        void SetExecutionDependencies(const c_vector[CObjectID] &dependencies)
-        void IncrementNumForwards()
-        const c_vector[CObjectID] &GetDependencies() const
-        void CopyTaskExecutionSpec(const CTask &task)
-
-    cdef c_string SerializeTaskAsString(
-        const c_vector[CObjectID] *dependencies,
-        const CTaskSpecification *task_spec)
+        CTask(CTaskSpec task_spec, CTaskExecutionSpec task_execution_spec)

python/ray/includes/task.pxi

@@ -5,18 +5,19 @@ from libcpp.memory cimport (
     static_pointer_cast,
 )
 from ray.includes.task cimport (
-    CTaskArgument,
-    CTaskArgumentByReference,
-    CTaskArgumentByValue,
-    CTaskSpecification,
-    SerializeTaskAsString,
+    CTask,
+    CTaskExecutionSpec,
+    CTaskSpec,
+    RpcTaskExecutionSpec,
+    TaskSpecBuilder,
+    TaskTableData,
 )
 
 
-cdef class Task:
+cdef class TaskSpec:
+    """Cython wrapper class of C++ `ray::raylet::TaskSpecification`."""
     cdef:
-        unique_ptr[CTaskSpecification] task_spec
-        unique_ptr[c_vector[CObjectID]] execution_dependencies
+        unique_ptr[CTaskSpec] task_spec
 
     def __init__(self, JobID job_id, function_descriptor, arguments,
                  int num_returns, TaskID parent_task_id, int parent_counter,
@@ -24,73 +25,78 @@ cdef class Task:
                  ObjectID actor_creation_dummy_object_id,
                  int32_t max_actor_reconstructions, ActorID actor_id,
                  ActorHandleID actor_handle_id, int actor_counter,
-                 new_actor_handles, execution_arguments, resource_map,
-                 placement_resource_map):
+                 new_actor_handles, resource_map, placement_resource_map):
         cdef:
+            TaskSpecBuilder builder
             unordered_map[c_string, double] required_resources
             unordered_map[c_string, double] required_placement_resources
-            c_vector[shared_ptr[CTaskArgument]] task_args
-            c_vector[CActorHandleID] task_new_actor_handles
             c_vector[c_string] c_function_descriptor
             c_string pickled_str
-            c_vector[CObjectID] references
+            c_vector[CActorHandleID] c_new_actor_handles
 
+        # Convert function descriptor to C++ vector.
         for item in function_descriptor:
             if not isinstance(item, bytes):
                 raise TypeError(
                     "'function_descriptor' takes a list of byte strings.")
             c_function_descriptor.push_back(item)
 
-        # Parse the resource map.
+        # Convert resource map to C++ unordered_map.
         if resource_map is not None:
             required_resources = resource_map_from_dict(resource_map)
         if placement_resource_map is not None:
             required_placement_resources = (
                 resource_map_from_dict(placement_resource_map))
 
-        # Parse the arguments from the list.
+        # Build common task spec.
+        builder.SetCommonTaskSpec(
+            LANGUAGE_PYTHON,
+            c_function_descriptor,
+            job_id.native(),
+            parent_task_id.native(),
+            parent_counter,
+            num_returns,
+            required_resources,
+            required_placement_resources,
+        )
+
+        # Build arguments.
         for arg in arguments:
             if isinstance(arg, ObjectID):
-                references = c_vector[CObjectID]()
-                references.push_back((<ObjectID>arg).native())
-                task_args.push_back(
-                    static_pointer_cast[CTaskArgument,
-                                        CTaskArgumentByReference](
-                        make_shared[CTaskArgumentByReference](references)))
+                builder.AddByRefArg((<ObjectID>arg).native())
             else:
                 pickled_str = pickle.dumps(
                     arg, protocol=pickle.HIGHEST_PROTOCOL)
-                task_args.push_back(
-                    static_pointer_cast[CTaskArgument,
-                                        CTaskArgumentByValue](
-                        make_shared[CTaskArgumentByValue](
-                            <uint8_t *>pickled_str.c_str(),
-                            pickled_str.size())))
+                builder.AddByValueArg(pickled_str)
 
-        for new_actor_handle in new_actor_handles:
-            task_new_actor_handles.push_back(
-                (<ActorHandleID?>new_actor_handle).native())
-
-        self.task_spec.reset(new CTaskSpecification(
-            job_id.native(), parent_task_id.native(), parent_counter, actor_creation_id.native(),
-            actor_creation_dummy_object_id.native(), max_actor_reconstructions, actor_id.native(),
-            actor_handle_id.native(), actor_counter, task_new_actor_handles, task_args, num_returns,
-            required_resources, required_placement_resources, LANGUAGE_PYTHON,
-            c_function_descriptor))
-
-        # Set the task's execution dependencies.
-        self.execution_dependencies.reset(new c_vector[CObjectID]())
-        if execution_arguments is not None:
-            for execution_arg in execution_arguments:
-                self.execution_dependencies.get().push_back(
-                    (<ObjectID?>execution_arg).native())
+        if not actor_creation_id.is_nil():
+            # Actor creation task.
+            builder.SetActorCreationTaskSpec(
+                actor_creation_id.native(),
+                max_actor_reconstructions,
+                [],
+            )
+        elif not actor_id.is_nil():
+            # Actor task.
+            for new_actor_handle in new_actor_handles:
+                c_new_actor_handles.push_back(
+                    (<ActorHandleID?>new_actor_handle).native())
+            builder.SetActorTaskSpec(
+                actor_id.native(),
+                actor_handle_id.native(),
+                actor_creation_dummy_object_id.native(),
+                actor_counter,
+                c_new_actor_handles,
+            )
+        else:
+            # Normal task.
+            pass
+        self.task_spec.reset(new CTaskSpec(builder.GetMessage()))
 
     @staticmethod
-    cdef make(unique_ptr[CTaskSpecification]& task_spec):
-        cdef Task self = Task.__new__(Task)
+    cdef make(unique_ptr[CTaskSpec]& task_spec):
+        cdef TaskSpec self = TaskSpec.__new__(TaskSpec)
         self.task_spec.reset(task_spec.release())
-        # The created task does not include any execution dependencies.
-        self.execution_dependencies.reset(new c_vector[CObjectID]())
         return self
 
     @staticmethod
@@ -103,11 +109,8 @@ cdef class Task:
         Returns:
             Python task specification object.
         """
-        cdef Task self = Task.__new__(Task)
-        # TODO(pcm): Use flatbuffers validation here.
-        self.task_spec.reset(new CTaskSpecification(task_spec_str))
-        # The created task does not include any execution dependencies.
-        self.execution_dependencies.reset(new c_vector[CObjectID]())
+        cdef TaskSpec self = TaskSpec.__new__(TaskSpec)
+        self.task_spec.reset(new CTaskSpec(task_spec_str))
         return self
 
     def to_string(self):
@@ -116,11 +119,7 @@ cdef class Task:
         Returns:
             String representing the task specification.
         """
-        return self.task_spec.get().SerializeAsString()
-
-    def _serialized_raylet_task(self):
-        return SerializeTaskAsString(
-            self.execution_dependencies.get(), self.task_spec.get())
+        return self.task_spec.get().Serialize()
 
     def job_id(self):
         """Return the job ID for this task."""
@@ -150,7 +149,7 @@ cdef class Task:
     def arguments(self):
         """Return the arguments for the task."""
         cdef:
-            CTaskSpecification *task_spec = self.task_spec.get()
+            CTaskSpec*task_spec = self.task_spec.get()
             int64_t num_args = task_spec.NumArgs()
             int32_t lang = <int32_t>task_spec.GetLanguage()
             int count
@@ -175,7 +174,7 @@ cdef class Task:
 
     def returns(self):
         """Return the object IDs for the return values of the task."""
-        cdef CTaskSpecification *task_spec = self.task_spec.get()
+        cdef CTaskSpec *task_spec = self.task_spec.get()
         return_id_list = []
         for i in range(task_spec.NumReturns()):
             return_id_list.append(ObjectID(task_spec.ReturnId(i).Binary()))
@@ -221,3 +220,59 @@ cdef class Task:
     def actor_counter(self):
         """Return the actor counter for this task."""
         return self.task_spec.get().ActorCounter()
+
+
+cdef class TaskExecutionSpec:
+    """Cython wrapper class of C++ `ray::raylet::TaskExecutionSpecification`."""
+    cdef:
+        unique_ptr[CTaskExecutionSpec] c_spec
+
+    def __init__(self, execution_dependencies):
+        cdef:
+            RpcTaskExecutionSpec message;
+
+        for dependency in execution_dependencies:
+            message.add_dependencies(
+                (<ObjectID?>dependency).binary())
+        self.c_spec.reset(new CTaskExecutionSpec(message))
+
+    @staticmethod
+    def from_string(const c_string& string):
+        """Convert a string to a Ray `TaskExecutionSpec` Python object.
+        """
+        cdef TaskExecutionSpec self = TaskExecutionSpec.__new__(TaskExecutionSpec)
+        self.c_spec.reset(new CTaskExecutionSpec(string))
+        return self
+
+    def dependencies(self):
+        cdef:
+            CObjectID c_id
+            c_vector[CObjectID] dependencies = (
+                self.c_spec.get().ExecutionDependencies())
+        ret = []
+        for c_id in dependencies:
+            ret.append(ObjectID(c_id.Binary()))
+        return ret
+
+    def num_forwards(self):
+        return self.c_spec.get().NumForwards()
+
+
+cdef class Task:
+    """Cython wrapper class of C++ `ray::raylet::Task`."""
+    cdef:
+        unique_ptr[CTask] c_task
+
+    def __init__(self, TaskSpec task_spec, TaskExecutionSpec task_execution_spec):
+        self.c_task.reset(new CTask(task_spec.task_spec.get()[0],
+                                    task_execution_spec.c_spec.get()[0]))
+
+
+def generate_gcs_task_table_data(TaskSpec task_spec):
+    """Converts a Python `TaskSpec` object to serialized GCS `TaskTableData`.
+    """
+    cdef:
+        TaskTableData task_table_data
+    task_table_data.mutable_task().mutable_task_spec().CopyFrom(
+        task_spec.task_spec.get().GetMessage())
+    return task_table_data.SerializeAsString()