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

python/ray/__init__.py

@@ -61,19 +61,32 @@ from ray._raylet import (
     ObjectID,
     TaskID,
     UniqueID,
-    _ID_TYPES,
 )  # noqa: E402
 
 _config = _Config()
 
 from ray.profiling import profile  # noqa: E402
-from ray.worker import (error_info, init, connect, disconnect, get, put, wait,
-                        remote, get_gpu_ids, get_resource_ids, get_webui_url,
-                        register_custom_serializer, shutdown,
-                        is_initialized)  # noqa: E402
-from ray.worker import (SCRIPT_MODE, WORKER_MODE, LOCAL_MODE,
-                        PYTHON_MODE)  # noqa: E402
-from ray.worker import global_state  # noqa: E402
+from ray.worker import (
+    LOCAL_MODE,
+    PYTHON_MODE,
+    SCRIPT_MODE,
+    WORKER_MODE,
+    connect,
+    disconnect,
+    error_info,
+    get,
+    get_gpu_ids,
+    get_resource_ids,
+    get_webui_url,
+    global_state,
+    init,
+    is_initialized,
+    put,
+    register_custom_serializer,
+    remote,
+    shutdown,
+    wait,
+)  # noqa: E402
 import ray.internal  # noqa: E402
 # We import ray.actor because some code is run in actor.py which initializes
 # some functions in the worker.
@@ -84,14 +97,34 @@ from ray.actor import method  # noqa: E402
 __version__ = "0.7.0.dev0"
 
 __all__ = [
-    "error_info", "init", "connect", "disconnect", "get", "put", "wait",
-    "remote", "profile", "actor", "method", "get_gpu_ids", "get_resource_ids",
-    "get_webui_url", "register_custom_serializer", "shutdown",
-    "is_initialized", "SCRIPT_MODE", "WORKER_MODE", "LOCAL_MODE",
-    "PYTHON_MODE", "global_state", "_config", "__version__", "internal",
-    "_ID_TYPES"
+    "LOCAL_MODE",
+    "PYTHON_MODE",
+    "SCRIPT_MODE",
+    "WORKER_MODE",
+    "__version__",
+    "_config",
+    "actor",
+    "connect",
+    "disconnect",
+    "error_info",
+    "get",
+    "get_gpu_ids",
+    "get_resource_ids",
+    "get_webui_url",
+    "global_state",
+    "init",
+    "internal",
+    "is_initialized",
+    "method",
+    "profile",
+    "put",
+    "register_custom_serializer",
+    "remote",
+    "shutdown",
+    "wait",
 ]
 
+# ID types
 __all__ += [
     "ActorCheckpointID",
     "ActorClassID",

python/ray/_raylet.pyx

@@ -19,20 +19,8 @@ include "includes/ray_config.pxi"
 include "includes/task.pxi"
 
 from ray.includes.common cimport (
-    CActorCheckpointID,
-    CActorClassID,
-    CActorHandleID,
-    CActorID,
-    CClientID,
-    CConfigID,
-    CDriverID,
-    CFunctionID,
     CLanguage,
-    CObjectID,
     CRayStatus,
-    CTaskID,
-    CUniqueID,
-    CWorkerID,
     LANGUAGE_CPP,
     LANGUAGE_JAVA,
     LANGUAGE_PYTHON,
@@ -44,6 +32,10 @@ from ray.includes.libraylet cimport (
     ResourceMappingType,
     WaitResultPair,
 )
+from ray.includes.unique_ids cimport (
+    CActorCheckpointID,
+    CObjectID,
+)
 from ray.includes.task cimport CTaskSpecification
 from ray.includes.ray_config cimport RayConfig
 from ray.utils import decode

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,
+]

python/ray/worker.py

@@ -1128,7 +1128,7 @@ def _initialize_serialization(driver_id, worker=global_worker):
     serialization_context.set_pickle(pickle.dumps, pickle.loads)
     pyarrow.register_torch_serialization_handlers(serialization_context)
 
-    for id_type in ray._ID_TYPES:
+    for id_type in ray._raylet._ID_TYPES:
         serialization_context.register_type(
             id_type,
             "{}.{}".format(id_type.__module__, id_type.__name__),