507 lines
21 KiB
Cython
507 lines
21 KiB
Cython
# cython: profile = False
|
|
# distutils: language = c++
|
|
# cython: embedsignature = True
|
|
|
|
from libc.stdint cimport int64_t, uint64_t
|
|
from libcpp cimport bool as c_bool
|
|
from libcpp.functional cimport function
|
|
from libcpp.memory cimport shared_ptr, unique_ptr
|
|
from libcpp.pair cimport pair as c_pair
|
|
from libcpp.string cimport string as c_string
|
|
from libcpp.unordered_map cimport unordered_map
|
|
from libcpp.utility cimport pair
|
|
from libcpp.vector cimport vector as c_vector
|
|
|
|
from ray.includes.unique_ids cimport (
|
|
CActorID,
|
|
CClusterID,
|
|
CNodeID,
|
|
CJobID,
|
|
CTaskID,
|
|
CObjectID,
|
|
CPlacementGroupID,
|
|
CWorkerID,
|
|
ObjectIDIndexType,
|
|
)
|
|
|
|
from ray.includes.common cimport (
|
|
CAddress,
|
|
CObjectReference,
|
|
CActorCreationOptions,
|
|
CBuffer,
|
|
CPlacementGroupCreationOptions,
|
|
CObjectLocation,
|
|
CObjectReference,
|
|
CRayFunction,
|
|
CRayObject,
|
|
CRayStatus,
|
|
CTaskArg,
|
|
CTaskOptions,
|
|
CTaskType,
|
|
CWorkerType,
|
|
CLanguage,
|
|
CGcsClientOptions,
|
|
LocalMemoryBuffer,
|
|
CJobConfig,
|
|
CConcurrencyGroup,
|
|
CSchedulingStrategy,
|
|
CWorkerExitType,
|
|
CLineageReconstructionTask,
|
|
)
|
|
from ray.includes.function_descriptor cimport (
|
|
CFunctionDescriptor,
|
|
)
|
|
|
|
from ray.includes.optional cimport (
|
|
optional,
|
|
)
|
|
|
|
ctypedef unordered_map[c_string, c_vector[pair[int64_t, double]]] \
|
|
ResourceMappingType
|
|
|
|
ctypedef void (*ray_callback_function) \
|
|
(shared_ptr[CRayObject] result_object,
|
|
CObjectID object_id, void* user_data)
|
|
|
|
ctypedef void (*plasma_callback_function) \
|
|
(CObjectID object_id, int64_t data_size, int64_t metadata_size)
|
|
|
|
# NOTE: This ctypedef is needed, because Cython doesn't compile
|
|
# "pair[shared_ptr[const CActorHandle], CRayStatus]".
|
|
# This is a bug of cython: https://github.com/cython/cython/issues/3967.
|
|
ctypedef shared_ptr[const CActorHandle] ActorHandleSharedPtr
|
|
|
|
cdef extern from "ray/core_worker/profile_event.h" nogil:
|
|
cdef cppclass CProfileEvent "ray::core::worker::ProfileEvent":
|
|
void SetExtraData(const c_string &extra_data)
|
|
|
|
cdef extern from "ray/core_worker/task_execution/fiber.h" nogil:
|
|
cdef cppclass CFiberEvent "ray::core::FiberEvent":
|
|
CFiberEvent()
|
|
void Wait()
|
|
void Notify()
|
|
|
|
cdef extern from "ray/core_worker/experimental_mutable_object_manager.h" nogil:
|
|
cdef cppclass CReaderRefInfo "ray::experimental::ReaderRefInfo":
|
|
CReaderRefInfo()
|
|
CObjectID reader_ref_id
|
|
CActorID owner_reader_actor_id
|
|
int64_t num_reader_actors
|
|
|
|
|
|
cdef extern from "ray/core_worker/context.h" nogil:
|
|
cdef cppclass CWorkerContext "ray::core::WorkerContext":
|
|
c_bool CurrentActorIsAsync()
|
|
void SetCurrentActorShouldExit()
|
|
c_bool GetCurrentActorShouldExit()
|
|
const c_string &GetCurrentSerializedRuntimeEnv()
|
|
int CurrentActorMaxConcurrency()
|
|
CActorID GetRootDetachedActorID()
|
|
|
|
cdef extern from "ray/core_worker/generator_waiter.h" nogil:
|
|
cdef cppclass CTaskGeneratorBackpressureWaiter "ray::core::TaskGeneratorBackpressureWaiter": # noqa
|
|
CTaskGeneratorBackpressureWaiter(
|
|
int64_t generator_backpressure_num_objects,
|
|
(CRayStatus() nogil) check_signals)
|
|
CRayStatus WaitAllObjectsReported()
|
|
CRayStatus WaitUntilObjectConsumed()
|
|
c_bool IsBackpressured()
|
|
c_bool NeedsObjectConsumedUpdates()
|
|
|
|
cdef cppclass CActorWideGeneratorBackpressureWaiter "ray::core::ActorWideGeneratorBackpressureWaiter": # noqa
|
|
pass
|
|
|
|
cdef cppclass CActorTaskBackpressureMetadata "ray::core::ActorTaskBackpressureMetadata": # noqa
|
|
CActorTaskBackpressureMetadata(
|
|
shared_ptr[CActorWideGeneratorBackpressureWaiter] actor_waiter)
|
|
CRayStatus ReserveSlot(int64_t num_objects)
|
|
c_bool TryReserveSlot(int64_t num_objects)
|
|
void ReleaseSlot(int64_t num_objects)
|
|
void Teardown()
|
|
|
|
cdef extern from "ray/core_worker/core_worker.h" nogil:
|
|
cdef cppclass CActorHandle "ray::core::ActorHandle":
|
|
CActorID GetActorID() const
|
|
CJobID CreationJobID() const
|
|
CLanguage ActorLanguage() const
|
|
CFunctionDescriptor ActorCreationTaskFunctionDescriptor() const
|
|
c_string ExtensionData() const
|
|
int MaxPendingCalls() const
|
|
int MaxTaskRetries() const
|
|
c_bool EnableTaskEvents() const
|
|
c_bool AllowOutOfOrderExecution() const
|
|
c_bool EnableTensorTransport() const
|
|
int64_t ActorGeneratorBackpressureNumObjects() const
|
|
|
|
cdef cppclass CCoreWorker "ray::core::CoreWorker":
|
|
CWorkerType GetWorkerType()
|
|
CLanguage GetLanguage()
|
|
|
|
c_vector[CObjectReference] SubmitTask(
|
|
const CRayFunction &function,
|
|
const c_vector[unique_ptr[CTaskArg]] &args,
|
|
const CTaskOptions &options,
|
|
int max_retries,
|
|
c_bool retry_exceptions,
|
|
const CSchedulingStrategy &scheduling_strategy,
|
|
c_string debugger_breakpoint,
|
|
c_string serialized_retry_exception_allowlist,
|
|
c_string call_site,
|
|
const CTaskID current_task_id)
|
|
CRayStatus CreateActor(
|
|
const CRayFunction &function,
|
|
const c_vector[unique_ptr[CTaskArg]] &args,
|
|
const CActorCreationOptions &options,
|
|
const c_string &extension_data,
|
|
c_string call_site,
|
|
CActorID *actor_id)
|
|
CRayStatus CreatePlacementGroup(
|
|
const CPlacementGroupCreationOptions &options,
|
|
CPlacementGroupID *placement_group_id)
|
|
CRayStatus RemovePlacementGroup(
|
|
const CPlacementGroupID &placement_group_id)
|
|
CRayStatus WaitPlacementGroupReady(
|
|
const CPlacementGroupID &placement_group_id, int64_t timeout_seconds)
|
|
CObjectID AsyncWaitPlacementGroupReady(
|
|
const CPlacementGroupID &placement_group_id,
|
|
const c_string &serialized_object_data,
|
|
const c_string &serialized_object_metadata)
|
|
CRayStatus SubmitActorTask(
|
|
const CActorID &actor_id, const CRayFunction &function,
|
|
const c_vector[unique_ptr[CTaskArg]] &args,
|
|
const CTaskOptions &options,
|
|
int max_retries,
|
|
c_bool retry_exceptions,
|
|
c_string serialized_retry_exception_allowlist,
|
|
c_string call_site,
|
|
c_vector[CObjectReference] &task_returns,
|
|
const CTaskID current_task_id)
|
|
CRayStatus KillActor(
|
|
const CActorID &actor_id, c_bool force_kill,
|
|
c_bool no_restart)
|
|
CRayStatus CancelTask(const CObjectID &object_id, c_bool force_kill,
|
|
c_bool recursive)
|
|
c_bool IsTaskCanceled(const CTaskID &task_id) const
|
|
c_bool ShouldInterruptTaskForCancellation() const
|
|
|
|
unique_ptr[CProfileEvent] CreateProfileEvent(
|
|
const c_string &event_type)
|
|
CRayStatus AllocateReturnObject(
|
|
const CObjectID &object_id,
|
|
const size_t &data_size,
|
|
const shared_ptr[CBuffer] &metadata,
|
|
const c_vector[CObjectID] &contained_object_id,
|
|
const CAddress &caller_address,
|
|
int64_t *task_output_inlined_bytes,
|
|
shared_ptr[CRayObject] *return_object)
|
|
CRayStatus SealReturnObject(
|
|
const CObjectID &return_id,
|
|
const shared_ptr[CRayObject] &return_object,
|
|
const CObjectID &generator_id,
|
|
const CAddress &caller_address
|
|
)
|
|
c_bool PinExistingReturnObject(
|
|
const CObjectID &return_id,
|
|
shared_ptr[CRayObject] *return_object,
|
|
const CObjectID &generator_id,
|
|
const CAddress &caller_address)
|
|
void AsyncDelObjectRefStream(const CObjectID &generator_id)
|
|
CRayStatus TryReadObjectRefStream(
|
|
const CObjectID &generator_id,
|
|
CObjectReference *object_ref_out)
|
|
CRayStatus TryReadObjectRefStreamN(
|
|
const CObjectID &generator_id,
|
|
int64_t num_items)
|
|
c_bool StreamingGeneratorIsFinished(const CObjectID &generator_id) const
|
|
pair[CObjectReference, c_bool] PeekObjectRefStream(
|
|
const CObjectID &generator_id)
|
|
c_vector[pair[CObjectReference, c_bool]] PeekObjectRefStreamN(
|
|
const CObjectID &generator_id, int64_t num_items)
|
|
CObjectID PeekObjectIdStream(
|
|
const CObjectID &generator_id)
|
|
CObjectID AllocateDynamicReturnId(
|
|
const CAddress &owner_address,
|
|
const CTaskID &task_id,
|
|
optional[ObjectIDIndexType] put_index)
|
|
|
|
CJobID GetCurrentJobId()
|
|
CTaskID GetCurrentTaskId()
|
|
const c_string GetCurrentTaskName()
|
|
const c_string GetCurrentTaskFunctionName()
|
|
void UpdateTaskIsDebuggerPaused(
|
|
const CTaskID &task_id,
|
|
const c_bool is_debugger_paused)
|
|
int64_t GetCurrentTaskAttemptNumber()
|
|
CNodeID GetCurrentNodeId()
|
|
int64_t GetTaskDepth()
|
|
c_bool GetCurrentTaskRetryExceptions()
|
|
CPlacementGroupID GetCurrentPlacementGroupId() const
|
|
CWorkerID GetWorkerID()
|
|
c_bool ShouldCaptureChildTasksInPlacementGroup()
|
|
CActorID GetActorId() const
|
|
const c_string GetActorName()
|
|
const ResourceMappingType &GetResourceIDs() const
|
|
void RemoveActorHandleReference(const CActorID &actor_id)
|
|
optional[int] GetLocalActorState(const CActorID &actor_id) const
|
|
CActorID DeserializeAndRegisterActorHandle(const c_string &bytes, const
|
|
CObjectID &outer_object_id,
|
|
c_bool add_local_ref)
|
|
CRayStatus SerializeActorHandle(const CActorID &actor_id, c_string
|
|
*bytes,
|
|
CObjectID *c_actor_handle_id)
|
|
ActorHandleSharedPtr GetActorHandle(const CActorID &actor_id) const
|
|
pair[ActorHandleSharedPtr, CRayStatus] GetNamedActorHandle(
|
|
const c_string &name, const c_string &ray_namespace)
|
|
pair[c_vector[c_pair[c_string, c_string]], CRayStatus] ListNamedActors(
|
|
c_bool all_namespaces)
|
|
void AddLocalReference(const CObjectID &object_id)
|
|
void RemoveLocalReference(const CObjectID &object_id)
|
|
c_bool AddObjectOutOfScopeOrFreedCallback(
|
|
const CObjectID &object_id,
|
|
void (*callback)(const CObjectID &, void *) nogil,
|
|
void *callback_context)
|
|
CRayStatus CheckObjectOwnedByUs(const CObjectID &object_id) const
|
|
void PutObjectIntoPlasma(const CRayObject &object,
|
|
const CObjectID &object_id)
|
|
const CAddress &GetRpcAddress() const
|
|
CRayStatus GetOwnerAddress(const CObjectID &object_id,
|
|
CAddress *owner_address) const
|
|
c_vector[CObjectReference] GetObjectRefs(
|
|
const c_vector[CObjectID] &object_ids) const
|
|
|
|
CRayStatus GetOwnershipInfo(const CObjectID &object_id,
|
|
CAddress *owner_address,
|
|
c_string *object_status)
|
|
void RegisterOwnershipInfoAndResolveFuture(
|
|
const CObjectID &object_id,
|
|
const CObjectID &outer_object_id,
|
|
const CAddress &owner_address,
|
|
const c_string &object_status)
|
|
|
|
CRayStatus Put(const CRayObject &object,
|
|
const c_vector[CObjectID] &contained_object_ids,
|
|
CObjectID *object_id)
|
|
CRayStatus Put(const CRayObject &object,
|
|
const c_vector[CObjectID] &contained_object_ids,
|
|
const CObjectID &object_id)
|
|
CRayStatus CreateOwnedAndIncrementLocalRef(
|
|
c_bool is_mutable,
|
|
const shared_ptr[CBuffer] &metadata,
|
|
const size_t data_size,
|
|
const c_vector[CObjectID] &contained_object_ids,
|
|
CObjectID *object_id, shared_ptr[CBuffer] *data,
|
|
c_bool inline_small_object,
|
|
const optional[c_string] &tensor_transport)
|
|
CRayStatus CreateExisting(const shared_ptr[CBuffer] &metadata,
|
|
const size_t data_size,
|
|
const CObjectID &object_id,
|
|
const CAddress &owner_address,
|
|
shared_ptr[CBuffer] *data,
|
|
c_bool created_by_worker)
|
|
CRayStatus ExperimentalChannelWriteAcquire(
|
|
const CObjectID &object_id,
|
|
const shared_ptr[CBuffer] &metadata,
|
|
uint64_t data_size,
|
|
int64_t num_readers,
|
|
int64_t timeout_ms,
|
|
shared_ptr[CBuffer] *data)
|
|
CRayStatus ExperimentalChannelWriteRelease(
|
|
const CObjectID &object_id)
|
|
CRayStatus ExperimentalChannelSetError(
|
|
const CObjectID &object_id)
|
|
void ExperimentalRegisterMutableObjectWriter(
|
|
const CObjectID &writer_object_id,
|
|
const c_vector[CNodeID] &remote_reader_node_ids)
|
|
CRayStatus ExperimentalRegisterMutableObjectReader(const CObjectID &object_id)
|
|
CRayStatus ExperimentalRegisterMutableObjectReaderRemote(
|
|
const CObjectID &object_id,
|
|
const c_vector[CReaderRefInfo] &remote_reader_ref_info)
|
|
CRayStatus SealOwned(const CObjectID &object_id, c_bool pin_object)
|
|
CRayStatus SealExisting(const CObjectID &object_id, c_bool pin_object,
|
|
const CObjectID &generator_id,
|
|
const unique_ptr[CAddress] &owner_address)
|
|
CRayStatus Get(const c_vector[CObjectID] &ids, int64_t timeout_ms,
|
|
c_vector[shared_ptr[CRayObject]] results)
|
|
CRayStatus GetIfLocal(
|
|
const c_vector[CObjectID] &ids,
|
|
c_vector[shared_ptr[CRayObject]] *results)
|
|
CRayStatus Contains(const CObjectID &object_id, c_bool *has_object,
|
|
c_bool *is_in_plasma)
|
|
CRayStatus Wait(const c_vector[CObjectID] &object_ids, int num_objects,
|
|
int64_t timeout_ms, c_vector[c_bool] *results,
|
|
c_bool fetch_local)
|
|
CRayStatus Delete(const c_vector[CObjectID] &object_ids,
|
|
c_bool local_only)
|
|
CRayStatus GetLocalObjectLocations(
|
|
const c_vector[CObjectID] &object_ids,
|
|
c_vector[optional[CObjectLocation]] *results)
|
|
CRayStatus GetLocationFromOwner(
|
|
const c_vector[CObjectID] &object_ids,
|
|
int64_t timeout_ms,
|
|
c_vector[shared_ptr[CObjectLocation]] *results)
|
|
CRayStatus TriggerGlobalGC()
|
|
CRayStatus ReportGeneratorItemReturns(
|
|
const c_vector[c_pair[CObjectID, shared_ptr[CRayObject]]] &dynamic_return_objects,
|
|
const CObjectID &generator_id,
|
|
const CAddress &caller_address,
|
|
int64_t item_index,
|
|
uint64_t attempt_number,
|
|
shared_ptr[CTaskGeneratorBackpressureWaiter] waiter,
|
|
shared_ptr[CActorTaskBackpressureMetadata] actor_metadata)
|
|
void MarkGeneratorBackpressureTaskFinished(
|
|
const CObjectID &generator_id)
|
|
c_bool TeardownGeneratorBackpressureTask(
|
|
const CObjectID &generator_id)
|
|
void RegisterGeneratorBackpressureState(
|
|
const CObjectID &generator_id,
|
|
shared_ptr[CTaskGeneratorBackpressureWaiter] waiter,
|
|
shared_ptr[CActorTaskBackpressureMetadata] actor_metadata,
|
|
const CAddress &owner_address)
|
|
void SetAsyncGeneratorBackpressureUnblockNotify(
|
|
const CObjectID &generator_id,
|
|
(void(void *) noexcept nogil) fn,
|
|
void *ctx)
|
|
void ClearAsyncGeneratorBackpressureUnblockNotify(
|
|
const CObjectID &generator_id)
|
|
void NotifyAsyncGeneratorBackpressureUnblock(
|
|
const CObjectID &generator_id,
|
|
c_bool notify_all)
|
|
shared_ptr[CActorWideGeneratorBackpressureWaiter] GetActorGeneratorWaiter() const
|
|
|
|
# Param output contains the usage string if successful.
|
|
# Returns an error status if unable to communicate with the plasma store.
|
|
CRayStatus GetPlasmaUsage(c_string &output)
|
|
|
|
int GetMemoryStoreSize()
|
|
|
|
CWorkerContext &GetWorkerContext()
|
|
void YieldCurrentFiber(CFiberEvent &coroutine_done)
|
|
|
|
unordered_map[CObjectID, pair[size_t, size_t]] GetAllReferenceCounts()
|
|
c_vector[CTaskID] GetPendingChildrenTasks(const CTaskID &task_id) const
|
|
|
|
void GetAsync(const CObjectID &object_id,
|
|
ray_callback_function success_callback,
|
|
void* python_user_callback)
|
|
|
|
CRayStatus PushError(const CJobID &job_id, const c_string &type,
|
|
const c_string &error_message, double timestamp)
|
|
CRayStatus SetResource(const c_string &resource_name,
|
|
const double capacity,
|
|
const CNodeID &client_Id)
|
|
|
|
CJobConfig GetJobConfig()
|
|
|
|
int64_t GetLocalMemoryStoreBytesUsed() const
|
|
|
|
void RecordTaskLogStart(
|
|
const CTaskID &task_id,
|
|
int attempt_number,
|
|
const c_string& stdout_path,
|
|
const c_string& stderr_path,
|
|
int64_t stdout_start_offset,
|
|
int64_t stderr_start_offset) const
|
|
|
|
void RecordTaskLogEnd(
|
|
const CTaskID &task_id,
|
|
int attempt_number,
|
|
int64_t stdout_end_offset,
|
|
int64_t stderr_end_offset) const
|
|
|
|
void Exit(const CWorkerExitType exit_type,
|
|
const c_string &detail,
|
|
const shared_ptr[LocalMemoryBuffer] &creation_task_exception_pb_bytes)
|
|
|
|
unordered_map[CLineageReconstructionTask, uint64_t] \
|
|
GetLocalOngoingLineageReconstructionTasks() const
|
|
|
|
cdef cppclass CCoreWorkerOptions "ray::core::CoreWorkerOptions":
|
|
CWorkerType worker_type
|
|
CLanguage language
|
|
c_string store_socket
|
|
c_string raylet_socket
|
|
CJobID job_id
|
|
CGcsClientOptions gcs_options
|
|
c_bool enable_logging
|
|
c_string log_dir
|
|
c_bool install_failure_signal_handler
|
|
c_bool interactive
|
|
c_string node_ip_address
|
|
int node_manager_port
|
|
c_string driver_name
|
|
(CRayStatus(
|
|
const CAddress &caller_address,
|
|
CTaskType task_type,
|
|
const c_string name,
|
|
const CRayFunction &ray_function,
|
|
const unordered_map[c_string, double] &resources,
|
|
const c_vector[shared_ptr[CRayObject]] &args,
|
|
const c_vector[CObjectReference] &arg_refs,
|
|
const c_string debugger_breakpoint,
|
|
const c_string serialized_retry_exception_allowlist,
|
|
c_vector[c_pair[CObjectID, shared_ptr[CRayObject]]] *returns,
|
|
c_vector[c_pair[CObjectID, shared_ptr[CRayObject]]] *dynamic_returns,
|
|
c_vector[c_pair[CObjectID, c_bool]] *streaming_generator_returns,
|
|
shared_ptr[LocalMemoryBuffer]
|
|
&creation_task_exception_pb_bytes,
|
|
c_bool *is_retryable_error,
|
|
c_string *actor_repr_name,
|
|
c_string *application_error,
|
|
const c_vector[CConcurrencyGroup] &defined_concurrency_groups,
|
|
const c_string name_of_concurrency_group_to_execute,
|
|
c_bool is_reattempt,
|
|
c_bool is_streaming_generator,
|
|
c_bool should_retry_exceptions,
|
|
int64_t generator_backpressure_num_objects,
|
|
int64_t num_objects_per_yield,
|
|
optional[c_string] tensor_transport
|
|
) nogil) task_execution_callback
|
|
(void(const CObjectID &) nogil) free_actor_object_callback
|
|
(void(const CObjectID &, const c_string &) nogil) set_direct_transport_metadata
|
|
(function[void()]() nogil) initialize_thread_callback
|
|
(CRayStatus() nogil) check_signals
|
|
(void() nogil) gc_collect
|
|
(c_vector[c_string](
|
|
const c_vector[CObjectReference] &) nogil) spill_objects
|
|
(int64_t(
|
|
const c_vector[CObjectReference] &,
|
|
const c_vector[c_string] &) nogil) restore_spilled_objects
|
|
(void(
|
|
const c_vector[c_string]&,
|
|
CWorkerType) nogil) delete_spilled_objects
|
|
(void(
|
|
const c_string&,
|
|
const c_vector[c_string]&) nogil) run_on_util_worker_handler
|
|
(void(const CRayObject&) nogil) unhandled_exception_handler
|
|
(c_bool(const CTaskID &c_task_id) nogil) cancel_async_actor_task
|
|
(void() noexcept nogil) actor_shutdown_callback
|
|
(void(c_string *stack_out) nogil) get_lang_stack
|
|
int num_workers
|
|
(c_bool(const CTaskID &) nogil) kill_main
|
|
CCoreWorkerOptions()
|
|
c_string serialized_job_config
|
|
int metrics_agent_port
|
|
int runtime_env_hash
|
|
CWorkerID worker_id
|
|
CClusterID cluster_id
|
|
c_string session_name
|
|
c_string entrypoint
|
|
int64_t worker_launch_time_ms
|
|
int64_t worker_launched_time_ms
|
|
c_string debug_source
|
|
|
|
cdef cppclass CCoreWorkerProcess "ray::core::CoreWorkerProcess":
|
|
@staticmethod
|
|
void Initialize(const CCoreWorkerOptions &options)
|
|
# Only call this in CoreWorker.__cinit__,
|
|
# use CoreWorker.core_worker to access C++ CoreWorker.
|
|
|
|
@staticmethod
|
|
CCoreWorker &GetCoreWorker()
|
|
|
|
@staticmethod
|
|
void Shutdown()
|
|
|
|
@staticmethod
|
|
void RunTaskExecutionLoop()
|