Simplify Cython ID types and fix bug of ActorCheckpointID (#4045)

python/ray/includes/common.pxd

@@ -6,18 +6,9 @@ from libcpp.unordered_map cimport unordered_map
 from libcpp.vector cimport vector as c_vector
 
 from ray.includes.unique_ids cimport (
-    ActorCheckpointID as CActorCheckpointID,
-    ActorClassID as CActorClassID,
-    ActorHandleID as CActorHandleID,
-    ActorID as CActorID,
-    CUniqueID,
-    ClientID as CClientID,
-    ConfigID as CConfigID,
-    DriverID as CDriverID,
-    FunctionID as CFunctionID,
-    ObjectID as CObjectID,
-    TaskID as CTaskID,
-    WorkerID as CWorkerID,
+    CDriverID,
+    CObjectID,
+    CTaskID,
 )
 
 

python/ray/includes/libraylet.pxd

@@ -6,22 +6,17 @@ from libcpp.utility cimport pair
 from libcpp.unordered_map cimport unordered_map
 from libcpp.vector cimport vector as c_vector
 
-
 from ray.includes.common cimport (
+    CLanguage,
+    CRayStatus,
+)
+from ray.includes.unique_ids cimport (
     CActorCheckpointID,
-    CActorClassID,
-    CActorHandleID,
     CActorID,
     CClientID,
-    CConfigID,
     CDriverID,
-    CFunctionID,
-    CLanguage,
     CObjectID,
-    CRayStatus,
     CTaskID,
-    CUniqueID,
-    CWorkerID,
 )
 from ray.includes.task cimport CTaskSpecification
 

python/ray/includes/task.pxd

@@ -6,9 +6,16 @@ from libcpp.unordered_map cimport unordered_map
 from libcpp.vector cimport vector as c_vector
 
 from ray.includes.common cimport (
-    CUniqueID, CTaskID, CObjectID, CFunctionID, CActorClassID, CActorID,
-    CActorHandleID, CWorkerID, CDriverID, CConfigID, CClientID,
-    CLanguage, ResourceSet)
+    CLanguage,
+    ResourceSet,
+)
+from ray.includes.unique_ids cimport (
+    CActorHandleID,
+    CActorID,
+    CDriverID,
+    CObjectID,
+    CTaskID,
+)
 
 
 cdef extern from "ray/raylet/task_execution_spec.h" namespace "ray::raylet" nogil:

python/ray/includes/task.pxi

@@ -1,6 +1,16 @@
 from libc.stdint cimport uint8_t
-from libcpp.memory cimport shared_ptr, make_shared, static_pointer_cast
-from ray.includes.task cimport CTaskSpecification, CTaskArgument, CTaskArgumentByValue, CTaskArgumentByReference, SerializeTaskAsString
+from libcpp.memory cimport (
+    make_shared,
+    shared_ptr,
+    static_pointer_cast,
+)
+from ray.includes.task cimport (
+    CTaskArgument,
+    CTaskArgumentByReference,
+    CTaskArgumentByValue,
+    CTaskSpecification,
+    SerializeTaskAsString,
+)
 
 from ray.utils import _random_string
 

python/ray/includes/unique_ids.pxd

@@ -22,14 +22,14 @@ cdef extern from "ray/id.h" namespace "ray" nogil:
         c_string binary() const;
         c_string hex() const;
 
-    ctypedef CUniqueID TaskID
-    ctypedef CUniqueID ObjectID
-    ctypedef CUniqueID FunctionID
-    ctypedef CUniqueID ActorID
-    ctypedef CUniqueID ActorClassID
-    ctypedef CUniqueID ActorHandleID
-    ctypedef CUniqueID ActorCheckpointID
-    ctypedef CUniqueID WorkerID
-    ctypedef CUniqueID DriverID
-    ctypedef CUniqueID ConfigID
-    ctypedef CUniqueID ClientID
+ctypedef CUniqueID CActorCheckpointID
+ctypedef CUniqueID CActorClassID
+ctypedef CUniqueID CActorHandleID
+ctypedef CUniqueID CActorID
+ctypedef CUniqueID CClientID
+ctypedef CUniqueID CConfigID
+ctypedef CUniqueID CDriverID
+ctypedef CUniqueID CFunctionID
+ctypedef CUniqueID CObjectID
+ctypedef CUniqueID CTaskID
+ctypedef CUniqueID CWorkerID

python/ray/includes/unique_ids.pxi

@@ -7,6 +7,10 @@ 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,
+)
+from ray.includes.unique_ids cimport (
     CActorCheckpointID,
     CActorClassID,
     CActorHandleID,
@@ -19,8 +23,6 @@ from ray.includes.common cimport (
     CTaskID,
     CUniqueID,
     CWorkerID,
-    ComputePutId,
-    ComputeTaskId,
 )
 
 from ray.utils import decode
@@ -317,5 +319,15 @@ cdef class ActorClassID(UniqueID):
         return "ActorClassID(" + self.hex() + ")"
 
 
-_ID_TYPES = [UniqueID, ObjectID, TaskID, ClientID, DriverID, ActorID,
-             ActorHandleID, FunctionID, ActorClassID]
+_ID_TYPES = [
+    ActorCheckpointID,
+    ActorClassID,
+    ActorHandleID,
+    ActorID,
+    ClientID,
+    DriverID,
+    FunctionID,
+    ObjectID,
+    TaskID,
+    UniqueID,
+]