Define common data structures with protobuf. (#5121)

2026-06-28 02:46:49 +08:00 · 2019-07-08 22:41:37 +08:00
parent b4e51c8aa1
commit 8a30b93e42
64 changed files with 1233 additions and 1561 deletions
@@ -24,8 +24,8 @@ from ray.includes.common cimport (
 )
 from ray.includes.libraylet cimport (
    CRayletClient,
-    GCSProfileEventT,
-    GCSProfileTableDataT,
+    GCSProfileEvent,
+    GCSProfileTableData,
    ResourceMappingType,
    WaitResultPair,
 )
@@ -34,7 +34,7 @@ from ray.includes.unique_ids cimport (
    CObjectID,
    CClientID,
 )
-from ray.includes.task cimport CTaskSpecification
+from ray.includes.task cimport CTaskSpec
 from ray.includes.ray_config cimport RayConfig
 from ray.utils import decode

@@ -232,18 +232,22 @@ cdef class RayletClient:
    def disconnect(self):
        check_status(self.client.get().Disconnect())

-    def submit_task(self, Task task_spec):
+    def submit_task(self, TaskSpec task_spec, execution_dependencies):
+        cdef:
+            CObjectID c_id
+            c_vector[CObjectID] c_dependencies
+        for dep in execution_dependencies:
+            c_dependencies.push_back((<ObjectID>dep).native())
        check_status(self.client.get().SubmitTask(
-            task_spec.execution_dependencies.get()[0],
-            task_spec.task_spec.get()[0]))
+            c_dependencies, task_spec.task_spec.get()[0]))

    def get_task(self):
        cdef:
-            unique_ptr[CTaskSpecification] task_spec
+            unique_ptr[CTaskSpec] task_spec

        with nogil:
            check_status(self.client.get().GetTask(&task_spec))
-        return Task.make(task_spec)
+        return TaskSpec.make(task_spec)

    def task_done(self):
        check_status(self.client.get().TaskDone())
@@ -303,19 +307,19 @@ cdef class RayletClient:
    def push_profile_events(self, component_type, UniqueID component_id,
                            node_ip_address, profile_data):
        cdef:
-            GCSProfileTableDataT profile_info
-            GCSProfileEventT *profile_event
+            GCSProfileTableData profile_info
+            GCSProfileEvent *profile_event
            c_string event_type

        if len(profile_data) == 0:
            return  # Short circuit if there are no profile events.

-        profile_info.component_type = component_type.encode("ascii")
-        profile_info.component_id = component_id.binary()
-        profile_info.node_ip_address = node_ip_address.encode("ascii")
+        profile_info.set_component_type(component_type.encode("ascii"))
+        profile_info.set_component_id(component_id.binary())
+        profile_info.set_node_ip_address(node_ip_address.encode("ascii"))

        for py_profile_event in profile_data:
-            profile_event = new GCSProfileEventT()
+            profile_event = profile_info.add_profile_events()
            if not isinstance(py_profile_event, dict):
                raise TypeError(
                    "Incorrect type for a profile event. Expected dict "
@@ -325,28 +329,22 @@ cdef class RayletClient:
            # that will cause segfaults in the node manager.
            for key_string, event_data in py_profile_event.items():
                if key_string == "event_type":
-                    profile_event.event_type = event_data.encode("ascii")
-                    if profile_event.event_type.length() == 0:
+                    if len(event_data) == 0:
                        raise ValueError(
                            "'event_type' should not be a null string.")
+                    profile_event.set_event_type(event_data.encode("ascii"))
                elif key_string == "start_time":
-                    profile_event.start_time = float(event_data)
+                    profile_event.set_start_time(float(event_data))
                elif key_string == "end_time":
-                    profile_event.end_time = float(event_data)
+                    profile_event.set_end_time(float(event_data))
                elif key_string == "extra_data":
-                    profile_event.extra_data = event_data.encode("ascii")
-                    if profile_event.extra_data.length() == 0:
+                    if len(event_data) == 0:
                        raise ValueError(
                            "'extra_data' should not be a null string.")
+                    profile_event.set_extra_data(event_data.encode("ascii"))
                else:
                    raise ValueError(
                        "Unknown profile event key '%s'" % key_string)
-            # 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. "emplace_back" of vector has not been
-            # supported by Cython
-            profile_info.profile_events.push_back(
-                unique_ptr[GCSProfileEventT](profile_event))

        check_status(self.client.get().PushProfileEvents(profile_info))

@@ -2,8 +2,6 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function

-from ray.core.generated.ray.protocol.Task import Task
-
 from ray.core.generated.gcs_pb2 import (
    ActorCheckpointIdData,
    ClientTableData,
@@ -33,7 +31,6 @@ __all__ = [
    "ProfileTableData",
    "TablePrefix",
    "TablePubsub",
-    "Task",
    "TaskTableData",
    "construct_error_message",
 ]
@@ -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"
@@ -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,
@@ -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)
@@ -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()
@@ -305,12 +305,9 @@ class GlobalState(object):
        assert len(gcs_entries.entries) == 1
        task_table_data = gcs_utils.TaskTableData.FromString(
            gcs_entries.entries[0])
-        task_table_message = gcs_utils.Task.GetRootAsTask(
-            task_table_data.task, 0)

-        execution_spec = task_table_message.TaskExecutionSpec()
-        task_spec = task_table_message.TaskSpecification()
-        task = ray._raylet.Task.from_string(task_spec)
+        task = ray._raylet.TaskSpec.from_string(
+            task_table_data.task.task_spec.SerializeToString())
        function_descriptor_list = task.function_descriptor_list()
        function_descriptor = FunctionDescriptor.from_bytes_list(
            function_descriptor_list)
@@ -335,14 +332,12 @@ class GlobalState(object):
            "FunctionName": function_descriptor.function_name,
        }

+        execution_spec = ray._raylet.TaskExecutionSpec.from_string(
+            task_table_data.task.task_execution_spec.SerializeToString())
        return {
            "ExecutionSpec": {
-                "Dependencies": [
-                    execution_spec.Dependencies(i)
-                    for i in range(execution_spec.DependenciesLength())
-                ],
-                "LastTimestamp": execution_spec.LastTimestamp(),
-                "NumForwards": execution_spec.NumForwards()
+                "Dependencies": execution_spec.dependencies(),
+                "NumForwards": execution_spec.num_forwards(),
            },
            "TaskSpec": task_spec_info
        }
@@ -688,7 +688,7 @@ class Worker(object):
            function_descriptor_list = (
                function_descriptor.get_function_descriptor_list())
            assert isinstance(job_id, JobID)
-            task = ray._raylet.Task(
+            task = ray._raylet.TaskSpec(
                job_id,
                function_descriptor_list,
                args_for_raylet,
@@ -702,11 +702,10 @@ class Worker(object):
                actor_handle_id,
                actor_counter,
                new_actor_handles,
-                execution_dependencies,
                resources,
                placement_resources,
            )
-            self.raylet_client.submit_task(task)
+            self.raylet_client.submit_task(task, execution_dependencies)

            return task.returns()

@@ -1864,7 +1863,7 @@ def connect(node,
        nil_actor_counter = 0

        function_descriptor = FunctionDescriptor.for_driver_task()
-        driver_task = ray._raylet.Task(
+        driver_task_spec = ray._raylet.TaskSpec(
            worker.current_job_id,
            function_descriptor.get_function_descriptor_list(),
            [],  # arguments.
@@ -1878,24 +1877,25 @@ def connect(node,
            ActorHandleID.nil(),  # actor_handle_id.
            nil_actor_counter,  # actor_counter.
            [],  # new_actor_handles.
-            [],  # execution_dependencies.
            {},  # resource_map.
            {},  # placement_resource_map.
        )
-        task_table_data = ray.gcs_utils.TaskTableData()
-        task_table_data.task = driver_task._serialized_raylet_task()
+        task_table_data = ray._raylet.generate_gcs_task_table_data(
+            driver_task_spec)

        # Add the driver task to the task table.
        ray.state.state._execute_command(
-            driver_task.task_id(), "RAY.TABLE_ADD",
+            driver_task_spec.task_id(),
+            "RAY.TABLE_ADD",
            ray.gcs_utils.TablePrefix.Value("RAYLET_TASK"),
            ray.gcs_utils.TablePubsub.Value("RAYLET_TASK_PUBSUB"),
-            driver_task.task_id().binary(),
-            task_table_data.SerializeToString())
+            driver_task_spec.task_id().binary(),
+            task_table_data,
+        )

        # Set the driver's current task ID to the task ID assigned to the
        # driver task.
-        worker.task_context.current_task_id = driver_task.task_id()
+        worker.task_context.current_task_id = driver_task_spec.task_id()

    worker.raylet_client = ray._raylet.RayletClient(
        node.raylet_socket_name,
@@ -28,11 +28,10 @@ ray_files = [
    "ray/dashboard/res/main.css", "ray/dashboard/res/main.js"
 ]

-# These are the directories where automatically generated Python flatbuffer
+# These are the directories where automatically generated Python protobuf
 # bindings are created.
 generated_python_directories = [
-    "ray/core/generated", "ray/core/generated/ray",
-    "ray/core/generated/ray/protocol"
+    "ray/core/generated",
 ]

 optional_ray_files = []
@@ -88,7 +87,7 @@ class build_ext(_build_ext.build_ext):

        files_to_include = ray_files + pyarrow_files + modin_files

-        # Copy over the autogenerated flatbuffer Python bindings.
+        # Copy over the autogenerated protobuf Python bindings.
        for directory in generated_python_directories:
            for filename in os.listdir(directory):
                if filename[-3:] == ".py":
@@ -148,7 +147,6 @@ requires = [
    # NOTE: Don't upgrade the version of six! Doing so causes installation
    # problems. See https://github.com/ray-project/ray/issues/4169.
    "six >= 1.0.0",
-    "flatbuffers",
    "faulthandler;python_version<'3.3'",
    "protobuf >= 3.8.0",
 ]