[core worker] Python core worker object interface (#5272)

python/ray/includes/buffer.pxi

@@ -0,0 +1,79 @@
+from cpython cimport Py_buffer, PyBytes_FromStringAndSize
+from libc.stdint cimport int64_t, uintptr_t
+from libc.stdio cimport printf
+from libcpp.memory cimport shared_ptr
+
+from ray.includes.common cimport CBuffer
+
+
+cdef class Buffer:
+    """Cython wrapper class of C++ `ray::Buffer`.
+
+    This class implements the Python 'buffer protocol', which allows
+    us to use it for calls into pyarrow (and other Python libraries
+    down the line) without having to copy the data.
+
+    See https://docs.python.org/3/c-api/buffer.html for details.
+    """
+    cdef:
+        shared_ptr[CBuffer] buffer
+        Py_ssize_t shape
+        Py_ssize_t strides
+
+    @staticmethod
+    cdef make(const shared_ptr[CBuffer]& buffer):
+        cdef Buffer self = Buffer.__new__(Buffer)
+        self.buffer = buffer
+        self.shape = <Py_ssize_t>self.size
+        self.strides = <Py_ssize_t>(1)
+        return self
+
+    def __len__(self):
+        return self.size
+
+    @property
+    def size(self):
+        """
+        The buffer size in bytes.
+        """
+        return self.buffer.get().Size()
+
+    def to_pybytes(self):
+        """
+        Return this buffer as a Python bytes object.  Memory is copied.
+        """
+        return PyBytes_FromStringAndSize(
+            <const char*>self.buffer.get().Data(),
+            self.buffer.get().Size())
+
+    def __getbuffer__(self, Py_buffer* buffer, int flags):
+        buffer.readonly = 0
+        buffer.buf = <char *>self.buffer.get().Data()
+        buffer.format = 'b'
+        buffer.internal = NULL
+        buffer.itemsize = 1
+        buffer.len = self.size
+        buffer.ndim = 1
+        buffer.obj = self
+        buffer.shape = &self.shape
+        buffer.strides = &self.strides
+        buffer.suboffsets = NULL
+
+    def __getsegcount__(self, Py_ssize_t *len_out):
+        if len_out != NULL:
+            len_out[0] = <Py_ssize_t>self.size
+        return 1
+
+    def __getreadbuffer__(self, Py_ssize_t idx, void **p):
+        if idx != 0:
+            raise SystemError("accessing non-existent buffer segment")
+        if p != NULL:
+            p[0] = <void*> self.buffer.get().Data()
+        return self.size
+
+    def __getwritebuffer__(self, Py_ssize_t idx, void **p):
+        if idx != 0:
+            raise SystemError("accessing non-existent buffer segment")
+        if p != NULL:
+            p[0] = <void*> self.buffer.get().Data()
+        return self.size

python/ray/includes/common.pxd

@@ -1,7 +1,8 @@
-from libcpp.string cimport string as c_string
 from libcpp cimport bool as c_bool
+from libcpp.memory cimport shared_ptr
+from libcpp.string cimport string as c_string
 
-from libc.stdint cimport int64_t
+from libc.stdint cimport uint8_t
 from libcpp.unordered_map cimport unordered_map
 from libcpp.vector cimport vector as c_vector
 
@@ -49,6 +50,9 @@ cdef extern from "ray/common/status.h" namespace "ray" nogil:
         @staticmethod
         CRayStatus RedisError()
 
+        @staticmethod
+        CRayStatus ObjectStoreFull()
+
         c_bool ok()
         c_bool IsOutOfMemory()
         c_bool IsKeyError()
@@ -58,6 +62,7 @@ cdef extern from "ray/common/status.h" namespace "ray" nogil:
         c_bool IsUnknownError()
         c_bool IsNotImplemented()
         c_bool IsRedisError()
+        c_bool IsObjectStoreFull()
 
         c_string ToString()
         c_string CodeAsString()
@@ -90,19 +95,24 @@ cdef extern from "ray/common/id.h" namespace "ray" nogil:
 cdef extern from "ray/protobuf/common.pb.h" nogil:
     cdef cppclass CLanguage "Language":
         pass
+    cdef cppclass CWorkerType "ray::WorkerType":
+        pass
 
 
 # This is a workaround for C++ enum class since Cython has no corresponding
 # representation.
-cdef extern from "ray/protobuf/common.pb.h" namespace "Language" nogil:
+cdef extern from "ray/protobuf/common.pb.h" nogil:
     cdef CLanguage LANGUAGE_PYTHON "Language::PYTHON"
     cdef CLanguage LANGUAGE_CPP "Language::CPP"
     cdef CLanguage LANGUAGE_JAVA "Language::JAVA"
 
+cdef extern from "ray/protobuf/common.pb.h" nogil:
+    cdef CWorkerType WORKER_TYPE_WORKER "ray::WorkerType::WORKER"
+    cdef CWorkerType WORKER_TYPE_DRIVER "ray::WorkerType::DRIVER"
 
-cdef extern from "ray/common/task/scheduling_resources.h" \
-        namespace "ray" nogil:
-    cdef cppclass ResourceSet "ResourceSet":
+
+cdef extern from "ray/common/task/scheduling_resources.h" nogil:
+    cdef cppclass ResourceSet "ray::ResourceSet":
         ResourceSet()
         ResourceSet(const unordered_map[c_string, double] &resource_map)
         ResourceSet(const c_vector[c_string] &resource_labels,
@@ -111,7 +121,8 @@ cdef extern from "ray/common/task/scheduling_resources.h" \
         c_bool IsEqual(const ResourceSet &other) const
         c_bool IsSubset(const ResourceSet &other) const
         c_bool IsSuperset(const ResourceSet &other) const
-        c_bool AddOrUpdateResource(const c_string &resource_name, double capacity)
+        c_bool AddOrUpdateResource(const c_string &resource_name,
+                                   double capacity)
         c_bool RemoveResource(const c_string &resource_name)
         void AddResources(const ResourceSet &other)
         c_bool SubtractResourcesStrict(const ResourceSet &other)
@@ -120,3 +131,25 @@ cdef extern from "ray/common/task/scheduling_resources.h" \
         c_bool IsEmpty() const
         const unordered_map[c_string, double] &GetResourceMap() const
         const c_string ToString() const
+
+cdef extern from "ray/common/buffer.h" namespace "ray" nogil:
+    cdef cppclass CBuffer "ray::Buffer":
+        uint8_t *Data() const
+        size_t Size() const
+
+    cdef cppclass LocalMemoryBuffer(CBuffer):
+        LocalMemoryBuffer(uint8_t *data, size_t size)
+
+cdef extern from "ray/core_worker/store_provider/store_provider.h" nogil:
+    cdef cppclass CRayObject "ray::RayObject":
+        const shared_ptr[CBuffer] &GetData()
+        const size_t DataSize() const
+        const shared_ptr[CBuffer] &GetMetadata() const
+        c_bool HasData() const
+        c_bool HasMetadata() const
+
+cdef extern from "ray/gcs/gcs_client_interface.h" nogil:
+    cdef cppclass CGcsClientOptions "ray::gcs::GcsClientOptions":
+        CGcsClientOptions(const c_string &ip, int port,
+                          const c_string &password,
+                          c_bool is_test_client)

python/ray/includes/common.pxi

@@ -0,0 +1,23 @@
+from libcpp cimport bool as c_bool
+from libcpp.string cimport string as c_string
+
+from ray.includes.common cimport CGcsClientOptions
+
+
+cdef class GcsClientOptions:
+    """Cython wrapper class of C++ `ray::gcs::GcsClientOptions`."""
+    cdef:
+        unique_ptr[CGcsClientOptions] gcs_client_options
+
+    def __init__(self, redis_ip, int redis_port,
+                 redis_password, c_bool is_test_client=False):
+        if not redis_password:
+            redis_password = ""
+        self.gcs_client_options.reset(
+            new CGcsClientOptions(redis_ip.encode("ascii"),
+                                  redis_port,
+                                  redis_password.encode("ascii"),
+                                  is_test_client))
+
+    cdef CGcsClientOptions* native(self):
+        return <CGcsClientOptions*>(self.gcs_client_options.get())

python/ray/includes/libcoreworker.pxd

@@ -0,0 +1,62 @@
+from libc.stdint cimport int64_t
+from libcpp cimport bool as c_bool
+from libcpp.memory cimport shared_ptr
+from libcpp.string cimport string as c_string
+from libcpp.vector cimport vector as c_vector
+
+from ray.includes.unique_ids cimport (
+    CJobID,
+    CTaskID,
+    CObjectID,
+)
+from ray.includes.common cimport (
+    CBuffer,
+    CRayStatus,
+    CRayObject,
+    CWorkerType,
+    CLanguage,
+    CGcsClientOptions,
+)
+from ray.includes.libraylet cimport CRayletClient
+
+
+cdef extern from "ray/core_worker/object_interface.h" nogil:
+    cdef cppclass CObjectInterface "ray::CoreWorkerObjectInterface":
+        CRayStatus SetClientOptions(c_string client_name, int64_t limit)
+        CRayStatus Put(const CRayObject &object, CObjectID *object_id)
+        CRayStatus Put(const CRayObject &object, const CObjectID &object_id)
+        CRayStatus Create(const shared_ptr[CBuffer] &metadata,
+                          const size_t data_size, const CObjectID &object_id,
+                          shared_ptr[CBuffer] *data)
+        CRayStatus Seal(const CObjectID &object_id)
+        CRayStatus Get(const c_vector[CObjectID] &ids, int64_t timeout_ms,
+                       c_vector[shared_ptr[CRayObject]] *results)
+        CRayStatus Contains(const CObjectID &object_id, c_bool *has_object)
+        CRayStatus Wait(const c_vector[CObjectID] &object_ids, int num_objects,
+                        int64_t timeout_ms, c_vector[c_bool] *results)
+        CRayStatus Delete(const c_vector[CObjectID] &object_ids,
+                          c_bool local_only, c_bool delete_creating_tasks)
+        c_string MemoryUsageString()
+
+cdef extern from "ray/core_worker/core_worker.h" nogil:
+    cdef cppclass CCoreWorker "ray::CoreWorker":
+        CCoreWorker(const CWorkerType worker_type, const CLanguage language,
+                    const c_string &store_socket,
+                    const c_string &raylet_socket, const CJobID &job_id,
+                    const CGcsClientOptions &gcs_options,
+                    const c_string log_dir, void* execution_callback,
+                    c_bool use_memory_store_)
+        void Disconnect()
+        CWorkerType &GetWorkerType()
+        CLanguage &GetLanguage()
+        CObjectInterface &Objects()
+        # CTaskSubmissionInterface &Tasks()
+        # CTaskExecutionInterface &Execution()
+
+        # TODO(edoakes): remove this once the raylet client is no longer used
+        # directly.
+        CRayletClient &GetRayletClient()
+        # TODO(edoakes): remove this once the Python core worker uses the task
+        # interfaces
+        void SetCurrentJobId(const CJobID &job_id)
+        void SetCurrentTaskId(const CTaskID &task_id)

python/ray/includes/libraylet.pxd

@@ -71,7 +71,9 @@ cdef extern from "ray/raylet/raylet_client.h" nogil:
                                           CActorCheckpointID &checkpoint_id)
         CRayStatus NotifyActorResumedFromCheckpoint(
             const CActorID &actor_id, const CActorCheckpointID &checkpoint_id)
-        CRayStatus SetResource(const c_string &resource_name, const double capacity, const CClientID &client_Id)
+        CRayStatus SetResource(const c_string &resource_name,
+                               const double capacity,
+                               const CClientID &client_Id)
         CLanguage GetLanguage() const
         CWorkerID GetWorkerID() const
         CJobID GetJobID() const

python/ray/includes/task.pxd

@@ -17,7 +17,7 @@ from ray.includes.unique_ids cimport (
     CTaskID,
 )
 
-cdef extern from "ray/protobuf/common.pb.h" namespace "ray::rpc" nogil:
+cdef extern from "ray/protobuf/common.pb.h" nogil:
     cdef cppclass RpcTaskSpec "ray::rpc::TaskSpec":
         void CopyFrom(const RpcTaskSpec &value)
 
@@ -29,13 +29,13 @@ cdef extern from "ray/protobuf/common.pb.h" namespace "ray::rpc" nogil:
         RpcTaskSpec *mutable_task_spec()
 
 
-cdef extern from "ray/protobuf/gcs.pb.h" namespace "ray::rpc" nogil:
+cdef extern from "ray/protobuf/gcs.pb.h" nogil:
     cdef cppclass TaskTableData "ray::rpc::TaskTableData":
         RpcTask *mutable_task()
         const c_string &SerializeAsString()
 
 
-cdef extern from "ray/common/task/task_spec.h" namespace "ray" nogil:
+cdef extern from "ray/common/task/task_spec.h" nogil:
     cdef cppclass CTaskSpec "ray::TaskSpecification":
         CTaskSpec(const RpcTaskSpec message)
         CTaskSpec(const c_string &serialized_binary)
@@ -77,18 +77,20 @@ cdef extern from "ray/common/task/task_spec.h" namespace "ray" nogil:
         c_vector[CActorHandleID] NewActorHandles() const
 
 
-cdef extern from "ray/common/task/task_util.h" namespace "ray" nogil:
+cdef extern from "ray/common/task/task_util.h" nogil:
     cdef cppclass TaskSpecBuilder "ray::TaskSpecBuilder":
         TaskSpecBuilder &SetCommonTaskSpec(
             const CTaskID &task_id, 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)
+            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)  # noqa: E501
 
         TaskSpecBuilder &AddByRefArg(const CObjectID &arg_id)
 
-        TaskSpecBuilder &AddByValueArg(const c_string &data, const c_string &metadata)
+        TaskSpecBuilder &AddByValueArg(const c_string &data,
+                                       const c_string &metadata)
 
         TaskSpecBuilder &SetActorCreationTaskSpec(
             const CActorID &actor_id, uint64_t max_reconstructions,
@@ -100,12 +102,12 @@ cdef extern from "ray/common/task/task_util.h" namespace "ray" nogil:
             const CObjectID &actor_creation_dummy_object_id,
             const CObjectID &previous_actor_task_dummy_object_id,
             uint64_t actor_counter,
-            const c_vector[CActorHandleID] &new_handle_ids);
+            const c_vector[CActorHandleID] &new_handle_ids)
 
         RpcTaskSpec GetMessage()
 
 
-cdef extern from "ray/common/task/task_execution_spec.h" namespace "ray" nogil:
+cdef extern from "ray/common/task/task_execution_spec.h" nogil:
     cdef cppclass CTaskExecutionSpec "ray::TaskExecutionSpecification":
         CTaskExecutionSpec(RpcTaskExecutionSpec message)
         CTaskExecutionSpec(const c_string &serialized_binary)
@@ -113,6 +115,6 @@ cdef extern from "ray/common/task/task_execution_spec.h" namespace "ray" nogil:
         c_vector[CObjectID] ExecutionDependencies()
         uint64_t NumForwards()
 
-cdef extern from "ray/common/task/task.h" namespace "ray" nogil:
+cdef extern from "ray/common/task/task.h" nogil:
     cdef cppclass CTask "ray::Task":
         CTask(CTaskSpec task_spec, CTaskExecutionSpec task_execution_spec)

python/ray/includes/task.pxi

@@ -1,9 +1,3 @@
-from libc.stdint cimport uint8_t
-from libcpp.memory cimport (
-    make_shared,
-    shared_ptr,
-    static_pointer_cast,
-)
 from ray.includes.task cimport (
     CTask,
     CTaskExecutionSpec,
@@ -184,8 +178,9 @@ cdef class TaskSpec:
                     arg_list.append(
                         ObjectID(task_spec.ArgId(i, 0).Binary()))
                 else:
-                    data = (task_spec.ArgData(i)[:task_spec.ArgDataSize(i)])
-                    metadata = (task_spec.ArgMetadata(i)[:task_spec.ArgMetadataSize(i)])
+                    data = task_spec.ArgData(i)[:task_spec.ArgDataSize(i)]
+                    metadata = task_spec.ArgMetadata(i)[
+                        :task_spec.ArgMetadataSize(i)]
                     if metadata == RAW_BUFFER_METADATA:
                         obj = data
                     else:

python/ray/includes/unique_ids.pxd

@@ -20,10 +20,10 @@ cdef extern from "ray/common/id.h" namespace "ray" nogil:
         c_bool IsNil() const
         c_bool operator==(const CBaseID &rhs) const
         c_bool operator!=(const CBaseID &rhs) const
-        const uint8_t *data() const;
+        const uint8_t *data() const
 
-        c_string Binary() const;
-        c_string Hex() const;
+        c_string Binary() const
+        c_string Hex() const
 
     cdef cppclass CUniqueID "ray::UniqueID"(CBaseID):
         CUniqueID()
@@ -65,8 +65,8 @@ cdef extern from "ray/common/id.h" namespace "ray" nogil:
         size_t Size()
 
         @staticmethod
-        CActorID Of(CJobID job_id, CTaskID parent_task_id, int64_t parent_task_counter)
-
+        CActorID Of(CJobID job_id, CTaskID parent_task_id,
+                    int64_t parent_task_counter)
 
     cdef cppclass CActorHandleID "ray::ActorHandleID"(CUniqueID):
 
@@ -123,10 +123,12 @@ cdef extern from "ray/common/id.h" namespace "ray" nogil:
         CTaskID ForActorCreationTask(CActorID actor_id)
 
         @staticmethod
-        CTaskID ForActorTask(CJobID job_id, CTaskID parent_task_id, int64_t parent_task_counter, CActorID actor_id)
+        CTaskID ForActorTask(CJobID job_id, CTaskID parent_task_id,
+                             int64_t parent_task_counter, CActorID actor_id)
 
         @staticmethod
-        CTaskID ForNormalTask(CJobID job_id, CTaskID parent_task_id, int64_t parent_task_counter)
+        CTaskID ForNormalTask(CJobID job_id, CTaskID parent_task_id,
+                              int64_t parent_task_counter)
 
     cdef cppclass CObjectID" ray::ObjectID"(CBaseID[CObjectID]):
 
@@ -140,10 +142,11 @@ cdef extern from "ray/common/id.h" namespace "ray" nogil:
         const CObjectID Nil()
 
         @staticmethod
-        CObjectID ForPut(const CTaskID &task_id, int64_t index, int64_t transport_type);
+        CObjectID ForPut(const CTaskID &task_id, int64_t index,
+                         int64_t transport_type)
 
         @staticmethod
-        CObjectID ForTaskReturn(const CTaskID &task_id, int64_t index);
+        CObjectID ForTaskReturn(const CTaskID &task_id, int64_t index)
 
         @staticmethod
         size_t Size()