Refactor ID Serial 1: Separate ObjectID and TaskID from UniqueID (#4776)

* Enable BaseId.

* Change TaskID and make python test pass

* Remove unnecessary functions and fix test failure and change TaskID to
16 bytes.

* Java code change draft

* Refine

* Lint

* Update java/api/src/main/java/org/ray/api/id/TaskId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Update java/api/src/main/java/org/ray/api/id/BaseId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Update java/api/src/main/java/org/ray/api/id/BaseId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Update java/api/src/main/java/org/ray/api/id/ObjectId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Address comment

* Lint

* Fix SINGLE_PROCESS

* Fix comments

* Refine code

* Refine test

* Resolve conflict

python/ray/includes/common.pxd

@@ -81,15 +81,9 @@ cdef extern from "ray/status.h" namespace "ray::StatusCode" nogil:
 
 
 cdef extern from "ray/id.h" namespace "ray" nogil:
-    const CTaskID FinishTaskId(const CTaskID &task_id)
-    const CObjectID ComputeReturnId(const CTaskID &task_id,
-                                    int64_t return_index)
-    const CObjectID ComputePutId(const CTaskID &task_id, int64_t put_index)
-    const CTaskID ComputeTaskId(const CObjectID &object_id)
     const CTaskID GenerateTaskId(const CDriverID &driver_id,
                                  const CTaskID &parent_task_id,
                                  int parent_task_counter)
-    int64_t ComputeObjectIndex(const CObjectID &object_id)
 
 
 cdef extern from "ray/gcs/format/gcs_generated.h" nogil:

python/ray/includes/unique_ids.pxd

@@ -1,12 +1,35 @@
 from libcpp cimport bool as c_bool
 from libcpp.string cimport string as c_string
-from libc.stdint cimport uint8_t
+from libc.stdint cimport uint8_t, int64_t
 
 cdef extern from "ray/id.h" namespace "ray" nogil:
-    cdef cppclass CUniqueID "ray::UniqueID":
+    cdef cppclass CBaseID[T]:
+        @staticmethod
+        T from_random()
+
+        @staticmethod
+        T from_binary(const c_string &binary)
+
+        @staticmethod
+        const T nil()
+
+        @staticmethod
+        size_t size()
+
+        size_t hash() const
+        c_bool is_nil() const
+        c_bool operator==(const CBaseID &rhs) const
+        c_bool operator!=(const CBaseID &rhs) const
+        const uint8_t *data() const;
+
+        c_string binary() const;
+        c_string hex() const;
+
+    cdef cppclass CUniqueID "ray::UniqueID"(CBaseID):
         CUniqueID()
-        CUniqueID(const c_string &binary)
-        CUniqueID(const CUniqueID &from_id)
+
+        @staticmethod
+        size_t size()
 
         @staticmethod
         CUniqueID from_random()
@@ -17,15 +40,8 @@ cdef extern from "ray/id.h" namespace "ray" nogil:
         @staticmethod
         const CUniqueID nil()
 
-        size_t hash() const
-        c_bool is_nil() const
-        c_bool operator==(const CUniqueID& rhs) const
-        c_bool operator!=(const CUniqueID& rhs) const
-        const uint8_t *data() const
-        uint8_t *mutable_data()
-        size_t size() const
-        c_string binary() const
-        c_string hex() const
+        @staticmethod
+        size_t size()
 
     cdef cppclass CActorCheckpointID "ray::ActorCheckpointID"(CUniqueID):
 
@@ -67,16 +83,40 @@ cdef extern from "ray/id.h" namespace "ray" nogil:
         @staticmethod
         CDriverID from_binary(const c_string &binary)
 
-    cdef cppclass CTaskID "ray::TaskID"(CUniqueID):
+    cdef cppclass CTaskID "ray::TaskID"(CBaseID[CTaskID]):
 
         @staticmethod
         CTaskID from_binary(const c_string &binary)
 
-    cdef cppclass CObjectID" ray::ObjectID"(CUniqueID):
+        @staticmethod
+        const CTaskID nil()
+
+        @staticmethod
+        size_t size()
+
+    cdef cppclass CObjectID" ray::ObjectID"(CBaseID[CObjectID]):
 
         @staticmethod
         CObjectID from_binary(const c_string &binary)
 
+        @staticmethod
+        const CObjectID nil()
+
+        @staticmethod
+        CObjectID for_put(const CTaskID &task_id, int64_t index);
+
+        @staticmethod
+        CObjectID for_task_return(const CTaskID &task_id, int64_t index);
+
+        @staticmethod
+        size_t size()
+
+        c_bool is_put()
+
+        int64_t object_index() const 
+
+        CTaskID task_id() const
+
     cdef cppclass CWorkerID "ray::WorkerID"(CUniqueID):
 
         @staticmethod

python/ray/includes/unique_ids.pxi

@@ -6,10 +6,8 @@ See https://github.com/ray-project/ray/issues/3721.
 
 # WARNING: Any additional ID types defined in this file must be added to the
 # _ID_TYPES list at the bottom of this file.
-from ray.includes.common cimport (
-    ComputePutId,
-    ComputeTaskId,
-)
+import os
+
 from ray.includes.unique_ids cimport (
     CActorCheckpointID,
     CActorClassID,
@@ -28,12 +26,12 @@ from ray.includes.unique_ids cimport (
 from ray.utils import decode
 
 
-def check_id(b):
+def check_id(b, size=kUniqueIDSize):
     if not isinstance(b, bytes):
         raise TypeError("Unsupported type: " + str(type(b)))
-    if len(b) != kUniqueIDSize:
+    if len(b) != size:
         raise ValueError("ID string needs to have length " +
-                         str(kUniqueIDSize))
+                         str(size))
 
 
 cdef extern from "ray/constants.h" nogil:
@@ -41,28 +39,27 @@ cdef extern from "ray/constants.h" nogil:
     cdef int64_t kMaxTaskPuts
 
 
-cdef class UniqueID:
-    cdef CUniqueID data
+cdef class BaseID:
 
-    def __init__(self, id):
-        check_id(id)
-        self.data = CUniqueID.from_binary(id)
+    # To avoid the error of "Python int too large to convert to C ssize_t",
+    # here `cdef size_t` is required.
+    cdef size_t hash(self):
+        pass
 
-    @classmethod
-    def from_binary(cls, id_bytes):
-        if not isinstance(id_bytes, bytes):
-            raise TypeError("Expect bytes, got " + str(type(id_bytes)))
-        return cls(id_bytes)
+    def binary(self):
+        pass
 
-    @classmethod
-    def nil(cls):
-        return cls(CUniqueID.nil().binary())
+    def size(self):
+        pass
 
-    def __hash__(self):
-        return self.data.hash()
+    def hex(self):
+        pass
 
     def is_nil(self):
-        return self.data.is_nil()
+        pass
+
+    def __hash__(self):
+        return self.hash()
 
     def __eq__(self, other):
         return type(self) == type(other) and self.binary() == other.binary()
@@ -70,18 +67,9 @@ cdef class UniqueID:
     def __ne__(self, other):
         return self.binary() != other.binary()
 
-    def size(self):
-        return self.data.size()
-
-    def binary(self):
-        return self.data.binary()
-
     def __bytes__(self):
         return self.binary()
 
-    def hex(self):
-        return decode(self.data.hex())
-
     def __hex__(self):
         return self.hex()
 
@@ -98,11 +86,52 @@ cdef class UniqueID:
         # NOTE: The hash function used here must match the one in
         # GetRedisContext in src/ray/gcs/tables.h. Changes to the
         # hash function should only be made through std::hash in
-        # src/common/common.h
+        # src/common/common.h.
+        # Do not use __hash__ that returns signed uint64_t, which
+        # is different from std::hash in c++ code.
+        return self.hash()
+
+
+cdef class UniqueID(BaseID):
+    cdef CUniqueID data
+
+    def __init__(self, id):
+        check_id(id)
+        self.data = CUniqueID.from_binary(id)
+
+    @classmethod
+    def from_binary(cls, id_bytes):
+        if not isinstance(id_bytes, bytes):
+            raise TypeError("Expect bytes, got " + str(type(id_bytes)))
+        return cls(id_bytes)
+
+    @classmethod
+    def nil(cls):
+        return cls(CUniqueID.nil().binary())
+
+    
+    @classmethod
+    def from_random(cls):
+        return cls(os.urandom(CUniqueID.size()))
+
+    def size(self):
+        return CUniqueID.size()
+
+    def binary(self):
+        return self.data.binary()
+
+    def hex(self):
+        return decode(self.data.hex())
+    
+    def is_nil(self):
+        return self.data.is_nil()
+
+    cdef size_t hash(self):
         return self.data.hash()
 
 
-cdef class ObjectID(UniqueID):
+cdef class ObjectID(BaseID):
+    cdef CObjectID data
 
     def __init__(self, id):
         check_id(id)
@@ -111,16 +140,67 @@ cdef class ObjectID(UniqueID):
     cdef CObjectID native(self):
         return <CObjectID>self.data
 
+    def size(self):
+        return CObjectID.size()
 
-cdef class TaskID(UniqueID):
+    def binary(self):
+        return self.data.binary()
+
+    def hex(self):
+        return decode(self.data.hex())
+    
+    def is_nil(self):
+        return self.data.is_nil()
+
+    cdef size_t hash(self):
+        return self.data.hash()
+
+    @classmethod
+    def nil(cls):
+        return cls(CObjectID.nil().binary())
+
+    @classmethod
+    def from_random(cls):
+        return cls(os.urandom(CObjectID.size()))
+
+
+cdef class TaskID(BaseID):
+    cdef CTaskID data
 
     def __init__(self, id):
-        check_id(id)
+        check_id(id, CTaskID.size())
         self.data = CTaskID.from_binary(<c_string>id)
 
     cdef CTaskID native(self):
         return <CTaskID>self.data
 
+    def size(self):
+        return CTaskID.size()
+
+    def binary(self):
+        return self.data.binary()
+
+    def hex(self):
+        return decode(self.data.hex())
+
+    def is_nil(self):
+        return self.data.is_nil()
+
+    cdef size_t hash(self):
+        return self.data.hash()
+
+    @classmethod
+    def nil(cls):
+        return cls(CTaskID.nil().binary())
+
+    @classmethod
+    def size(cla):
+        return CTaskID.size()
+
+    @classmethod
+    def from_random(cls):
+        return cls(os.urandom(CTaskID.size()))
+
 
 cdef class ClientID(UniqueID):