From 042ad84573b07c6da399c9a33469341aacf5622a Mon Sep 17 00:00:00 2001 From: Hao Chen Date: Fri, 15 Feb 2019 20:15:16 +0800 Subject: [PATCH] Simplify Cython ID types and fix bug of ActorCheckpointID (#4045) --- python/ray/__init__.py | 61 +++++++++++++++++++++++------- python/ray/_raylet.pyx | 16 ++------ python/ray/includes/common.pxd | 15 ++------ python/ray/includes/libraylet.pxd | 13 ++----- python/ray/includes/task.pxd | 13 +++++-- python/ray/includes/task.pxi | 14 ++++++- python/ray/includes/unique_ids.pxd | 22 +++++------ python/ray/includes/unique_ids.pxi | 20 ++++++++-- python/ray/worker.py | 2 +- 9 files changed, 108 insertions(+), 68 deletions(-) diff --git a/python/ray/__init__.py b/python/ray/__init__.py index 6857827dc..a38797c68 100644 --- a/python/ray/__init__.py +++ b/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", diff --git a/python/ray/_raylet.pyx b/python/ray/_raylet.pyx index b37ace64f..ae1abfc88 100644 --- a/python/ray/_raylet.pyx +++ b/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 diff --git a/python/ray/includes/common.pxd b/python/ray/includes/common.pxd index bdabd87e4..c25cab321 100644 --- a/python/ray/includes/common.pxd +++ b/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, ) diff --git a/python/ray/includes/libraylet.pxd b/python/ray/includes/libraylet.pxd index dd5e6c0ca..29e444a50 100644 --- a/python/ray/includes/libraylet.pxd +++ b/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 diff --git a/python/ray/includes/task.pxd b/python/ray/includes/task.pxd index 80b02c7d4..9e487a2ce 100644 --- a/python/ray/includes/task.pxd +++ b/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: diff --git a/python/ray/includes/task.pxi b/python/ray/includes/task.pxi index aa1e8f66c..71ba91283 100644 --- a/python/ray/includes/task.pxi +++ b/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 diff --git a/python/ray/includes/unique_ids.pxd b/python/ray/includes/unique_ids.pxd index d29e5d398..da4abb3ac 100644 --- a/python/ray/includes/unique_ids.pxd +++ b/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 diff --git a/python/ray/includes/unique_ids.pxi b/python/ray/includes/unique_ids.pxi index 47d31ab09..be0bdd3f5 100644 --- a/python/ray/includes/unique_ids.pxi +++ b/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, +] diff --git a/python/ray/worker.py b/python/ray/worker.py index 1cd99ee57..8a7fe7b5a 100644 --- a/python/ray/worker.py +++ b/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__),