mirror of
https://github.com/wassname/ray.git
synced 2026-07-26 13:37:24 +08:00
Migrate Python C extension to Cython (#3541)
This commit is contained in:
Si-Yuan
@@ -0,0 +1,120 @@
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp cimport bool as c_bool
|
||||
|
||||
from libc.stdint cimport int64_t
|
||||
from libcpp.unordered_map cimport unordered_map
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
|
||||
from ray.includes.unique_ids cimport (
|
||||
CUniqueID, TaskID as CTaskID, ObjectID as CObjectID,
|
||||
FunctionID as CFunctionID, ActorClassID as CActorClassID, ActorID as CActorID,
|
||||
ActorHandleID as CActorHandleID, WorkerID as CWorkerID,
|
||||
DriverID as CDriverID, ConfigID as CConfigID, ClientID as CClientID)
|
||||
|
||||
|
||||
cdef extern from "ray/status.h" namespace "ray" nogil:
|
||||
cdef cppclass StatusCode:
|
||||
pass
|
||||
|
||||
cdef cppclass CRayStatus "ray::Status":
|
||||
RayStatus()
|
||||
RayStatus(StatusCode code, const c_string &msg)
|
||||
RayStatus(const CRayStatus &s);
|
||||
|
||||
@staticmethod
|
||||
CRayStatus OK()
|
||||
@staticmethod
|
||||
CRayStatus OutOfMemory()
|
||||
@staticmethod
|
||||
CRayStatus KeyError()
|
||||
@staticmethod
|
||||
CRayStatus Invalid()
|
||||
@staticmethod
|
||||
CRayStatus IOError()
|
||||
@staticmethod
|
||||
CRayStatus TypeError()
|
||||
@staticmethod
|
||||
CRayStatus UnknownError()
|
||||
@staticmethod
|
||||
CRayStatus NotImplemented()
|
||||
@staticmethod
|
||||
CRayStatus RedisError()
|
||||
|
||||
c_bool ok()
|
||||
c_bool IsOutOfMemory()
|
||||
c_bool IsKeyError()
|
||||
c_bool IsInvalid()
|
||||
c_bool IsIOError()
|
||||
c_bool IsTypeError()
|
||||
c_bool IsUnknownError()
|
||||
c_bool IsNotImplemented()
|
||||
c_bool IsRedisError()
|
||||
|
||||
c_string ToString()
|
||||
c_string CodeAsString()
|
||||
StatusCode code()
|
||||
c_string message()
|
||||
|
||||
# We can later add more of the common status factory methods as needed
|
||||
cdef CRayStatus RayStatus_OK "Status::OK"()
|
||||
cdef CRayStatus RayStatus_Invalid "Status::Invalid"()
|
||||
|
||||
|
||||
cdef extern from "ray/status.h" namespace "ray::StatusCode" nogil:
|
||||
cdef StatusCode StatusCode_OK "OK"
|
||||
cdef StatusCode StatusCode_OutOfMemory "OutOfMemory"
|
||||
cdef StatusCode StatusCode_KeyError "KeyError"
|
||||
cdef StatusCode StatusCode_TypeError "TypeError"
|
||||
cdef StatusCode StatusCode_Invalid "Invalid"
|
||||
cdef StatusCode StatusCode_IOError "IOError"
|
||||
cdef StatusCode StatusCode_UnknownError "UnknownError"
|
||||
cdef StatusCode StatusCode_NotImplemented "NotImplemented"
|
||||
cdef StatusCode StatusCode_RedisError "RedisError"
|
||||
|
||||
|
||||
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:
|
||||
cdef cppclass GCSArg "Arg":
|
||||
pass
|
||||
|
||||
cdef cppclass CLanguage "Language":
|
||||
pass
|
||||
|
||||
|
||||
# This is a workaround for C++ enum class since Cython has no corresponding representation.
|
||||
cdef extern from "ray/gcs/format/gcs_generated.h" namespace "Language" nogil:
|
||||
cdef CLanguage LANGUAGE_PYTHON "Language::PYTHON"
|
||||
cdef CLanguage LANGUAGE_CPP "Language::CPP"
|
||||
cdef CLanguage LANGUAGE_JAVA "Language::JAVA"
|
||||
|
||||
|
||||
cdef extern from "ray/raylet/scheduling_resources.h" namespace "ray::raylet" nogil:
|
||||
cdef cppclass ResourceSet "ResourceSet":
|
||||
ResourceSet()
|
||||
ResourceSet(const unordered_map[c_string, double] &resource_map)
|
||||
ResourceSet(const c_vector[c_string] &resource_labels, const c_vector[double] resource_capacity)
|
||||
c_bool operator==(const ResourceSet &rhs) const
|
||||
c_bool IsEqual(const ResourceSet &other) const
|
||||
c_bool IsSubset(const ResourceSet &other) const
|
||||
c_bool IsSuperset(const ResourceSet &other) const
|
||||
c_bool AddResource(const c_string &resource_name, double capacity)
|
||||
c_bool RemoveResource(const c_string &resource_name)
|
||||
c_bool AddResourcesStrict(const ResourceSet &other)
|
||||
void AddResources(const ResourceSet &other)
|
||||
c_bool SubtractResourcesStrict(const ResourceSet &other)
|
||||
c_bool GetResource(const c_string &resource_name, double *value) const
|
||||
double GetNumCpus() const
|
||||
c_bool IsEmpty() const
|
||||
const unordered_map[c_string, double] &GetResourceMap() const
|
||||
const c_string ToString() const
|
||||
@@ -0,0 +1,64 @@
|
||||
from libc.stdint cimport int64_t
|
||||
from libcpp cimport bool as c_bool
|
||||
from libcpp.memory cimport unique_ptr
|
||||
from libcpp.string cimport string as c_string
|
||||
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 (
|
||||
CUniqueID, CTaskID, CObjectID, CFunctionID, CActorClassID, CActorID,
|
||||
CActorHandleID, CWorkerID, CDriverID, CConfigID, CClientID,
|
||||
CLanguage, CRayStatus)
|
||||
from ray.includes.task cimport CTaskSpecification
|
||||
|
||||
|
||||
cdef extern from "ray/gcs/format/gcs_generated.h" nogil:
|
||||
cdef cppclass GCSProfileEventT "ProfileEventT":
|
||||
c_string event_type
|
||||
double start_time
|
||||
double end_time
|
||||
c_string extra_data
|
||||
GCSProfileEventT()
|
||||
|
||||
cdef cppclass GCSProfileTableDataT "ProfileTableDataT":
|
||||
c_string component_type
|
||||
c_string component_id
|
||||
c_string node_ip_address
|
||||
c_vector[unique_ptr[GCSProfileEventT]] profile_events
|
||||
GCSProfileTableDataT()
|
||||
|
||||
|
||||
ctypedef unordered_map[c_string, c_vector[pair[int64_t, double]]] ResourceMappingType
|
||||
ctypedef pair[c_vector[CObjectID], c_vector[CObjectID]] WaitResultPair
|
||||
|
||||
|
||||
cdef extern from "ray/raylet/raylet_client.h" nogil:
|
||||
cdef cppclass CRayletClient "RayletClient":
|
||||
CRayletClient(const c_string &raylet_socket,
|
||||
const CClientID &client_id,
|
||||
c_bool is_worker, const CDriverID &driver_id,
|
||||
const CLanguage &language)
|
||||
CRayStatus Disconnect()
|
||||
CRayStatus SubmitTask(const c_vector[CObjectID] &execution_dependencies,
|
||||
const CTaskSpecification &task_spec)
|
||||
CRayStatus GetTask(unique_ptr[CTaskSpecification] *task_spec)
|
||||
CRayStatus TaskDone()
|
||||
CRayStatus FetchOrReconstruct(c_vector[CObjectID] &object_ids,
|
||||
c_bool fetch_only,
|
||||
const CTaskID ¤t_task_id)
|
||||
CRayStatus NotifyUnblocked(const CTaskID ¤t_task_id)
|
||||
CRayStatus Wait(const c_vector[CObjectID] &object_ids, int num_returns,
|
||||
int64_t timeout_milliseconds, c_bool wait_local,
|
||||
const CTaskID ¤t_task_id, WaitResultPair *result)
|
||||
CRayStatus PushError(const CDriverID &job_id, const c_string &type,
|
||||
const c_string &error_message, double timestamp)
|
||||
CRayStatus PushProfileEvents(const GCSProfileTableDataT &profile_events)
|
||||
CRayStatus FreeObjects(const c_vector[CObjectID] &object_ids,
|
||||
c_bool local_only)
|
||||
CLanguage GetLanguage() const
|
||||
CClientID GetClientID() const
|
||||
CDriverID GetDriverID() const
|
||||
c_bool IsWorker() const
|
||||
const ResourceMappingType &GetResourceIDs() const
|
||||
@@ -0,0 +1,83 @@
|
||||
from libc.stdint cimport int64_t, uint64_t
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp.unordered_map cimport unordered_map
|
||||
|
||||
|
||||
cdef extern from "ray/ray_config.h" nogil:
|
||||
cdef cppclass RayConfig "RayConfig":
|
||||
@staticmethod
|
||||
RayConfig &instance()
|
||||
|
||||
int64_t ray_protocol_version() const
|
||||
|
||||
int64_t handler_warning_timeout_ms() const
|
||||
|
||||
int64_t heartbeat_timeout_milliseconds() const
|
||||
|
||||
int64_t debug_dump_period_milliseconds() const
|
||||
|
||||
int64_t num_heartbeats_timeout() const
|
||||
|
||||
uint64_t num_heartbeats_warning() const
|
||||
|
||||
int64_t initial_reconstruction_timeout_milliseconds() const
|
||||
|
||||
int64_t get_timeout_milliseconds() const
|
||||
|
||||
uint64_t max_lineage_size() const
|
||||
|
||||
int64_t worker_get_request_size() const
|
||||
|
||||
int64_t worker_fetch_request_size() const
|
||||
|
||||
int64_t actor_max_dummy_objects() const
|
||||
|
||||
int64_t num_connect_attempts() const
|
||||
|
||||
int64_t connect_timeout_milliseconds() const
|
||||
|
||||
int64_t local_scheduler_fetch_timeout_milliseconds() const
|
||||
|
||||
int64_t local_scheduler_reconstruction_timeout_milliseconds() const
|
||||
|
||||
int64_t max_num_to_reconstruct() const
|
||||
|
||||
int64_t local_scheduler_fetch_request_size() const
|
||||
|
||||
int64_t kill_worker_timeout_milliseconds() const
|
||||
|
||||
int64_t max_time_for_handler_milliseconds() const
|
||||
|
||||
int64_t size_limit() const
|
||||
|
||||
int64_t num_elements_limit() const
|
||||
|
||||
int64_t max_time_for_loop() const
|
||||
|
||||
int64_t redis_db_connect_retries()
|
||||
|
||||
int64_t redis_db_connect_wait_milliseconds() const
|
||||
|
||||
int64_t plasma_default_release_delay() const
|
||||
|
||||
int64_t L3_cache_size_bytes() const
|
||||
|
||||
int64_t max_tasks_to_spillback() const
|
||||
|
||||
int64_t actor_creation_num_spillbacks_warning() const
|
||||
|
||||
int node_manager_forward_task_retry_timeout_milliseconds() const
|
||||
|
||||
int object_manager_pull_timeout_ms() const
|
||||
|
||||
int object_manager_push_timeout_ms() const
|
||||
|
||||
int object_manager_repeated_push_delay_ms() const
|
||||
|
||||
uint64_t object_manager_default_chunk_size() const
|
||||
|
||||
int num_workers_per_process() const
|
||||
|
||||
int64_t max_task_lease_timeout_ms() const
|
||||
|
||||
void initialize(const unordered_map[c_string, int] &config_map)
|
||||
@@ -0,0 +1,146 @@
|
||||
from ray.includes.ray_config cimport RayConfig
|
||||
|
||||
cdef class Config:
|
||||
@staticmethod
|
||||
def ray_protocol_version():
|
||||
return RayConfig.instance().ray_protocol_version()
|
||||
|
||||
@staticmethod
|
||||
def handler_warning_timeout_ms():
|
||||
return RayConfig.instance().handler_warning_timeout_ms()
|
||||
|
||||
@staticmethod
|
||||
def heartbeat_timeout_milliseconds():
|
||||
return RayConfig.instance().heartbeat_timeout_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def debug_dump_period_milliseconds():
|
||||
return RayConfig.instance().debug_dump_period_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def num_heartbeats_timeout():
|
||||
return RayConfig.instance().num_heartbeats_timeout()
|
||||
|
||||
@staticmethod
|
||||
def num_heartbeats_warning():
|
||||
return RayConfig.instance().num_heartbeats_warning()
|
||||
|
||||
@staticmethod
|
||||
def initial_reconstruction_timeout_milliseconds():
|
||||
return RayConfig.instance().initial_reconstruction_timeout_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def get_timeout_milliseconds():
|
||||
return RayConfig.instance().get_timeout_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def max_lineage_size():
|
||||
return RayConfig.instance().max_lineage_size()
|
||||
|
||||
@staticmethod
|
||||
def worker_get_request_size():
|
||||
return RayConfig.instance().worker_get_request_size()
|
||||
|
||||
@staticmethod
|
||||
def worker_fetch_request_size():
|
||||
return RayConfig.instance().worker_fetch_request_size()
|
||||
|
||||
@staticmethod
|
||||
def actor_max_dummy_objects():
|
||||
return RayConfig.instance().actor_max_dummy_objects()
|
||||
|
||||
@staticmethod
|
||||
def num_connect_attempts():
|
||||
return RayConfig.instance().num_connect_attempts()
|
||||
|
||||
@staticmethod
|
||||
def connect_timeout_milliseconds():
|
||||
return RayConfig.instance().connect_timeout_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def local_scheduler_fetch_timeout_milliseconds():
|
||||
return RayConfig.instance().local_scheduler_fetch_timeout_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def local_scheduler_reconstruction_timeout_milliseconds():
|
||||
return RayConfig.instance().local_scheduler_reconstruction_timeout_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def max_num_to_reconstruct():
|
||||
return RayConfig.instance().max_num_to_reconstruct()
|
||||
|
||||
@staticmethod
|
||||
def local_scheduler_fetch_request_size():
|
||||
return RayConfig.instance().local_scheduler_fetch_request_size()
|
||||
|
||||
@staticmethod
|
||||
def kill_worker_timeout_milliseconds():
|
||||
return RayConfig.instance().kill_worker_timeout_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def max_time_for_handler_milliseconds():
|
||||
return RayConfig.instance().max_time_for_handler_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def size_limit():
|
||||
return RayConfig.instance().size_limit()
|
||||
|
||||
@staticmethod
|
||||
def num_elements_limit():
|
||||
return RayConfig.instance().num_elements_limit()
|
||||
|
||||
@staticmethod
|
||||
def max_time_for_loop():
|
||||
return RayConfig.instance().max_time_for_loop()
|
||||
|
||||
@staticmethod
|
||||
def redis_db_connect_retries():
|
||||
return RayConfig.instance().redis_db_connect_retries()
|
||||
|
||||
@staticmethod
|
||||
def redis_db_connect_wait_milliseconds():
|
||||
return RayConfig.instance().redis_db_connect_wait_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def plasma_default_release_delay():
|
||||
return RayConfig.instance().plasma_default_release_delay()
|
||||
|
||||
@staticmethod
|
||||
def L3_cache_size_bytes():
|
||||
return RayConfig.instance().L3_cache_size_bytes()
|
||||
|
||||
@staticmethod
|
||||
def max_tasks_to_spillback():
|
||||
return RayConfig.instance().max_tasks_to_spillback()
|
||||
|
||||
@staticmethod
|
||||
def actor_creation_num_spillbacks_warning():
|
||||
return RayConfig.instance().actor_creation_num_spillbacks_warning()
|
||||
|
||||
@staticmethod
|
||||
def node_manager_forward_task_retry_timeout_milliseconds():
|
||||
return RayConfig.instance().node_manager_forward_task_retry_timeout_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def object_manager_pull_timeout_ms():
|
||||
return RayConfig.instance().object_manager_pull_timeout_ms()
|
||||
|
||||
@staticmethod
|
||||
def object_manager_push_timeout_ms():
|
||||
return RayConfig.instance().object_manager_push_timeout_ms()
|
||||
|
||||
@staticmethod
|
||||
def object_manager_repeated_push_delay_ms():
|
||||
return RayConfig.instance().object_manager_repeated_push_delay_ms()
|
||||
|
||||
@staticmethod
|
||||
def object_manager_default_chunk_size():
|
||||
return RayConfig.instance().object_manager_default_chunk_size()
|
||||
|
||||
@staticmethod
|
||||
def num_workers_per_process():
|
||||
return RayConfig.instance().num_workers_per_process()
|
||||
|
||||
@staticmethod
|
||||
def max_task_lease_timeout_ms():
|
||||
return RayConfig.instance().max_task_lease_timeout_ms()
|
||||
@@ -0,0 +1,102 @@
|
||||
from libc.stdint cimport int64_t, uint8_t
|
||||
from libcpp cimport bool as c_bool
|
||||
from libcpp.memory cimport unique_ptr, shared_ptr
|
||||
from libcpp.string cimport string as c_string
|
||||
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)
|
||||
|
||||
|
||||
cdef extern from "ray/raylet/task_execution_spec.h" namespace "ray::raylet" nogil:
|
||||
cdef cppclass CTaskExecutionSpecification "ray::raylet::TaskExecutionSpecification":
|
||||
CTaskExecutionSpecification(const c_vector[CObjectID] &&dependencies)
|
||||
CTaskExecutionSpecification(const c_vector[CObjectID] &&dependencies, int num_forwards)
|
||||
c_vector[CObjectID] ExecutionDependencies() const
|
||||
void SetExecutionDependencies(const c_vector[CObjectID] &dependencies)
|
||||
int NumForwards() const
|
||||
void IncrementNumForwards()
|
||||
int64_t LastTimestamp() const
|
||||
void SetLastTimestamp(int64_t new_timestamp)
|
||||
|
||||
|
||||
cdef extern from "ray/raylet/task_spec.h" namespace "ray::raylet" nogil:
|
||||
cdef cppclass CTaskArgument "ray::raylet::TaskArgument":
|
||||
pass
|
||||
|
||||
cdef cppclass CTaskArgumentByReference "ray::raylet::TaskArgumentByReference":
|
||||
CTaskArgumentByReference(const c_vector[CObjectID] &references);
|
||||
|
||||
cdef cppclass CTaskArgumentByValue "ray::raylet::TaskArgumentByValue":
|
||||
CTaskArgumentByValue(const uint8_t *value, size_t length);
|
||||
|
||||
cdef cppclass CTaskSpecification "ray::raylet::TaskSpecification":
|
||||
CTaskSpecification(const CDriverID &driver_id, const CTaskID &parent_task_id,
|
||||
int64_t parent_counter,
|
||||
const c_vector[shared_ptr[CTaskArgument]] &task_arguments,
|
||||
int64_t num_returns,
|
||||
const unordered_map[c_string, double] &required_resources,
|
||||
const CLanguage &language,
|
||||
const c_vector[c_string] &function_descriptor)
|
||||
CTaskSpecification(
|
||||
const CDriverID &driver_id, const CTaskID &parent_task_id, int64_t parent_counter,
|
||||
const CActorID &actor_creation_id, const CObjectID &actor_creation_dummy_object_id,
|
||||
int64_t max_actor_reconstructions, const CActorID &actor_id,
|
||||
const CActorHandleID &actor_handle_id, int64_t actor_counter,
|
||||
const c_vector[CActorHandleID] &new_actor_handles,
|
||||
const c_vector[shared_ptr[CTaskArgument]] &task_arguments,
|
||||
int64_t num_returns,
|
||||
const unordered_map[c_string, double] &required_resources,
|
||||
const unordered_map[c_string, double] &required_placement_resources,
|
||||
const CLanguage &language, const c_vector[c_string] &function_descriptor)
|
||||
CTaskSpecification(const c_string &string)
|
||||
c_string SerializeAsString() const
|
||||
|
||||
CTaskID TaskId() const
|
||||
CDriverID DriverId() const
|
||||
CTaskID ParentTaskId() const
|
||||
int64_t ParentCounter() const
|
||||
c_vector[c_string] FunctionDescriptor() const
|
||||
c_string FunctionDescriptorString() const
|
||||
int64_t NumArgs() const
|
||||
int64_t NumReturns() const
|
||||
c_bool ArgByRef(int64_t arg_index) const
|
||||
int ArgIdCount(int64_t arg_index) const
|
||||
CObjectID ArgId(int64_t arg_index, int64_t id_index) const
|
||||
CObjectID ReturnId(int64_t return_index) const
|
||||
const uint8_t *ArgVal(int64_t arg_index) const
|
||||
size_t ArgValLength(int64_t arg_index) const
|
||||
double GetRequiredResource(const c_string &resource_name) const
|
||||
const ResourceSet GetRequiredResources() const
|
||||
const ResourceSet GetRequiredPlacementResources() const
|
||||
c_bool IsDriverTask() const
|
||||
CLanguage GetLanguage() const
|
||||
|
||||
c_bool IsActorCreationTask() const
|
||||
c_bool IsActorTask() const
|
||||
CActorID ActorCreationId() const
|
||||
CObjectID ActorCreationDummyObjectId() const
|
||||
int64_t MaxActorReconstructions() const
|
||||
CActorID ActorId() const
|
||||
CActorHandleID ActorHandleId() const
|
||||
int64_t ActorCounter() const
|
||||
CObjectID ActorDummyObject() const
|
||||
c_vector[CActorHandleID] NewActorHandles() const
|
||||
|
||||
|
||||
cdef extern from "ray/raylet/task.h" namespace "ray::raylet" nogil:
|
||||
cdef cppclass CTask "ray::raylet::Task":
|
||||
CTask(const CTaskExecutionSpecification &execution_spec,
|
||||
const CTaskSpecification &task_spec)
|
||||
const CTaskExecutionSpecification &GetTaskExecutionSpec() const
|
||||
const CTaskSpecification &GetTaskSpecification() const
|
||||
void SetExecutionDependencies(const c_vector[CObjectID] &dependencies)
|
||||
void IncrementNumForwards()
|
||||
const c_vector[CObjectID] &GetDependencies() const
|
||||
void CopyTaskExecutionSpec(const CTask &task)
|
||||
|
||||
cdef c_string SerializeTaskAsString(const c_vector[CObjectID] *dependencies,
|
||||
const CTaskSpecification *task_spec)
|
||||
@@ -0,0 +1,185 @@
|
||||
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 ray.utils import _random_string
|
||||
|
||||
cdef class Task:
|
||||
cdef:
|
||||
unique_ptr[CTaskSpecification] task_spec
|
||||
unique_ptr[c_vector[CObjectID]] execution_dependencies
|
||||
|
||||
def __init__(self, DriverID driver_id, function_descriptor, arguments,
|
||||
int num_returns, TaskID parent_task_id, int parent_counter,
|
||||
ActorID actor_creation_id,
|
||||
ObjectID actor_creation_dummy_object_id,
|
||||
int32_t max_actor_reconstructions, ActorID actor_id,
|
||||
ActorHandleID actor_handle_id, int actor_counter,
|
||||
new_actor_handles, execution_arguments, resource_map,
|
||||
placement_resource_map):
|
||||
cdef:
|
||||
unordered_map[c_string, double] required_resources
|
||||
unordered_map[c_string, double] required_placement_resources
|
||||
c_vector[shared_ptr[CTaskArgument]] task_args
|
||||
c_vector[CActorHandleID] task_new_actor_handles
|
||||
c_vector[c_string] c_function_descriptor
|
||||
c_string pickled_str
|
||||
c_vector[CObjectID] references
|
||||
|
||||
for item in function_descriptor:
|
||||
if not isinstance(item, bytes):
|
||||
raise TypeError("'function_descriptor' takes a list of byte strings.")
|
||||
c_function_descriptor.push_back(item)
|
||||
|
||||
# Parse the resource map.
|
||||
if resource_map is not None:
|
||||
required_resources = resource_map_from_python_dict(resource_map)
|
||||
if required_resources.count(b"CPU") == 0:
|
||||
required_resources[b"CPU"] = 1.0
|
||||
if placement_resource_map is not None:
|
||||
required_placement_resources = resource_map_from_python_dict(placement_resource_map)
|
||||
|
||||
# Parse the arguments from the list.
|
||||
for arg in arguments:
|
||||
if isinstance(arg, ObjectID):
|
||||
references = c_vector[CObjectID]()
|
||||
references.push_back((<ObjectID>arg).data)
|
||||
task_args.push_back(static_pointer_cast[CTaskArgument, CTaskArgumentByReference](make_shared[CTaskArgumentByReference](references)))
|
||||
else:
|
||||
pickled_str = pickle.dumps(arg, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
task_args.push_back(static_pointer_cast[CTaskArgument, CTaskArgumentByValue](make_shared[CTaskArgumentByValue](<uint8_t *>pickled_str.c_str(), pickled_str.size())))
|
||||
|
||||
for new_actor_handle in new_actor_handles:
|
||||
task_new_actor_handles.push_back((<ActorHandleID?>new_actor_handle).data)
|
||||
|
||||
self.task_spec.reset(new CTaskSpecification(
|
||||
CUniqueID(driver_id.data), parent_task_id.data, parent_counter, actor_creation_id.data,
|
||||
actor_creation_dummy_object_id.data, max_actor_reconstructions, CUniqueID(actor_id.data),
|
||||
CUniqueID(actor_handle_id.data), actor_counter, task_new_actor_handles, task_args, num_returns,
|
||||
required_resources, required_placement_resources, LANGUAGE_PYTHON,
|
||||
c_function_descriptor))
|
||||
|
||||
# Set the task's execution dependencies.
|
||||
self.execution_dependencies.reset(new c_vector[CObjectID]())
|
||||
if execution_arguments is not None:
|
||||
for execution_arg in execution_arguments:
|
||||
self.execution_dependencies.get().push_back((<ObjectID?>execution_arg).data)
|
||||
|
||||
@staticmethod
|
||||
cdef make(unique_ptr[CTaskSpecification]& task_spec):
|
||||
cdef Task self = Task.__new__(Task)
|
||||
self.task_spec.reset(task_spec.release())
|
||||
# The created task does not include any execution dependencies.
|
||||
self.execution_dependencies.reset(new c_vector[CObjectID]())
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def from_string(const c_string& task_spec_str):
|
||||
"""Convert a string to a Ray task specification Python object.
|
||||
|
||||
Args:
|
||||
task_spec_str: String representation of the task specification.
|
||||
|
||||
Returns:
|
||||
Python task specification object.
|
||||
"""
|
||||
cdef Task self = Task.__new__(Task)
|
||||
# TODO(pcm): Use flatbuffers validation here.
|
||||
self.task_spec.reset(new CTaskSpecification(task_spec_str))
|
||||
# The created task does not include any execution dependencies.
|
||||
self.execution_dependencies.reset(new c_vector[CObjectID]())
|
||||
return self
|
||||
|
||||
def to_string(self):
|
||||
"""Convert a Ray task specification Python object to a string.
|
||||
|
||||
Returns:
|
||||
String representing the task specification.
|
||||
"""
|
||||
return self.task_spec.get().SerializeAsString()
|
||||
|
||||
def _serialized_raylet_task(self):
|
||||
return SerializeTaskAsString(self.execution_dependencies.get(), self.task_spec.get())
|
||||
|
||||
def driver_id(self):
|
||||
"""Return the driver ID for this task."""
|
||||
return DriverID.from_native(self.task_spec.get().DriverId())
|
||||
|
||||
def task_id(self):
|
||||
"""Return the task ID for this task."""
|
||||
return TaskID.from_native(self.task_spec.get().TaskId())
|
||||
|
||||
def parent_task_id(self):
|
||||
"""Return the task ID of the parent task."""
|
||||
return TaskID.from_native(self.task_spec.get().ParentTaskId())
|
||||
|
||||
def parent_counter(self):
|
||||
"""Return the parent counter of this task."""
|
||||
return self.task_spec.get().ParentCounter()
|
||||
|
||||
def function_descriptor_list(self):
|
||||
"""Return the function descriptor for this task."""
|
||||
cdef c_vector[c_string] function_descriptor = self.task_spec.get().FunctionDescriptor()
|
||||
results = []
|
||||
for i in range(function_descriptor.size()):
|
||||
results.append(function_descriptor[i])
|
||||
return results
|
||||
|
||||
def arguments(self):
|
||||
"""Return the arguments for the task."""
|
||||
cdef:
|
||||
CTaskSpecification *task_spec = self.task_spec.get()
|
||||
int64_t num_args = task_spec.NumArgs()
|
||||
int count
|
||||
arg_list = []
|
||||
for i in range(num_args):
|
||||
count = task_spec.ArgIdCount(i)
|
||||
if count > 0:
|
||||
assert count == 1
|
||||
arg_list.append(ObjectID.from_native(task_spec.ArgId(i, 0)))
|
||||
else:
|
||||
serialized_str = task_spec.ArgVal(i)[:task_spec.ArgValLength(i)]
|
||||
obj = pickle.loads(serialized_str)
|
||||
arg_list.append(obj)
|
||||
return arg_list
|
||||
|
||||
def returns(self):
|
||||
"""Return the object IDs for the return values of the task."""
|
||||
cdef CTaskSpecification *task_spec = self.task_spec.get()
|
||||
return_id_list = []
|
||||
for i in range(task_spec.NumReturns()):
|
||||
return_id_list.append(ObjectID.from_native(task_spec.ReturnId(i)))
|
||||
return return_id_list
|
||||
|
||||
def required_resources(self):
|
||||
"""Return the resource dictionary of the task."""
|
||||
cdef:
|
||||
unordered_map[c_string, double] resource_map = self.task_spec.get().GetRequiredResources().GetResourceMap()
|
||||
c_string resource_name
|
||||
double resource_value
|
||||
unordered_map[c_string, double].iterator iterator = resource_map.begin()
|
||||
|
||||
required_resources = {}
|
||||
while iterator != resource_map.end():
|
||||
resource_name = dereference(iterator).first
|
||||
py_resource_name = str(resource_name) # bytes for Py2, unicode for Py3
|
||||
resource_value = dereference(iterator).second
|
||||
required_resources[py_resource_name] = resource_value
|
||||
postincrement(iterator)
|
||||
return required_resources
|
||||
|
||||
def actor_creation_id(self):
|
||||
"""Return the actor creation ID for the task."""
|
||||
return ActorID.from_native(self.task_spec.get().ActorCreationId())
|
||||
|
||||
def actor_creation_dummy_object_id(self):
|
||||
"""Return the actor creation dummy object ID for the task."""
|
||||
return ObjectID.from_native(self.task_spec.get().ActorCreationDummyObjectId())
|
||||
|
||||
def actor_id(self):
|
||||
"""Return the actor ID for this task."""
|
||||
return ActorID.from_native(self.task_spec.get().ActorId())
|
||||
|
||||
def actor_counter(self):
|
||||
"""Return the actor counter for this task."""
|
||||
return self.task_spec.get().ActorCounter()
|
||||
@@ -0,0 +1,34 @@
|
||||
from libcpp cimport bool as c_bool
|
||||
from libcpp.string cimport string as c_string
|
||||
from libc.stdint cimport uint8_t
|
||||
|
||||
cdef extern from "ray/id.h" namespace "ray" nogil:
|
||||
cdef cppclass CUniqueID "ray::UniqueID":
|
||||
CUniqueID()
|
||||
CUniqueID(const CUniqueID &from_id)
|
||||
@staticmethod
|
||||
CUniqueID from_random()
|
||||
@staticmethod
|
||||
CUniqueID from_binary(const c_string & binary)
|
||||
@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;
|
||||
|
||||
ctypedef CUniqueID TaskID
|
||||
ctypedef CUniqueID ObjectID
|
||||
ctypedef CUniqueID FunctionID
|
||||
ctypedef CUniqueID ActorID
|
||||
ctypedef CUniqueID ActorClassID
|
||||
ctypedef CUniqueID ActorHandleID
|
||||
ctypedef CUniqueID WorkerID
|
||||
ctypedef CUniqueID DriverID
|
||||
ctypedef CUniqueID ConfigID
|
||||
ctypedef CUniqueID ClientID
|
||||
@@ -0,0 +1,280 @@
|
||||
"""This is a module for unique IDs in Ray.
|
||||
We define different types for different IDs for type safety.
|
||||
|
||||
See https://github.com/ray-project/ray/issues/3721.
|
||||
"""
|
||||
|
||||
from ray.includes.common cimport (
|
||||
CUniqueID, CTaskID, CObjectID, CFunctionID, CActorClassID, CActorID,
|
||||
CActorHandleID, CWorkerID, CDriverID, CConfigID, CClientID,
|
||||
ComputePutId, ComputeTaskId)
|
||||
|
||||
from ray.utils import decode
|
||||
|
||||
|
||||
def check_id(b):
|
||||
if not isinstance(b, bytes):
|
||||
raise TypeError("Unsupported type: " + str(type(b)))
|
||||
if len(b) != kUniqueIDSize:
|
||||
raise ValueError("ID string needs to have length " + str(kUniqueIDSize))
|
||||
|
||||
|
||||
cdef extern from "ray/constants.h" nogil:
|
||||
cdef int64_t kUniqueIDSize
|
||||
cdef int64_t kMaxTaskPuts
|
||||
|
||||
|
||||
cdef class UniqueID:
|
||||
cdef CUniqueID data
|
||||
|
||||
def __init__(self, id):
|
||||
if not id:
|
||||
self.data = CUniqueID()
|
||||
else:
|
||||
check_id(id)
|
||||
self.data = CUniqueID.from_binary(id)
|
||||
|
||||
@staticmethod
|
||||
cdef from_native(const CUniqueID& cpp_id):
|
||||
cdef UniqueID self = UniqueID.__new__(UniqueID)
|
||||
self.data = cpp_id
|
||||
return self
|
||||
|
||||
@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)
|
||||
|
||||
@staticmethod
|
||||
def nil():
|
||||
return UniqueID.from_native(CUniqueID.nil())
|
||||
|
||||
def __hash__(self):
|
||||
return self.data.hash()
|
||||
|
||||
def is_nil(self):
|
||||
return self.data.is_nil()
|
||||
|
||||
def __eq__(self, other):
|
||||
return self.binary() == other.binary()
|
||||
|
||||
def __ne__(self, other):
|
||||
return self.binary() != other.binary()
|
||||
|
||||
def size(self):
|
||||
return self.data.size()
|
||||
|
||||
def __len__(self):
|
||||
return self.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()
|
||||
|
||||
def __repr__(self):
|
||||
return "UniqueID(" + self.hex() + ")"
|
||||
|
||||
def __str__(self):
|
||||
return self.__repr__()
|
||||
|
||||
def __reduce__(self):
|
||||
return type(self), (self.binary(),)
|
||||
|
||||
def redis_shard_hash(self):
|
||||
# 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
|
||||
return self.data.hash()
|
||||
|
||||
|
||||
cdef class ObjectID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
if not id:
|
||||
self.data = CUniqueID()
|
||||
else:
|
||||
check_id(id)
|
||||
self.data = CUniqueID.from_binary(id)
|
||||
|
||||
@staticmethod
|
||||
cdef from_native(const CObjectID& cpp_id):
|
||||
cdef ObjectID self = ObjectID.__new__(ObjectID)
|
||||
self.data = cpp_id
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def nil():
|
||||
return ObjectID.from_native(CObjectID.nil())
|
||||
|
||||
def __repr__(self):
|
||||
return "ObjectID(" + self.hex() + ")"
|
||||
|
||||
|
||||
cdef class TaskID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
if not id:
|
||||
self.data = CUniqueID()
|
||||
else:
|
||||
check_id(id)
|
||||
self.data = CUniqueID.from_binary(id)
|
||||
|
||||
@staticmethod
|
||||
cdef from_native(const CTaskID& cpp_id):
|
||||
cdef TaskID self = TaskID.__new__(TaskID)
|
||||
self.data = cpp_id
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def nil():
|
||||
return TaskID.from_native(CTaskID.nil())
|
||||
|
||||
def __repr__(self):
|
||||
return "TaskID(" + self.hex() + ")"
|
||||
|
||||
|
||||
cdef class ClientID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
if not id:
|
||||
self.data = CUniqueID()
|
||||
else:
|
||||
check_id(id)
|
||||
self.data = CUniqueID.from_binary(id)
|
||||
|
||||
@staticmethod
|
||||
cdef from_native(const CClientID& cpp_id):
|
||||
cdef ClientID self = ClientID.__new__(ClientID)
|
||||
self.data = cpp_id
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def nil():
|
||||
return ClientID.from_native(CClientID.nil())
|
||||
|
||||
def __repr__(self):
|
||||
return "ClientID(" + self.hex() + ")"
|
||||
|
||||
|
||||
cdef class DriverID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
if not id:
|
||||
self.data = CUniqueID()
|
||||
else:
|
||||
check_id(id)
|
||||
self.data = CUniqueID.from_binary(id)
|
||||
|
||||
@staticmethod
|
||||
cdef from_native(const CDriverID& cpp_id):
|
||||
cdef DriverID self = DriverID.__new__(DriverID)
|
||||
self.data = cpp_id
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def nil():
|
||||
return DriverID.from_native(CDriverID.nil())
|
||||
|
||||
def __repr__(self):
|
||||
return "DriverID(" + self.hex() + ")"
|
||||
|
||||
|
||||
cdef class ActorID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
if not id:
|
||||
self.data = CUniqueID()
|
||||
else:
|
||||
check_id(id)
|
||||
self.data = CUniqueID.from_binary(id)
|
||||
|
||||
@staticmethod
|
||||
cdef from_native(const CActorID& cpp_id):
|
||||
cdef ActorID self = ActorID.__new__(ActorID)
|
||||
self.data = cpp_id
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def nil():
|
||||
return ActorID.from_native(CActorID.nil())
|
||||
|
||||
def __repr__(self):
|
||||
return "ActorID(" + self.hex() + ")"
|
||||
|
||||
|
||||
cdef class ActorHandleID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
if not id:
|
||||
self.data = CUniqueID()
|
||||
else:
|
||||
check_id(id)
|
||||
self.data = CUniqueID.from_binary(id)
|
||||
|
||||
@staticmethod
|
||||
cdef from_native(const CActorHandleID& cpp_id):
|
||||
cdef ActorHandleID self = ActorHandleID.__new__(ActorHandleID)
|
||||
self.data = cpp_id
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def nil():
|
||||
return ActorHandleID.from_native(CActorHandleID.nil())
|
||||
|
||||
def __repr__(self):
|
||||
return "ActorHandleID(" + self.hex() + ")"
|
||||
|
||||
|
||||
cdef class FunctionID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
if not id:
|
||||
self.data = CUniqueID()
|
||||
else:
|
||||
check_id(id)
|
||||
self.data = CUniqueID.from_binary(id)
|
||||
|
||||
@staticmethod
|
||||
cdef from_native(const CFunctionID& cpp_id):
|
||||
cdef FunctionID self = FunctionID.__new__(FunctionID)
|
||||
self.data = cpp_id
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def nil():
|
||||
return FunctionID.from_native(CFunctionID.nil())
|
||||
|
||||
def __repr__(self):
|
||||
return "FunctionID(" + self.hex() + ")"
|
||||
|
||||
|
||||
cdef class ActorClassID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
if not id:
|
||||
self.data = CUniqueID()
|
||||
else:
|
||||
check_id(id)
|
||||
self.data = CUniqueID.from_binary(id)
|
||||
|
||||
@staticmethod
|
||||
cdef from_native(const CActorClassID& cpp_id):
|
||||
cdef ActorClassID self = ActorClassID.__new__(ActorClassID)
|
||||
self.data = cpp_id
|
||||
return self
|
||||
|
||||
@staticmethod
|
||||
def nil():
|
||||
return ActorClassID.from_native(CActorClassID.nil())
|
||||
|
||||
def __repr__(self):
|
||||
return "ActorClassID(" + self.hex() + ")"
|
||||
Reference in New Issue
Block a user