chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,6 @@
|
||||
from libc.stddef cimport size_t
|
||||
from libcpp.string cimport string
|
||||
|
||||
cdef extern from "<array>" namespace "std":
|
||||
cdef cppclass array_string_2 "std::array<std::string, 2>":
|
||||
string& operator[](size_t) except +
|
||||
@@ -0,0 +1,74 @@
|
||||
from cpython cimport Py_buffer, PyBytes_FromStringAndSize
|
||||
from libc.stdint cimport int64_t, uintptr_t
|
||||
from libc.stdio cimport printf
|
||||
from libcpp.memory cimport shared_ptr
|
||||
|
||||
from ray.includes.common cimport CBuffer
|
||||
|
||||
|
||||
cdef class Buffer:
|
||||
"""Cython wrapper class of C++ `ray::Buffer`.
|
||||
|
||||
This class implements the Python 'buffer protocol', which allows
|
||||
us to use it for calls into Python libraries without having to
|
||||
copy the data.
|
||||
|
||||
See https://docs.python.org/3/c-api/buffer.html for details.
|
||||
"""
|
||||
@staticmethod
|
||||
cdef make(const shared_ptr[CBuffer]& buffer):
|
||||
cdef Buffer self = Buffer.__new__(Buffer)
|
||||
self.buffer = buffer
|
||||
self.shape = <Py_ssize_t>self.size
|
||||
self.strides = <Py_ssize_t>(1)
|
||||
return self
|
||||
|
||||
def __len__(self):
|
||||
return self.size
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
"""
|
||||
The buffer size in bytes.
|
||||
"""
|
||||
return self.buffer.get().Size()
|
||||
|
||||
def to_pybytes(self):
|
||||
"""
|
||||
Return this buffer as a Python bytes object. Memory is copied.
|
||||
"""
|
||||
return PyBytes_FromStringAndSize(
|
||||
<const char*>self.buffer.get().Data(),
|
||||
self.buffer.get().Size())
|
||||
|
||||
def __getbuffer__(self, Py_buffer* buffer, int flags):
|
||||
buffer.readonly = 0
|
||||
buffer.buf = <char *>self.buffer.get().Data()
|
||||
buffer.format = 'B'
|
||||
buffer.internal = NULL
|
||||
buffer.itemsize = 1
|
||||
buffer.len = self.size
|
||||
buffer.ndim = 1
|
||||
buffer.obj = self
|
||||
buffer.shape = &self.shape
|
||||
buffer.strides = &self.strides
|
||||
buffer.suboffsets = NULL
|
||||
|
||||
def __getsegcount__(self, Py_ssize_t *len_out):
|
||||
if len_out != NULL:
|
||||
len_out[0] = <Py_ssize_t>self.size
|
||||
return 1
|
||||
|
||||
def __getreadbuffer__(self, Py_ssize_t idx, void **p):
|
||||
if idx != 0:
|
||||
raise SystemError("accessing non-existent buffer segment")
|
||||
if p != NULL:
|
||||
p[0] = <void*> self.buffer.get().Data()
|
||||
return self.size
|
||||
|
||||
def __getwritebuffer__(self, Py_ssize_t idx, void **p):
|
||||
if idx != 0:
|
||||
raise SystemError("accessing non-existent buffer segment")
|
||||
if p != NULL:
|
||||
p[0] = <void*> self.buffer.get().Data()
|
||||
return self.size
|
||||
@@ -0,0 +1,900 @@
|
||||
from libcpp cimport bool as c_bool
|
||||
from libcpp.memory cimport shared_ptr, unique_ptr
|
||||
from libcpp.string cimport string as c_string
|
||||
|
||||
from libc.stdint cimport uint8_t, int32_t, uint64_t, int64_t, uint32_t
|
||||
from libcpp.unordered_map cimport unordered_map
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
from libcpp.pair cimport pair as c_pair
|
||||
from ray.includes.optional cimport (
|
||||
optional,
|
||||
)
|
||||
from ray.includes.unique_ids cimport (
|
||||
CActorID,
|
||||
CJobID,
|
||||
CClusterID,
|
||||
CWorkerID,
|
||||
CObjectID,
|
||||
CTaskID,
|
||||
CPlacementGroupID,
|
||||
CNodeID,
|
||||
)
|
||||
from ray.includes.function_descriptor cimport (
|
||||
CFunctionDescriptor,
|
||||
)
|
||||
|
||||
|
||||
cdef extern from * namespace "polyfill" nogil:
|
||||
"""
|
||||
namespace polyfill {
|
||||
|
||||
template <typename T>
|
||||
inline typename std::remove_reference<T>::type&& move(T& t) {
|
||||
return std::move(t);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline typename std::remove_reference<T>::type&& move(T&& t) {
|
||||
return std::move(t);
|
||||
}
|
||||
|
||||
} // namespace polyfill
|
||||
"""
|
||||
cdef T move[T](T)
|
||||
|
||||
|
||||
cdef extern from "ray/common/status.h" namespace "ray" nogil:
|
||||
cdef enum class CStatusCode "ray::StatusCode":
|
||||
pass
|
||||
c_bool operator==(CStatusCode lhs, CStatusCode rhs)
|
||||
|
||||
cdef cppclass CRayStatus "ray::Status":
|
||||
CRayStatus()
|
||||
CRayStatus(CStatusCode code, const c_string &msg)
|
||||
CRayStatus(CStatusCode code, const c_string &msg, int rpc_code)
|
||||
CRayStatus(const CRayStatus &s)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus OK()
|
||||
|
||||
@staticmethod
|
||||
CRayStatus OutOfMemory(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus KeyError(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus Invalid(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus IOError(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus TypeError(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus UnknownError(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus NotImplemented(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus ObjectStoreFull(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus RedisError(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus TimedOut(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus InvalidArgument(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus Interrupted(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus IntentionalSystemExit(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus UnexpectedSystemExit(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus CreationTaskError(const c_string &msg)
|
||||
|
||||
@staticmethod
|
||||
CRayStatus NotFound()
|
||||
|
||||
@staticmethod
|
||||
CRayStatus ObjectRefEndOfStream()
|
||||
|
||||
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 IsObjectStoreFull()
|
||||
c_bool IsAlreadyExists()
|
||||
c_bool IsOutOfDisk()
|
||||
c_bool IsRedisError()
|
||||
c_bool IsTimedOut()
|
||||
c_bool IsInvalidArgument()
|
||||
c_bool IsInterrupted()
|
||||
c_bool IsObjectNotFound()
|
||||
c_bool IsNotFound()
|
||||
c_bool IsObjectUnknownOwner()
|
||||
c_bool IsRpcError()
|
||||
c_bool IsOutOfResource()
|
||||
c_bool IsObjectRefEndOfStream()
|
||||
c_bool IsIntentionalSystemExit()
|
||||
c_bool IsUnexpectedSystemExit()
|
||||
c_bool IsChannelError()
|
||||
c_bool IsChannelTimeoutError()
|
||||
c_bool IsUnauthenticated()
|
||||
|
||||
c_string ToString()
|
||||
c_string CodeAsString()
|
||||
CStatusCode code()
|
||||
c_string message()
|
||||
int rpc_code()
|
||||
|
||||
# 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 CRayStatus RayStatus_NotImplemented "Status::NotImplemented"()
|
||||
|
||||
cdef extern from "ray/common/status_or.h" namespace "ray" nogil:
|
||||
cdef cppclass CStatusOr "ray::StatusOr"[T]:
|
||||
c_bool ok()
|
||||
const CRayStatus &status() const
|
||||
T &value()
|
||||
|
||||
cdef extern from "ray/common/status.h" namespace "ray::StatusT" nogil:
|
||||
cdef cppclass CStatusTIOError "ray::StatusT::IOError":
|
||||
CStatusTIOError(const c_string &msg)
|
||||
c_string message() const
|
||||
|
||||
cdef cppclass CStatusTTimedOut "ray::StatusT::TimedOut":
|
||||
CStatusTTimedOut(const c_string &msg)
|
||||
c_string message() const
|
||||
|
||||
cdef cppclass CStatusTInvalid "ray::StatusT::Invalid":
|
||||
CStatusTInvalid(const c_string &msg)
|
||||
c_string message() const
|
||||
|
||||
cdef extern from "ray/common/status.h" namespace "ray" nogil:
|
||||
cdef cppclass CWaitForPersistedPortResult "ray::StatusSetOr<int, ray::StatusT::IOError, ray::StatusT::TimedOut, ray::StatusT::Invalid>":
|
||||
c_bool has_value()
|
||||
c_bool has_error()
|
||||
int &value()
|
||||
c_string message()
|
||||
|
||||
cdef extern from "ray/util/port_persistence.h" namespace "ray" nogil:
|
||||
c_string GetPortFileName "ray::GetPortFileName"(
|
||||
const CNodeID &node_id,
|
||||
const c_string &port_name)
|
||||
CRayStatus PersistPort "ray::PersistPort"(
|
||||
const c_string &dir,
|
||||
const CNodeID &node_id,
|
||||
const c_string &port_name,
|
||||
int port)
|
||||
CWaitForPersistedPortResult WaitForPersistedPort "ray::WaitForPersistedPort"(
|
||||
const c_string &dir,
|
||||
const CNodeID &node_id,
|
||||
const c_string &port_name,
|
||||
int timeout_ms,
|
||||
int poll_interval_ms)
|
||||
|
||||
cdef extern from "ray/common/id.h" namespace "ray" nogil:
|
||||
const CTaskID GenerateTaskId(const CJobID &job_id,
|
||||
const CTaskID &parent_task_id,
|
||||
int parent_task_counter)
|
||||
|
||||
|
||||
cdef extern from "src/ray/protobuf/common.pb.h" nogil:
|
||||
cdef cppclass CLanguage "ray::Language":
|
||||
pass
|
||||
cdef cppclass CWorkerType "ray::core::WorkerType":
|
||||
pass
|
||||
cdef cppclass CWorkerExitType "ray::rpc::WorkerExitType":
|
||||
pass
|
||||
cdef cppclass CTaskType "ray::TaskType":
|
||||
pass
|
||||
cdef cppclass CPlacementStrategy "ray::core::PlacementStrategy":
|
||||
pass
|
||||
cdef cppclass CDefaultSchedulingStrategy "ray::rpc::DefaultSchedulingStrategy": # noqa: E501
|
||||
CDefaultSchedulingStrategy()
|
||||
cdef cppclass CSpreadSchedulingStrategy "ray::rpc::SpreadSchedulingStrategy": # noqa: E501
|
||||
CSpreadSchedulingStrategy()
|
||||
cdef cppclass CPlacementGroupSchedulingStrategy "ray::rpc::PlacementGroupSchedulingStrategy": # noqa: E501
|
||||
CPlacementGroupSchedulingStrategy()
|
||||
void set_placement_group_id(const c_string& placement_group_id)
|
||||
void set_placement_group_bundle_index(int64_t placement_group_bundle_index) # noqa: E501
|
||||
void set_placement_group_capture_child_tasks(c_bool placement_group_capture_child_tasks) # noqa: E501
|
||||
cdef cppclass CNodeAffinitySchedulingStrategy "ray::rpc::NodeAffinitySchedulingStrategy": # noqa: E501
|
||||
CNodeAffinitySchedulingStrategy()
|
||||
void set_node_id(const c_string& node_id)
|
||||
void set_soft(c_bool soft)
|
||||
void set_spill_on_unavailable(c_bool spill_on_unavailable)
|
||||
void set_fail_on_unavailable(c_bool fail_on_unavailable)
|
||||
cdef cppclass CSchedulingStrategy "ray::rpc::SchedulingStrategy":
|
||||
CSchedulingStrategy()
|
||||
void clear_scheduling_strategy()
|
||||
CSpreadSchedulingStrategy* mutable_spread_scheduling_strategy()
|
||||
CDefaultSchedulingStrategy* mutable_default_scheduling_strategy()
|
||||
CPlacementGroupSchedulingStrategy* mutable_placement_group_scheduling_strategy() # noqa: E501
|
||||
CNodeAffinitySchedulingStrategy* mutable_node_affinity_scheduling_strategy()
|
||||
CNodeLabelSchedulingStrategy* mutable_node_label_scheduling_strategy()
|
||||
cdef cppclass CAddress "ray::rpc::Address":
|
||||
CAddress()
|
||||
const c_string &SerializeAsString() const
|
||||
void ParseFromString(const c_string &serialized)
|
||||
void CopyFrom(const CAddress& address)
|
||||
const c_string &worker_id()
|
||||
cdef cppclass CObjectReference "ray::rpc::ObjectReference":
|
||||
CObjectReference()
|
||||
CAddress owner_address() const
|
||||
const c_string &object_id() const
|
||||
const c_string &call_site() const
|
||||
c_bool has_tensor_transport() const
|
||||
const c_string &tensor_transport() const
|
||||
cdef cppclass CNodeLabelSchedulingStrategy "ray::rpc::NodeLabelSchedulingStrategy": # noqa: E501
|
||||
CNodeLabelSchedulingStrategy()
|
||||
CLabelMatchExpressions* mutable_hard()
|
||||
CLabelMatchExpressions* mutable_soft()
|
||||
cdef cppclass CLabelMatchExpressions "ray::rpc::LabelMatchExpressions": # noqa: E501
|
||||
CLabelMatchExpressions()
|
||||
CLabelMatchExpression* add_expressions()
|
||||
cdef cppclass CLabelMatchExpression "ray::rpc::LabelMatchExpression": # noqa: E501
|
||||
CLabelMatchExpression()
|
||||
void set_key(const c_string &key)
|
||||
CLabelOperator* mutable_operator_()
|
||||
cdef cppclass CLabelIn "ray::rpc::LabelIn": # noqa: E501
|
||||
CLabelIn()
|
||||
void add_values(const c_string &value)
|
||||
cdef cppclass CLabelNotIn "ray::rpc::LabelNotIn": # noqa: E501
|
||||
CLabelNotIn()
|
||||
void add_values(const c_string &value)
|
||||
cdef cppclass CLabelExists "ray::rpc::LabelExists": # noqa: E501
|
||||
CLabelExists()
|
||||
cdef cppclass CLabelDoesNotExist "ray::rpc::LabelDoesNotExist": # noqa: E501
|
||||
CLabelDoesNotExist()
|
||||
cdef cppclass CLabelNotIn "ray::rpc::LabelNotIn": # noqa: E501
|
||||
CLabelNotIn()
|
||||
void add_values(const c_string &value)
|
||||
cdef cppclass CLabelOperator "ray::rpc::LabelOperator": # noqa: E501
|
||||
CLabelOperator()
|
||||
CLabelIn* mutable_label_in()
|
||||
CLabelNotIn* mutable_label_not_in()
|
||||
CLabelExists* mutable_label_exists()
|
||||
CLabelDoesNotExist* mutable_label_does_not_exist()
|
||||
cdef cppclass CLineageReconstructionTask "ray::rpc::LineageReconstructionTask":
|
||||
CLineageReconstructionTask()
|
||||
const c_string &SerializeAsString() const
|
||||
|
||||
cdef extern from "ray/common/scheduling/cluster_resource_data.h" namespace "ray" nogil:
|
||||
cdef cppclass CNodeResources "ray::NodeResources":
|
||||
CNodeResources()
|
||||
unordered_map[c_string, c_string] labels
|
||||
c_bool HasRequiredLabels(const CLabelSelector &label_selector) const
|
||||
|
||||
void SetNodeResourcesLabels(CNodeResources& resources, const unordered_map[c_string, c_string]& labels)
|
||||
|
||||
cdef extern from "ray/common/scheduling/label_selector.h" namespace "ray":
|
||||
cdef cppclass CLabelSelector "ray::LabelSelector":
|
||||
CLabelSelector() nogil except +
|
||||
void AddConstraint(const c_string& key, const c_string& value) nogil except +
|
||||
|
||||
cdef extern from "ray/common/scheduling/fallback_strategy.h" namespace "ray":
|
||||
cdef cppclass CFallbackOption "ray::FallbackOption":
|
||||
CLabelSelector label_selector
|
||||
|
||||
CFallbackOption() nogil except +
|
||||
CFallbackOption(CLabelSelector) nogil except +
|
||||
|
||||
# This is a workaround for C++ enum class since Cython has no corresponding
|
||||
# representation.
|
||||
cdef extern from "src/ray/protobuf/common.pb.h" nogil:
|
||||
cdef CLanguage LANGUAGE_PYTHON "ray::Language::PYTHON"
|
||||
cdef CLanguage LANGUAGE_CPP "ray::Language::CPP"
|
||||
cdef CLanguage LANGUAGE_JAVA "ray::Language::JAVA"
|
||||
|
||||
cdef extern from "src/ray/protobuf/common.pb.h" nogil:
|
||||
cdef CWorkerType WORKER_TYPE_WORKER "ray::core::WorkerType::WORKER"
|
||||
cdef CWorkerType WORKER_TYPE_DRIVER "ray::core::WorkerType::DRIVER"
|
||||
cdef CWorkerType WORKER_TYPE_SPILL_WORKER "ray::core::WorkerType::SPILL_WORKER" # noqa: E501
|
||||
cdef CWorkerType WORKER_TYPE_RESTORE_WORKER "ray::core::WorkerType::RESTORE_WORKER" # noqa: E501
|
||||
cdef CWorkerType WORKER_TYPE_UTIL_WORKER "ray::core::WorkerType::UTIL_WORKER" # noqa: E501
|
||||
cdef CWorkerExitType WORKER_EXIT_TYPE_USER_ERROR "ray::rpc::WorkerExitType::USER_ERROR" # noqa: E501
|
||||
cdef CWorkerExitType WORKER_EXIT_TYPE_SYSTEM_ERROR "ray::rpc::WorkerExitType::SYSTEM_ERROR" # noqa: E501
|
||||
cdef CWorkerExitType WORKER_EXIT_TYPE_INTENTIONAL_SYSTEM_ERROR "ray::rpc::WorkerExitType::INTENDED_SYSTEM_EXIT" # noqa: E501
|
||||
|
||||
cdef extern from "src/ray/protobuf/common.pb.h" nogil:
|
||||
cdef CTaskType TASK_TYPE_NORMAL_TASK "ray::TaskType::NORMAL_TASK"
|
||||
cdef CTaskType TASK_TYPE_ACTOR_CREATION_TASK "ray::TaskType::ACTOR_CREATION_TASK" # noqa: E501
|
||||
cdef CTaskType TASK_TYPE_ACTOR_TASK "ray::TaskType::ACTOR_TASK"
|
||||
|
||||
cdef extern from "src/ray/protobuf/common.pb.h" nogil:
|
||||
cdef CPlacementStrategy PLACEMENT_STRATEGY_PACK \
|
||||
"ray::core::PlacementStrategy::PACK"
|
||||
cdef CPlacementStrategy PLACEMENT_STRATEGY_SPREAD \
|
||||
"ray::core::PlacementStrategy::SPREAD"
|
||||
cdef CPlacementStrategy PLACEMENT_STRATEGY_STRICT_PACK \
|
||||
"ray::core::PlacementStrategy::STRICT_PACK"
|
||||
cdef CPlacementStrategy PLACEMENT_STRATEGY_STRICT_SPREAD \
|
||||
"ray::core::PlacementStrategy::STRICT_SPREAD"
|
||||
|
||||
cdef extern from "ray/common/buffer.h" namespace "ray" nogil:
|
||||
cdef cppclass CBuffer "ray::Buffer":
|
||||
uint8_t *Data() const
|
||||
size_t Size() const
|
||||
c_bool IsPlasmaBuffer() const
|
||||
|
||||
cdef cppclass LocalMemoryBuffer(CBuffer):
|
||||
LocalMemoryBuffer(uint8_t *data, size_t size, c_bool copy_data)
|
||||
LocalMemoryBuffer(size_t size)
|
||||
|
||||
cdef cppclass SharedMemoryBuffer(CBuffer):
|
||||
SharedMemoryBuffer(
|
||||
const shared_ptr[CBuffer] &buffer,
|
||||
int64_t offset,
|
||||
int64_t size)
|
||||
c_bool IsPlasmaBuffer() const
|
||||
|
||||
cdef extern from "ray/common/ray_object.h" nogil:
|
||||
cdef cppclass CRayObject "ray::RayObject":
|
||||
CRayObject(const shared_ptr[CBuffer] &data,
|
||||
const shared_ptr[CBuffer] &metadata,
|
||||
const c_vector[CObjectReference] &nested_refs)
|
||||
c_bool HasData() const
|
||||
c_bool HasMetadata() const
|
||||
const size_t DataSize() const
|
||||
const shared_ptr[CBuffer] &GetData()
|
||||
const shared_ptr[CBuffer] &GetMetadata() const
|
||||
c_bool IsInPlasmaError() const
|
||||
optional[c_string] GetTensorTransport() const
|
||||
void SetDirectTransportMetadata(c_string direct_transport_metadata)
|
||||
|
||||
cdef extern from "ray/core_worker/common.h" nogil:
|
||||
cdef cppclass CRayFunction "ray::core::RayFunction":
|
||||
CRayFunction()
|
||||
CRayFunction(CLanguage language,
|
||||
const CFunctionDescriptor &function_descriptor)
|
||||
CLanguage GetLanguage()
|
||||
const CFunctionDescriptor GetFunctionDescriptor()
|
||||
|
||||
cdef cppclass CTaskArg "ray::TaskArg":
|
||||
pass
|
||||
|
||||
cdef cppclass CTaskArgByReference "ray::TaskArgByReference":
|
||||
CTaskArgByReference(const CObjectID &object_id,
|
||||
const CAddress &owner_address,
|
||||
const c_string &call_site,
|
||||
optional[c_string] tensor_transport)
|
||||
|
||||
cdef cppclass CTaskArgByValue "ray::TaskArgByValue":
|
||||
CTaskArgByValue(const shared_ptr[CRayObject] &data)
|
||||
|
||||
cdef cppclass CTaskOptions "ray::core::TaskOptions":
|
||||
CTaskOptions()
|
||||
CTaskOptions(c_string name, int num_returns,
|
||||
unordered_map[c_string, double] &resources,
|
||||
c_string concurrency_group_name,
|
||||
int64_t generator_backpressure_num_objects,
|
||||
int64_t num_objects_per_yield)
|
||||
CTaskOptions(c_string name, int num_returns,
|
||||
unordered_map[c_string, double] &resources,
|
||||
c_string concurrency_group_name,
|
||||
int64_t generator_backpressure_num_objects,
|
||||
int64_t num_objects_per_yield,
|
||||
c_string serialized_runtime_env)
|
||||
CTaskOptions(c_string name, int num_returns,
|
||||
unordered_map[c_string, double] &resources,
|
||||
c_string concurrency_group_name,
|
||||
int64_t generator_backpressure_num_objects,
|
||||
int64_t num_objects_per_yield,
|
||||
c_string serialized_runtime_env,
|
||||
c_bool enable_task_events,
|
||||
const unordered_map[c_string, c_string] &labels,
|
||||
CLabelSelector label_selector,
|
||||
optional[c_string] tensor_transport,
|
||||
c_vector[CFallbackOption] fallback_strategy)
|
||||
|
||||
cdef cppclass CActorCreationOptions "ray::core::ActorCreationOptions":
|
||||
CActorCreationOptions()
|
||||
CActorCreationOptions(
|
||||
int64_t max_restarts,
|
||||
int64_t max_task_retries,
|
||||
int32_t max_concurrency,
|
||||
const unordered_map[c_string, double] &resources,
|
||||
const unordered_map[c_string, double] &placement_resources,
|
||||
const c_vector[c_string] &dynamic_worker_options,
|
||||
optional[c_bool] is_detached, c_string &name, c_string &ray_namespace,
|
||||
c_bool is_asyncio,
|
||||
const CSchedulingStrategy &scheduling_strategy,
|
||||
c_string serialized_runtime_env,
|
||||
const c_vector[CConcurrencyGroup] &concurrency_groups,
|
||||
c_bool allow_out_of_order_execution,
|
||||
int32_t max_pending_calls,
|
||||
c_bool enable_tensor_transport,
|
||||
c_bool enable_task_events,
|
||||
const unordered_map[c_string, c_string] &labels,
|
||||
CLabelSelector label_selector,
|
||||
c_vector[CFallbackOption] fallback_strategy,
|
||||
int64_t actor_generator_backpressure_num_objects)
|
||||
|
||||
cdef cppclass CPlacementGroupCreationOptions \
|
||||
"ray::core::PlacementGroupCreationOptions":
|
||||
CPlacementGroupCreationOptions()
|
||||
CPlacementGroupCreationOptions(
|
||||
const c_string &name,
|
||||
CPlacementStrategy strategy,
|
||||
const c_vector[unordered_map[c_string, double]] &bundles,
|
||||
c_bool is_detached,
|
||||
CNodeID soft_target_node_id,
|
||||
const c_vector[unordered_map[c_string, c_string]] &bundle_label_selector,
|
||||
const unordered_map[c_string, CPlacementStrategy] &topology_strategy,
|
||||
)
|
||||
|
||||
cdef cppclass CObjectLocation "ray::core::ObjectLocation":
|
||||
const int64_t GetObjectSize() const
|
||||
const c_vector[CNodeID] &GetNodeIDs() const
|
||||
c_bool IsSpilled() const
|
||||
const c_string &GetSpilledURL() const
|
||||
const CNodeID &GetSpilledNodeID() const
|
||||
const c_bool GetDidSpill() const
|
||||
|
||||
cdef extern from "ray/common/python_callbacks.h" namespace "ray":
|
||||
cdef cppclass MultiItemPyCallback[T]:
|
||||
MultiItemPyCallback(
|
||||
object (*)(CRayStatus, c_vector[T]) nogil,
|
||||
void (object, object) nogil,
|
||||
object) nogil
|
||||
|
||||
cdef cppclass OptionalItemPyCallback[T]:
|
||||
OptionalItemPyCallback(
|
||||
object (*)(CRayStatus, optional[T]) nogil,
|
||||
void (object, object) nogil,
|
||||
object) nogil
|
||||
|
||||
cdef cppclass StatusPyCallback:
|
||||
StatusPyCallback(
|
||||
object (*)(CRayStatus) nogil,
|
||||
void (object, object) nogil,
|
||||
object) nogil
|
||||
|
||||
cdef extern from "ray/gcs_rpc_client/accessors/actor_info_accessor_interface.h" nogil:
|
||||
cdef cppclass CActorInfoAccessorInterface "ray::gcs::ActorInfoAccessorInterface":
|
||||
void AsyncGetAllByFilter(
|
||||
const optional[CActorID] &actor_id,
|
||||
const optional[CJobID] &job_id,
|
||||
const optional[c_string] &actor_state_name,
|
||||
const MultiItemPyCallback[CActorTableData] &callback,
|
||||
int64_t timeout_ms)
|
||||
|
||||
cdef extern from "ray/gcs_rpc_client/accessor.h" nogil:
|
||||
cdef cppclass CJobInfoAccessor "ray::gcs::JobInfoAccessor":
|
||||
CRayStatus GetAll(
|
||||
const optional[c_string] &job_or_submission_id,
|
||||
c_bool skip_submission_job_info_field,
|
||||
c_bool skip_is_running_tasks_field,
|
||||
c_vector[CJobTableData] &result,
|
||||
int64_t timeout_ms)
|
||||
|
||||
void AsyncGetAll(
|
||||
const optional[c_string] &job_or_submission_id,
|
||||
c_bool skip_submission_job_info_field,
|
||||
c_bool skip_is_running_tasks_field,
|
||||
const MultiItemPyCallback[CJobTableData] &callback,
|
||||
int64_t timeout_ms)
|
||||
|
||||
cdef cppclass CNodeInfoAccessor "ray::gcs::NodeInfoAccessor":
|
||||
CRayStatus CheckAlive(
|
||||
const c_vector[CNodeID] &node_ids,
|
||||
int64_t timeout_ms,
|
||||
c_vector[c_bool] &result)
|
||||
|
||||
void AsyncCheckAlive(
|
||||
const c_vector[CNodeID] &node_ids,
|
||||
int64_t timeout_ms,
|
||||
const MultiItemPyCallback[c_bool] &callback)
|
||||
|
||||
CRayStatus DrainNodes(
|
||||
const c_vector[CNodeID] &node_ids,
|
||||
int64_t timeout_ms,
|
||||
c_vector[c_string] &drained_node_ids)
|
||||
|
||||
CStatusOr[c_vector[CGcsNodeInfo]] GetAllNoCache(
|
||||
int64_t timeout_ms,
|
||||
optional[CGcsNodeState] state_filter,
|
||||
const c_vector[CNodeSelector] &node_selectors)
|
||||
|
||||
void AsyncGetAll(
|
||||
const OptionalItemPyCallback[c_pair[c_vector[CGcsNodeInfo], int64_t]] &callback,
|
||||
int64_t timeout_ms,
|
||||
optional[CGcsNodeState] state_filter,
|
||||
const c_vector[CNodeSelector] &node_selectors,
|
||||
optional[int64_t] limit) const
|
||||
|
||||
cdef cppclass CNodeResourceInfoAccessor "ray::gcs::NodeResourceInfoAccessor":
|
||||
CRayStatus GetAllResourceUsage(
|
||||
int64_t timeout_ms,
|
||||
CGetAllResourceUsageReply &serialized_reply)
|
||||
|
||||
cdef cppclass CInternalKVAccessor "ray::gcs::InternalKVAccessor":
|
||||
CRayStatus Keys(
|
||||
const c_string &ns,
|
||||
const c_string &prefix,
|
||||
int64_t timeout_ms,
|
||||
c_vector[c_string] &value)
|
||||
|
||||
CRayStatus Put(
|
||||
const c_string &ns,
|
||||
const c_string &key,
|
||||
const c_string &value,
|
||||
c_bool overwrite,
|
||||
int64_t timeout_ms,
|
||||
c_bool &added)
|
||||
|
||||
CRayStatus Get(
|
||||
const c_string &ns,
|
||||
const c_string &key,
|
||||
int64_t timeout_ms,
|
||||
c_string &value)
|
||||
|
||||
CRayStatus MultiGet(
|
||||
const c_string &ns,
|
||||
const c_vector[c_string] &keys,
|
||||
int64_t timeout_ms,
|
||||
unordered_map[c_string, c_string] &values)
|
||||
|
||||
CRayStatus Del(
|
||||
const c_string &ns,
|
||||
const c_string &key,
|
||||
c_bool del_by_prefix,
|
||||
int64_t timeout_ms,
|
||||
int& num_deleted)
|
||||
|
||||
CRayStatus Exists(
|
||||
const c_string &ns,
|
||||
const c_string &key,
|
||||
int64_t timeout_ms,
|
||||
c_bool &exists)
|
||||
|
||||
void AsyncInternalKVKeys(
|
||||
const c_string &ns,
|
||||
const c_string &prefix,
|
||||
int64_t timeout_ms,
|
||||
const OptionalItemPyCallback[c_vector[c_string]] &callback)
|
||||
|
||||
void AsyncInternalKVGet(
|
||||
const c_string &ns,
|
||||
const c_string &key,
|
||||
int64_t timeout_ms,
|
||||
const OptionalItemPyCallback[c_string] &callback)
|
||||
|
||||
void AsyncInternalKVMultiGet(
|
||||
const c_string &ns,
|
||||
const c_vector[c_string] &keys,
|
||||
int64_t timeout_ms,
|
||||
const OptionalItemPyCallback[unordered_map[c_string, c_string]] &callback)
|
||||
|
||||
void AsyncInternalKVPut(
|
||||
const c_string &ns,
|
||||
const c_string &key,
|
||||
const c_string &value,
|
||||
c_bool overwrite,
|
||||
int64_t timeout_ms,
|
||||
const OptionalItemPyCallback[c_bool] &callback)
|
||||
|
||||
void AsyncInternalKVExists(
|
||||
const c_string &ns,
|
||||
const c_string &key,
|
||||
int64_t timeout_ms,
|
||||
const OptionalItemPyCallback[c_bool] &callback)
|
||||
|
||||
void AsyncInternalKVDel(
|
||||
const c_string &ns,
|
||||
const c_string &key,
|
||||
c_bool del_by_prefix,
|
||||
int64_t timeout_ms,
|
||||
const OptionalItemPyCallback[int] &callback)
|
||||
|
||||
cdef cppclass CRuntimeEnvAccessor "ray::gcs::RuntimeEnvAccessor":
|
||||
CRayStatus PinRuntimeEnvUri(
|
||||
const c_string &uri,
|
||||
int expiration_s,
|
||||
int64_t timeout_ms)
|
||||
|
||||
cdef cppclass CAutoscalerStateAccessor "ray::gcs::AutoscalerStateAccessor":
|
||||
|
||||
CRayStatus RequestClusterResourceConstraint(
|
||||
int64_t timeout_ms,
|
||||
const c_vector[unordered_map[c_string, double]] &bundles,
|
||||
const c_vector[unordered_map[c_string, c_string]] &label_selectors,
|
||||
const c_vector[int64_t] &count_array,
|
||||
)
|
||||
|
||||
CRayStatus GetClusterResourceState(
|
||||
int64_t timeout_ms,
|
||||
c_string &serialized_reply
|
||||
)
|
||||
|
||||
CRayStatus GetClusterStatus(
|
||||
int64_t timeout_ms,
|
||||
c_string &serialized_reply
|
||||
)
|
||||
|
||||
void AsyncGetClusterStatus(
|
||||
int64_t timeout_ms,
|
||||
const OptionalItemPyCallback[CGetClusterStatusReply] &callback)
|
||||
|
||||
CRayStatus ReportAutoscalingState(
|
||||
int64_t timeout_ms,
|
||||
const c_string &serialized_state
|
||||
)
|
||||
|
||||
CRayStatus ReportClusterConfig(
|
||||
int64_t timeout_ms,
|
||||
const c_string &serialized_cluster_config
|
||||
)
|
||||
|
||||
CRayStatus DrainNode(
|
||||
const c_string &node_id,
|
||||
int32_t reason,
|
||||
const c_string &reason_message,
|
||||
int64_t deadline_timestamp_ms,
|
||||
int64_t timeout_ms,
|
||||
c_bool &is_accepted,
|
||||
c_string &rejection_reason_message
|
||||
)
|
||||
|
||||
CRayStatus ResizeRayletResourceInstances(
|
||||
const c_string &node_id,
|
||||
const unordered_map[c_string, double] &resources,
|
||||
int64_t timeout_ms,
|
||||
unordered_map[c_string, double] &total_resources
|
||||
)
|
||||
|
||||
cdef cppclass CPublisherAccessor "ray::gcs::PublisherAccessor":
|
||||
CRayStatus PublishError(
|
||||
c_string key_id,
|
||||
CErrorTableData data,
|
||||
int64_t timeout_ms)
|
||||
|
||||
CRayStatus PublishLogs(
|
||||
c_string key_id,
|
||||
CLogBatch data,
|
||||
int64_t timeout_ms)
|
||||
|
||||
void AsyncPublishNodeResourceUsage(
|
||||
c_string key_id,
|
||||
c_string node_resource_usage,
|
||||
const StatusPyCallback &callback
|
||||
)
|
||||
|
||||
cdef cppclass CTaskInfoAccessor "ray::gcs::TaskInfoAccessor":
|
||||
void AsyncAddEvents(
|
||||
CAddEventsRequest &&request,
|
||||
const StatusPyCallback &callback,
|
||||
int64_t timeout_ms)
|
||||
|
||||
|
||||
cdef extern from "ray/gcs_rpc_client/gcs_client.h" nogil:
|
||||
cdef enum CGrpcStatusCode "grpc::StatusCode":
|
||||
UNAVAILABLE "grpc::StatusCode::UNAVAILABLE",
|
||||
UNKNOWN "grpc::StatusCode::UNKNOWN",
|
||||
DEADLINE_EXCEEDED "grpc::StatusCode::DEADLINE_EXCEEDED",
|
||||
RESOURCE_EXHAUSTED "grpc::StatusCode::RESOURCE_EXHAUSTED",
|
||||
UNIMPLEMENTED "grpc::StatusCode::UNIMPLEMENTED",
|
||||
|
||||
cdef cppclass CGcsClientOptions "ray::gcs::GcsClientOptions":
|
||||
CGcsClientOptions(
|
||||
c_string gcs_address, int port, CClusterID cluster_id,
|
||||
c_bool allow_cluster_id_nil, c_bool fetch_cluster_id_if_nil)
|
||||
|
||||
cdef cppclass CGcsClient "ray::gcs::GcsClient":
|
||||
CGcsClient(CGcsClientOptions options)
|
||||
|
||||
c_pair[c_string, int] GetGcsServerAddress() const
|
||||
CClusterID GetClusterId() const
|
||||
|
||||
CActorInfoAccessorInterface& Actors()
|
||||
CJobInfoAccessor& Jobs()
|
||||
CInternalKVAccessor& InternalKV()
|
||||
CNodeInfoAccessor& Nodes()
|
||||
CNodeResourceInfoAccessor& NodeResources()
|
||||
CRuntimeEnvAccessor& RuntimeEnvs()
|
||||
CAutoscalerStateAccessor& Autoscaler()
|
||||
CPublisherAccessor& Publisher()
|
||||
CTaskInfoAccessor& Tasks()
|
||||
CGcsRpcClient& GetGcsRpcClient()
|
||||
|
||||
cdef CRayStatus ConnectOnSingletonIoContext(CGcsClient &gcs_client, int timeout_ms)
|
||||
|
||||
cdef extern from "ray/gcs_rpc_client/rpc_client.h" namespace "ray::rpc::events" nogil:
|
||||
cdef cppclass CAddEventsRequest "ray::rpc::events::AddEventsRequest":
|
||||
bint ParseFromString(const c_string &data)
|
||||
cdef cppclass CAddEventsReply "ray::rpc::events::AddEventsReply":
|
||||
pass
|
||||
|
||||
cdef extern from "ray/gcs_rpc_client/rpc_client.h" namespace "ray::rpc" nogil:
|
||||
cdef cppclass CGcsRpcClient "ray::rpc::GcsRpcClient":
|
||||
pass
|
||||
|
||||
cdef extern from "ray/gcs_rpc_client/gcs_client.h" namespace "ray::gcs" nogil:
|
||||
unordered_map[c_string, double] PythonGetResourcesTotal(
|
||||
const CGcsNodeInfo& node_info)
|
||||
|
||||
cdef extern from "ray/pubsub/python_gcs_subscriber.h" nogil:
|
||||
cdef cppclass CPythonGcsSubscriber "ray::pubsub::PythonGcsSubscriber":
|
||||
|
||||
CPythonGcsSubscriber(
|
||||
const c_string& gcs_address, int gcs_port, CChannelType channel_type,
|
||||
const c_string& subscriber_id, const c_string& worker_id)
|
||||
|
||||
CRayStatus Subscribe()
|
||||
|
||||
int64_t last_batch_size()
|
||||
|
||||
CRayStatus PollError(
|
||||
c_string* key_id, int64_t timeout_ms, CErrorTableData* data)
|
||||
|
||||
CRayStatus PollLogs(
|
||||
c_string* key_id, int64_t timeout_ms, CLogBatch* data)
|
||||
|
||||
CRayStatus Close()
|
||||
|
||||
cdef extern from "ray/pubsub/python_gcs_subscriber.h" namespace "ray::pubsub" nogil:
|
||||
c_vector[c_string] PythonGetLogBatchLines(CLogBatch log_batch)
|
||||
|
||||
cdef extern from "ray/gcs_rpc_client/gcs_client.h" namespace "ray::gcs" nogil:
|
||||
unordered_map[c_string, c_string] PythonGetNodeLabels(
|
||||
const CGcsNodeInfo& node_info)
|
||||
|
||||
cdef extern from "src/ray/protobuf/gcs.pb.h" nogil:
|
||||
cdef enum CChannelType "ray::rpc::ChannelType":
|
||||
RAY_ERROR_INFO_CHANNEL "ray::rpc::ChannelType::RAY_ERROR_INFO_CHANNEL",
|
||||
RAY_LOG_CHANNEL "ray::rpc::ChannelType::RAY_LOG_CHANNEL",
|
||||
GCS_ACTOR_CHANNEL "ray::rpc::ChannelType::GCS_ACTOR_CHANNEL",
|
||||
|
||||
cdef cppclass CJobConfig "ray::rpc::JobConfig":
|
||||
c_string ray_namespace() const
|
||||
const c_string &SerializeAsString() const
|
||||
|
||||
cdef cppclass CNodeDeathInfo "ray::rpc::NodeDeathInfo":
|
||||
int reason() const
|
||||
c_string reason_message() const
|
||||
|
||||
cdef cppclass CGcsNodeInfo "ray::rpc::GcsNodeInfo":
|
||||
c_string node_id() const
|
||||
c_string node_name() const
|
||||
int state() const
|
||||
c_string node_manager_address() const
|
||||
c_string node_manager_hostname() const
|
||||
int node_manager_port() const
|
||||
int object_manager_port() const
|
||||
c_string object_store_socket_name() const
|
||||
c_string raylet_socket_name() const
|
||||
int metrics_export_port() const
|
||||
int metrics_agent_port() const
|
||||
int dashboard_agent_listen_port() const
|
||||
int runtime_env_agent_port() const
|
||||
CNodeDeathInfo death_info() const
|
||||
void ParseFromString(const c_string &serialized)
|
||||
const c_string& SerializeAsString() const
|
||||
|
||||
cdef enum CGcsNodeState "ray::rpc::GcsNodeInfo_GcsNodeState":
|
||||
ALIVE "ray::rpc::GcsNodeInfo_GcsNodeState_ALIVE",
|
||||
|
||||
cdef cppclass CNodeSelector "ray::rpc::GetAllNodeInfoRequest::NodeSelector":
|
||||
CNodeSelector()
|
||||
void set_node_id(const c_string &node_id)
|
||||
void set_node_name(const c_string &node_name)
|
||||
void set_node_ip_address(const c_string &node_ip_address)
|
||||
void set_is_head_node(c_bool is_head_node)
|
||||
void ParseFromString(const c_string &serialized)
|
||||
|
||||
cdef cppclass CJobTableData "ray::rpc::JobTableData":
|
||||
c_string job_id() const
|
||||
c_bool is_dead() const
|
||||
CJobConfig config() const
|
||||
const c_string &SerializeAsString() const
|
||||
|
||||
cdef cppclass CGetAllResourceUsageReply "ray::rpc::GetAllResourceUsageReply":
|
||||
const c_string& SerializeAsString() const
|
||||
|
||||
cdef cppclass CPythonFunction "ray::rpc::PythonFunction":
|
||||
void set_key(const c_string &key)
|
||||
c_string key() const
|
||||
|
||||
cdef cppclass CErrorTableData "ray::rpc::ErrorTableData":
|
||||
c_string job_id() const
|
||||
c_string type() const
|
||||
c_string error_message() const
|
||||
double timestamp() const
|
||||
|
||||
void set_job_id(const c_string &job_id)
|
||||
void set_type(const c_string &type)
|
||||
void set_error_message(const c_string &error_message)
|
||||
void set_timestamp(double timestamp)
|
||||
|
||||
cdef cppclass CLogBatch "ray::rpc::LogBatch":
|
||||
c_string ip() const
|
||||
c_string pid() const
|
||||
c_string job_id() const
|
||||
c_bool is_error() const
|
||||
c_string actor_name() const
|
||||
c_string task_name() const
|
||||
|
||||
void set_ip(const c_string &ip)
|
||||
void set_pid(const c_string &pid)
|
||||
void set_job_id(const c_string &job_id)
|
||||
void set_is_error(c_bool is_error)
|
||||
void add_lines(const c_string &line)
|
||||
void set_actor_name(const c_string &actor_name)
|
||||
void set_task_name(const c_string &task_name)
|
||||
|
||||
cdef cppclass CActorTableData "ray::rpc::ActorTableData":
|
||||
CAddress address() const
|
||||
void ParseFromString(const c_string &serialized)
|
||||
const c_string &SerializeAsString() const
|
||||
|
||||
cdef extern from "src/ray/protobuf/autoscaler.pb.h" nogil:
|
||||
cdef cppclass CGetClusterStatusReply "ray::rpc::autoscaler::GetClusterStatusReply":
|
||||
c_string serialized_cluster_status() const
|
||||
void ParseFromString(const c_string &serialized)
|
||||
const c_string &SerializeAsString() const
|
||||
|
||||
cdef extern from "ray/raylet_rpc_client/raylet_client_with_io_context.h" nogil:
|
||||
cdef cppclass CRayletClientWithIoContext "ray::rpc::RayletClientWithIoContext":
|
||||
CRayletClientWithIoContext(const c_string &ip_address, int port)
|
||||
CRayStatus GetWorkerPIDs(const OptionalItemPyCallback[c_vector[int32_t]] &callback,
|
||||
int64_t timeout_ms)
|
||||
CRayStatus GetAgentPIDs(const OptionalItemPyCallback[c_vector[int32_t]] &callback,
|
||||
int64_t timeout_ms)
|
||||
|
||||
cdef extern from "ray/common/task/task_spec.h" nogil:
|
||||
cdef cppclass CConcurrencyGroup "ray::ConcurrencyGroup":
|
||||
CConcurrencyGroup(
|
||||
c_string name,
|
||||
uint32_t max_concurrency,
|
||||
c_vector[CFunctionDescriptor] c_fds)
|
||||
CConcurrencyGroup()
|
||||
c_string GetName() const
|
||||
uint32_t GetMaxConcurrency() const
|
||||
c_vector[CFunctionDescriptor] GetFunctionDescriptors() const
|
||||
|
||||
cdef extern from "ray/common/constants.h" nogil:
|
||||
cdef const char[] kWorkerSetupHookKeyName
|
||||
cdef int kResourceUnitScaling
|
||||
cdef const char[] kImplicitResourcePrefix
|
||||
cdef int kStreamingGeneratorReturn
|
||||
cdef const char[] kGcsAutoscalerStateNamespace
|
||||
cdef const char[] kGcsAutoscalerV2EnabledKey
|
||||
cdef const char[] kGcsAutoscalerClusterConfigKey
|
||||
cdef const char[] kGcsPidKey
|
||||
cdef const char[] kNodeTypeNameEnv
|
||||
cdef const char[] kNodeMarketTypeEnv
|
||||
cdef const char[] kNodeRegionEnv
|
||||
cdef const char[] kNodeZoneEnv
|
||||
cdef const char[] kLabelKeyNodeID
|
||||
cdef const char[] kLabelKeyNodeAcceleratorType
|
||||
cdef const char[] kLabelKeyNodeMarketType
|
||||
cdef const char[] kLabelKeyNodeRegion
|
||||
cdef const char[] kLabelKeyNodeZone
|
||||
cdef const char[] kLabelKeyNodeGroup
|
||||
cdef const char[] kLabelKeyTpuTopology
|
||||
# Port names for local port discovery
|
||||
cdef const char[] kRuntimeEnvAgentPortName
|
||||
cdef const char[] kMetricsAgentPortName
|
||||
cdef const char[] kMetricsExportPortName
|
||||
cdef const char[] kDashboardAgentListenPortName
|
||||
cdef const char[] kGcsServerPortName
|
||||
cdef const char[] kLabelKeyTpuSliceName
|
||||
cdef const char[] kLabelKeyTpuWorkerId
|
||||
cdef const char[] kLabelKeyTpuPodType
|
||||
cdef const char[] kRayInternalNamespacePrefix
|
||||
@@ -0,0 +1,216 @@
|
||||
from libcpp cimport bool as c_bool
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
|
||||
from ray.includes.common cimport (
|
||||
CObjectLocation,
|
||||
CGcsClientOptions,
|
||||
CPythonGcsSubscriber,
|
||||
kWorkerSetupHookKeyName,
|
||||
kResourceUnitScaling,
|
||||
kImplicitResourcePrefix,
|
||||
kStreamingGeneratorReturn,
|
||||
kGcsAutoscalerStateNamespace,
|
||||
kGcsAutoscalerV2EnabledKey,
|
||||
kGcsAutoscalerClusterConfigKey,
|
||||
kGcsPidKey,
|
||||
kNodeTypeNameEnv,
|
||||
kNodeMarketTypeEnv,
|
||||
kNodeRegionEnv,
|
||||
kNodeZoneEnv,
|
||||
kLabelKeyNodeAcceleratorType,
|
||||
kLabelKeyNodeMarketType,
|
||||
kLabelKeyNodeRegion,
|
||||
kLabelKeyNodeZone,
|
||||
kLabelKeyNodeID,
|
||||
kLabelKeyNodeGroup,
|
||||
kLabelKeyTpuTopology,
|
||||
kLabelKeyTpuSliceName,
|
||||
kLabelKeyTpuWorkerId,
|
||||
kLabelKeyTpuPodType,
|
||||
kRayInternalNamespacePrefix,
|
||||
kRuntimeEnvAgentPortName,
|
||||
kMetricsAgentPortName,
|
||||
kMetricsExportPortName,
|
||||
kDashboardAgentListenPortName,
|
||||
kGcsServerPortName,
|
||||
)
|
||||
|
||||
from ray.exceptions import (
|
||||
RayActorError,
|
||||
ActorAlreadyExistsError,
|
||||
ActorDiedError,
|
||||
RayError,
|
||||
RaySystemError,
|
||||
AuthenticationError,
|
||||
RayTaskError,
|
||||
ObjectStoreFullError,
|
||||
OutOfDiskError,
|
||||
GetTimeoutError,
|
||||
TaskCancelledError,
|
||||
AsyncioActorExit,
|
||||
PendingCallsLimitExceeded,
|
||||
RpcError,
|
||||
ObjectRefStreamEndOfStreamError,
|
||||
)
|
||||
|
||||
|
||||
cdef class GcsClientOptions:
|
||||
"""Cython wrapper class of C++ `ray::gcs::GcsClientOptions`."""
|
||||
cdef:
|
||||
unique_ptr[CGcsClientOptions] inner
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls, gcs_address, cluster_id_hex, allow_cluster_id_nil, fetch_cluster_id_if_nil
|
||||
):
|
||||
"""
|
||||
Creates a GcsClientOption with a maybe-Nil cluster_id, and may fetch from GCS.
|
||||
"""
|
||||
cdef CClusterID c_cluster_id = CClusterID.Nil()
|
||||
if cluster_id_hex:
|
||||
c_cluster_id = CClusterID.FromHex(cluster_id_hex)
|
||||
self = GcsClientOptions()
|
||||
try:
|
||||
ip, port_str = parse_address(gcs_address)
|
||||
port = int(port_str)
|
||||
self.inner.reset(
|
||||
new CGcsClientOptions(
|
||||
ip, port, c_cluster_id, allow_cluster_id_nil, allow_cluster_id_nil))
|
||||
except Exception:
|
||||
raise ValueError(f"Invalid gcs_address: {gcs_address}")
|
||||
return self
|
||||
|
||||
cdef CGcsClientOptions* native(self):
|
||||
return <CGcsClientOptions*>(self.inner.get())
|
||||
|
||||
cdef int check_status(const CRayStatus& status) except -1 nogil:
|
||||
if status.ok():
|
||||
return 0
|
||||
|
||||
with gil:
|
||||
message = status.message().decode()
|
||||
|
||||
if status.IsObjectStoreFull():
|
||||
raise ObjectStoreFullError(message)
|
||||
elif status.IsInvalidArgument():
|
||||
raise ValueError(message)
|
||||
elif status.IsAlreadyExists():
|
||||
raise ActorAlreadyExistsError(message)
|
||||
elif status.IsOutOfDisk():
|
||||
raise OutOfDiskError(message)
|
||||
elif status.IsObjectRefEndOfStream():
|
||||
raise ObjectRefStreamEndOfStreamError(message)
|
||||
elif status.IsInterrupted():
|
||||
raise KeyboardInterrupt()
|
||||
elif status.IsTimedOut():
|
||||
raise GetTimeoutError(message)
|
||||
elif status.IsNotFound():
|
||||
raise ValueError(message)
|
||||
elif status.IsObjectNotFound():
|
||||
raise ValueError(message)
|
||||
elif status.IsObjectUnknownOwner():
|
||||
raise ValueError(message)
|
||||
elif status.IsIOError():
|
||||
raise IOError(message)
|
||||
elif status.IsUnauthenticated():
|
||||
raise AuthenticationError(message)
|
||||
elif status.IsRpcError():
|
||||
raise RpcError(message, rpc_code=status.rpc_code())
|
||||
elif status.IsIntentionalSystemExit():
|
||||
with gil:
|
||||
raise_sys_exit_with_custom_error_message(message)
|
||||
elif status.IsUnexpectedSystemExit():
|
||||
with gil:
|
||||
raise_sys_exit_with_custom_error_message(
|
||||
message, exit_code=1)
|
||||
elif status.IsChannelError():
|
||||
raise RayChannelError(message)
|
||||
elif status.IsChannelTimeoutError():
|
||||
raise RayChannelTimeoutError(message)
|
||||
else:
|
||||
raise RaySystemError(message)
|
||||
|
||||
cdef int check_status_timeout_as_rpc_error(const CRayStatus& status) except -1 nogil:
|
||||
"""
|
||||
Same as check_status, except that it raises RpcError for timeout. This is for
|
||||
backward compatibility: on timeout, `ray.get` raises GetTimeoutError, while
|
||||
GcsClient methods raise RpcError. So in the binding, `get_objects` use check_status
|
||||
and GcsClient methods use check_status_timeout_as_rpc_error.
|
||||
"""
|
||||
if status.IsTimedOut():
|
||||
raise RpcError(status.message().decode(),
|
||||
rpc_code=CGrpcStatusCode.DEADLINE_EXCEEDED)
|
||||
return check_status(status)
|
||||
|
||||
|
||||
WORKER_PROCESS_SETUP_HOOK_KEY_NAME_GCS = str(kWorkerSetupHookKeyName)
|
||||
RESOURCE_UNIT_SCALING = kResourceUnitScaling
|
||||
IMPLICIT_RESOURCE_PREFIX = kImplicitResourcePrefix.decode()
|
||||
STREAMING_GENERATOR_RETURN = kStreamingGeneratorReturn
|
||||
GCS_AUTOSCALER_STATE_NAMESPACE = kGcsAutoscalerStateNamespace.decode()
|
||||
GCS_AUTOSCALER_V2_ENABLED_KEY = kGcsAutoscalerV2EnabledKey.decode()
|
||||
GCS_AUTOSCALER_CLUSTER_CONFIG_KEY = kGcsAutoscalerClusterConfigKey.decode()
|
||||
GCS_PID_KEY = kGcsPidKey.decode()
|
||||
|
||||
# Ray node label related constants from src/ray/common/constants.h
|
||||
NODE_TYPE_NAME_ENV = kNodeTypeNameEnv.decode()
|
||||
NODE_MARKET_TYPE_ENV = kNodeMarketTypeEnv.decode()
|
||||
NODE_REGION_ENV = kNodeRegionEnv.decode()
|
||||
NODE_ZONE_ENV = kNodeZoneEnv.decode()
|
||||
|
||||
RAY_NODE_ACCELERATOR_TYPE_KEY = kLabelKeyNodeAcceleratorType.decode()
|
||||
RAY_NODE_MARKET_TYPE_KEY = kLabelKeyNodeMarketType.decode()
|
||||
RAY_NODE_REGION_KEY = kLabelKeyNodeRegion.decode()
|
||||
RAY_NODE_ZONE_KEY = kLabelKeyNodeZone.decode()
|
||||
# Keep this in sync with NODE_ID_LABEL_KEY in ray.util.placement_group. That
|
||||
# module cannot import this exported value because it forms a circular dependency
|
||||
RAY_NODE_ID_KEY = kLabelKeyNodeID.decode()
|
||||
RAY_NODE_GROUP_KEY = kLabelKeyNodeGroup.decode()
|
||||
|
||||
# TPU specifc Ray node label related constants
|
||||
RAY_NODE_TPU_TOPOLOGY_KEY = kLabelKeyTpuTopology.decode()
|
||||
RAY_NODE_TPU_SLICE_NAME_KEY = kLabelKeyTpuSliceName.decode()
|
||||
RAY_NODE_TPU_WORKER_ID_KEY = kLabelKeyTpuWorkerId.decode()
|
||||
RAY_NODE_TPU_POD_TYPE_KEY = kLabelKeyTpuPodType.decode()
|
||||
|
||||
RAY_INTERNAL_NAMESPACE_PREFIX = kRayInternalNamespacePrefix.decode()
|
||||
# Prefix for namespaces which are used internally by ray.
|
||||
# Jobs within these namespaces should be hidden from users
|
||||
|
||||
# Port names for local port discovery
|
||||
RUNTIME_ENV_AGENT_PORT_NAME = kRuntimeEnvAgentPortName.decode()
|
||||
METRICS_AGENT_PORT_NAME = kMetricsAgentPortName.decode()
|
||||
METRICS_EXPORT_PORT_NAME = kMetricsExportPortName.decode()
|
||||
DASHBOARD_AGENT_LISTEN_PORT_NAME = kDashboardAgentListenPortName.decode()
|
||||
GCS_SERVER_PORT_NAME = kGcsServerPortName.decode()
|
||||
# and should not be considered user activity.
|
||||
RAY_INTERNAL_DASHBOARD_NAMESPACE = f"{RAY_INTERNAL_NAMESPACE_PREFIX}dashboard"
|
||||
|
||||
# Util functions for async handling
|
||||
|
||||
cdef incremented_fut():
|
||||
fut = concurrent.futures.Future()
|
||||
cpython.Py_INCREF(fut)
|
||||
return fut
|
||||
|
||||
cdef void assign_and_decrement_fut(result, fut) noexcept with gil:
|
||||
assert isinstance(fut, concurrent.futures.Future)
|
||||
|
||||
assert not fut.done()
|
||||
try:
|
||||
ret, exc = result
|
||||
if exc:
|
||||
fut.set_exception(exc)
|
||||
else:
|
||||
fut.set_result(ret)
|
||||
finally:
|
||||
# We INCREFed it in `incremented_fut` to keep it alive during the async wait,
|
||||
# and we DECREF it here to balance it.
|
||||
cpython.Py_DECREF(fut)
|
||||
|
||||
cdef raise_or_return(tup):
|
||||
ret, exc = tup
|
||||
if exc:
|
||||
raise exc
|
||||
return ret
|
||||
@@ -0,0 +1,30 @@
|
||||
from libc.stdint cimport int64_t
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp.memory cimport unique_ptr
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
|
||||
cdef extern from "ray/observability/ray_event_interface.h" namespace "ray::observability" nogil:
|
||||
cdef cppclass CRayEventInterface "ray::observability::RayEventInterface":
|
||||
pass
|
||||
|
||||
cdef extern from "ray/observability/python_event_interface.h" namespace "ray::observability" nogil:
|
||||
unique_ptr[CRayEventInterface] CreatePythonRayEvent(
|
||||
int source_type,
|
||||
int event_type,
|
||||
int severity,
|
||||
const c_string &entity_id,
|
||||
const c_string &message,
|
||||
const c_string &session_name,
|
||||
const c_string &serialized_event_data,
|
||||
int nested_event_field_number,
|
||||
const c_string &event_id,
|
||||
int64_t timestamp_ns)
|
||||
|
||||
cdef cppclass CPythonEventRecorder "ray::observability::PythonEventRecorder":
|
||||
CPythonEventRecorder(int aggregator_port,
|
||||
const c_string &node_ip,
|
||||
const c_string &node_id_hex,
|
||||
size_t max_buffer_size,
|
||||
const c_string &metric_source)
|
||||
void AddEvents(c_vector[unique_ptr[CRayEventInterface]] &&data_list)
|
||||
void Shutdown()
|
||||
@@ -0,0 +1,239 @@
|
||||
"""Cython bindings for RayEventRecorder.
|
||||
|
||||
This module provides Python access to the C++ RayEventRecorder for emitting
|
||||
internal Ray events from Python code (e.g., submission job events).
|
||||
"""
|
||||
|
||||
from ray.includes.event_recorder cimport (
|
||||
CRayEventInterface,
|
||||
CPythonEventRecorder,
|
||||
CreatePythonRayEvent,
|
||||
)
|
||||
from ray.includes.common cimport move
|
||||
from libc.stdint cimport int64_t
|
||||
from libcpp.memory cimport unique_ptr
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
from libcpp.string cimport string as c_string
|
||||
|
||||
import logging
|
||||
import threading
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
cdef class RayEvent:
|
||||
"""Python wrapper holding event data for transfer to C++.
|
||||
|
||||
This class stores event metadata and serialized protobuf data. When added
|
||||
to the EventRecorder, it creates the underlying C++ RayEvent object.
|
||||
|
||||
Args:
|
||||
source_type: Integer value of RayEvent.SourceType enum.
|
||||
event_type: Integer value of RayEvent.EventType enum.
|
||||
severity: Integer value of RayEvent.Severity enum.
|
||||
entity_id: Unique identifier for the event entity (e.g., submission_id).
|
||||
message: Optional message associated with the event.
|
||||
session_name: The Ray session name.
|
||||
serialized_data: Serialized protobuf bytes of the nested event message.
|
||||
nested_event_field_number: The field number in RayEvent proto for the
|
||||
nested event message. Use RayEventProto.<FIELD>_FIELD_NUMBER constants.
|
||||
event_id: Optional explicit event id bytes. Pass empty bytes to let the
|
||||
C++ layer generate a random id (matching the convention used by the
|
||||
other RayEventInterface subclasses). A non-empty value is preserved
|
||||
as-is — use this to reuse an id from an upstream source
|
||||
(e.g., a Kubernetes event uid).
|
||||
timestamp_ns: Optional explicit event timestamp in nanoseconds since the
|
||||
unix epoch. Pass 0 to let the C++ layer capture the time at
|
||||
construction via absl::Now().
|
||||
"""
|
||||
cdef:
|
||||
int _source_type
|
||||
int _event_type
|
||||
int _severity
|
||||
str _entity_id
|
||||
str _message
|
||||
str _session_name
|
||||
bytes _serialized_data
|
||||
int _nested_event_field_number
|
||||
bytes _event_id
|
||||
int64_t _timestamp_ns
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
int source_type,
|
||||
int event_type,
|
||||
int severity,
|
||||
str entity_id,
|
||||
str message,
|
||||
str session_name,
|
||||
bytes serialized_data,
|
||||
int nested_event_field_number,
|
||||
bytes event_id = b"",
|
||||
int64_t timestamp_ns = 0,
|
||||
):
|
||||
self._source_type = source_type
|
||||
self._event_type = event_type
|
||||
self._severity = severity
|
||||
self._entity_id = entity_id
|
||||
self._message = message
|
||||
self._session_name = session_name
|
||||
self._serialized_data = serialized_data
|
||||
self._nested_event_field_number = nested_event_field_number
|
||||
self._event_id = event_id
|
||||
self._timestamp_ns = timestamp_ns
|
||||
|
||||
cdef unique_ptr[CRayEventInterface] to_cpp_event(self):
|
||||
"""Create the underlying C++ event. Ownership is transferred to caller."""
|
||||
return CreatePythonRayEvent(
|
||||
self._source_type,
|
||||
self._event_type,
|
||||
self._severity,
|
||||
self._entity_id.encode("utf-8"),
|
||||
self._message.encode("utf-8"),
|
||||
self._session_name.encode("utf-8"),
|
||||
self._serialized_data,
|
||||
self._nested_event_field_number,
|
||||
self._event_id,
|
||||
self._timestamp_ns,
|
||||
)
|
||||
|
||||
@property
|
||||
def entity_id(self):
|
||||
return self._entity_id
|
||||
|
||||
@property
|
||||
def event_type(self):
|
||||
return self._event_type
|
||||
|
||||
|
||||
# module-level Singleton instance, lazily created by EventRecorder.initialize().
|
||||
_event_recorder_instance = None
|
||||
# Guards singleton lifecycle and emission against concurrent shutdown/initialize.
|
||||
_event_recorder_lock = threading.RLock()
|
||||
|
||||
|
||||
cdef class EventRecorder:
|
||||
"""Per-process singleton for recording Ray events."""
|
||||
cdef unique_ptr[CPythonEventRecorder] _recorder
|
||||
|
||||
def __dealloc__(self):
|
||||
"""Safety-net cleanup. C++ destructor also calls Shutdown()."""
|
||||
if self._recorder.get() != NULL:
|
||||
self._recorder.get().Shutdown()
|
||||
self._recorder.reset()
|
||||
|
||||
@staticmethod
|
||||
def initialize(
|
||||
int aggregator_port,
|
||||
str node_ip,
|
||||
str node_id_hex,
|
||||
size_t max_buffer_size,
|
||||
str metric_source = "python",
|
||||
):
|
||||
"""Initialize the per-process event recorder.
|
||||
|
||||
Creates the underlying C++ PythonEventRecorder with a background I/O
|
||||
thread and gRPC client. No-op if already initialized.
|
||||
|
||||
Args:
|
||||
aggregator_port: Port of the event aggregator server (bound on 127.0.0.1).
|
||||
node_ip: IP address of the current node.
|
||||
node_id_hex: Hex-encoded node ID.
|
||||
max_buffer_size: Maximum number of events to buffer.
|
||||
metric_source: Label for the "Source" tag on dropped-events metrics
|
||||
(default "python").
|
||||
"""
|
||||
global _event_recorder_instance
|
||||
cdef EventRecorder rec
|
||||
with _event_recorder_lock:
|
||||
if _event_recorder_instance is not None:
|
||||
return
|
||||
|
||||
rec = EventRecorder()
|
||||
rec._recorder.reset(
|
||||
new CPythonEventRecorder(
|
||||
aggregator_port,
|
||||
node_ip.encode("utf-8"),
|
||||
node_id_hex.encode("utf-8"),
|
||||
max_buffer_size,
|
||||
metric_source.encode("utf-8"),
|
||||
)
|
||||
)
|
||||
_event_recorder_instance = rec
|
||||
|
||||
@staticmethod
|
||||
def instance():
|
||||
"""Get the per-process EventRecorder singleton.
|
||||
|
||||
Returns:
|
||||
The EventRecorder instance if initialized, None otherwise.
|
||||
"""
|
||||
with _event_recorder_lock:
|
||||
return _event_recorder_instance
|
||||
|
||||
@staticmethod
|
||||
def shutdown():
|
||||
"""Shutdown the event recorder.
|
||||
|
||||
Stops exporting events, performs a final flush, and releases all
|
||||
C++ resources. After this call, emit() and emit_batch() will be
|
||||
no-ops until initialize() is called again.
|
||||
"""
|
||||
global _event_recorder_instance
|
||||
cdef EventRecorder rec
|
||||
with _event_recorder_lock:
|
||||
if _event_recorder_instance is None:
|
||||
return
|
||||
|
||||
rec = <EventRecorder>_event_recorder_instance
|
||||
if rec._recorder.get() != NULL:
|
||||
with nogil:
|
||||
rec._recorder.get().Shutdown()
|
||||
rec._recorder.reset()
|
||||
_event_recorder_instance = None
|
||||
|
||||
@staticmethod
|
||||
def emit(RayEvent event):
|
||||
"""Emit a single event. No-op if not initialized.
|
||||
|
||||
Args:
|
||||
event: A RayEvent object (created via InternalEventBuilder.build()).
|
||||
|
||||
Returns:
|
||||
True if the event was successfully queued, False otherwise.
|
||||
"""
|
||||
return EventRecorder.emit_batch([event])
|
||||
|
||||
@staticmethod
|
||||
def emit_batch(list events):
|
||||
"""Emit multiple events. No-op if not initialized.
|
||||
|
||||
Args:
|
||||
events: List of RayEvent objects to emit.
|
||||
|
||||
Returns:
|
||||
True if events were successfully queued, False otherwise.
|
||||
"""
|
||||
cdef EventRecorder rec
|
||||
cdef c_vector[unique_ptr[CRayEventInterface]] cpp_events
|
||||
cdef RayEvent ev
|
||||
|
||||
if not events:
|
||||
return True
|
||||
|
||||
for ev in events:
|
||||
cpp_events.push_back(move(ev.to_cpp_event()))
|
||||
|
||||
with _event_recorder_lock:
|
||||
if _event_recorder_instance is None:
|
||||
logger.debug(
|
||||
"Event recorder not initialized, dropping %d events",
|
||||
len(events),
|
||||
)
|
||||
return False
|
||||
|
||||
rec = <EventRecorder>_event_recorder_instance
|
||||
|
||||
with nogil:
|
||||
rec._recorder.get().AddEvents(move(cpp_events))
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,80 @@
|
||||
from libc.stdint cimport uint8_t, uint64_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 (
|
||||
CLanguage,
|
||||
)
|
||||
from ray.includes.unique_ids cimport (
|
||||
CActorID,
|
||||
CJobID,
|
||||
CObjectID,
|
||||
CTaskID,
|
||||
)
|
||||
|
||||
cdef extern from "src/ray/protobuf/common.pb.h" nogil:
|
||||
cdef cppclass CFunctionDescriptorType \
|
||||
"ray::FunctionDescriptorType":
|
||||
pass
|
||||
|
||||
cdef CFunctionDescriptorType EmptyFunctionDescriptorType \
|
||||
"ray::FunctionDescriptorType::FUNCTION_DESCRIPTOR_NOT_SET"
|
||||
cdef CFunctionDescriptorType JavaFunctionDescriptorType \
|
||||
"ray::FunctionDescriptorType::kJavaFunctionDescriptor"
|
||||
cdef CFunctionDescriptorType PythonFunctionDescriptorType \
|
||||
"ray::FunctionDescriptorType::kPythonFunctionDescriptor"
|
||||
cdef CFunctionDescriptorType CppFunctionDescriptorType \
|
||||
"ray::FunctionDescriptorType::kCppFunctionDescriptor"
|
||||
|
||||
|
||||
cdef extern from "ray/common/function_descriptor.h" nogil:
|
||||
cdef cppclass CFunctionDescriptorInterface \
|
||||
"ray::FunctionDescriptorInterface":
|
||||
CFunctionDescriptorType Type()
|
||||
c_string ToString()
|
||||
c_string Serialize()
|
||||
|
||||
ctypedef shared_ptr[CFunctionDescriptorInterface] CFunctionDescriptor \
|
||||
"ray::FunctionDescriptor"
|
||||
|
||||
cdef cppclass CFunctionDescriptorBuilder "ray::FunctionDescriptorBuilder":
|
||||
@staticmethod
|
||||
CFunctionDescriptor Empty()
|
||||
|
||||
@staticmethod
|
||||
CFunctionDescriptor BuildJava(const c_string &class_name,
|
||||
const c_string &function_name,
|
||||
const c_string &signature)
|
||||
|
||||
@staticmethod
|
||||
CFunctionDescriptor BuildPython(const c_string &module_name,
|
||||
const c_string &class_name,
|
||||
const c_string &function_name,
|
||||
const c_string &function_source_hash)
|
||||
|
||||
@staticmethod
|
||||
CFunctionDescriptor BuildCpp(const c_string &function_name,
|
||||
const c_string &caller,
|
||||
const c_string &class_name)
|
||||
|
||||
@staticmethod
|
||||
CFunctionDescriptor Deserialize(const c_string &serialized_binary)
|
||||
|
||||
cdef cppclass CJavaFunctionDescriptor "ray::JavaFunctionDescriptor":
|
||||
c_string ClassName()
|
||||
c_string FunctionName()
|
||||
c_string Signature()
|
||||
|
||||
cdef cppclass CPythonFunctionDescriptor "ray::PythonFunctionDescriptor":
|
||||
c_string ModuleName()
|
||||
c_string ClassName()
|
||||
c_string FunctionName()
|
||||
c_string FunctionHash()
|
||||
|
||||
cdef cppclass CCppFunctionDescriptor "ray::CppFunctionDescriptor":
|
||||
c_string FunctionName()
|
||||
c_string Caller()
|
||||
c_string ClassName()
|
||||
@@ -0,0 +1,396 @@
|
||||
from ray.includes.function_descriptor cimport (
|
||||
CFunctionDescriptor,
|
||||
CFunctionDescriptorBuilder,
|
||||
CPythonFunctionDescriptor,
|
||||
CJavaFunctionDescriptor,
|
||||
CCppFunctionDescriptor,
|
||||
EmptyFunctionDescriptorType,
|
||||
JavaFunctionDescriptorType,
|
||||
PythonFunctionDescriptorType,
|
||||
CppFunctionDescriptorType,
|
||||
)
|
||||
|
||||
import hashlib
|
||||
import cython
|
||||
import inspect
|
||||
import uuid
|
||||
import ray._private.ray_constants as ray_constants
|
||||
|
||||
|
||||
ctypedef object (*FunctionDescriptor_from_cpp)(const CFunctionDescriptor &)
|
||||
cdef unordered_map[int, FunctionDescriptor_from_cpp] \
|
||||
FunctionDescriptor_constructor_map
|
||||
cdef CFunctionDescriptorToPython(CFunctionDescriptor function_descriptor):
|
||||
cdef int function_descriptor_type = <int>function_descriptor.get().Type()
|
||||
it = FunctionDescriptor_constructor_map.find(function_descriptor_type)
|
||||
if it == FunctionDescriptor_constructor_map.end():
|
||||
raise Exception("Can't construct FunctionDescriptor from type {}"
|
||||
.format(function_descriptor_type))
|
||||
else:
|
||||
constructor = dereference(it).second
|
||||
return constructor(function_descriptor)
|
||||
|
||||
|
||||
@cython.auto_pickle(False)
|
||||
cdef class FunctionDescriptor:
|
||||
def __cinit__(self, *args, **kwargs):
|
||||
if type(self) == FunctionDescriptor:
|
||||
raise Exception("type {} is abstract".format(type(self).__name__))
|
||||
|
||||
def __hash__(self):
|
||||
return hash(self.descriptor.get().ToString())
|
||||
|
||||
def __eq__(self, other):
|
||||
return (type(self) == type(other) and
|
||||
self.descriptor.get().ToString() ==
|
||||
(<FunctionDescriptor>other).descriptor.get().ToString())
|
||||
|
||||
def __repr__(self):
|
||||
return <str>self.descriptor.get().ToString()
|
||||
|
||||
def to_dict(self):
|
||||
d = {"type": type(self).__name__}
|
||||
for k, v in vars(type(self)).items():
|
||||
if inspect.isgetsetdescriptor(v):
|
||||
d[k] = v.__get__(self)
|
||||
return d
|
||||
|
||||
@property
|
||||
def repr(self):
|
||||
return self.__repr__()
|
||||
|
||||
|
||||
FunctionDescriptor_constructor_map[<int>EmptyFunctionDescriptorType] = \
|
||||
EmptyFunctionDescriptor.from_cpp
|
||||
|
||||
|
||||
@cython.auto_pickle(False)
|
||||
cdef class EmptyFunctionDescriptor(FunctionDescriptor):
|
||||
def __cinit__(self):
|
||||
self.descriptor = CFunctionDescriptorBuilder.Empty()
|
||||
|
||||
def __reduce__(self):
|
||||
return EmptyFunctionDescriptor, ()
|
||||
|
||||
@staticmethod
|
||||
cdef from_cpp(const CFunctionDescriptor &c_function_descriptor):
|
||||
return EmptyFunctionDescriptor()
|
||||
|
||||
@property
|
||||
def repr(self):
|
||||
return "FunctionDescriptor(empty)"
|
||||
|
||||
|
||||
FunctionDescriptor_constructor_map[<int>JavaFunctionDescriptorType] = \
|
||||
JavaFunctionDescriptor.from_cpp
|
||||
|
||||
|
||||
@cython.auto_pickle(False)
|
||||
cdef class JavaFunctionDescriptor(FunctionDescriptor):
|
||||
cdef:
|
||||
CJavaFunctionDescriptor *typed_descriptor
|
||||
|
||||
def __cinit__(self,
|
||||
class_name,
|
||||
function_name,
|
||||
signature):
|
||||
self.descriptor = CFunctionDescriptorBuilder.BuildJava(
|
||||
class_name, function_name, signature)
|
||||
self.typed_descriptor = <CJavaFunctionDescriptor*>(
|
||||
self.descriptor.get())
|
||||
|
||||
def __reduce__(self):
|
||||
return JavaFunctionDescriptor, (self.typed_descriptor.ClassName(),
|
||||
self.typed_descriptor.FunctionName(),
|
||||
self.typed_descriptor.Signature())
|
||||
|
||||
@staticmethod
|
||||
cdef from_cpp(const CFunctionDescriptor &c_function_descriptor):
|
||||
cdef CJavaFunctionDescriptor *typed_descriptor = \
|
||||
<CJavaFunctionDescriptor*>(c_function_descriptor.get())
|
||||
return JavaFunctionDescriptor(typed_descriptor.ClassName(),
|
||||
typed_descriptor.FunctionName(),
|
||||
typed_descriptor.Signature())
|
||||
|
||||
@property
|
||||
def class_name(self):
|
||||
"""Get the class name of current function descriptor.
|
||||
|
||||
Returns:
|
||||
The class name of the function descriptor. It could be
|
||||
empty if the function is not a class method.
|
||||
"""
|
||||
return <str>self.typed_descriptor.ClassName()
|
||||
|
||||
@property
|
||||
def function_name(self):
|
||||
"""Get the function name of current function descriptor.
|
||||
|
||||
Returns:
|
||||
The function name of the function descriptor.
|
||||
"""
|
||||
return <str>self.typed_descriptor.FunctionName()
|
||||
|
||||
@property
|
||||
def signature(self):
|
||||
"""Get the signature of current function descriptor.
|
||||
|
||||
Returns:
|
||||
The signature of the function descriptor.
|
||||
"""
|
||||
return <str>self.typed_descriptor.Signature()
|
||||
|
||||
@property
|
||||
def repr(self):
|
||||
return f"{self.class_name}.{self.function_name}"
|
||||
|
||||
|
||||
FunctionDescriptor_constructor_map[<int>PythonFunctionDescriptorType] = \
|
||||
PythonFunctionDescriptor.from_cpp
|
||||
|
||||
|
||||
@cython.auto_pickle(False)
|
||||
cdef class PythonFunctionDescriptor(FunctionDescriptor):
|
||||
cdef:
|
||||
CPythonFunctionDescriptor *typed_descriptor
|
||||
object _function_id
|
||||
|
||||
def __cinit__(self,
|
||||
module_name,
|
||||
function_name,
|
||||
class_name="",
|
||||
function_source_hash=""):
|
||||
self.descriptor = CFunctionDescriptorBuilder.BuildPython(
|
||||
module_name, class_name, function_name, function_source_hash)
|
||||
self.typed_descriptor = <CPythonFunctionDescriptor*>(
|
||||
self.descriptor.get())
|
||||
|
||||
def __reduce__(self):
|
||||
return PythonFunctionDescriptor, (self.typed_descriptor.ModuleName(),
|
||||
self.typed_descriptor.FunctionName(),
|
||||
self.typed_descriptor.ClassName(),
|
||||
self.typed_descriptor.FunctionHash())
|
||||
|
||||
@staticmethod
|
||||
cdef from_cpp(const CFunctionDescriptor &c_function_descriptor):
|
||||
cdef CPythonFunctionDescriptor *typed_descriptor = \
|
||||
<CPythonFunctionDescriptor*>(c_function_descriptor.get())
|
||||
return PythonFunctionDescriptor(typed_descriptor.ModuleName(),
|
||||
typed_descriptor.FunctionName(),
|
||||
typed_descriptor.ClassName(),
|
||||
typed_descriptor.FunctionHash())
|
||||
|
||||
@classmethod
|
||||
def from_function(cls, function, function_uuid):
|
||||
"""Create a FunctionDescriptor from a function instance.
|
||||
|
||||
This function is used to create the function descriptor from
|
||||
a python function. If a function is a class function, it should
|
||||
not be used by this function.
|
||||
|
||||
Args:
|
||||
cls: Current class which is required argument for classmethod.
|
||||
function: the python function used to create the function
|
||||
descriptor.
|
||||
function_uuid: Used to uniquely identify a function.
|
||||
Ideally we can use the pickled function bytes
|
||||
but cloudpickle isn't stable in some cases
|
||||
for the same function.
|
||||
|
||||
Returns:
|
||||
The FunctionDescriptor instance created according to the function.
|
||||
"""
|
||||
module_name = cls._get_module_name(function)
|
||||
function_name = function.__qualname__
|
||||
class_name = ""
|
||||
|
||||
return cls(module_name, function_name, class_name, function_uuid.hex)
|
||||
|
||||
@classmethod
|
||||
def from_class(cls, target_class):
|
||||
"""Create a FunctionDescriptor from a class.
|
||||
|
||||
Args:
|
||||
cls: Current class which is required argument for classmethod.
|
||||
target_class: the python class used to create the function
|
||||
descriptor.
|
||||
|
||||
Returns:
|
||||
The FunctionDescriptor instance created according to the class.
|
||||
"""
|
||||
module_name = cls._get_module_name(target_class)
|
||||
class_name = target_class.__qualname__
|
||||
# Use a random uuid as function hash to solve actor name conflict.
|
||||
return cls(module_name, "__init__", class_name, uuid.uuid4().hex)
|
||||
|
||||
@property
|
||||
def module_name(self):
|
||||
"""Get the module name of current function descriptor.
|
||||
|
||||
Returns:
|
||||
The module name of the function descriptor.
|
||||
"""
|
||||
return <str>self.typed_descriptor.ModuleName()
|
||||
|
||||
@property
|
||||
def class_name(self):
|
||||
"""Get the class name of current function descriptor.
|
||||
|
||||
Returns:
|
||||
The class name of the function descriptor. It could be
|
||||
empty if the function is not a class method.
|
||||
"""
|
||||
return <str>self.typed_descriptor.ClassName()
|
||||
|
||||
@property
|
||||
def function_name(self):
|
||||
"""Get the function name of current function descriptor.
|
||||
|
||||
Returns:
|
||||
The function name of the function descriptor.
|
||||
"""
|
||||
return <str>self.typed_descriptor.FunctionName()
|
||||
|
||||
@property
|
||||
def function_hash(self):
|
||||
"""Get the hash string of the function source code.
|
||||
|
||||
Returns:
|
||||
The hex of function hash if the source code is available.
|
||||
Otherwise, it will be an empty string.
|
||||
"""
|
||||
return <str>self.typed_descriptor.FunctionHash()
|
||||
|
||||
@property
|
||||
def function_id(self):
|
||||
"""Get the function id calculated from this descriptor.
|
||||
|
||||
Returns:
|
||||
The value of ray.ObjectRef that represents the function id.
|
||||
"""
|
||||
if not self._function_id:
|
||||
self._function_id = self._get_function_id()
|
||||
return self._function_id
|
||||
|
||||
@property
|
||||
def repr(self):
|
||||
"""Get the module_name.Optional[class_name].function_name
|
||||
of the descriptor.
|
||||
|
||||
Returns:
|
||||
The value of module_name.Optional[class_name].function_name
|
||||
"""
|
||||
if self.is_actor_method():
|
||||
return ".".join(
|
||||
[self.module_name, self.class_name, self.function_name])
|
||||
else:
|
||||
return ".".join(
|
||||
[self.module_name, self.function_name])
|
||||
|
||||
def _get_function_id(self):
|
||||
"""Calculate the function id of current function descriptor.
|
||||
|
||||
This function id is calculated from all the fields of function
|
||||
descriptor.
|
||||
|
||||
Returns:
|
||||
ray.ObjectRef to represent the function descriptor.
|
||||
"""
|
||||
function_id_hash = hashlib.shake_128()
|
||||
# Include the function module and name in the hash.
|
||||
function_id_hash.update(self.typed_descriptor.ModuleName())
|
||||
function_id_hash.update(self.typed_descriptor.FunctionName())
|
||||
function_id_hash.update(self.typed_descriptor.ClassName())
|
||||
function_id_hash.update(self.typed_descriptor.FunctionHash())
|
||||
# Compute the function ID.
|
||||
function_id = function_id_hash.digest(ray_constants.ID_SIZE)
|
||||
return ray.FunctionID(function_id)
|
||||
|
||||
@staticmethod
|
||||
def _get_module_name(object):
|
||||
"""Get the module name from object. If the module is __main__,
|
||||
get the module name from file.
|
||||
|
||||
Returns:
|
||||
Module name of object.
|
||||
"""
|
||||
module_name = object.__module__
|
||||
if module_name == "__main__":
|
||||
try:
|
||||
file_path = inspect.getfile(object)
|
||||
n = inspect.getmodulename(file_path)
|
||||
if n:
|
||||
module_name = n
|
||||
except (TypeError, OSError):
|
||||
pass
|
||||
return module_name
|
||||
|
||||
def is_actor_method(self):
|
||||
"""Whether this function descriptor is an actor method.
|
||||
|
||||
Returns:
|
||||
True if it's an actor method, False if it's a normal function.
|
||||
"""
|
||||
return not self.typed_descriptor.ClassName().empty()
|
||||
|
||||
|
||||
FunctionDescriptor_constructor_map[<int>CppFunctionDescriptorType] = \
|
||||
CppFunctionDescriptor.from_cpp
|
||||
|
||||
|
||||
@cython.auto_pickle(False)
|
||||
cdef class CppFunctionDescriptor(FunctionDescriptor):
|
||||
cdef:
|
||||
CCppFunctionDescriptor *typed_descriptor
|
||||
|
||||
def __cinit__(self,
|
||||
function_name, caller, class_name=""):
|
||||
self.descriptor = CFunctionDescriptorBuilder.BuildCpp(
|
||||
function_name, caller, class_name)
|
||||
self.typed_descriptor = <CCppFunctionDescriptor*>(
|
||||
self.descriptor.get())
|
||||
|
||||
def __reduce__(self):
|
||||
return CppFunctionDescriptor, (self.typed_descriptor.FunctionName(),
|
||||
self.typed_descriptor.Caller(),
|
||||
self.typed_descriptor.ClassName())
|
||||
|
||||
@staticmethod
|
||||
cdef from_cpp(const CFunctionDescriptor &c_function_descriptor):
|
||||
cdef CCppFunctionDescriptor *typed_descriptor = \
|
||||
<CCppFunctionDescriptor*>(c_function_descriptor.get())
|
||||
return CppFunctionDescriptor(typed_descriptor.FunctionName(),
|
||||
typed_descriptor.Caller(),
|
||||
typed_descriptor.ClassName())
|
||||
|
||||
@property
|
||||
def function_name(self):
|
||||
"""Get the function name of current function descriptor.
|
||||
|
||||
Returns:
|
||||
The function name of the function descriptor.
|
||||
"""
|
||||
return <str>self.typed_descriptor.FunctionName()
|
||||
|
||||
@property
|
||||
def caller(self):
|
||||
"""Get the caller of current function descriptor.
|
||||
|
||||
Returns:
|
||||
The caller of the function descriptor.
|
||||
"""
|
||||
return <str>self.typed_descriptor.Caller()
|
||||
|
||||
@property
|
||||
def class_name(self):
|
||||
"""Get the class name of current function descriptor,
|
||||
when it is empty, it is a non-member function.
|
||||
|
||||
Returns:
|
||||
The class name of the function descriptor.
|
||||
"""
|
||||
return <str>self.typed_descriptor.ClassName()
|
||||
|
||||
@property
|
||||
def repr(self):
|
||||
return f"{self.class_name}::{self.function_name}"
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,150 @@
|
||||
import random
|
||||
|
||||
from libcpp.memory cimport shared_ptr
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
from libcpp.utility cimport move
|
||||
|
||||
from ray.includes.common cimport(
|
||||
CPythonGcsSubscriber,
|
||||
CErrorTableData,
|
||||
CLogBatch,
|
||||
PythonGetLogBatchLines,
|
||||
RAY_ERROR_INFO_CHANNEL,
|
||||
RAY_LOG_CHANNEL,
|
||||
)
|
||||
|
||||
cdef class _GcsSubscriber:
|
||||
"""Cython wrapper class of C++ `ray::pubsub::PythonGcsSubscriber`."""
|
||||
cdef:
|
||||
shared_ptr[CPythonGcsSubscriber] inner
|
||||
|
||||
def _construct(self, address, channel, worker_id):
|
||||
cdef:
|
||||
c_worker_id = worker_id or b""
|
||||
# subscriber_id needs to match the binary format of a random
|
||||
# SubscriberID / UniqueID, which is 28 (kUniqueIDSize) random bytes.
|
||||
subscriber_id = bytes(bytearray(random.getrandbits(8) for _ in range(28)))
|
||||
gcs_address, gcs_port = parse_address(address)
|
||||
self.inner.reset(new CPythonGcsSubscriber(
|
||||
gcs_address, int(gcs_port), channel, subscriber_id, c_worker_id))
|
||||
|
||||
def subscribe(self):
|
||||
"""Registers a subscription for the subscriber's channel type.
|
||||
|
||||
Before the registration, published messages in the channel will not be
|
||||
saved for the subscriber.
|
||||
"""
|
||||
with nogil:
|
||||
check_status(self.inner.get().Subscribe())
|
||||
|
||||
@property
|
||||
def last_batch_size(self):
|
||||
"""Batch size of the result from last poll.
|
||||
|
||||
Used to indicate whether the subscriber can keep up.
|
||||
"""
|
||||
return self.inner.get().last_batch_size()
|
||||
|
||||
def close(self):
|
||||
"""Closes the subscriber and its active subscription."""
|
||||
with nogil:
|
||||
check_status(self.inner.get().Close())
|
||||
|
||||
|
||||
cdef class GcsErrorSubscriber(_GcsSubscriber):
|
||||
"""Subscriber to error info. Thread safe.
|
||||
|
||||
Usage example:
|
||||
subscriber = GcsErrorSubscriber()
|
||||
# Subscribe to the error channel.
|
||||
subscriber.subscribe()
|
||||
...
|
||||
while running:
|
||||
error_id, error_data = subscriber.poll()
|
||||
......
|
||||
# Unsubscribe from the error channels.
|
||||
subscriber.close()
|
||||
"""
|
||||
|
||||
def __init__(self, address, worker_id=None):
|
||||
self._construct(address, RAY_ERROR_INFO_CHANNEL, worker_id)
|
||||
|
||||
def poll(self, timeout=None):
|
||||
"""Polls for new error messages.
|
||||
|
||||
Returns:
|
||||
A tuple of error message ID and dict describing the error,
|
||||
or None, None if polling times out or subscriber closed.
|
||||
"""
|
||||
cdef:
|
||||
CErrorTableData error_data
|
||||
c_string key_id
|
||||
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
||||
|
||||
with nogil:
|
||||
check_status(self.inner.get().PollError(&key_id, timeout_ms, &error_data))
|
||||
|
||||
if key_id == b"":
|
||||
return None, None
|
||||
|
||||
return (bytes(key_id), {
|
||||
"job_id": error_data.job_id(),
|
||||
"type": error_data.type().decode(),
|
||||
"error_message": error_data.error_message().decode(),
|
||||
"timestamp": error_data.timestamp(),
|
||||
})
|
||||
|
||||
|
||||
cdef class GcsLogSubscriber(_GcsSubscriber):
|
||||
"""Subscriber to logs. Thread safe.
|
||||
|
||||
Usage example:
|
||||
subscriber = GcsLogSubscriber()
|
||||
# Subscribe to the log channel.
|
||||
subscriber.subscribe()
|
||||
...
|
||||
while running:
|
||||
log = subscriber.poll()
|
||||
......
|
||||
# Unsubscribe from the log channel.
|
||||
subscriber.close()
|
||||
"""
|
||||
|
||||
def __init__(self, address, worker_id=None):
|
||||
self._construct(address, RAY_LOG_CHANNEL, worker_id)
|
||||
|
||||
def poll(self, timeout=None):
|
||||
"""Polls for new log messages.
|
||||
|
||||
Returns:
|
||||
A dict containing a batch of log lines and their metadata.
|
||||
"""
|
||||
cdef:
|
||||
CLogBatch log_batch
|
||||
c_string key_id
|
||||
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
||||
c_vector[c_string] c_log_lines
|
||||
c_string c_log_line
|
||||
|
||||
with nogil:
|
||||
check_status(self.inner.get().PollLogs(&key_id, timeout_ms, &log_batch))
|
||||
|
||||
result = {
|
||||
"ip": log_batch.ip().decode(),
|
||||
"pid": log_batch.pid().decode(),
|
||||
"job": log_batch.job_id().decode(),
|
||||
"is_err": log_batch.is_error(),
|
||||
"actor_name": log_batch.actor_name().decode(),
|
||||
"task_name": log_batch.task_name().decode(),
|
||||
}
|
||||
|
||||
with nogil:
|
||||
c_log_lines = PythonGetLogBatchLines(move(log_batch))
|
||||
|
||||
log_lines = []
|
||||
for c_log_line in c_log_lines:
|
||||
log_lines.append(c_log_line.decode())
|
||||
|
||||
result["lines"] = log_lines
|
||||
return result
|
||||
@@ -0,0 +1,140 @@
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp cimport bool as c_bool
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
from libcpp.unordered_map cimport unordered_map
|
||||
from libcpp.memory cimport unique_ptr
|
||||
from libc.stdint cimport (
|
||||
int32_t as c_int32_t,
|
||||
uint32_t as c_uint32_t,
|
||||
int64_t as c_int64_t,
|
||||
)
|
||||
from ray.includes.unique_ids cimport (
|
||||
CActorID,
|
||||
CJobID,
|
||||
CNodeID,
|
||||
CObjectID,
|
||||
CWorkerID,
|
||||
CPlacementGroupID,
|
||||
)
|
||||
from ray.includes.common cimport (
|
||||
CRayStatus,
|
||||
CGcsClientOptions,
|
||||
)
|
||||
from ray.includes.optional cimport (
|
||||
optional
|
||||
)
|
||||
|
||||
cdef extern from "ray/gcs_rpc_client/global_state_accessor.h" nogil:
|
||||
cdef cppclass CGlobalStateAccessor "ray::gcs::GlobalStateAccessor":
|
||||
CGlobalStateAccessor(const CGcsClientOptions&)
|
||||
c_bool Connect()
|
||||
void Disconnect()
|
||||
c_vector[c_string] GetAllJobInfo(
|
||||
c_bool skip_submission_job_info_field, c_bool skip_is_running_tasks_field)
|
||||
CJobID GetNextJobID()
|
||||
c_vector[c_string] GetAllNodeInfo()
|
||||
c_vector[c_string] GetAllAvailableResources()
|
||||
c_vector[c_string] GetAllTotalResources()
|
||||
unordered_map[CNodeID, c_int64_t] GetDrainingNodes()
|
||||
unique_ptr[c_string] GetInternalKV(
|
||||
const c_string &namespace, const c_string &key)
|
||||
c_vector[c_string] GetAllTaskEvents()
|
||||
unique_ptr[c_string] GetObjectInfo(const CObjectID &object_id)
|
||||
unique_ptr[c_string] GetAllResourceUsage()
|
||||
c_vector[c_string] GetAllActorInfo(
|
||||
optional[CActorID], optional[CJobID], optional[c_string])
|
||||
unique_ptr[c_string] GetActorInfo(const CActorID &actor_id)
|
||||
unique_ptr[c_string] GetWorkerInfo(const CWorkerID &worker_id)
|
||||
c_vector[c_string] GetAllWorkerInfo()
|
||||
c_bool AddWorkerInfo(const c_string &serialized_string)
|
||||
c_bool UpdateWorkerDebuggerPort(const CWorkerID &worker_id,
|
||||
const c_uint32_t debuger_port)
|
||||
c_bool UpdateWorkerNumPausedThreads(const CWorkerID &worker_id,
|
||||
const c_int32_t num_paused_threads_delta)
|
||||
c_uint32_t GetWorkerDebuggerPort(const CWorkerID &worker_id)
|
||||
unique_ptr[c_string] GetPlacementGroupInfo(
|
||||
const CPlacementGroupID &placement_group_id)
|
||||
unique_ptr[c_string] GetPlacementGroupByName(
|
||||
const c_string &placement_group_name,
|
||||
const c_string &ray_namespace,
|
||||
)
|
||||
c_vector[c_string] GetAllPlacementGroupInfo()
|
||||
c_string GetSystemConfig()
|
||||
CRayStatus GetNode(
|
||||
const c_string &node_id_hex_str,
|
||||
c_string *node_info)
|
||||
|
||||
cdef extern from * namespace "ray::gcs" nogil:
|
||||
"""
|
||||
#include <thread>
|
||||
#include "ray/gcs/store_client_kv.h"
|
||||
#include "ray/gcs/store_client/redis_store_client.h"
|
||||
#include "ray/util/clock.h"
|
||||
#include "ray/util/raii.h"
|
||||
namespace ray {
|
||||
namespace gcs {
|
||||
|
||||
bool RedisGetKeySync(const std::string& host,
|
||||
int32_t port,
|
||||
const std::string& username,
|
||||
const std::string& password,
|
||||
bool use_ssl,
|
||||
const std::string& config,
|
||||
const std::string& key,
|
||||
std::string* data) {
|
||||
// Logging default value see class `RayLog`.
|
||||
InitShutdownRAII ray_log_shutdown_raii(ray::RayLog::StartRayLog,
|
||||
ray::RayLog::ShutDownRayLog,
|
||||
"ray_init",
|
||||
ray::RayLogLevel::WARNING,
|
||||
/*log_filepath=*/"",
|
||||
/*err_log_filepath=*/"",
|
||||
/*log_rotation_max_size=*/1ULL << 29,
|
||||
/*log_rotation_file_num=*/10);
|
||||
|
||||
std::string config_list;
|
||||
RAY_CHECK(absl::Base64Unescape(config, &config_list));
|
||||
RayConfig::instance().initialize(config_list);
|
||||
|
||||
instrumented_io_context io_service{/*enable_lag_probe=*/false, /*running_on_single_thread=*/true};
|
||||
RedisClientOptions options{host, port, username, password, use_ssl};
|
||||
Clock clock;
|
||||
auto client = std::make_unique<StoreClientInternalKV>(
|
||||
std::make_unique<RedisStoreClient>(io_service, options, clock));
|
||||
|
||||
bool ret_val = false;
|
||||
client->Get("session", key, {[&](std::optional<std::string> result) {
|
||||
if (result.has_value()) {
|
||||
*data = result.value();
|
||||
ret_val = true;
|
||||
} else {
|
||||
RAY_LOG(INFO) << "Failed to retrieve the key " << key
|
||||
<< " from persistent storage.";
|
||||
ret_val = false;
|
||||
}
|
||||
}, io_service});
|
||||
io_service.run_for(std::chrono::milliseconds(1000));
|
||||
|
||||
return ret_val;
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
"""
|
||||
c_bool RedisGetKeySync(const c_string& host,
|
||||
c_int32_t port,
|
||||
const c_string& username,
|
||||
const c_string& password,
|
||||
c_bool use_ssl,
|
||||
const c_string& config,
|
||||
const c_string& key,
|
||||
c_string* data)
|
||||
|
||||
|
||||
cdef extern from "ray/gcs/store_client/redis_store_client.h" namespace "ray::gcs" nogil:
|
||||
c_bool RedisDelKeyPrefixSync(const c_string& host,
|
||||
c_int32_t port,
|
||||
const c_string& username,
|
||||
const c_string& password,
|
||||
c_bool use_ssl,
|
||||
const c_string& key_prefix)
|
||||
@@ -0,0 +1,289 @@
|
||||
from ray.includes.common cimport (
|
||||
CGcsClientOptions,
|
||||
CGcsNodeState,
|
||||
PythonGetResourcesTotal,
|
||||
PythonGetNodeLabels
|
||||
)
|
||||
|
||||
from ray.includes.unique_ids cimport (
|
||||
CActorID,
|
||||
CNodeID,
|
||||
CObjectID,
|
||||
CWorkerID,
|
||||
CPlacementGroupID
|
||||
)
|
||||
|
||||
from ray.includes.global_state_accessor cimport (
|
||||
CGlobalStateAccessor,
|
||||
RedisDelKeyPrefixSync,
|
||||
)
|
||||
|
||||
from ray.includes.optional cimport (
|
||||
optional,
|
||||
nullopt,
|
||||
make_optional
|
||||
)
|
||||
|
||||
from libc.stdint cimport uint32_t as c_uint32_t, int32_t as c_int32_t
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp.memory cimport make_unique as c_make_unique
|
||||
|
||||
cdef class GlobalStateAccessor:
|
||||
"""Cython wrapper class of C++ `ray::gcs::GlobalStateAccessor`."""
|
||||
cdef:
|
||||
unique_ptr[CGlobalStateAccessor] inner
|
||||
|
||||
def __cinit__(self, GcsClientOptions gcs_options):
|
||||
cdef CGcsClientOptions *opts
|
||||
opts = gcs_options.native()
|
||||
self.inner = c_make_unique[CGlobalStateAccessor](opts[0])
|
||||
|
||||
def connect(self):
|
||||
cdef c_bool result
|
||||
with nogil:
|
||||
result = self.inner.get().Connect()
|
||||
return result
|
||||
|
||||
def get_job_table(
|
||||
self, *, skip_submission_job_info_field=False, skip_is_running_tasks_field=False
|
||||
):
|
||||
cdef c_vector[c_string] result
|
||||
cdef c_bool c_skip_submission_job_info_field = skip_submission_job_info_field
|
||||
cdef c_bool c_skip_is_running_tasks_field = skip_is_running_tasks_field
|
||||
|
||||
with nogil:
|
||||
result = self.inner.get().GetAllJobInfo(
|
||||
c_skip_submission_job_info_field, c_skip_is_running_tasks_field)
|
||||
return result
|
||||
|
||||
def get_next_job_id(self):
|
||||
cdef CJobID cjob_id
|
||||
with nogil:
|
||||
cjob_id = self.inner.get().GetNextJobID()
|
||||
return cjob_id.ToInt()
|
||||
|
||||
def get_node_table(self):
|
||||
cdef:
|
||||
c_vector[c_string] items
|
||||
c_string item
|
||||
CGcsNodeInfo c_node_info
|
||||
unordered_map[c_string, double] c_resources
|
||||
with nogil:
|
||||
items = self.inner.get().GetAllNodeInfo()
|
||||
results = []
|
||||
for item in items:
|
||||
c_node_info.ParseFromString(item)
|
||||
node_info = {
|
||||
"NodeID": ray._common.utils.binary_to_hex(c_node_info.node_id()),
|
||||
"Alive": c_node_info.state() == CGcsNodeState.ALIVE,
|
||||
"NodeManagerAddress": c_node_info.node_manager_address().decode(),
|
||||
"NodeManagerHostname": c_node_info.node_manager_hostname().decode(),
|
||||
"NodeManagerPort": c_node_info.node_manager_port(),
|
||||
"ObjectManagerPort": c_node_info.object_manager_port(),
|
||||
"ObjectStoreSocketName":
|
||||
c_node_info.object_store_socket_name().decode(),
|
||||
"RayletSocketName": c_node_info.raylet_socket_name().decode(),
|
||||
"MetricsExportPort": c_node_info.metrics_export_port(),
|
||||
"MetricsAgentPort": c_node_info.metrics_agent_port(),
|
||||
"DashboardAgentListenPort": c_node_info.dashboard_agent_listen_port(),
|
||||
"NodeName": c_node_info.node_name().decode(),
|
||||
"RuntimeEnvAgentPort": c_node_info.runtime_env_agent_port(),
|
||||
"DeathReason": c_node_info.death_info().reason(),
|
||||
"DeathReasonMessage":
|
||||
c_node_info.death_info().reason_message().decode(),
|
||||
}
|
||||
node_info["alive"] = node_info["Alive"]
|
||||
c_resources = PythonGetResourcesTotal(c_node_info)
|
||||
node_info["Resources"] = (
|
||||
{key.decode(): value for key, value in c_resources}
|
||||
if node_info["Alive"]
|
||||
else {}
|
||||
)
|
||||
c_labels = PythonGetNodeLabels(c_node_info)
|
||||
node_info["Labels"] = \
|
||||
{key.decode(): value.decode() for key, value in c_labels}
|
||||
results.append(node_info)
|
||||
return results
|
||||
|
||||
def get_draining_nodes(self):
|
||||
cdef:
|
||||
unordered_map[CNodeID, int64_t] draining_nodes
|
||||
unordered_map[CNodeID, int64_t].iterator draining_nodes_it
|
||||
|
||||
with nogil:
|
||||
draining_nodes = self.inner.get().GetDrainingNodes()
|
||||
draining_nodes_it = draining_nodes.begin()
|
||||
results = {}
|
||||
while draining_nodes_it != draining_nodes.end():
|
||||
draining_node_id = dereference(draining_nodes_it).first
|
||||
results[ray._common.utils.binary_to_hex(
|
||||
draining_node_id.Binary())] = dereference(draining_nodes_it).second
|
||||
postincrement(draining_nodes_it)
|
||||
|
||||
return results
|
||||
|
||||
def get_internal_kv(self, namespace, key):
|
||||
cdef:
|
||||
c_string c_namespace = namespace
|
||||
c_string c_key = key
|
||||
unique_ptr[c_string] result
|
||||
with nogil:
|
||||
result = self.inner.get().GetInternalKV(c_namespace, c_key)
|
||||
if result:
|
||||
return c_string(result.get().data(), result.get().size())
|
||||
return None
|
||||
|
||||
def get_all_available_resources(self):
|
||||
cdef c_vector[c_string] result
|
||||
with nogil:
|
||||
result = self.inner.get().GetAllAvailableResources()
|
||||
return result
|
||||
|
||||
def get_all_total_resources(self):
|
||||
cdef c_vector[c_string] result
|
||||
with nogil:
|
||||
result = self.inner.get().GetAllTotalResources()
|
||||
return result
|
||||
|
||||
def get_task_events(self):
|
||||
cdef c_vector[c_string] result
|
||||
with nogil:
|
||||
result = self.inner.get().GetAllTaskEvents()
|
||||
return result
|
||||
|
||||
def get_all_resource_usage(self):
|
||||
"""Get newest resource usage of all nodes from GCS service."""
|
||||
cdef unique_ptr[c_string] result
|
||||
with nogil:
|
||||
result = self.inner.get().GetAllResourceUsage()
|
||||
if result:
|
||||
return c_string(result.get().data(), result.get().size())
|
||||
return None
|
||||
|
||||
def get_actor_table(self, job_id, actor_state_name):
|
||||
cdef c_vector[c_string] result
|
||||
cdef optional[CActorID] cactor_id = nullopt
|
||||
cdef optional[CJobID] cjob_id
|
||||
cdef optional[c_string] cactor_state_name
|
||||
cdef c_string c_name
|
||||
if job_id is not None:
|
||||
cjob_id = make_optional[CJobID](CJobID.FromBinary(job_id.binary()))
|
||||
if actor_state_name is not None:
|
||||
c_name = actor_state_name
|
||||
cactor_state_name = make_optional[c_string](c_name)
|
||||
with nogil:
|
||||
result = self.inner.get().GetAllActorInfo(
|
||||
cactor_id, cjob_id, cactor_state_name)
|
||||
return result
|
||||
|
||||
def get_actor_info(self, actor_id):
|
||||
cdef unique_ptr[c_string] actor_info
|
||||
cdef CActorID cactor_id = CActorID.FromBinary(actor_id.binary())
|
||||
with nogil:
|
||||
actor_info = self.inner.get().GetActorInfo(cactor_id)
|
||||
if actor_info:
|
||||
return c_string(actor_info.get().data(), actor_info.get().size())
|
||||
return None
|
||||
|
||||
def get_worker_table(self):
|
||||
cdef c_vector[c_string] result
|
||||
with nogil:
|
||||
result = self.inner.get().GetAllWorkerInfo()
|
||||
return result
|
||||
|
||||
def get_worker_info(self, worker_id):
|
||||
cdef unique_ptr[c_string] worker_info
|
||||
cdef CWorkerID cworker_id = <CWorkerID>CUniqueID.FromBinary(worker_id.binary())
|
||||
with nogil:
|
||||
worker_info = self.inner.get().GetWorkerInfo(cworker_id)
|
||||
if worker_info:
|
||||
return c_string(worker_info.get().data(), worker_info.get().size())
|
||||
return None
|
||||
|
||||
def add_worker_info(self, serialized_string):
|
||||
cdef c_bool result
|
||||
cdef c_string cserialized_string = serialized_string
|
||||
with nogil:
|
||||
result = self.inner.get().AddWorkerInfo(cserialized_string)
|
||||
return result
|
||||
|
||||
def get_worker_debugger_port(self, worker_id):
|
||||
cdef c_uint32_t result
|
||||
cdef CWorkerID cworker_id = <CWorkerID>CUniqueID.FromBinary(worker_id.binary())
|
||||
with nogil:
|
||||
result = self.inner.get().GetWorkerDebuggerPort(cworker_id)
|
||||
return result
|
||||
|
||||
def update_worker_debugger_port(self, worker_id, debugger_port):
|
||||
cdef c_bool result
|
||||
cdef CWorkerID cworker_id = <CWorkerID>CUniqueID.FromBinary(worker_id.binary())
|
||||
cdef c_uint32_t cdebugger_port = debugger_port
|
||||
with nogil:
|
||||
result = self.inner.get().UpdateWorkerDebuggerPort(
|
||||
cworker_id,
|
||||
cdebugger_port)
|
||||
return result
|
||||
|
||||
def update_worker_num_paused_threads(self, worker_id, num_paused_threads_delta):
|
||||
cdef c_bool result
|
||||
cdef CWorkerID cworker_id = <CWorkerID>CUniqueID.FromBinary(worker_id.binary())
|
||||
cdef c_int32_t cnum_paused_threads_delta = num_paused_threads_delta
|
||||
|
||||
with nogil:
|
||||
result = self.inner.get().UpdateWorkerNumPausedThreads(
|
||||
cworker_id, cnum_paused_threads_delta)
|
||||
return result
|
||||
|
||||
def get_placement_group_table(self):
|
||||
cdef c_vector[c_string] result
|
||||
with nogil:
|
||||
result = self.inner.get().GetAllPlacementGroupInfo()
|
||||
return result
|
||||
|
||||
def get_placement_group_info(self, placement_group_id):
|
||||
cdef unique_ptr[c_string] result
|
||||
cdef CPlacementGroupID cplacement_group_id = (
|
||||
CPlacementGroupID.FromBinary(placement_group_id.binary()))
|
||||
with nogil:
|
||||
result = self.inner.get().GetPlacementGroupInfo(
|
||||
cplacement_group_id)
|
||||
if result:
|
||||
return c_string(result.get().data(), result.get().size())
|
||||
return None
|
||||
|
||||
def get_placement_group_by_name(self, placement_group_name, ray_namespace):
|
||||
cdef unique_ptr[c_string] result
|
||||
cdef c_string cplacement_group_name = placement_group_name
|
||||
cdef c_string cray_namespace = ray_namespace
|
||||
with nogil:
|
||||
result = self.inner.get().GetPlacementGroupByName(
|
||||
cplacement_group_name, cray_namespace)
|
||||
if result:
|
||||
return c_string(result.get().data(), result.get().size())
|
||||
return None
|
||||
|
||||
def get_system_config(self):
|
||||
return self.inner.get().GetSystemConfig()
|
||||
|
||||
def get_node(self, node_id):
|
||||
cdef CRayStatus status
|
||||
cdef c_string cnode_id = node_id
|
||||
cdef c_string cnode_info_str
|
||||
cdef CGcsNodeInfo c_node_info
|
||||
with nogil:
|
||||
status = self.inner.get().GetNode(cnode_id, &cnode_info_str)
|
||||
if not status.ok():
|
||||
raise RuntimeError(status.message())
|
||||
c_node_info.ParseFromString(cnode_info_str)
|
||||
c_labels = PythonGetNodeLabels(c_node_info)
|
||||
return {
|
||||
"object_store_socket_name": c_node_info.object_store_socket_name().decode(),
|
||||
"raylet_socket_name": c_node_info.raylet_socket_name().decode(),
|
||||
"node_manager_port": c_node_info.node_manager_port(),
|
||||
"node_id": c_node_info.node_id().hex(),
|
||||
"runtime_env_agent_port": c_node_info.runtime_env_agent_port(),
|
||||
"metrics_agent_port": c_node_info.metrics_agent_port(),
|
||||
"metrics_export_port": c_node_info.metrics_export_port(),
|
||||
"dashboard_agent_listen_port": c_node_info.dashboard_agent_listen_port(),
|
||||
"labels": {key.decode(): value.decode() for key, value in c_labels},
|
||||
}
|
||||
@@ -0,0 +1,506 @@
|
||||
# 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()
|
||||
@@ -0,0 +1,48 @@
|
||||
from libcpp.string cimport string as c_string
|
||||
|
||||
from ray.includes.libcoreworker cimport CProfileEvent
|
||||
|
||||
import json
|
||||
import traceback
|
||||
|
||||
cdef class ProfileEvent:
|
||||
"""Cython wrapper class of C++ `ray::core::worker::ProfileEvent`."""
|
||||
cdef:
|
||||
unique_ptr[CProfileEvent] inner
|
||||
object extra_data
|
||||
|
||||
@staticmethod
|
||||
cdef make(unique_ptr[CProfileEvent] event, object extra_data):
|
||||
cdef ProfileEvent self = ProfileEvent.__new__(ProfileEvent)
|
||||
self.inner = move(event)
|
||||
self.extra_data = extra_data
|
||||
return self
|
||||
|
||||
def set_extra_data(self, c_string extra_data):
|
||||
self.inner.get().SetExtraData(extra_data)
|
||||
|
||||
def __enter__(self):
|
||||
pass
|
||||
|
||||
def __exit__(self, type, value, tb):
|
||||
extra_data = None
|
||||
if type is not None:
|
||||
extra_data = {
|
||||
"type": str(type),
|
||||
"value": str(value),
|
||||
"traceback": str(traceback.format_exc()),
|
||||
}
|
||||
elif self.extra_data is not None:
|
||||
extra_data = self.extra_data
|
||||
|
||||
if not extra_data:
|
||||
self.inner.get().SetExtraData(b"{}")
|
||||
elif isinstance(extra_data, dict):
|
||||
self.inner.get().SetExtraData(
|
||||
json.dumps(extra_data).encode("ascii"))
|
||||
else:
|
||||
self.inner.get().SetExtraData(extra_data)
|
||||
|
||||
# Deleting the CProfileEvent will add it to a queue to be pushed to
|
||||
# the driver.
|
||||
self.inner.reset()
|
||||
@@ -0,0 +1,39 @@
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
from libcpp.pair cimport pair as c_pair
|
||||
|
||||
cdef extern from "ray/stats/metric.h" nogil:
|
||||
cdef cppclass CMetric "ray::stats::Metric":
|
||||
CMetric(const c_string &name,
|
||||
const c_string &description,
|
||||
const c_string &unit,
|
||||
const c_vector[c_string] &tag_keys)
|
||||
c_string GetName() const
|
||||
void Record(double value)
|
||||
void RecordForCython(double value,
|
||||
c_vector[c_pair[c_string, c_string]] tags)
|
||||
|
||||
cdef cppclass CGauge "ray::stats::Gauge":
|
||||
CGauge(const c_string &name,
|
||||
const c_string &description,
|
||||
const c_string &unit,
|
||||
const c_vector[c_string] &tag_keys)
|
||||
|
||||
cdef cppclass CCount "ray::stats::Count":
|
||||
CCount(const c_string &name,
|
||||
const c_string &description,
|
||||
const c_string &unit,
|
||||
const c_vector[c_string] &tag_keys)
|
||||
|
||||
cdef cppclass CSum "ray::stats::Sum":
|
||||
CSum(const c_string &name,
|
||||
const c_string &description,
|
||||
const c_string &unit,
|
||||
const c_vector[c_string] &tag_keys)
|
||||
|
||||
cdef cppclass CHistogram "ray::stats::Histogram":
|
||||
CHistogram(const c_string &name,
|
||||
const c_string &description,
|
||||
const c_string &unit,
|
||||
const c_vector[double] &boundaries,
|
||||
const c_vector[c_string] &tag_keys)
|
||||
@@ -0,0 +1,205 @@
|
||||
from ray.includes.metric cimport (
|
||||
CCount,
|
||||
CGauge,
|
||||
CHistogram,
|
||||
CSum,
|
||||
CMetric,
|
||||
)
|
||||
from libcpp.utility cimport move
|
||||
from libcpp.memory cimport unique_ptr
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
from libcpp.pair cimport pair as c_pair
|
||||
|
||||
cdef class Metric:
|
||||
"""Cython wrapper class of C++ `ray::stats::Metric`.
|
||||
|
||||
It's an abstract class of all metric types.
|
||||
"""
|
||||
cdef:
|
||||
unique_ptr[CMetric] metric
|
||||
c_vector[c_string] c_tag_keys
|
||||
|
||||
def __init__(self, tag_keys):
|
||||
for tag_key in tag_keys:
|
||||
self.c_tag_keys.push_back(tag_key.encode("ascii"))
|
||||
|
||||
def record(self, value, tags=None):
|
||||
"""Record a measurement of metric.
|
||||
|
||||
Flush a metric raw point to stats module with a key-value dict tags.
|
||||
Args:
|
||||
value (double): metric name.
|
||||
tags (dict): default none.
|
||||
"""
|
||||
cdef c_vector[c_pair[c_string, c_string]] c_tags
|
||||
cdef double c_value
|
||||
# Default tags will be exported if it's empty map.
|
||||
if tags:
|
||||
c_tags.reserve(len(tags))
|
||||
for tag_k, tag_v in tags.items():
|
||||
if tag_v is not None:
|
||||
c_tags.push_back(c_pair[c_string, c_string](
|
||||
tag_k.encode("ascii"),
|
||||
tag_v.encode("ascii")))
|
||||
c_value = value
|
||||
with nogil:
|
||||
self.metric.get().RecordForCython(c_value, move(c_tags))
|
||||
|
||||
def get_name(self):
|
||||
return self.metric.get().GetName()
|
||||
|
||||
|
||||
cdef class Gauge(Metric):
|
||||
"""Cython wrapper class of C++ `ray::stats::Gauge`.
|
||||
|
||||
Gauge: Keeps the last recorded value, drops everything before.
|
||||
|
||||
Example:
|
||||
|
||||
>>> gauge = Gauge(
|
||||
"ray.worker.metric",
|
||||
"description",
|
||||
["tagk1", "tagk2"]).
|
||||
value = 5
|
||||
key1= "key1"
|
||||
key2 = "key2"
|
||||
gauge.record(value, {"tagk1": key1, "tagk2": key2})
|
||||
"""
|
||||
def __init__(self, name, description, tag_keys):
|
||||
"""Create a gauge metric
|
||||
|
||||
Args:
|
||||
name (string): metric name.
|
||||
description (string): description of this metric.
|
||||
tag_keys (list): a list of tay keys in string format.
|
||||
"""
|
||||
super().__init__(tag_keys)
|
||||
|
||||
self.metric.reset(
|
||||
new CGauge(
|
||||
name.encode("ascii"),
|
||||
description.encode("ascii"),
|
||||
b"", # Unit, unused.
|
||||
self.c_tag_keys
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
cdef class Count(Metric):
|
||||
"""Cython wrapper class of C++ `ray::stats::Count`.
|
||||
|
||||
Example:
|
||||
|
||||
>>> count = Count(
|
||||
"ray.worker.metric",
|
||||
"description",
|
||||
["tagk1", "tagk2"]).
|
||||
value = 5
|
||||
key1= "key1"
|
||||
key2 = "key2"
|
||||
|
||||
count.record(value, {"tagk1": key1, "tagk2": key2})
|
||||
|
||||
Count: The count of the number of metric points.
|
||||
"""
|
||||
def __init__(self, name, description, tag_keys):
|
||||
"""Create a count metric
|
||||
|
||||
Args:
|
||||
name (string): metric name.
|
||||
description (string): description of this metric.
|
||||
tag_keys (list): a list of tay keys in string format.
|
||||
"""
|
||||
super().__init__(tag_keys)
|
||||
|
||||
self.metric.reset(
|
||||
new CCount(
|
||||
name.encode("ascii"),
|
||||
description.encode("ascii"),
|
||||
b"", # Unit, unused.
|
||||
self.c_tag_keys
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
cdef class Sum(Metric):
|
||||
"""Cython wrapper class of C++ `ray::stats::Sum`.
|
||||
|
||||
Example:
|
||||
|
||||
>>> metric_sum = Sum(
|
||||
"ray.worker.metric",
|
||||
"description",
|
||||
["tagk1", "tagk2"]).
|
||||
value = 5
|
||||
key1= "key1"
|
||||
key2 = "key2"
|
||||
|
||||
metric_sum.record(value, {"tagk1": key1, "tagk2": key2})
|
||||
|
||||
Sum: A sum up of the metric points.
|
||||
"""
|
||||
def __init__(self, name, description, tag_keys):
|
||||
"""Create a sum metric
|
||||
|
||||
Args:
|
||||
name (string): metric name.
|
||||
description (string): description of this metric.
|
||||
tag_keys (list): a list of tay keys in string format.
|
||||
"""
|
||||
|
||||
super().__init__(tag_keys)
|
||||
|
||||
self.metric.reset(
|
||||
new CSum(
|
||||
name.encode("ascii"),
|
||||
description.encode("ascii"),
|
||||
b"", # Unit, unused.
|
||||
self.c_tag_keys
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
cdef class Histogram(Metric):
|
||||
"""Cython wrapper class of C++ `ray::stats::Histogram`.
|
||||
|
||||
Example:
|
||||
|
||||
>>> histogram = Histogram(
|
||||
"ray.worker.histogram1",
|
||||
"description",
|
||||
[1.0, 2.0], # boundaries.
|
||||
["tagk1"])
|
||||
value = 5
|
||||
key1= "key1"
|
||||
|
||||
histogram.record(value, {"tagk1": key1})
|
||||
|
||||
Histogram: Histogram distribution of metric points.
|
||||
"""
|
||||
def __init__(self, name, description, boundaries, tag_keys):
|
||||
"""Create a sum metric
|
||||
|
||||
Args:
|
||||
name (string): metric name.
|
||||
description (string): description of this metric.
|
||||
boundaries (list): a double type list boundaries of histogram.
|
||||
tag_keys (list): a list of tay key in string format.
|
||||
"""
|
||||
|
||||
super().__init__(tag_keys)
|
||||
|
||||
cdef c_vector[double] c_boundaries
|
||||
for value in boundaries:
|
||||
c_boundaries.push_back(value)
|
||||
|
||||
self.metric.reset(
|
||||
new CHistogram(
|
||||
name.encode("ascii"),
|
||||
description.encode("ascii"),
|
||||
b"", # Unit, unused.
|
||||
c_boundaries,
|
||||
self.c_tag_keys
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,15 @@
|
||||
from libc.stddef cimport size_t
|
||||
from libcpp.string cimport string
|
||||
from libcpp cimport bool
|
||||
from ray.includes.array cimport array_string_2
|
||||
from ray.includes.optional cimport optional
|
||||
|
||||
cdef extern from "ray/util/network_util.h" namespace "ray":
|
||||
string BuildAddress(const string &host, int port)
|
||||
string BuildAddress(const string &host, const string &port)
|
||||
optional[array_string_2] ParseAddress(const string &address)
|
||||
string GetNodeIpAddressFromPerspective(const optional[string] &address)
|
||||
bool IsIPv6(const string &host)
|
||||
string GetLocalhostIP()
|
||||
string GetAllInterfacesIP()
|
||||
bool IsLocalhost(const string &address)
|
||||
@@ -0,0 +1,123 @@
|
||||
from ray.includes.network_util cimport (
|
||||
BuildAddress,
|
||||
ParseAddress,
|
||||
GetNodeIpAddressFromPerspective,
|
||||
IsIPv6,
|
||||
GetLocalhostIP,
|
||||
GetAllInterfacesIP,
|
||||
IsLocalhost,
|
||||
array_string_2,
|
||||
optional,
|
||||
)
|
||||
from libcpp.string cimport string
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
|
||||
def parse_address(address: str) -> Optional[Tuple[str, str]]:
|
||||
"""Parse a network address string into host and port.
|
||||
|
||||
Args:
|
||||
address: The address string to parse (e.g., "localhost:8000", "[::1]:8000").
|
||||
|
||||
Returns:
|
||||
Tuple with (host, port) if port found, None if no colon separator.
|
||||
"""
|
||||
cdef optional[array_string_2] res = ParseAddress(address.encode('utf-8'))
|
||||
if not res.has_value():
|
||||
return None
|
||||
|
||||
cdef array_string_2 ip_port = res.value()
|
||||
return (ip_port[0].decode('utf-8'), ip_port[1].decode('utf-8'))
|
||||
|
||||
|
||||
def build_address(host: str, port: Union[int, str]) -> str:
|
||||
"""Build a network address string from host and port.
|
||||
|
||||
Args:
|
||||
host: The hostname or IP address.
|
||||
port: The port number (int or string).
|
||||
|
||||
Returns:
|
||||
Formatted address string (e.g., "localhost:8000" or "[::1]:8000").
|
||||
"""
|
||||
cdef string host_c = host.encode('utf-8')
|
||||
cdef string result
|
||||
cdef string port_c
|
||||
|
||||
if isinstance(port, int):
|
||||
result = BuildAddress(host_c, <int>port)
|
||||
else:
|
||||
port_c = str(port).encode('utf-8')
|
||||
result = BuildAddress(host_c, port_c)
|
||||
|
||||
return result.decode('utf-8')
|
||||
|
||||
|
||||
def node_ip_address_from_perspective(address=None) -> str:
|
||||
"""IP address by which the local node can be reached from the address.
|
||||
|
||||
If no address is given, defaults to public DNS servers for detection.
|
||||
|
||||
Args:
|
||||
address: The IP address and port of any known live service on the
|
||||
network you care about.
|
||||
|
||||
Returns:
|
||||
The IP address by which the local node can be reached from the address.
|
||||
"""
|
||||
cdef string node_ip
|
||||
cdef optional[string] address_c
|
||||
cdef string address_str
|
||||
if address is not None:
|
||||
address_str = address.encode('utf-8')
|
||||
address_c = optional[string](address_str)
|
||||
else:
|
||||
address_c = optional[string]()
|
||||
node_ip = GetNodeIpAddressFromPerspective(address_c)
|
||||
return node_ip.decode('utf-8')
|
||||
|
||||
|
||||
def is_ipv6(host: str) -> bool:
|
||||
"""Check if a host is resolved to IPv6.
|
||||
|
||||
Args:
|
||||
host: The IP or domain name to check (must be without port).
|
||||
|
||||
Returns:
|
||||
True if the host is resolved to IPv6, False if IPv4.
|
||||
"""
|
||||
cdef string host_c = host.encode('utf-8')
|
||||
return IsIPv6(host_c)
|
||||
|
||||
|
||||
def get_localhost_ip() -> str:
|
||||
"""Get localhost loopback IP with IPv4/IPv6 support.
|
||||
|
||||
Returns:
|
||||
"127.0.0.1" for IPv4 or "::1" for IPv6.
|
||||
"""
|
||||
cdef string result = GetLocalhostIP()
|
||||
return result.decode('utf-8')
|
||||
|
||||
|
||||
def get_all_interfaces_ip() -> str:
|
||||
"""Get the IP address to bind to all network interfaces.
|
||||
|
||||
Returns "0.0.0.0" for IPv4 or "::" for IPv6, depending on the system's
|
||||
localhost resolution.
|
||||
"""
|
||||
cdef string result = GetAllInterfacesIP()
|
||||
return result.decode('utf-8')
|
||||
|
||||
|
||||
def is_localhost(host: str) -> bool:
|
||||
"""Check if the given host string represents a localhost address.
|
||||
|
||||
Args:
|
||||
host: The hostname or IP address to check.
|
||||
|
||||
Returns:
|
||||
True if the host is a localhost address (127.0.0.1, ::1, or "localhost").
|
||||
"""
|
||||
cdef string host_c = host.encode('utf-8')
|
||||
return IsLocalhost(host_c)
|
||||
@@ -0,0 +1,182 @@
|
||||
from ray.includes.unique_ids cimport CObjectID
|
||||
from ray.includes.optional cimport (
|
||||
optional,
|
||||
nullopt,
|
||||
)
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import functools
|
||||
import logging
|
||||
import threading
|
||||
from typing import Callable, Any, Union, Optional
|
||||
from _collections_abc import GenericAlias
|
||||
from builtins import StopAsyncIteration
|
||||
|
||||
import ray
|
||||
import cython
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _set_future_helper(
|
||||
result: Any,
|
||||
*,
|
||||
py_future: Union[asyncio.Future, concurrent.futures.Future],
|
||||
):
|
||||
# Issue #11030, #8841
|
||||
# If this future has result set already, we just need to
|
||||
# skip the set result/exception procedure.
|
||||
if py_future.done():
|
||||
return
|
||||
|
||||
if isinstance(result, RayTaskError):
|
||||
exc = result.as_instanceof_cause()
|
||||
# Convert StopIteration to RuntimeError to prevent segfaults due to
|
||||
# Cpython's behavior w.r.t. PEP479
|
||||
if isinstance(exc, StopIteration) or isinstance(exc, StopAsyncIteration):
|
||||
runtime_error = RuntimeError(f"generator raised {type(exc).__name__}")
|
||||
runtime_error.__cause__ = exc
|
||||
py_future.set_exception(runtime_error)
|
||||
else:
|
||||
py_future.set_exception(exc)
|
||||
elif isinstance(result, RayError):
|
||||
# Directly raise exception for RayActorError
|
||||
py_future.set_exception(result)
|
||||
else:
|
||||
py_future.set_result(result)
|
||||
|
||||
|
||||
cdef class ObjectRef(BaseID):
|
||||
__class_getitem__ = classmethod(GenericAlias) # should match how typing.Generic works
|
||||
|
||||
def __cinit__(self):
|
||||
self.in_core_worker = False
|
||||
|
||||
def __init__(
|
||||
self, id, owner_addr="", call_site_data="",
|
||||
skip_adding_local_ref=False, tensor_transport: Optional[str] = None):
|
||||
self._set_id(id)
|
||||
self.owner_addr = owner_addr
|
||||
self.in_core_worker = False
|
||||
self.call_site_data = call_site_data
|
||||
self._tensor_transport = tensor_transport
|
||||
worker = ray._private.worker.global_worker
|
||||
# TODO(edoakes): We should be able to remove the in_core_worker flag.
|
||||
# But there are still some dummy object refs being created outside the
|
||||
# context of a core worker.
|
||||
if hasattr(worker, "core_worker"):
|
||||
if not skip_adding_local_ref:
|
||||
worker.core_worker.add_object_ref_reference(self)
|
||||
self.in_core_worker = True
|
||||
|
||||
def __dealloc__(self):
|
||||
if self.in_core_worker:
|
||||
try:
|
||||
worker = ray._private.worker.global_worker
|
||||
worker.core_worker.remove_object_ref_reference(self)
|
||||
except Exception as e:
|
||||
# There is a strange error in rllib that causes the above to
|
||||
# fail. Somehow the global 'ray' variable corresponding to the
|
||||
# imported package is None when this gets called. Unfortunately
|
||||
# this is hard to debug because __dealloc__ is called during
|
||||
# garbage collection so we can't get a good stack trace. In any
|
||||
# case, there's not much we can do besides ignore it
|
||||
# (re-importing ray won't help).
|
||||
pass
|
||||
|
||||
cdef CObjectID native(self):
|
||||
return <CObjectID>self.data
|
||||
|
||||
def binary(self):
|
||||
return self.data.Binary()
|
||||
|
||||
def hex(self):
|
||||
return decode(self.data.Hex())
|
||||
|
||||
def is_nil(self):
|
||||
return self.data.IsNil()
|
||||
|
||||
cdef size_t hash(self):
|
||||
return self.data.Hash()
|
||||
|
||||
def task_id(self):
|
||||
return TaskID(self.data.TaskId().Binary())
|
||||
|
||||
def job_id(self):
|
||||
return self.task_id().job_id()
|
||||
|
||||
def owner_address(self):
|
||||
return self.owner_addr
|
||||
|
||||
def call_site(self):
|
||||
return decode(self.call_site_data)
|
||||
|
||||
@classmethod
|
||||
def size(cls):
|
||||
return CObjectID.Size()
|
||||
|
||||
def _set_id(self, id):
|
||||
check_id(id)
|
||||
self.data = CObjectID.FromBinary(<c_string>id)
|
||||
|
||||
@classmethod
|
||||
def nil(cls):
|
||||
return cls(CObjectID.Nil().Binary())
|
||||
|
||||
@classmethod
|
||||
def from_random(cls):
|
||||
return cls(CObjectID.FromRandom().Binary())
|
||||
|
||||
def future(self) -> concurrent.futures.Future:
|
||||
"""Wrap ObjectRef with a concurrent.futures.Future
|
||||
|
||||
Note that the future cancellation will not cancel the correspoding
|
||||
task when the ObjectRef representing return object of a task.
|
||||
Additionally, future.running() will always be ``False`` even if the
|
||||
underlying task is running.
|
||||
"""
|
||||
py_future = concurrent.futures.Future()
|
||||
|
||||
self._on_completed(
|
||||
functools.partial(_set_future_helper, py_future=py_future))
|
||||
|
||||
# A hack to keep a reference to the object ref for ref counting.
|
||||
py_future.object_ref = self
|
||||
return py_future
|
||||
|
||||
def __await__(self):
|
||||
return self.as_future(_internal=True).__await__()
|
||||
|
||||
def as_future(self, _internal=False) -> asyncio.Future:
|
||||
"""Wrap ObjectRef with an asyncio.Future.
|
||||
|
||||
Note that the future cancellation will not cancel the correspoding
|
||||
task when the ObjectRef representing return object of a task.
|
||||
"""
|
||||
if not _internal:
|
||||
logger.warning("ref.as_future() is deprecated in favor of "
|
||||
"asyncio.wrap_future(ref.future()).")
|
||||
return asyncio.wrap_future(self.future())
|
||||
|
||||
def _on_completed(self, py_callback: Callable[[Any], None]):
|
||||
"""Register a callback that will be called after Object is ready.
|
||||
If the ObjectRef is already ready, the callback will be called soon.
|
||||
The callback should take the result as the only argument. The result
|
||||
can be an exception object in case of task error.
|
||||
"""
|
||||
core_worker = ray._private.worker.global_worker.core_worker
|
||||
core_worker.set_get_async_callback(self, py_callback)
|
||||
return self
|
||||
|
||||
def tensor_transport(self):
|
||||
return self._tensor_transport
|
||||
|
||||
cdef optional[c_string] c_tensor_transport(self):
|
||||
cdef:
|
||||
optional[c_string] c_tensor_transport = nullopt
|
||||
c_string c_tensor_transport_str
|
||||
if self._tensor_transport is not None:
|
||||
c_tensor_transport_str = self._tensor_transport.encode("utf-8")
|
||||
c_tensor_transport.emplace(move(c_tensor_transport_str))
|
||||
return c_tensor_transport
|
||||
@@ -0,0 +1,74 @@
|
||||
# source: object_ref.pxi
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
from typing import Any, Awaitable, Callable, Generator, Optional, TypeVar, Union
|
||||
|
||||
from ray.includes.unique_ids import BaseID, JobID, TaskID
|
||||
|
||||
_T = TypeVar("_T")
|
||||
def _set_future_helper(
|
||||
result: _T,
|
||||
*,
|
||||
py_future: Union[asyncio.Future[_T], concurrent.futures.Future[_T]],
|
||||
) -> None: ...
|
||||
|
||||
|
||||
_OR = TypeVar("_OR", bound=ObjectRef)
|
||||
class ObjectRef(BaseID, Awaitable[_T]):
|
||||
|
||||
|
||||
def __init__(
|
||||
self, id: bytes, owner_addr: str = "", call_site_data: str = "",
|
||||
skip_adding_local_ref: bool = False, tensor_transport: Optional[str] = None) -> None: ...
|
||||
|
||||
def __dealloc__(self) -> None: ...
|
||||
|
||||
|
||||
def task_id(self) -> TaskID: ...
|
||||
|
||||
def job_id(self) -> JobID: ...
|
||||
|
||||
def owner_address(self) -> str: ...
|
||||
|
||||
def call_site(self) -> str: ...
|
||||
|
||||
@classmethod
|
||||
def size(cls) -> int: ...
|
||||
|
||||
def _set_id(self, id: bytes) -> None: ...
|
||||
|
||||
@classmethod
|
||||
def nil(cls: type[_OR]) -> _OR: ...
|
||||
|
||||
@classmethod
|
||||
def from_random(cls: type[_OR]) -> _OR: ...
|
||||
|
||||
def future(self) -> concurrent.futures.Future[_T]:
|
||||
"""Wrap ObjectRef with a concurrent.futures.Future
|
||||
|
||||
Note that the future cancellation will not cancel the correspoding
|
||||
task when the ObjectRef representing return object of a task.
|
||||
Additionally, future.running() will always be ``False`` even if the
|
||||
underlying task is running.
|
||||
"""
|
||||
...
|
||||
|
||||
def __await__(self) -> Generator[Any, None, _T]: ...
|
||||
|
||||
def as_future(self, _internal=False) -> asyncio.Future[_T]:
|
||||
"""Wrap ObjectRef with an asyncio.Future.
|
||||
|
||||
Note that the future cancellation will not cancel the correspoding
|
||||
task when the ObjectRef representing return object of a task.
|
||||
"""
|
||||
...
|
||||
|
||||
def _on_completed(self, py_callback: Callable[[_T], None]):
|
||||
"""Register a callback that will be called after Object is ready.
|
||||
If the ObjectRef is already ready, the callback will be called soon.
|
||||
The callback should take the result as the only argument. The result
|
||||
can be an exception object in case of task error.
|
||||
"""
|
||||
...
|
||||
|
||||
def tensor_transport(self) -> int: ...
|
||||
@@ -0,0 +1,36 @@
|
||||
# Currently Cython does not support std::optional.
|
||||
# See: https://github.com/cython/cython/pull/3294
|
||||
from libcpp cimport bool
|
||||
|
||||
cdef extern from "<optional>" namespace "std" nogil:
|
||||
cdef cppclass nullopt_t:
|
||||
nullopt_t()
|
||||
|
||||
cdef nullopt_t nullopt
|
||||
|
||||
cdef cppclass optional[T]:
|
||||
ctypedef T value_type
|
||||
optional()
|
||||
optional(nullopt_t)
|
||||
optional(optional&) except +
|
||||
optional(T&) except +
|
||||
bool has_value()
|
||||
T& value()
|
||||
T& value_or[U](U& default_value)
|
||||
void swap(optional&)
|
||||
void reset()
|
||||
T& emplace(...)
|
||||
T& operator*()
|
||||
# T* operator->() # Not Supported
|
||||
optional& operator=(optional&)
|
||||
optional& operator=[U](U&)
|
||||
bool operator bool()
|
||||
bool operator!()
|
||||
bool operator==[U](optional&, U&)
|
||||
bool operator!=[U](optional&, U&)
|
||||
bool operator<[U](optional&, U&)
|
||||
bool operator>[U](optional&, U&)
|
||||
bool operator<=[U](optional&, U&)
|
||||
bool operator>=[U](optional&, U&)
|
||||
|
||||
optional[T] make_optional[T](...) except +
|
||||
@@ -0,0 +1,88 @@
|
||||
from libcpp cimport bool as c_bool
|
||||
from libc.stdint cimport int64_t, uint64_t, uint32_t
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp.unordered_map cimport unordered_map
|
||||
|
||||
|
||||
cdef extern from "ray/common/ray_config.h" nogil:
|
||||
cdef cppclass RayConfig "RayConfig":
|
||||
@staticmethod
|
||||
RayConfig &instance()
|
||||
|
||||
void initialize(const c_string& config_list)
|
||||
|
||||
int64_t ray_cookie() const
|
||||
|
||||
int64_t handler_warning_timeout_ms() const
|
||||
|
||||
int64_t debug_dump_period_milliseconds() const
|
||||
|
||||
int64_t object_timeout_milliseconds() const
|
||||
|
||||
int64_t raylet_client_num_connect_attempts() const
|
||||
|
||||
int64_t raylet_client_connect_timeout_milliseconds() const
|
||||
|
||||
int64_t kill_worker_timeout_milliseconds() const
|
||||
|
||||
int64_t worker_register_timeout_seconds() const
|
||||
|
||||
int64_t redis_db_connect_retries()
|
||||
|
||||
int64_t redis_db_connect_wait_milliseconds() const
|
||||
|
||||
int object_manager_pull_timeout_ms() const
|
||||
|
||||
int object_manager_push_timeout_ms() const
|
||||
|
||||
uint32_t maximum_gcs_deletion_batch_size() const
|
||||
|
||||
int64_t max_direct_call_object_size() const
|
||||
|
||||
int64_t task_rpc_inlined_bytes_limit() const
|
||||
|
||||
uint64_t metrics_report_interval_ms() const
|
||||
|
||||
c_bool enable_timeline() const
|
||||
|
||||
uint32_t max_grpc_message_size() const
|
||||
|
||||
c_bool record_ref_creation_sites() const
|
||||
|
||||
c_string REDIS_CA_CERT() const
|
||||
|
||||
c_string REDIS_CA_PATH() const
|
||||
|
||||
c_string REDIS_CLIENT_CERT() const
|
||||
|
||||
c_string REDIS_CLIENT_KEY() const
|
||||
|
||||
c_string REDIS_SERVER_NAME() const
|
||||
|
||||
int64_t health_check_initial_delay_ms() const
|
||||
|
||||
int64_t health_check_period_ms() const
|
||||
|
||||
int64_t health_check_timeout_ms() const
|
||||
|
||||
int64_t health_check_failure_threshold() const
|
||||
|
||||
int64_t grpc_client_keepalive_time_ms() const
|
||||
|
||||
int64_t grpc_client_keepalive_timeout_ms() const
|
||||
|
||||
c_bool enable_autoscaler_v2() const
|
||||
|
||||
c_string predefined_unit_instance_resources() const
|
||||
|
||||
c_string custom_unit_instance_resources() const
|
||||
|
||||
int64_t nums_py_gcs_reconnect_retry() const
|
||||
|
||||
int64_t py_gcs_connect_timeout_s() const
|
||||
|
||||
int maximum_gcs_destroyed_actor_cached_count() const
|
||||
|
||||
c_bool record_task_actor_creation_sites() const
|
||||
|
||||
c_bool start_python_gc_manager_thread() const
|
||||
@@ -0,0 +1,141 @@
|
||||
from libcpp.string cimport string as c_string
|
||||
from ray.includes.ray_config cimport RayConfig
|
||||
|
||||
cdef class Config:
|
||||
@staticmethod
|
||||
def initialize(c_string config_list):
|
||||
return RayConfig.instance().initialize(config_list)
|
||||
|
||||
@staticmethod
|
||||
def ray_cookie():
|
||||
return RayConfig.instance().ray_cookie()
|
||||
|
||||
@staticmethod
|
||||
def handler_warning_timeout_ms():
|
||||
return RayConfig.instance().handler_warning_timeout_ms()
|
||||
|
||||
@staticmethod
|
||||
def debug_dump_period_milliseconds():
|
||||
return RayConfig.instance().debug_dump_period_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def object_timeout_milliseconds():
|
||||
return (RayConfig.instance()
|
||||
.object_timeout_milliseconds())
|
||||
|
||||
@staticmethod
|
||||
def raylet_client_num_connect_attempts():
|
||||
return RayConfig.instance().raylet_client_num_connect_attempts()
|
||||
|
||||
@staticmethod
|
||||
def raylet_client_connect_timeout_milliseconds():
|
||||
return (RayConfig.instance()
|
||||
.raylet_client_connect_timeout_milliseconds())
|
||||
|
||||
@staticmethod
|
||||
def kill_worker_timeout_milliseconds():
|
||||
return RayConfig.instance().kill_worker_timeout_milliseconds()
|
||||
|
||||
@staticmethod
|
||||
def worker_register_timeout_seconds():
|
||||
return RayConfig.instance().worker_register_timeout_seconds()
|
||||
|
||||
@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 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 maximum_gcs_deletion_batch_size():
|
||||
return RayConfig.instance().maximum_gcs_deletion_batch_size()
|
||||
|
||||
@staticmethod
|
||||
def metrics_report_interval_ms():
|
||||
return RayConfig.instance().metrics_report_interval_ms()
|
||||
|
||||
@staticmethod
|
||||
def enable_timeline():
|
||||
return RayConfig.instance().enable_timeline()
|
||||
|
||||
@staticmethod
|
||||
def max_grpc_message_size():
|
||||
return RayConfig.instance().max_grpc_message_size()
|
||||
|
||||
@staticmethod
|
||||
def record_ref_creation_sites():
|
||||
return RayConfig.instance().record_ref_creation_sites()
|
||||
|
||||
@staticmethod
|
||||
def REDIS_CA_CERT():
|
||||
return RayConfig.instance().REDIS_CA_CERT()
|
||||
|
||||
@staticmethod
|
||||
def REDIS_CA_PATH():
|
||||
return RayConfig.instance().REDIS_CA_PATH()
|
||||
|
||||
@staticmethod
|
||||
def REDIS_CLIENT_CERT():
|
||||
return RayConfig.instance().REDIS_CLIENT_CERT()
|
||||
|
||||
@staticmethod
|
||||
def REDIS_CLIENT_KEY():
|
||||
return RayConfig.instance().REDIS_CLIENT_KEY()
|
||||
|
||||
@staticmethod
|
||||
def REDIS_SERVER_NAME():
|
||||
return RayConfig.instance().REDIS_SERVER_NAME()
|
||||
|
||||
@staticmethod
|
||||
def health_check_initial_delay_ms():
|
||||
return RayConfig.instance().health_check_initial_delay_ms()
|
||||
|
||||
@staticmethod
|
||||
def health_check_period_ms():
|
||||
return RayConfig.instance().health_check_period_ms()
|
||||
|
||||
@staticmethod
|
||||
def health_check_timeout_ms():
|
||||
return RayConfig.instance().health_check_timeout_ms()
|
||||
|
||||
@staticmethod
|
||||
def health_check_failure_threshold():
|
||||
return RayConfig.instance().health_check_failure_threshold()
|
||||
|
||||
@staticmethod
|
||||
def grpc_client_keepalive_time_ms():
|
||||
return RayConfig.instance().grpc_client_keepalive_time_ms()
|
||||
|
||||
@staticmethod
|
||||
def grpc_client_keepalive_timeout_ms():
|
||||
return RayConfig.instance().grpc_client_keepalive_timeout_ms()
|
||||
|
||||
@staticmethod
|
||||
def enable_autoscaler_v2():
|
||||
return RayConfig.instance().enable_autoscaler_v2()
|
||||
|
||||
@staticmethod
|
||||
def nums_py_gcs_reconnect_retry():
|
||||
return RayConfig.instance().nums_py_gcs_reconnect_retry()
|
||||
|
||||
@staticmethod
|
||||
def py_gcs_connect_timeout_s():
|
||||
return RayConfig.instance().py_gcs_connect_timeout_s()
|
||||
|
||||
@staticmethod
|
||||
def maximum_gcs_destroyed_actor_cached_count():
|
||||
return RayConfig.instance().maximum_gcs_destroyed_actor_cached_count()
|
||||
|
||||
@staticmethod
|
||||
def start_python_gc_manager_thread():
|
||||
return RayConfig.instance().start_python_gc_manager_thread()
|
||||
@@ -0,0 +1,67 @@
|
||||
from asyncio import Future
|
||||
import concurrent.futures
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
from libcpp.string cimport string as c_string
|
||||
from libc.stdint cimport int32_t
|
||||
from libcpp.utility cimport move
|
||||
from libcpp.memory cimport unique_ptr, make_unique, shared_ptr
|
||||
from ray.includes.common cimport (
|
||||
CRayletClientWithIoContext,
|
||||
CRayStatus,
|
||||
CAddress,
|
||||
OptionalItemPyCallback,
|
||||
)
|
||||
from ray.includes.optional cimport optional
|
||||
|
||||
|
||||
cdef convert_optional_vector_int32(
|
||||
CRayStatus status, optional[c_vector[int32_t]] vec) with gil:
|
||||
try:
|
||||
check_status_timeout_as_rpc_error(status)
|
||||
assert vec.has_value()
|
||||
return move(vec.value()), None
|
||||
except Exception as e:
|
||||
return None, e
|
||||
|
||||
|
||||
cdef class RayletClient:
|
||||
cdef:
|
||||
unique_ptr[CRayletClientWithIoContext] inner
|
||||
|
||||
def __cinit__(self, ip_address: str, port: int):
|
||||
cdef:
|
||||
c_string c_ip_address
|
||||
int32_t c_port
|
||||
c_ip_address = ip_address.encode('utf-8')
|
||||
c_port = <int32_t>port
|
||||
self.inner = make_unique[CRayletClientWithIoContext](c_ip_address, c_port)
|
||||
|
||||
def async_get_worker_pids(self, timeout_ms: int = 1000) -> Future[list[int]]:
|
||||
"""Get the PIDs of all workers registered with the raylet."""
|
||||
cdef:
|
||||
fut = incremented_fut()
|
||||
int32_t timeout = <int32_t>timeout_ms
|
||||
assert self.inner.get() is not NULL
|
||||
with nogil:
|
||||
self.inner.get().GetWorkerPIDs(
|
||||
OptionalItemPyCallback[c_vector[int32_t]](
|
||||
&convert_optional_vector_int32,
|
||||
assign_and_decrement_fut,
|
||||
fut),
|
||||
timeout)
|
||||
return asyncio.wrap_future(fut)
|
||||
|
||||
def async_get_agent_pids(self, timeout_ms: int = 1000) -> Future[list[int]]:
|
||||
"""Get the PIDs of dashboard and runtime_env agent."""
|
||||
cdef:
|
||||
fut = incremented_fut()
|
||||
int32_t timeout = <int32_t>timeout_ms
|
||||
assert self.inner.get() is not NULL
|
||||
with nogil:
|
||||
self.inner.get().GetAgentPIDs(
|
||||
OptionalItemPyCallback[c_vector[int32_t]](
|
||||
&convert_optional_vector_int32,
|
||||
assign_and_decrement_fut,
|
||||
fut),
|
||||
timeout)
|
||||
return asyncio.wrap_future(fut)
|
||||
@@ -0,0 +1,43 @@
|
||||
from libcpp cimport bool as c_bool
|
||||
from libcpp.memory cimport shared_ptr
|
||||
from libcpp.string cimport string
|
||||
from ray.includes.optional cimport optional
|
||||
|
||||
|
||||
cdef extern from "ray/rpc/authentication/authentication_mode.h" namespace "ray::rpc" nogil:
|
||||
cdef enum CAuthenticationMode "ray::rpc::AuthenticationMode":
|
||||
DISABLED "ray::rpc::AuthenticationMode::DISABLED"
|
||||
TOKEN "ray::rpc::AuthenticationMode::TOKEN"
|
||||
|
||||
CAuthenticationMode GetAuthenticationMode()
|
||||
c_bool IsK8sTokenAuthEnabled()
|
||||
|
||||
cdef extern from "ray/rpc/authentication/authentication_token.h" namespace "ray::rpc" nogil:
|
||||
cdef cppclass CAuthenticationToken "ray::rpc::AuthenticationToken":
|
||||
CAuthenticationToken()
|
||||
CAuthenticationToken(string value)
|
||||
c_bool empty()
|
||||
c_bool Equals(const CAuthenticationToken& other)
|
||||
string ToAuthorizationHeaderValue()
|
||||
string GetRawValue()
|
||||
@staticmethod
|
||||
CAuthenticationToken FromMetadata(string metadata_value)
|
||||
|
||||
cdef extern from "ray/rpc/authentication/authentication_token_loader.h" namespace "ray::rpc" nogil:
|
||||
cdef cppclass CTokenLoadResult "ray::rpc::TokenLoadResult":
|
||||
optional[CAuthenticationToken] token
|
||||
string error_message
|
||||
c_bool hasError()
|
||||
|
||||
cdef cppclass CAuthenticationTokenLoader "ray::rpc::AuthenticationTokenLoader":
|
||||
@staticmethod
|
||||
CAuthenticationTokenLoader& instance()
|
||||
void ResetCache()
|
||||
shared_ptr[const CAuthenticationToken] GetToken(c_bool ignore_auth_mode)
|
||||
CTokenLoadResult TryLoadToken(c_bool ignore_auth_mode)
|
||||
|
||||
cdef extern from "ray/rpc/authentication/authentication_token_validator.h" namespace "ray::rpc" nogil:
|
||||
cdef cppclass CAuthenticationTokenValidator "ray::rpc::AuthenticationTokenValidator":
|
||||
@staticmethod
|
||||
CAuthenticationTokenValidator& instance()
|
||||
c_bool ValidateToken(const shared_ptr[const CAuthenticationToken]& expected_token, const string& provided_metadata)
|
||||
@@ -0,0 +1,156 @@
|
||||
from libcpp cimport bool as c_bool
|
||||
from libcpp.memory cimport shared_ptr
|
||||
from ray.includes.rpc_token_authentication cimport (
|
||||
CAuthenticationMode,
|
||||
GetAuthenticationMode,
|
||||
IsK8sTokenAuthEnabled,
|
||||
CAuthenticationToken,
|
||||
CAuthenticationTokenLoader,
|
||||
CAuthenticationTokenValidator,
|
||||
CTokenLoadResult,
|
||||
)
|
||||
from ray._private.authentication.authentication_constants import AUTHORIZATION_HEADER_NAME
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Authentication mode enum exposed to Python
|
||||
class AuthenticationMode:
|
||||
DISABLED = CAuthenticationMode.DISABLED
|
||||
TOKEN = CAuthenticationMode.TOKEN
|
||||
|
||||
|
||||
def get_authentication_mode():
|
||||
"""Get the current authentication mode.
|
||||
|
||||
Returns:
|
||||
AuthenticationMode enum value (DISABLED or TOKEN)
|
||||
"""
|
||||
return GetAuthenticationMode()
|
||||
|
||||
|
||||
def is_k8s_token_auth_enabled():
|
||||
"""Returns whether Kubernetes token authentication is enabled.
|
||||
|
||||
Returns:
|
||||
bool: True if Kubernetes token auth is enabled, false otherwise.
|
||||
"""
|
||||
return IsK8sTokenAuthEnabled()
|
||||
|
||||
|
||||
def validate_authentication_token(provided_metadata: str) -> bool:
|
||||
"""Validate provided authentication token.
|
||||
|
||||
For TOKEN mode, compares against the expected token.
|
||||
If K8S auth is enabled, validates against the Kubernetes API.
|
||||
|
||||
Args:
|
||||
provided_metadata: Full authorization header value (e.g., "Bearer <token>")
|
||||
|
||||
Returns:
|
||||
bool: True if token is valid, False otherwise
|
||||
"""
|
||||
cdef shared_ptr[const CAuthenticationToken] expected_ptr
|
||||
|
||||
if get_authentication_mode() == CAuthenticationMode.TOKEN and not is_k8s_token_auth_enabled():
|
||||
expected_ptr = CAuthenticationTokenLoader.instance().GetToken(False)
|
||||
if not expected_ptr:
|
||||
return False
|
||||
|
||||
return CAuthenticationTokenValidator.instance().ValidateToken(
|
||||
expected_ptr, provided_metadata.encode())
|
||||
|
||||
|
||||
class AuthenticationTokenLoader:
|
||||
"""Python wrapper for C++ AuthenticationTokenLoader singleton."""
|
||||
|
||||
@staticmethod
|
||||
def instance():
|
||||
"""Get the singleton instance (returns a wrapper for convenience)."""
|
||||
return AuthenticationTokenLoader()
|
||||
|
||||
def has_token(self, ignore_auth_mode=False):
|
||||
"""Check if an authentication token exists without crashing.
|
||||
|
||||
Args:
|
||||
ignore_auth_mode: If True, bypass auth mode check and attempt to load token
|
||||
regardless of RAY_AUTH_MODE setting.
|
||||
|
||||
Returns:
|
||||
bool: True if a token exists, False otherwise
|
||||
|
||||
Raises:
|
||||
AuthenticationError: If any issues loading the token
|
||||
"""
|
||||
cdef CTokenLoadResult result
|
||||
cdef c_bool c_ignore_auth_mode = ignore_auth_mode
|
||||
|
||||
with nogil:
|
||||
result = CAuthenticationTokenLoader.instance().TryLoadToken(c_ignore_auth_mode)
|
||||
|
||||
if result.hasError():
|
||||
from ray.exceptions import AuthenticationError
|
||||
raise AuthenticationError(result.error_message.decode('utf-8'))
|
||||
|
||||
if not result.token.has_value() or result.token.value().empty():
|
||||
return False
|
||||
return True
|
||||
|
||||
def reset_cache(self):
|
||||
"""Reset the C++ authentication token cache.
|
||||
|
||||
This forces the token loader to reload the token from environment
|
||||
variables or files on the next request.
|
||||
"""
|
||||
CAuthenticationTokenLoader.instance().ResetCache()
|
||||
|
||||
def get_token_for_http_header(self, ignore_auth_mode=False) -> dict:
|
||||
"""Get authentication token as a dictionary for HTTP headers.
|
||||
|
||||
This method loads the token from C++ AuthenticationTokenLoader and returns it
|
||||
as a dictionary that can be merged with existing headers. It returns an empty
|
||||
dictionary if:
|
||||
- A token does not exist
|
||||
- The token is empty
|
||||
|
||||
Args:
|
||||
ignore_auth_mode: If True, bypass auth mode check and attempt to load token
|
||||
regardless of RAY_AUTH_MODE setting.
|
||||
|
||||
Returns:
|
||||
dict: Empty dict or {"authorization": "Bearer <token>"}
|
||||
"""
|
||||
if not self.has_token(ignore_auth_mode):
|
||||
return {}
|
||||
|
||||
# Get the token from C++ layer (returns shared_ptr)
|
||||
cdef shared_ptr[const CAuthenticationToken] token_ptr = \
|
||||
CAuthenticationTokenLoader.instance().GetToken(ignore_auth_mode)
|
||||
|
||||
if not token_ptr or token_ptr.get().empty():
|
||||
return {}
|
||||
|
||||
return {AUTHORIZATION_HEADER_NAME: token_ptr.get().ToAuthorizationHeaderValue().decode('utf-8')}
|
||||
|
||||
def get_raw_token(self, ignore_auth_mode=False) -> str:
|
||||
"""Get the raw authentication token value.
|
||||
|
||||
Args:
|
||||
ignore_auth_mode: If True, bypass auth mode check and attempt to load token
|
||||
regardless of RAY_AUTH_MODE setting.
|
||||
|
||||
Returns:
|
||||
str: The raw token string, or empty string if no token exists
|
||||
"""
|
||||
if not self.has_token(ignore_auth_mode):
|
||||
return ""
|
||||
|
||||
# Get the token from C++ layer (returns shared_ptr)
|
||||
cdef shared_ptr[const CAuthenticationToken] token_ptr = \
|
||||
CAuthenticationTokenLoader.instance().GetToken(ignore_auth_mode)
|
||||
|
||||
if not token_ptr or token_ptr.get().empty():
|
||||
return ""
|
||||
|
||||
return token_ptr.get().GetRawValue().decode('utf-8')
|
||||
@@ -0,0 +1,537 @@
|
||||
from libc.string cimport memcpy
|
||||
from libc.stdint cimport uintptr_t, uint64_t, INT32_MAX
|
||||
import contextlib
|
||||
import cython
|
||||
|
||||
DEF MEMCOPY_THREADS = 6
|
||||
|
||||
# This is the default alignment value for len(buffer) < 2048.
|
||||
DEF kMinorBufferAlign = 8
|
||||
# This is the default alignment value for len(buffer) >= 2048.
|
||||
# Some projects like Arrow use it for possible SIMD acceleration.
|
||||
DEF kMajorBufferAlign = 64
|
||||
DEF kMajorBufferSize = 2048
|
||||
DEF kMemcopyDefaultBlocksize = 64
|
||||
DEF kMemcopyDefaultThreshold = 1024 * 1024
|
||||
DEF kLanguageSpecificTypeExtensionId = 101
|
||||
DEF kMessagePackOffset = 9
|
||||
|
||||
cdef extern from "ray/util/memory.h" namespace "ray" nogil:
|
||||
void parallel_memcopy(uint8_t* dst, const uint8_t* src, int64_t nbytes,
|
||||
uintptr_t block_size, int num_threads)
|
||||
|
||||
cdef extern from "google/protobuf/repeated_field.h" nogil:
|
||||
cdef cppclass RepeatedField[Element]:
|
||||
const Element* data() const
|
||||
|
||||
cdef extern from "src/ray/protobuf/serialization.pb.h" nogil:
|
||||
cdef cppclass CPythonBuffer "ray::serialization::PythonBuffer":
|
||||
void set_address(uint64_t value)
|
||||
uint64_t address() const
|
||||
void set_length(int64_t value)
|
||||
int64_t length() const
|
||||
void set_itemsize(int64_t value)
|
||||
int64_t itemsize()
|
||||
void set_ndim(int32_t value)
|
||||
int32_t ndim()
|
||||
void set_readonly(c_bool value)
|
||||
c_bool readonly()
|
||||
void set_format(const c_string& value)
|
||||
const c_string &format()
|
||||
c_string* release_format()
|
||||
void add_shape(int64_t value)
|
||||
int64_t shape(int index)
|
||||
const RepeatedField[int64_t] &shape() const
|
||||
int shape_size()
|
||||
void add_strides(int64_t value)
|
||||
int64_t strides(int index)
|
||||
const RepeatedField[int64_t] &strides() const
|
||||
int strides_size()
|
||||
|
||||
cdef cppclass CPythonObject "ray::serialization::PythonObject":
|
||||
uint64_t inband_data_size() const
|
||||
void set_inband_data_size(uint64_t value)
|
||||
uint64_t raw_buffers_size() const
|
||||
void set_raw_buffers_size(uint64_t value)
|
||||
CPythonBuffer* add_buffer()
|
||||
CPythonBuffer& buffer(int index) const
|
||||
int buffer_size() const
|
||||
size_t ByteSizeLong() const
|
||||
int GetCachedSize() const
|
||||
uint8_t *SerializeWithCachedSizesToArray(uint8_t *target)
|
||||
c_bool ParseFromArray(void* data, int size)
|
||||
|
||||
|
||||
cdef int64_t padded_length(int64_t offset, int64_t alignment):
|
||||
return ((offset + alignment - 1) // alignment) * alignment
|
||||
|
||||
|
||||
cdef uint8_t* aligned_address(uint8_t* addr, uint64_t alignment) nogil:
|
||||
cdef uintptr_t u_addr = <uintptr_t>addr
|
||||
return <uint8_t*>(((u_addr + alignment - 1) // alignment) * alignment)
|
||||
|
||||
|
||||
cdef class SubBuffer:
|
||||
cdef:
|
||||
void *buf
|
||||
Py_ssize_t len
|
||||
int readonly
|
||||
c_string _format
|
||||
int ndim
|
||||
c_vector[Py_ssize_t] _shape
|
||||
c_vector[Py_ssize_t] _strides
|
||||
Py_ssize_t *suboffsets
|
||||
Py_ssize_t itemsize
|
||||
void *internal
|
||||
object buffer
|
||||
|
||||
def __cinit__(self, object buffer):
|
||||
# Increase ref count.
|
||||
self.buffer = buffer
|
||||
self.suboffsets = NULL
|
||||
self.internal = NULL
|
||||
|
||||
def __len__(self):
|
||||
return self.len // self.itemsize
|
||||
|
||||
@property
|
||||
def nbytes(self):
|
||||
"""
|
||||
The buffer size in bytes.
|
||||
"""
|
||||
return self.len
|
||||
|
||||
@property
|
||||
def readonly(self):
|
||||
return self.readonly
|
||||
|
||||
def tobytes(self):
|
||||
"""
|
||||
Return this buffer as a Python bytes object. Memory is copied.
|
||||
"""
|
||||
return PyBytes_FromStringAndSize(
|
||||
<const char*> self.buf, self.len)
|
||||
|
||||
def __getbuffer__(self, Py_buffer* buffer, int flags):
|
||||
if flags & cpython.PyBUF_WRITABLE:
|
||||
# Ray ensures all buffers are immutable.
|
||||
raise BufferError
|
||||
buffer.readonly = self.readonly
|
||||
buffer.buf = self.buf
|
||||
buffer.format = <char *>self._format.c_str()
|
||||
buffer.internal = self.internal
|
||||
buffer.itemsize = self.itemsize
|
||||
buffer.len = self.len
|
||||
buffer.ndim = self.ndim
|
||||
buffer.obj = self # This is important for GC.
|
||||
buffer.shape = self._shape.data()
|
||||
buffer.strides = self._strides.data()
|
||||
buffer.suboffsets = self.suboffsets
|
||||
|
||||
def __getsegcount__(self, Py_ssize_t *len_out):
|
||||
if len_out != NULL:
|
||||
len_out[0] = <Py_ssize_t> self.size
|
||||
return 1
|
||||
|
||||
def __getreadbuffer__(self, Py_ssize_t idx, void ** p):
|
||||
if idx != 0:
|
||||
raise SystemError("accessing non-existent buffer segment")
|
||||
if p != NULL:
|
||||
p[0] = self.buf
|
||||
return self.size
|
||||
|
||||
def __getwritebuffer__(self, Py_ssize_t idx, void ** p):
|
||||
if idx != 0:
|
||||
raise SystemError("accessing non-existent buffer segment")
|
||||
if p != NULL:
|
||||
p[0] = self.buf
|
||||
return self.size
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _temporarily_disable_gc():
|
||||
gc_enabled = gc.isenabled()
|
||||
try:
|
||||
if gc_enabled:
|
||||
gc.disable()
|
||||
yield
|
||||
finally:
|
||||
if gc_enabled:
|
||||
gc.enable()
|
||||
|
||||
|
||||
cdef class MessagePackSerializer(object):
|
||||
@staticmethod
|
||||
def dumps(o, python_serializer=None):
|
||||
def _default(obj):
|
||||
if python_serializer is not None:
|
||||
return msgpack.ExtType(kLanguageSpecificTypeExtensionId,
|
||||
msgpack.dumps(python_serializer(obj)))
|
||||
return obj
|
||||
try:
|
||||
# If we let strict_types is False, then whether list or tuple will
|
||||
# be packed to a message pack array. So, they can't be
|
||||
# distinguished when unpacking.
|
||||
return msgpack.dumps(o, default=_default,
|
||||
use_bin_type=True, strict_types=True)
|
||||
except ValueError as ex:
|
||||
# msgpack can't handle recursive objects, so we serialize them by
|
||||
# python serializer, e.g. pickle.
|
||||
return msgpack.dumps(_default(o), default=_default,
|
||||
use_bin_type=True, strict_types=True)
|
||||
|
||||
@classmethod
|
||||
def loads(cls, s, python_deserializer=None):
|
||||
def _ext_hook(code, data):
|
||||
if code == kLanguageSpecificTypeExtensionId:
|
||||
if python_deserializer is not None:
|
||||
return python_deserializer(msgpack.loads(data))
|
||||
raise Exception('Unrecognized ext type id: {}'.format(code))
|
||||
|
||||
with _temporarily_disable_gc(): # Performance optimization for msgpack
|
||||
return msgpack.loads(s, ext_hook=_ext_hook, raw=False,
|
||||
strict_map_key=False)
|
||||
|
||||
|
||||
@cython.boundscheck(False)
|
||||
@cython.wraparound(False)
|
||||
def split_buffer(Buffer buf):
|
||||
cdef:
|
||||
size_t size = buf.buffer.get().Size()
|
||||
uint8_t[:] bufferview = buf
|
||||
int64_t msgpack_bytes_length
|
||||
|
||||
assert kMessagePackOffset <= size
|
||||
header_unpacker = msgpack.Unpacker()
|
||||
header_unpacker.feed(bufferview[:kMessagePackOffset])
|
||||
msgpack_bytes_length = header_unpacker.unpack()
|
||||
assert kMessagePackOffset + msgpack_bytes_length <= <int64_t>size
|
||||
return (bufferview[kMessagePackOffset:
|
||||
kMessagePackOffset + msgpack_bytes_length],
|
||||
bufferview[kMessagePackOffset + msgpack_bytes_length:])
|
||||
|
||||
|
||||
# Note [Pickle5 serialization layout & alignment]
|
||||
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
# To ensure efficient data access, our serialize enforces alignment
|
||||
# when writing data to a buffer. See 'serialization.proto' for
|
||||
# the detail memory layout and alignment.
|
||||
|
||||
|
||||
@cython.boundscheck(False)
|
||||
@cython.wraparound(False)
|
||||
def unpack_pickle5_buffers(uint8_t[:] bufferview):
|
||||
cdef:
|
||||
const uint8_t *data = &bufferview[0]
|
||||
CPythonObject python_object
|
||||
CPythonBuffer *buffer_meta
|
||||
int inband_offset = sizeof(int64_t) * 2
|
||||
int64_t inband_size
|
||||
int64_t protobuf_size
|
||||
int32_t i
|
||||
const uint8_t *buffers_segment
|
||||
inband_size = (<int64_t*>data)[0]
|
||||
if inband_size < 0:
|
||||
raise ValueError("The inband data size should be positive."
|
||||
"Got negative instead. "
|
||||
"Maybe the buffer has been corrupted.")
|
||||
protobuf_size = (<int64_t*>data)[1]
|
||||
if protobuf_size > INT32_MAX or protobuf_size < 0:
|
||||
raise ValueError("Incorrect protobuf size. "
|
||||
"Maybe the buffer has been corrupted.")
|
||||
inband_data = bufferview[inband_offset:inband_offset + inband_size]
|
||||
if not python_object.ParseFromArray(
|
||||
data + inband_offset + inband_size, <int32_t>protobuf_size):
|
||||
raise ValueError("Protobuf object is corrupted.")
|
||||
buffers_segment = aligned_address(
|
||||
<uint8_t*>data + inband_offset + inband_size + protobuf_size,
|
||||
kMajorBufferAlign)
|
||||
pickled_buffers = []
|
||||
# Now read buffer meta
|
||||
for i in range(python_object.buffer_size()):
|
||||
buffer_meta = <CPythonBuffer *>&python_object.buffer(i)
|
||||
buffer = SubBuffer(bufferview)
|
||||
buffer.buf = <void*>(buffers_segment + buffer_meta.address())
|
||||
buffer.len = buffer_meta.length()
|
||||
buffer.itemsize = buffer_meta.itemsize()
|
||||
buffer.readonly = buffer_meta.readonly()
|
||||
buffer.ndim = buffer_meta.ndim()
|
||||
buffer._format = buffer_meta.format()
|
||||
buffer._shape.assign(
|
||||
buffer_meta.shape().data(),
|
||||
buffer_meta.shape().data() + buffer_meta.ndim())
|
||||
buffer._strides.assign(
|
||||
buffer_meta.strides().data(),
|
||||
buffer_meta.strides().data() + buffer_meta.ndim())
|
||||
buffer.internal = NULL
|
||||
buffer.suboffsets = NULL
|
||||
pickled_buffers.append(buffer)
|
||||
return inband_data, pickled_buffers
|
||||
|
||||
|
||||
cdef class Pickle5Writer:
|
||||
cdef:
|
||||
CPythonObject python_object
|
||||
c_vector[Py_buffer] buffers
|
||||
# Address of end of the current buffer, relative to the
|
||||
# begin offset of our buffers.
|
||||
uint64_t _curr_buffer_addr
|
||||
uint64_t _protobuf_offset
|
||||
int64_t _total_bytes
|
||||
|
||||
def __cinit__(self):
|
||||
self._curr_buffer_addr = 0
|
||||
self._total_bytes = -1
|
||||
|
||||
def __dealloc__(self):
|
||||
# We must release the buffer, or we could experience memory leaks.
|
||||
for i in range(self.buffers.size()):
|
||||
cpython.PyBuffer_Release(&self.buffers[i])
|
||||
|
||||
def buffer_callback(self, pickle_buffer):
|
||||
cdef:
|
||||
Py_buffer view
|
||||
int32_t i
|
||||
CPythonBuffer* buffer = self.python_object.add_buffer()
|
||||
cpython.PyObject_GetBuffer(pickle_buffer, &view,
|
||||
cpython.PyBUF_FULL_RO)
|
||||
buffer.set_length(view.len)
|
||||
buffer.set_ndim(view.ndim)
|
||||
# It should be 'view.readonly'. But for the sake of shared memory,
|
||||
# we have to make it immutable.
|
||||
buffer.set_readonly(1)
|
||||
buffer.set_itemsize(view.itemsize)
|
||||
if view.format:
|
||||
buffer.set_format(view.format)
|
||||
if view.shape:
|
||||
for i in range(view.ndim):
|
||||
buffer.add_shape(view.shape[i])
|
||||
if view.strides:
|
||||
for i in range(view.ndim):
|
||||
buffer.add_strides(view.strides[i])
|
||||
|
||||
# Increase buffer address.
|
||||
if view.len < kMajorBufferSize:
|
||||
self._curr_buffer_addr = padded_length(
|
||||
self._curr_buffer_addr, kMinorBufferAlign)
|
||||
else:
|
||||
self._curr_buffer_addr = padded_length(
|
||||
self._curr_buffer_addr, kMajorBufferAlign)
|
||||
buffer.set_address(self._curr_buffer_addr)
|
||||
self._curr_buffer_addr += view.len
|
||||
self.buffers.push_back(view)
|
||||
|
||||
def get_total_bytes(self, const uint8_t[:] inband):
|
||||
cdef:
|
||||
size_t protobuf_bytes = 0
|
||||
uint64_t inband_data_offset = sizeof(int64_t) * 2
|
||||
self.python_object.set_inband_data_size(len(inband))
|
||||
self.python_object.set_raw_buffers_size(self._curr_buffer_addr)
|
||||
# Since calculating the output size is expensive, we will
|
||||
# reuse the cached size.
|
||||
# However, protobuf could change the output size according to
|
||||
# different values, so we MUST NOT change 'python_object' afterwards.
|
||||
protobuf_bytes = self.python_object.ByteSizeLong()
|
||||
if protobuf_bytes > INT32_MAX:
|
||||
raise ValueError("Total buffer metadata size is bigger than %d. "
|
||||
"Consider reduce the number of buffers "
|
||||
"(number of numpy arrays, etc)." % INT32_MAX)
|
||||
self._protobuf_offset = inband_data_offset + len(inband)
|
||||
self._total_bytes = self._protobuf_offset + protobuf_bytes
|
||||
if self._curr_buffer_addr > 0:
|
||||
# reserve 'kMajorBufferAlign' bytes for possible buffer alignment
|
||||
self._total_bytes += kMajorBufferAlign + self._curr_buffer_addr
|
||||
return self._total_bytes
|
||||
|
||||
@cython.boundscheck(False)
|
||||
@cython.wraparound(False)
|
||||
cdef void write_to(self, const uint8_t[:] inband, uint8_t[:] data,
|
||||
int memcopy_threads) nogil:
|
||||
cdef:
|
||||
uint8_t *ptr = &data[0]
|
||||
uint64_t buffer_addr
|
||||
uint64_t buffer_len
|
||||
int i
|
||||
int64_t protobuf_size = self.python_object.GetCachedSize()
|
||||
if self._total_bytes < 0:
|
||||
raise ValueError("Must call 'get_total_bytes()' first "
|
||||
"to get the actual size")
|
||||
# Write inband data & protobuf size for deserialization.
|
||||
(<int64_t*>ptr)[0] = len(inband)
|
||||
(<int64_t*>ptr)[1] = protobuf_size
|
||||
# Write inband data.
|
||||
ptr += sizeof(int64_t) * 2
|
||||
with nogil:
|
||||
memcpy(ptr, &inband[0], len(inband))
|
||||
# Write protobuf data.
|
||||
ptr += len(inband)
|
||||
self.python_object.SerializeWithCachedSizesToArray(ptr)
|
||||
ptr += protobuf_size
|
||||
if self._curr_buffer_addr <= 0:
|
||||
# End of serialization. Writing more stuff will corrupt the memory.
|
||||
return
|
||||
# aligned to 64 bytes
|
||||
ptr = aligned_address(ptr, kMajorBufferAlign)
|
||||
for i in range(self.python_object.buffer_size()):
|
||||
buffer_addr = self.python_object.buffer(i).address()
|
||||
buffer_len = self.python_object.buffer(i).length()
|
||||
with nogil:
|
||||
if (memcopy_threads > 1 and
|
||||
buffer_len > kMemcopyDefaultThreshold):
|
||||
parallel_memcopy(ptr + buffer_addr,
|
||||
<const uint8_t*> self.buffers[i].buf,
|
||||
buffer_len,
|
||||
kMemcopyDefaultBlocksize, memcopy_threads)
|
||||
else:
|
||||
memcpy(ptr + buffer_addr, self.buffers[i].buf, buffer_len)
|
||||
|
||||
|
||||
cdef class SerializedObject(object):
|
||||
cdef:
|
||||
object _metadata
|
||||
object _contained_object_refs
|
||||
|
||||
def __init__(self, metadata, contained_object_refs=None):
|
||||
self._metadata = metadata
|
||||
self._contained_object_refs = contained_object_refs or []
|
||||
|
||||
@property
|
||||
def total_bytes(self):
|
||||
raise NotImplementedError("{}.total_bytes not implemented.".format(
|
||||
type(self).__name__))
|
||||
|
||||
@property
|
||||
def metadata(self):
|
||||
return self._metadata
|
||||
|
||||
@property
|
||||
def contained_object_refs(self):
|
||||
return self._contained_object_refs
|
||||
|
||||
@cython.boundscheck(False)
|
||||
@cython.wraparound(False)
|
||||
cdef void write_to(self, uint8_t[:] buffer) nogil:
|
||||
raise NotImplementedError("{}.write_to not implemented.".format(
|
||||
type(self).__name__))
|
||||
|
||||
|
||||
cdef class Pickle5SerializedObject(SerializedObject):
|
||||
cdef:
|
||||
const uint8_t[:] inband
|
||||
Pickle5Writer writer
|
||||
object _total_bytes
|
||||
|
||||
def __init__(self, metadata, inband, Pickle5Writer writer,
|
||||
contained_object_refs):
|
||||
super(Pickle5SerializedObject, self).__init__(metadata,
|
||||
contained_object_refs)
|
||||
self.inband = inband
|
||||
self.writer = writer
|
||||
# cached total bytes
|
||||
self._total_bytes = None
|
||||
|
||||
@property
|
||||
def total_bytes(self):
|
||||
if self._total_bytes is None:
|
||||
self._total_bytes = self.writer.get_total_bytes(self.inband)
|
||||
return self._total_bytes
|
||||
|
||||
@cython.boundscheck(False)
|
||||
@cython.wraparound(False)
|
||||
cdef void write_to(self, uint8_t[:] buffer) nogil:
|
||||
self.writer.write_to(self.inband, buffer, MEMCOPY_THREADS)
|
||||
|
||||
|
||||
cdef class MessagePackSerializedObject(SerializedObject):
|
||||
cdef:
|
||||
SerializedObject nest_serialized_object
|
||||
object msgpack_header
|
||||
object msgpack_data
|
||||
int64_t _msgpack_header_bytes
|
||||
int64_t _msgpack_data_bytes
|
||||
int64_t _total_bytes
|
||||
const uint8_t *msgpack_header_ptr
|
||||
const uint8_t *msgpack_data_ptr
|
||||
|
||||
def __init__(self, metadata, msgpack_data, contained_object_refs,
|
||||
SerializedObject nest_serialized_object=None):
|
||||
if nest_serialized_object:
|
||||
contained_object_refs.extend(
|
||||
nest_serialized_object.contained_object_refs
|
||||
)
|
||||
total_bytes = nest_serialized_object.total_bytes
|
||||
else:
|
||||
total_bytes = 0
|
||||
super(MessagePackSerializedObject, self).__init__(
|
||||
metadata,
|
||||
contained_object_refs,
|
||||
)
|
||||
self.nest_serialized_object = nest_serialized_object
|
||||
self.msgpack_header = msgpack_header = msgpack.dumps(len(msgpack_data))
|
||||
self.msgpack_data = msgpack_data
|
||||
self._msgpack_header_bytes = len(msgpack_header)
|
||||
self._msgpack_data_bytes = len(msgpack_data)
|
||||
self._total_bytes = (kMessagePackOffset +
|
||||
self._msgpack_data_bytes +
|
||||
total_bytes)
|
||||
self.msgpack_header_ptr = <const uint8_t*>msgpack_header
|
||||
self.msgpack_data_ptr = <const uint8_t*>msgpack_data
|
||||
assert self._msgpack_header_bytes <= kMessagePackOffset
|
||||
|
||||
@property
|
||||
def total_bytes(self):
|
||||
return self._total_bytes
|
||||
|
||||
def to_bytes(self) -> bytes:
|
||||
cdef shared_ptr[CBuffer] data = \
|
||||
dynamic_pointer_cast[CBuffer, LocalMemoryBuffer](
|
||||
make_shared[LocalMemoryBuffer](self._total_bytes))
|
||||
buffer = Buffer.make(data)
|
||||
self.write_to(buffer)
|
||||
return buffer.to_pybytes()
|
||||
|
||||
@cython.boundscheck(False)
|
||||
@cython.wraparound(False)
|
||||
cdef void write_to(self, uint8_t[:] buffer) nogil:
|
||||
cdef uint8_t *ptr = &buffer[0]
|
||||
|
||||
# Write msgpack data first.
|
||||
with nogil:
|
||||
memcpy(ptr, self.msgpack_header_ptr, self._msgpack_header_bytes)
|
||||
memcpy(ptr + kMessagePackOffset,
|
||||
self.msgpack_data_ptr, self._msgpack_data_bytes)
|
||||
|
||||
if self.nest_serialized_object is not None:
|
||||
self.nest_serialized_object.write_to(
|
||||
buffer[kMessagePackOffset + self._msgpack_data_bytes:])
|
||||
|
||||
|
||||
cdef class RawSerializedObject(SerializedObject):
|
||||
cdef:
|
||||
object value
|
||||
const uint8_t *value_ptr
|
||||
int64_t _total_bytes
|
||||
|
||||
def __init__(self, value):
|
||||
super(RawSerializedObject,
|
||||
self).__init__(ray_constants.OBJECT_METADATA_TYPE_RAW)
|
||||
self.value = value
|
||||
self.value_ptr = <const uint8_t*> value
|
||||
self._total_bytes = len(value)
|
||||
|
||||
@property
|
||||
def total_bytes(self):
|
||||
return self._total_bytes
|
||||
|
||||
@cython.boundscheck(False)
|
||||
@cython.wraparound(False)
|
||||
cdef void write_to(self, uint8_t[:] buffer) nogil:
|
||||
with nogil:
|
||||
if (MEMCOPY_THREADS > 1 and
|
||||
self._total_bytes > kMemcopyDefaultThreshold):
|
||||
parallel_memcopy(&buffer[0],
|
||||
self.value_ptr,
|
||||
self._total_bytes, kMemcopyDefaultBlocksize,
|
||||
MEMCOPY_THREADS)
|
||||
else:
|
||||
memcpy(&buffer[0], self.value_ptr, self._total_bytes)
|
||||
@@ -0,0 +1,17 @@
|
||||
from libcpp.string cimport string as c_string
|
||||
|
||||
cdef extern from *:
|
||||
"""
|
||||
extern "C" {
|
||||
#include "ray/thirdparty/setproctitle/spt_setup.h"
|
||||
}
|
||||
"""
|
||||
int spt_setup()
|
||||
|
||||
cdef extern from *:
|
||||
"""
|
||||
extern "C" {
|
||||
#include "ray/thirdparty/setproctitle/spt_status.h"
|
||||
}
|
||||
"""
|
||||
void set_ps_display(const char *activity, bint force)
|
||||
@@ -0,0 +1,33 @@
|
||||
import sys
|
||||
import psutil
|
||||
import subprocess
|
||||
import threading
|
||||
|
||||
from libcpp.string cimport string as c_string
|
||||
from ray.includes.setproctitle cimport (
|
||||
spt_setup,
|
||||
set_ps_display
|
||||
)
|
||||
|
||||
_current_proctitle = None
|
||||
_current_proctitle_lock = threading.Lock()
|
||||
|
||||
def setproctitle(title: str):
|
||||
global _current_proctitle
|
||||
cdef c_string c_title = title.encode("utf-8")
|
||||
|
||||
with _current_proctitle_lock:
|
||||
spt_setup()
|
||||
set_ps_display(c_title.c_str(), True)
|
||||
|
||||
_current_proctitle = title
|
||||
|
||||
def getproctitle() -> str:
|
||||
global _current_proctitle
|
||||
|
||||
with _current_proctitle_lock:
|
||||
if _current_proctitle is None:
|
||||
# The process title is not change so getting the process cmdline as the
|
||||
# initial title.
|
||||
_current_proctitle = subprocess.list2cmdline(psutil.Process().cmdline())
|
||||
return _current_proctitle
|
||||
@@ -0,0 +1,16 @@
|
||||
from libcpp.string cimport string as c_string
|
||||
from libc.stdint cimport uint64_t
|
||||
from libcpp cimport bool as c_bool
|
||||
|
||||
cdef extern from "ray/util/stream_redirection_options.h" nogil:
|
||||
cdef cppclass CStreamRedirectionOptions "ray::StreamRedirectionOption":
|
||||
CStreamRedirectionOptions()
|
||||
c_string file_path
|
||||
uint64_t rotation_max_size
|
||||
uint64_t rotation_max_file_count
|
||||
c_bool tee_to_stdout
|
||||
c_bool tee_to_stderr
|
||||
|
||||
cdef extern from "ray/util/stream_redirection.h" namespace "ray" nogil:
|
||||
void RedirectStdoutOncePerProcess(const CStreamRedirectionOptions& opt)
|
||||
void RedirectStderrOncePerProcess(const CStreamRedirectionOptions& opt)
|
||||
@@ -0,0 +1,168 @@
|
||||
from typing import Dict
|
||||
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp.unordered_map cimport unordered_map
|
||||
from libcpp.utility cimport move
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
|
||||
from ray.includes.common cimport (
|
||||
CConcurrencyGroup,
|
||||
CFallbackOption,
|
||||
CLabelSelector,
|
||||
CNodeResources,
|
||||
SetNodeResourcesLabels,
|
||||
)
|
||||
from ray.includes.function_descriptor cimport (
|
||||
CFunctionDescriptor,
|
||||
CFunctionDescriptorBuilder,
|
||||
)
|
||||
from ray.includes.ray_config cimport RayConfig
|
||||
|
||||
|
||||
cdef int prepare_labels(
|
||||
dict label_dict,
|
||||
unordered_map[c_string, c_string] *label_map) except -1:
|
||||
|
||||
if label_dict is None:
|
||||
return 0
|
||||
|
||||
label_map[0].reserve(len(label_dict))
|
||||
for key, value in label_dict.items():
|
||||
if not isinstance(key, str):
|
||||
raise ValueError(f"Label key must be string, but got {type(key)}")
|
||||
if not isinstance(value, str):
|
||||
raise ValueError(f"Label value must be string, but got {type(value)}")
|
||||
label_map[0][key.encode("utf-8")] = value.encode("utf-8")
|
||||
|
||||
return 0
|
||||
|
||||
cdef int prepare_label_selector(
|
||||
dict label_selector_dict,
|
||||
CLabelSelector *c_label_selector) except -1:
|
||||
|
||||
c_label_selector[0] = CLabelSelector()
|
||||
|
||||
if label_selector_dict is None:
|
||||
return 0
|
||||
|
||||
for key, value in label_selector_dict.items():
|
||||
if not isinstance(key, str):
|
||||
raise ValueError(f"Label selector key type must be string, but got {type(key)}")
|
||||
if not isinstance(value, str):
|
||||
raise ValueError(f"Label selector value must be string, but got {type(value)}")
|
||||
if key == "":
|
||||
raise ValueError("Label selector key must be a non-empty string.")
|
||||
if (value.startswith("in(") and value.endswith(")")) or \
|
||||
(value.startswith("!in(") and value.endswith(")")):
|
||||
inner = value[value.index("(")+1:-1].strip()
|
||||
if not inner:
|
||||
raise ValueError(f"No values provided for Label Selector '{value[:value.index('(')]}' operator on key '{key}'.")
|
||||
# Add key-value constraint to the LabelSelector object.
|
||||
c_label_selector[0].AddConstraint(key.encode("utf-8"), value.encode("utf-8"))
|
||||
|
||||
return 0
|
||||
|
||||
def node_labels_match_selector(node_labels: Dict[str, str], selector: Dict[str, str]) -> bool:
|
||||
"""
|
||||
Checks if the given node labels satisfy the label selector. This helper function exposes
|
||||
the C++ logic for determining if a node satisfies a label selector to the Python layer.
|
||||
"""
|
||||
cdef:
|
||||
CNodeResources c_node_resources
|
||||
CLabelSelector c_label_selector
|
||||
unordered_map[c_string, c_string] c_labels_map
|
||||
|
||||
prepare_labels(node_labels, &c_labels_map)
|
||||
SetNodeResourcesLabels(c_node_resources, c_labels_map)
|
||||
prepare_label_selector(selector, &c_label_selector)
|
||||
|
||||
# Return whether the node resources satisfy the label constraint.
|
||||
return c_node_resources.HasRequiredLabels(c_label_selector)
|
||||
|
||||
cdef int prepare_fallback_strategy(
|
||||
list fallback_strategy,
|
||||
c_vector[CFallbackOption] *fallback_strategy_vector) except -1:
|
||||
|
||||
cdef dict label_selector_dict
|
||||
cdef CLabelSelector c_label_selector
|
||||
|
||||
if fallback_strategy is None:
|
||||
return 0
|
||||
|
||||
for strategy_dict in fallback_strategy:
|
||||
if not isinstance(strategy_dict, dict):
|
||||
raise ValueError(
|
||||
"Fallback strategy must be a list of dicts, "
|
||||
f"but got list containing {type(strategy_dict)}")
|
||||
|
||||
label_selector_dict = strategy_dict.get("label_selector")
|
||||
|
||||
if label_selector_dict is not None and not isinstance(label_selector_dict, dict):
|
||||
raise ValueError("Invalid fallback strategy element: invalid 'label_selector'.")
|
||||
|
||||
prepare_label_selector(label_selector_dict, &c_label_selector)
|
||||
|
||||
fallback_strategy_vector.push_back(
|
||||
CFallbackOption(c_label_selector)
|
||||
)
|
||||
|
||||
return 0
|
||||
|
||||
cdef int prepare_resources(
|
||||
dict resource_dict,
|
||||
unordered_map[c_string, double] *resource_map) except -1:
|
||||
cdef:
|
||||
list unit_resources
|
||||
|
||||
if resource_dict is None:
|
||||
raise ValueError("Must provide resource map.")
|
||||
|
||||
resource_map[0].reserve(len(resource_dict))
|
||||
for key, value in resource_dict.items():
|
||||
if not (isinstance(value, int) or isinstance(value, float)):
|
||||
raise ValueError("Resource quantities may only be ints or floats.")
|
||||
if value < 0:
|
||||
raise ValueError("Resource quantities may not be negative.")
|
||||
if value > 0:
|
||||
unit_resources = (
|
||||
f"{RayConfig.instance().predefined_unit_instance_resources()\
|
||||
.decode('utf-8')},"
|
||||
f"{RayConfig.instance().custom_unit_instance_resources()\
|
||||
.decode('utf-8')}"
|
||||
).split(",")
|
||||
|
||||
if (value >= 1 and isinstance(value, float)
|
||||
and not value.is_integer() and str(key) in unit_resources):
|
||||
raise ValueError(
|
||||
f"{key} resource quantities >1 must",
|
||||
f" be whole numbers. The specified quantity {value} is invalid.")
|
||||
resource_map[0][key.encode("ascii")] = float(value)
|
||||
return 0
|
||||
|
||||
cdef c_vector[CFunctionDescriptor] prepare_function_descriptors(pyfd_list):
|
||||
cdef:
|
||||
c_vector[CFunctionDescriptor] fd_list
|
||||
|
||||
fd_list.reserve(len(pyfd_list))
|
||||
for pyfd in pyfd_list:
|
||||
fd_list.push_back(CFunctionDescriptorBuilder.BuildPython(
|
||||
pyfd.module_name, pyfd.class_name, pyfd.function_name, b""))
|
||||
return fd_list
|
||||
|
||||
|
||||
cdef int prepare_actor_concurrency_groups(
|
||||
dict concurrency_groups_dict,
|
||||
c_vector[CConcurrencyGroup] *concurrency_groups):
|
||||
|
||||
cdef:
|
||||
c_vector[CFunctionDescriptor] c_fd_list
|
||||
|
||||
if concurrency_groups_dict is None:
|
||||
raise ValueError("Must provide it...")
|
||||
|
||||
concurrency_groups.reserve(len(concurrency_groups_dict))
|
||||
for key, value in concurrency_groups_dict.items():
|
||||
c_fd_list = prepare_function_descriptors(value["function_descriptors"])
|
||||
concurrency_groups.push_back(CConcurrencyGroup(
|
||||
key.encode("ascii"), value["max_concurrency"], move(c_fd_list)))
|
||||
return 1
|
||||
@@ -0,0 +1,462 @@
|
||||
# cython: profile=False
|
||||
# cython: boundscheck=False
|
||||
# cython: wraparound=False
|
||||
# cython: cdivision=True
|
||||
# cython: initializedcheck=False
|
||||
|
||||
# =============================================================================
|
||||
# WARNING: This file is used EXCLUSIVELY by Ray Serve.
|
||||
# =============================================================================
|
||||
#
|
||||
# These Cython-optimized timeseries utilities exist solely to support the
|
||||
# Ray Serve controller's autoscaling metrics pipeline (specifically the
|
||||
# per-deployment timeseries aggregation in the Serve replica scheduler).
|
||||
#
|
||||
# This code lives in `ray/includes/` and is compiled into `_raylet.so` because
|
||||
# Rather than introducing a new top-level `.pyx` / `.so` for a single
|
||||
# Serve-internal optimization, we include it here alongside the other
|
||||
# `.pxi` helpers that ship in `_raylet`.
|
||||
#
|
||||
# If you are not working on Ray Serve autoscaling, you almost certainly do not
|
||||
# need to modify this file.
|
||||
# =============================================================================
|
||||
|
||||
|
||||
|
||||
# C library imports
|
||||
from libc.stdlib cimport malloc, free
|
||||
from libc.math cimport round as c_round, isnan, nan, isinf
|
||||
|
||||
# Heap node for k-way merge
|
||||
cdef struct _TsHeapNode:
|
||||
double timestamp
|
||||
int series_idx
|
||||
double value
|
||||
int position_in_series # Current position within the series
|
||||
|
||||
|
||||
cdef inline void _ts_heap_sift_down(_TsHeapNode* heap, int size, int pos) noexcept nogil:
|
||||
"""Sift down operation for min-heap (inline for performance)."""
|
||||
cdef int smallest, left, right
|
||||
cdef _TsHeapNode temp
|
||||
|
||||
while True:
|
||||
smallest = pos
|
||||
left = 2 * pos + 1
|
||||
right = 2 * pos + 2
|
||||
|
||||
if left < size and heap[left].timestamp < heap[smallest].timestamp:
|
||||
smallest = left
|
||||
|
||||
if right < size and heap[right].timestamp < heap[smallest].timestamp:
|
||||
smallest = right
|
||||
|
||||
if smallest == pos:
|
||||
break
|
||||
|
||||
# Swap
|
||||
temp = heap[pos]
|
||||
heap[pos] = heap[smallest]
|
||||
heap[smallest] = temp
|
||||
pos = smallest
|
||||
|
||||
|
||||
cdef inline void _ts_heap_sift_up(_TsHeapNode* heap, int pos) noexcept nogil:
|
||||
"""Sift up operation for min-heap (inline for performance)."""
|
||||
cdef int parent
|
||||
cdef _TsHeapNode temp
|
||||
|
||||
while pos > 0:
|
||||
parent = (pos - 1) // 2
|
||||
if heap[parent].timestamp <= heap[pos].timestamp:
|
||||
break
|
||||
|
||||
# Swap
|
||||
temp = heap[pos]
|
||||
heap[pos] = heap[parent]
|
||||
heap[parent] = temp
|
||||
pos = parent
|
||||
|
||||
|
||||
cdef inline void _ts_heap_pop(_TsHeapNode* heap, int* size) noexcept nogil:
|
||||
"""Remove minimum element from heap."""
|
||||
if size[0] <= 0:
|
||||
return
|
||||
|
||||
heap[0] = heap[size[0] - 1]
|
||||
size[0] -= 1
|
||||
if size[0] > 0:
|
||||
_ts_heap_sift_down(heap, size[0], 0)
|
||||
|
||||
|
||||
cdef inline void _ts_heap_push(_TsHeapNode* heap, int* size, _TsHeapNode node) noexcept nogil:
|
||||
"""Add element to heap."""
|
||||
heap[size[0]] = node
|
||||
_ts_heap_sift_up(heap, size[0])
|
||||
size[0] += 1
|
||||
|
||||
|
||||
cdef int _kway_merge_timeseries_nogil(double** timestamps_arrays, double** values_arrays,
|
||||
int* series_lengths, int num_series,
|
||||
int result_capacity,
|
||||
double** out_timestamps, double** out_values) noexcept nogil:
|
||||
"""
|
||||
Fully nogil k-way merge operating on C arrays.
|
||||
|
||||
Assumptions:
|
||||
- Each input series is sorted by timestamp in ascending order
|
||||
- Values represent instantaneous gauge measurements (non-negative)
|
||||
|
||||
Args:
|
||||
timestamps_arrays: Array of pointers to timestamp arrays for each series
|
||||
values_arrays: Array of pointers to value arrays for each series
|
||||
series_lengths: Array of lengths for each series
|
||||
num_series: Number of series to merge
|
||||
result_capacity: Pre-allocated capacity (should be >= sum of all series lengths)
|
||||
out_timestamps: Output pointer for result timestamps
|
||||
out_values: Output pointer for result values
|
||||
|
||||
Returns: Number of points in merged result, or -1 on error
|
||||
"""
|
||||
cdef:
|
||||
int i, pos_in_series, series_idx
|
||||
int heap_size = 0
|
||||
double timestamp, value, old_value
|
||||
double running_total = 0.0
|
||||
double rounded_timestamp, last_rounded_timestamp = -1.0
|
||||
_TsHeapNode new_node
|
||||
int result_count = 0
|
||||
# C arrays for performance
|
||||
double* current_values = <double*>malloc(num_series * sizeof(double))
|
||||
int* series_positions = <int*>malloc(num_series * sizeof(int))
|
||||
_TsHeapNode* merge_heap = <_TsHeapNode*>malloc(num_series * sizeof(_TsHeapNode))
|
||||
double* result_timestamps = <double*>malloc(result_capacity * sizeof(double))
|
||||
double* result_values = <double*>malloc(result_capacity * sizeof(double))
|
||||
|
||||
if not current_values or not series_positions or not merge_heap or not result_timestamps or not result_values:
|
||||
# Memory allocation failed
|
||||
if current_values:
|
||||
free(current_values)
|
||||
if series_positions:
|
||||
free(series_positions)
|
||||
if merge_heap:
|
||||
free(merge_heap)
|
||||
if result_timestamps:
|
||||
free(result_timestamps)
|
||||
if result_values:
|
||||
free(result_values)
|
||||
return -1
|
||||
|
||||
# Initialize arrays
|
||||
for i in range(num_series):
|
||||
current_values[i] = 0.0
|
||||
series_positions[i] = 0
|
||||
|
||||
# Push first element from each series to heap
|
||||
if series_lengths[i] > 0:
|
||||
merge_heap[heap_size].timestamp = timestamps_arrays[i][0]
|
||||
merge_heap[heap_size].series_idx = i
|
||||
merge_heap[heap_size].value = values_arrays[i][0]
|
||||
merge_heap[heap_size].position_in_series = 0
|
||||
heap_size += 1
|
||||
|
||||
# Build initial heap
|
||||
for i in range(heap_size // 2 - 1, -1, -1):
|
||||
_ts_heap_sift_down(merge_heap, heap_size, i)
|
||||
|
||||
# K-way merge
|
||||
while heap_size > 0:
|
||||
# Get minimum element
|
||||
timestamp = merge_heap[0].timestamp
|
||||
series_idx = merge_heap[0].series_idx
|
||||
value = merge_heap[0].value
|
||||
pos_in_series = merge_heap[0].position_in_series
|
||||
|
||||
# Update running total
|
||||
old_value = current_values[series_idx]
|
||||
current_values[series_idx] = value
|
||||
running_total += value - old_value
|
||||
|
||||
# Remove from heap
|
||||
_ts_heap_pop(merge_heap, &heap_size)
|
||||
|
||||
# Push next element from same series if available
|
||||
series_positions[series_idx] = pos_in_series + 1
|
||||
if series_positions[series_idx] < series_lengths[series_idx]:
|
||||
new_node.timestamp = timestamps_arrays[series_idx][series_positions[series_idx]]
|
||||
new_node.series_idx = series_idx
|
||||
new_node.value = values_arrays[series_idx][series_positions[series_idx]]
|
||||
new_node.position_in_series = series_positions[series_idx]
|
||||
|
||||
_ts_heap_push(merge_heap, &heap_size, new_node)
|
||||
|
||||
# Only add point if value changed
|
||||
if value != old_value:
|
||||
# Round to 10ms precision
|
||||
rounded_timestamp = c_round(timestamp * 100.0) / 100.0
|
||||
|
||||
# Check if we can merge with last point
|
||||
if result_count > 0 and last_rounded_timestamp == rounded_timestamp:
|
||||
# Update last point's value
|
||||
result_values[result_count - 1] = running_total
|
||||
else:
|
||||
# Add new point (capacity is pre-allocated to be large enough)
|
||||
result_timestamps[result_count] = rounded_timestamp
|
||||
result_values[result_count] = running_total
|
||||
result_count += 1
|
||||
last_rounded_timestamp = rounded_timestamp
|
||||
|
||||
# Clean up
|
||||
free(current_values)
|
||||
free(series_positions)
|
||||
free(merge_heap)
|
||||
|
||||
# Return results
|
||||
out_timestamps[0] = result_timestamps
|
||||
out_values[0] = result_values
|
||||
return result_count
|
||||
|
||||
|
||||
def merge_instantaneous_total_cython(list replicas_timeseries):
|
||||
"""
|
||||
Cython-optimized k-way merge for timeseries.
|
||||
|
||||
This is a drop-in replacement for the Python version with 5-10x speedup.
|
||||
|
||||
Assumptions:
|
||||
- Each input timeseries is sorted by timestamp in ascending order
|
||||
- Values represent instantaneous gauge measurements
|
||||
|
||||
Args:
|
||||
replicas_timeseries: List of timeseries. Each timeseries is a list of
|
||||
objects with .timestamp and .value attributes.
|
||||
|
||||
Returns:
|
||||
List of (timestamp, value) tuples representing the merged timeseries.
|
||||
"""
|
||||
# Filter empty series
|
||||
cdef list active_series = [series for series in replicas_timeseries if series]
|
||||
|
||||
if not active_series:
|
||||
return []
|
||||
|
||||
if len(active_series) == 1:
|
||||
# Convert to tuples for consistent return type
|
||||
return [(point.timestamp, point.value) for point in active_series[0]]
|
||||
|
||||
cdef:
|
||||
int num_series = len(active_series)
|
||||
int i, j
|
||||
int total_points = 0
|
||||
object point, series
|
||||
bint alloc_failed = False
|
||||
# C arrays for all timestamps and values
|
||||
double** timestamps_arrays = <double**>malloc(num_series * sizeof(double*))
|
||||
double** values_arrays = <double**>malloc(num_series * sizeof(double*))
|
||||
int* series_lengths = <int*>malloc(num_series * sizeof(int))
|
||||
double* result_timestamps = NULL
|
||||
double* result_values = NULL
|
||||
int result_count
|
||||
|
||||
if not timestamps_arrays or not values_arrays or not series_lengths:
|
||||
# Memory allocation failed
|
||||
if timestamps_arrays:
|
||||
free(timestamps_arrays)
|
||||
if values_arrays:
|
||||
free(values_arrays)
|
||||
if series_lengths:
|
||||
free(series_lengths)
|
||||
raise MemoryError("Failed to allocate memory for merge operation")
|
||||
|
||||
# Initialize pointers to NULL for safe cleanup
|
||||
for i in range(num_series):
|
||||
timestamps_arrays[i] = NULL
|
||||
values_arrays[i] = NULL
|
||||
|
||||
try:
|
||||
# Extract all data from Python objects into C arrays
|
||||
for i in range(num_series):
|
||||
series = active_series[i]
|
||||
series_lengths[i] = len(series)
|
||||
total_points += series_lengths[i]
|
||||
|
||||
timestamps_arrays[i] = <double*>malloc(series_lengths[i] * sizeof(double))
|
||||
values_arrays[i] = <double*>malloc(series_lengths[i] * sizeof(double))
|
||||
|
||||
if not timestamps_arrays[i] or not values_arrays[i]:
|
||||
alloc_failed = True
|
||||
break
|
||||
|
||||
# Copy data from Python objects to C arrays
|
||||
for j in range(series_lengths[i]):
|
||||
point = series[j]
|
||||
timestamps_arrays[i][j] = point.timestamp
|
||||
values_arrays[i][j] = point.value
|
||||
|
||||
if alloc_failed:
|
||||
raise MemoryError("Failed to allocate memory for series data")
|
||||
|
||||
# Perform merge with full nogil
|
||||
# Pass total_points as capacity (worst case: all points output)
|
||||
with nogil:
|
||||
result_count = _kway_merge_timeseries_nogil(timestamps_arrays, values_arrays,
|
||||
series_lengths, num_series,
|
||||
total_points,
|
||||
&result_timestamps, &result_values)
|
||||
|
||||
if result_count < 0:
|
||||
# Note: _kway_merge_timeseries_nogil frees all memory on error
|
||||
raise MemoryError("Failed during merge operation")
|
||||
|
||||
# Convert C arrays back to Python tuples
|
||||
merged = [None] * result_count
|
||||
for i in range(result_count):
|
||||
merged[i] = (result_timestamps[i], result_values[i])
|
||||
|
||||
return merged
|
||||
|
||||
finally:
|
||||
# Centralized cleanup: safe even if some allocations failed
|
||||
# Free result arrays (allocated by _kway_merge_timeseries_nogil)
|
||||
if result_timestamps:
|
||||
free(result_timestamps)
|
||||
if result_values:
|
||||
free(result_values)
|
||||
if timestamps_arrays:
|
||||
for i in range(num_series):
|
||||
if timestamps_arrays[i]:
|
||||
free(timestamps_arrays[i])
|
||||
free(timestamps_arrays)
|
||||
if values_arrays:
|
||||
for i in range(num_series):
|
||||
if values_arrays[i]:
|
||||
free(values_arrays[i])
|
||||
free(values_arrays)
|
||||
if series_lengths:
|
||||
free(series_lengths)
|
||||
|
||||
|
||||
cdef double _compute_time_weighted_average_nogil(double* timestamps, double* values, int n,
|
||||
double window_start, double window_end) noexcept nogil:
|
||||
"""
|
||||
Fully nogil time-weighted average computation on C arrays.
|
||||
|
||||
Returns: Time-weighted average or NaN to indicate None (invalid result)
|
||||
"""
|
||||
cdef:
|
||||
int i
|
||||
double total_weighted_value = 0.0
|
||||
double total_duration = 0.0
|
||||
double current_value = 0.0
|
||||
double current_time
|
||||
double timestamp, value, duration
|
||||
|
||||
if window_end <= window_start:
|
||||
return nan("")
|
||||
|
||||
current_time = window_start
|
||||
|
||||
# Find value at window_start (LOCF)
|
||||
for i in range(n):
|
||||
timestamp = timestamps[i]
|
||||
|
||||
if timestamp <= window_start:
|
||||
current_value = values[i]
|
||||
else:
|
||||
break
|
||||
|
||||
# Process segments
|
||||
for i in range(n):
|
||||
timestamp = timestamps[i]
|
||||
value = values[i]
|
||||
|
||||
if timestamp <= window_start:
|
||||
continue
|
||||
|
||||
if timestamp >= window_end:
|
||||
break
|
||||
|
||||
# Add contribution of current segment
|
||||
# Note: timestamp < window_end is guaranteed here due to the break above
|
||||
duration = timestamp - current_time
|
||||
|
||||
if duration > 0:
|
||||
total_weighted_value += current_value * duration
|
||||
total_duration += duration
|
||||
|
||||
current_value = value
|
||||
current_time = timestamp
|
||||
|
||||
# Add final segment
|
||||
if current_time < window_end:
|
||||
duration = window_end - current_time
|
||||
total_weighted_value += current_value * duration
|
||||
total_duration += duration
|
||||
|
||||
if total_duration > 0:
|
||||
return total_weighted_value / total_duration
|
||||
|
||||
return nan("")
|
||||
|
||||
|
||||
def time_weighted_average_cython(list timeseries, double window_start,
|
||||
double window_end, double last_window_s=1.0):
|
||||
"""
|
||||
Cython-optimized time-weighted average calculation.
|
||||
|
||||
Assumptions:
|
||||
- Input timeseries is sorted by timestamp in ascending order
|
||||
- Values are treated as a step function (LOCF - Last Observation Carried Forward)
|
||||
|
||||
Args:
|
||||
timeseries: List of objects with .timestamp and .value attributes
|
||||
window_start: Start of window (negative infinity means use first timestamp)
|
||||
window_end: End of window (negative infinity means use last timestamp + last_window_s)
|
||||
last_window_s: Window size for last segment
|
||||
|
||||
Returns:
|
||||
Time-weighted average or None (returned as None when result would be invalid)
|
||||
"""
|
||||
if not timeseries:
|
||||
return None
|
||||
|
||||
cdef:
|
||||
int n = len(timeseries)
|
||||
int i
|
||||
double result
|
||||
object point
|
||||
double* timestamps = <double*>malloc(n * sizeof(double))
|
||||
double* values = <double*>malloc(n * sizeof(double))
|
||||
|
||||
if not timestamps or not values:
|
||||
if timestamps:
|
||||
free(timestamps)
|
||||
if values:
|
||||
free(values)
|
||||
raise MemoryError("Failed to allocate memory for time weighted average")
|
||||
|
||||
try:
|
||||
# Extract data from Python objects into C arrays
|
||||
for i in range(n):
|
||||
point = timeseries[i]
|
||||
timestamps[i] = point.timestamp
|
||||
values[i] = point.value
|
||||
|
||||
# Handle window boundaries
|
||||
# Use negative infinity as sentinel for None (any valid float including -1.0 works)
|
||||
if isinf(window_start) and window_start < 0:
|
||||
window_start = timestamps[0]
|
||||
|
||||
if isinf(window_end) and window_end < 0:
|
||||
window_end = timestamps[n - 1] + last_window_s
|
||||
|
||||
# Compute with full nogil
|
||||
with nogil:
|
||||
result = _compute_time_weighted_average_nogil(timestamps, values, n,
|
||||
window_start, window_end)
|
||||
|
||||
return None if isnan(result) else result
|
||||
|
||||
finally:
|
||||
free(timestamps)
|
||||
free(values)
|
||||
@@ -0,0 +1,191 @@
|
||||
from libcpp cimport bool as c_bool
|
||||
from libcpp.string cimport string as c_string
|
||||
from libc.stdint cimport uint8_t, uint32_t, int64_t
|
||||
|
||||
# Note: we removed the staticmethod declarations in
|
||||
# https://github.com/ray-project/ray/pull/47984 due
|
||||
# to a compiler bug in Cython 3.0.x -- we should see
|
||||
# if we can bring them back in Cython 3.1.x if the
|
||||
# bug is fixed.
|
||||
cdef extern from "ray/common/id.h" namespace "ray" nogil:
|
||||
cdef cppclass CBaseID[T]:
|
||||
size_t Hash() const
|
||||
c_bool IsNil() const
|
||||
c_bool operator==(const CBaseID &rhs) const
|
||||
c_bool operator!=(const CBaseID &rhs) const
|
||||
const uint8_t *data() const
|
||||
|
||||
c_string Binary() const
|
||||
c_string Hex() const
|
||||
|
||||
cdef cppclass CUniqueID "ray::UniqueID"(CBaseID[CUniqueID]):
|
||||
CUniqueID()
|
||||
|
||||
@staticmethod
|
||||
size_t Size()
|
||||
|
||||
@staticmethod
|
||||
CUniqueID FromRandom()
|
||||
|
||||
@staticmethod
|
||||
CUniqueID FromBinary(const c_string &binary)
|
||||
|
||||
@staticmethod
|
||||
const CUniqueID Nil()
|
||||
|
||||
cdef cppclass CActorClassID "ray::ActorClassID"(CBaseID[CActorClassID]):
|
||||
|
||||
@staticmethod
|
||||
CActorClassID FromHex(const c_string &hex_str)
|
||||
|
||||
cdef cppclass CActorID "ray::ActorID"(CBaseID[CActorID]):
|
||||
|
||||
@staticmethod
|
||||
CActorID FromBinary(const c_string &binary)
|
||||
|
||||
@staticmethod
|
||||
CActorID FromHex(const c_string &hex_str)
|
||||
|
||||
@staticmethod
|
||||
const CActorID Nil()
|
||||
|
||||
@staticmethod
|
||||
size_t Size()
|
||||
|
||||
@staticmethod
|
||||
CActorID Of(CJobID job_id, CTaskID parent_task_id,
|
||||
int64_t parent_task_counter)
|
||||
|
||||
CJobID JobId()
|
||||
|
||||
cdef cppclass CNodeID "ray::NodeID"(CBaseID[CNodeID]):
|
||||
|
||||
@staticmethod
|
||||
CNodeID FromHex(const c_string &hex_str)
|
||||
|
||||
@staticmethod
|
||||
const CNodeID Nil()
|
||||
|
||||
cdef cppclass CConfigID "ray::ConfigID"(CBaseID[CConfigID]):
|
||||
pass
|
||||
|
||||
cdef cppclass CFunctionID "ray::FunctionID"(CBaseID[CFunctionID]):
|
||||
|
||||
@staticmethod
|
||||
CFunctionID FromHex(const c_string &hex_str)
|
||||
|
||||
cdef cppclass CJobID "ray::JobID"(CBaseID[CJobID]):
|
||||
|
||||
@staticmethod
|
||||
CJobID FromBinary(const c_string &binary)
|
||||
|
||||
@staticmethod
|
||||
CJobID FromHex(const c_string &hex_str)
|
||||
|
||||
@staticmethod
|
||||
const CJobID Nil()
|
||||
|
||||
@staticmethod
|
||||
size_t Size()
|
||||
|
||||
@staticmethod
|
||||
CJobID FromInt(uint32_t value)
|
||||
|
||||
uint32_t ToInt()
|
||||
|
||||
cdef cppclass CTaskID "ray::TaskID"(CBaseID[CTaskID]):
|
||||
|
||||
@staticmethod
|
||||
CTaskID FromBinary(const c_string &binary)
|
||||
|
||||
@staticmethod
|
||||
CTaskID FromHex(const c_string &hex_str)
|
||||
|
||||
@staticmethod
|
||||
const CTaskID Nil()
|
||||
|
||||
@staticmethod
|
||||
size_t Size()
|
||||
|
||||
@staticmethod
|
||||
CTaskID ForDriverTask(const CJobID &job_id)
|
||||
|
||||
@staticmethod
|
||||
CTaskID FromRandom(const CJobID &job_id)
|
||||
|
||||
@staticmethod
|
||||
CTaskID ForActorCreationTask(CActorID actor_id)
|
||||
|
||||
@staticmethod
|
||||
CTaskID ForActorTask(CJobID job_id, CTaskID parent_task_id,
|
||||
int64_t parent_task_counter, CActorID actor_id)
|
||||
|
||||
@staticmethod
|
||||
CTaskID ForNormalTask(CJobID job_id, CTaskID parent_task_id,
|
||||
int64_t parent_task_counter)
|
||||
|
||||
CActorID ActorId() const
|
||||
|
||||
CJobID JobId() const
|
||||
|
||||
cdef cppclass CObjectID" ray::ObjectID"(CBaseID[CObjectID]):
|
||||
|
||||
@staticmethod
|
||||
int64_t MaxObjectIndex()
|
||||
|
||||
@staticmethod
|
||||
CObjectID FromBinary(const c_string &binary)
|
||||
|
||||
@staticmethod
|
||||
CObjectID FromRandom()
|
||||
|
||||
@staticmethod
|
||||
const CObjectID Nil()
|
||||
|
||||
@staticmethod
|
||||
CObjectID FromIndex(const CTaskID &task_id, int64_t index)
|
||||
|
||||
@staticmethod
|
||||
size_t Size()
|
||||
|
||||
c_bool is_put()
|
||||
|
||||
int64_t ObjectIndex() const
|
||||
|
||||
CTaskID TaskId() const
|
||||
|
||||
cdef cppclass CClusterID "ray::ClusterID"(CBaseID[CClusterID]):
|
||||
|
||||
@staticmethod
|
||||
CClusterID FromHex(const c_string &hex_str)
|
||||
|
||||
@staticmethod
|
||||
CClusterID FromRandom()
|
||||
|
||||
@staticmethod
|
||||
const CClusterID Nil()
|
||||
|
||||
cdef cppclass CWorkerID "ray::WorkerID"(CBaseID[CWorkerID]):
|
||||
|
||||
@staticmethod
|
||||
CWorkerID FromHex(const c_string &hex_str)
|
||||
|
||||
cdef cppclass CPlacementGroupID "ray::PlacementGroupID" \
|
||||
(CBaseID[CPlacementGroupID]):
|
||||
|
||||
@staticmethod
|
||||
CPlacementGroupID FromBinary(const c_string &binary)
|
||||
|
||||
@staticmethod
|
||||
CPlacementGroupID FromHex(const c_string &hex_str)
|
||||
|
||||
@staticmethod
|
||||
const CPlacementGroupID Nil()
|
||||
|
||||
@staticmethod
|
||||
size_t Size()
|
||||
|
||||
@staticmethod
|
||||
CPlacementGroupID Of(CJobID job_id)
|
||||
|
||||
ctypedef uint32_t ObjectIDIndexType
|
||||
@@ -0,0 +1,436 @@
|
||||
"""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.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from ray.includes.unique_ids cimport (
|
||||
CActorClassID,
|
||||
CActorID,
|
||||
CNodeID,
|
||||
CConfigID,
|
||||
CJobID,
|
||||
CFunctionID,
|
||||
CObjectID,
|
||||
CTaskID,
|
||||
CUniqueID,
|
||||
CWorkerID,
|
||||
CPlacementGroupID,
|
||||
CClusterID,
|
||||
)
|
||||
|
||||
|
||||
import ray
|
||||
from ray._common.utils import decode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def check_id(b, size=kUniqueIDSize):
|
||||
if not isinstance(b, bytes):
|
||||
raise TypeError(f"Unsupported type: {type(b)}, expected bytes")
|
||||
if len(b) != size:
|
||||
raise ValueError("ID string needs to have length " +
|
||||
str(size) + ", got " + str(len(b)))
|
||||
|
||||
|
||||
def check_nil(id_obj):
|
||||
if id_obj.is_nil():
|
||||
raise ValueError(f"{id_obj.__class__.__name__} is nil")
|
||||
|
||||
|
||||
cdef extern from "ray/common/constants.h" nogil:
|
||||
cdef int64_t kUniqueIDSize
|
||||
|
||||
|
||||
cdef class BaseID:
|
||||
|
||||
@classmethod
|
||||
def from_binary(cls, id_bytes):
|
||||
return cls(id_bytes)
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, hex_id):
|
||||
raise NotImplementedError
|
||||
|
||||
cdef size_t hash(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def binary(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def size(cls):
|
||||
raise NotImplementedError
|
||||
|
||||
def hex(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def is_nil(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def __hash__(self):
|
||||
return self.hash()
|
||||
|
||||
def __eq__(self, other):
|
||||
return type(self) == type(other) and self.binary() == other.binary()
|
||||
|
||||
def __ne__(self, other):
|
||||
return type(self) != type(other) or self.binary() != other.binary()
|
||||
|
||||
def __bytes__(self):
|
||||
return self.binary()
|
||||
|
||||
def __hex__(self):
|
||||
return self.hex()
|
||||
|
||||
def __repr__(self):
|
||||
return self.__class__.__name__ + "(" + 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.
|
||||
# Do not use __hash__ that returns signed uint64_t, which
|
||||
# is different from std::hash in c++ code.
|
||||
return self.hash()
|
||||
|
||||
|
||||
cdef class UniqueID(BaseID):
|
||||
cdef CUniqueID data
|
||||
|
||||
def __init__(self, id):
|
||||
check_id(id)
|
||||
self.data = CUniqueID.FromBinary(id)
|
||||
|
||||
@classmethod
|
||||
def nil(cls):
|
||||
return cls(CUniqueID.Nil().Binary())
|
||||
|
||||
@classmethod
|
||||
def from_random(cls):
|
||||
return cls(CUniqueID.FromRandom().Binary())
|
||||
|
||||
@classmethod
|
||||
def size(cls):
|
||||
return CUniqueID.Size()
|
||||
|
||||
def binary(self):
|
||||
return self.data.Binary()
|
||||
|
||||
def hex(self):
|
||||
return decode(self.data.Hex())
|
||||
|
||||
def is_nil(self):
|
||||
return self.data.IsNil()
|
||||
|
||||
cdef size_t hash(self):
|
||||
return self.data.Hash()
|
||||
|
||||
|
||||
cdef class TaskID(BaseID):
|
||||
cdef CTaskID data
|
||||
|
||||
def __init__(self, id):
|
||||
check_id(id, CTaskID.Size())
|
||||
self.data = CTaskID.FromBinary(<c_string>id)
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, hex_id):
|
||||
binary_id = CTaskID.FromHex(<c_string>hex_id).Binary()
|
||||
return cls(binary_id)
|
||||
|
||||
cdef CTaskID native(self):
|
||||
return <CTaskID>self.data
|
||||
|
||||
@classmethod
|
||||
def size(cls):
|
||||
return CTaskID.Size()
|
||||
|
||||
def binary(self):
|
||||
return self.data.Binary()
|
||||
|
||||
def hex(self):
|
||||
return decode(self.data.Hex())
|
||||
|
||||
def is_nil(self):
|
||||
return self.data.IsNil()
|
||||
|
||||
def actor_id(self):
|
||||
return ActorID(self.data.ActorId().Binary())
|
||||
|
||||
def job_id(self):
|
||||
check_nil(self)
|
||||
return JobID(self.data.JobId().Binary())
|
||||
|
||||
cdef size_t hash(self):
|
||||
return self.data.Hash()
|
||||
|
||||
@classmethod
|
||||
def nil(cls):
|
||||
return cls(CTaskID.Nil().Binary())
|
||||
|
||||
@classmethod
|
||||
def for_fake_task(cls, job_id):
|
||||
return cls(CTaskID.FromRandom(
|
||||
CJobID.FromBinary(job_id.binary())).Binary())
|
||||
|
||||
@classmethod
|
||||
def for_driver_task(cls, job_id):
|
||||
return cls(CTaskID.ForDriverTask(
|
||||
CJobID.FromBinary(job_id.binary())).Binary())
|
||||
|
||||
@classmethod
|
||||
def for_actor_creation_task(cls, actor_id):
|
||||
assert isinstance(actor_id, ActorID)
|
||||
return cls(CTaskID.ForActorCreationTask(
|
||||
CActorID.FromBinary(actor_id.binary())).Binary())
|
||||
|
||||
@classmethod
|
||||
def for_actor_task(cls, job_id, parent_task_id,
|
||||
parent_task_counter, actor_id):
|
||||
assert isinstance(job_id, JobID)
|
||||
assert isinstance(parent_task_id, TaskID)
|
||||
assert isinstance(actor_id, ActorID)
|
||||
return cls(CTaskID.ForActorTask(
|
||||
CJobID.FromBinary(job_id.binary()),
|
||||
CTaskID.FromBinary(parent_task_id.binary()),
|
||||
parent_task_counter,
|
||||
CActorID.FromBinary(actor_id.binary())).Binary())
|
||||
|
||||
@classmethod
|
||||
def for_normal_task(cls, job_id, parent_task_id, parent_task_counter):
|
||||
assert isinstance(job_id, JobID)
|
||||
assert isinstance(parent_task_id, TaskID)
|
||||
return cls(CTaskID.ForNormalTask(
|
||||
CJobID.FromBinary(job_id.binary()),
|
||||
CTaskID.FromBinary(parent_task_id.binary()),
|
||||
parent_task_counter).Binary())
|
||||
|
||||
cdef class NodeID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
check_id(id)
|
||||
self.data = CUniqueID.FromBinary(<c_string>id)
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, hex_id):
|
||||
binary_id = CNodeID.FromHex(<c_string>hex_id).Binary()
|
||||
return cls(binary_id)
|
||||
|
||||
cdef CNodeID native(self):
|
||||
return <CNodeID>self.data
|
||||
|
||||
|
||||
cdef class JobID(BaseID):
|
||||
cdef CJobID data
|
||||
|
||||
def __init__(self, id):
|
||||
check_id(id, CJobID.Size())
|
||||
self.data = CJobID.FromBinary(<c_string>id)
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, hex_id):
|
||||
binary_id = CJobID.FromHex(<c_string>hex_id).Binary()
|
||||
return cls(binary_id)
|
||||
|
||||
cdef CJobID native(self):
|
||||
return <CJobID>self.data
|
||||
|
||||
@classmethod
|
||||
def from_int(cls, value):
|
||||
assert value < 2**32, "Maximum JobID integer is 2**32 - 1."
|
||||
return cls(CJobID.FromInt(value).Binary())
|
||||
|
||||
@classmethod
|
||||
def nil(cls):
|
||||
return cls(CJobID.Nil().Binary())
|
||||
|
||||
@classmethod
|
||||
def size(cls):
|
||||
return CJobID.Size()
|
||||
|
||||
def int(self):
|
||||
return self.data.ToInt()
|
||||
|
||||
def binary(self):
|
||||
return self.data.Binary()
|
||||
|
||||
def hex(self):
|
||||
return decode(self.data.Hex())
|
||||
|
||||
def is_nil(self):
|
||||
return self.data.IsNil()
|
||||
|
||||
cdef size_t hash(self):
|
||||
return self.data.Hash()
|
||||
|
||||
cdef class WorkerID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
check_id(id)
|
||||
self.data = CUniqueID.FromBinary(<c_string>id)
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, hex_id):
|
||||
binary_id = CWorkerID.FromHex(<c_string>hex_id).Binary()
|
||||
return cls(binary_id)
|
||||
|
||||
cdef CWorkerID native(self):
|
||||
return <CWorkerID>self.data
|
||||
|
||||
cdef class ActorID(BaseID):
|
||||
|
||||
def __init__(self, id):
|
||||
self._set_id(id)
|
||||
|
||||
@classmethod
|
||||
def of(cls, job_id, parent_task_id, parent_task_counter):
|
||||
assert isinstance(job_id, JobID)
|
||||
assert isinstance(parent_task_id, TaskID)
|
||||
return cls(CActorID.Of(CJobID.FromBinary(job_id.binary()),
|
||||
CTaskID.FromBinary(parent_task_id.binary()),
|
||||
parent_task_counter).Binary())
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, hex_id):
|
||||
binary_id = CActorID.FromHex(<c_string>hex_id).Binary()
|
||||
return cls(binary_id)
|
||||
|
||||
@classmethod
|
||||
def nil(cls):
|
||||
return cls(CActorID.Nil().Binary())
|
||||
|
||||
@classmethod
|
||||
def from_random(cls):
|
||||
return cls(os.urandom(CActorID.Size()))
|
||||
|
||||
@classmethod
|
||||
def size(cls):
|
||||
return CActorID.Size()
|
||||
|
||||
def _set_id(self, id):
|
||||
check_id(id, CActorID.Size())
|
||||
self.data = CActorID.FromBinary(<c_string>id)
|
||||
|
||||
@property
|
||||
def job_id(self):
|
||||
check_nil(self)
|
||||
return JobID(self.data.JobId().Binary())
|
||||
|
||||
def binary(self):
|
||||
return self.data.Binary()
|
||||
|
||||
def hex(self):
|
||||
return decode(self.data.Hex())
|
||||
|
||||
def is_nil(self):
|
||||
return self.data.IsNil()
|
||||
|
||||
cdef size_t hash(self):
|
||||
return self.data.Hash()
|
||||
|
||||
cdef CActorID native(self):
|
||||
return <CActorID>self.data
|
||||
|
||||
|
||||
cdef class FunctionID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
check_id(id)
|
||||
self.data = CUniqueID.FromBinary(<c_string>id)
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, hex_id):
|
||||
binary_id = CFunctionID.FromHex(<c_string>hex_id).Binary()
|
||||
return cls(binary_id)
|
||||
|
||||
cdef CFunctionID native(self):
|
||||
return <CFunctionID>self.data
|
||||
|
||||
|
||||
cdef class ActorClassID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
check_id(id)
|
||||
self.data = CUniqueID.FromBinary(<c_string>id)
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, hex_id):
|
||||
binary_id = CActorClassID.FromHex(<c_string>hex_id).Binary()
|
||||
return cls(binary_id)
|
||||
|
||||
cdef CActorClassID native(self):
|
||||
return <CActorClassID>self.data
|
||||
|
||||
cdef class ClusterID(UniqueID):
|
||||
|
||||
def __init__(self, id):
|
||||
check_id(id)
|
||||
self.data = CUniqueID.FromBinary(<c_string>id)
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, hex_id):
|
||||
binary_id = CClusterID.FromHex(<c_string>hex_id).Binary()
|
||||
return cls(binary_id)
|
||||
|
||||
cdef CClusterID native(self):
|
||||
return <CClusterID>self.data
|
||||
|
||||
|
||||
# This type alias is for backward compatibility.
|
||||
ObjectID = ObjectRef
|
||||
|
||||
cdef class PlacementGroupID(BaseID):
|
||||
cdef CPlacementGroupID data
|
||||
|
||||
def __init__(self, id):
|
||||
check_id(id, CPlacementGroupID.Size())
|
||||
self.data = CPlacementGroupID.FromBinary(<c_string>id)
|
||||
|
||||
cdef CPlacementGroupID native(self):
|
||||
return <CPlacementGroupID>self.data
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, hex_id):
|
||||
binary_id = CPlacementGroupID.FromHex(<c_string>hex_id).Binary()
|
||||
return cls(binary_id)
|
||||
|
||||
@classmethod
|
||||
def from_random(cls):
|
||||
return cls(os.urandom(CPlacementGroupID.Size()))
|
||||
|
||||
@classmethod
|
||||
def of(cls, job_id):
|
||||
assert isinstance(job_id, JobID)
|
||||
return cls(CPlacementGroupID.Of(CJobID.FromBinary(job_id.binary())).Binary())
|
||||
|
||||
@classmethod
|
||||
def nil(cls):
|
||||
return cls(CPlacementGroupID.Nil().Binary())
|
||||
|
||||
@classmethod
|
||||
def size(cls):
|
||||
return CPlacementGroupID.Size()
|
||||
|
||||
def binary(self):
|
||||
return self.data.Binary()
|
||||
|
||||
def hex(self):
|
||||
return decode(self.data.Hex())
|
||||
|
||||
def is_nil(self):
|
||||
return self.data.IsNil()
|
||||
|
||||
cdef size_t hash(self):
|
||||
return self.data.Hash()
|
||||
@@ -0,0 +1,146 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Tuple, TypeVar
|
||||
|
||||
# backwards compatibility. Luckily circular references are fine in type stubs
|
||||
from ray._raylet import ObjectRef
|
||||
|
||||
ObjectID = ObjectRef
|
||||
|
||||
# implementations are in unique_ids.pxi
|
||||
def check_id(b: bytes, size: int = ...) -> None: ...
|
||||
|
||||
_BID = TypeVar("_BID", bound=BaseID)
|
||||
class BaseID:
|
||||
|
||||
@classmethod
|
||||
def from_binary(cls: type[_BID], id_bytes: bytes) -> _BID: ...
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls: type[_BID], hex_id: str | bytes) -> _BID: ...
|
||||
|
||||
def binary(self) -> bytes: ...
|
||||
|
||||
@classmethod
|
||||
def size(cls) -> int: ...
|
||||
|
||||
def hex(self) -> str: ...
|
||||
|
||||
def is_nil(self) -> bool: ...
|
||||
|
||||
def __hash__(self) -> int: ...
|
||||
|
||||
def __eq__(self, other: object) -> bool: ...
|
||||
|
||||
def __ne__(self, other: object) -> bool: ...
|
||||
|
||||
def __bytes__(self) -> bytes: ...
|
||||
|
||||
def __hex__(self) -> str: ...
|
||||
|
||||
def __repr__(self) -> str: ...
|
||||
|
||||
def __str__(self) -> str: ...
|
||||
|
||||
def __reduce__(self: _BID) -> Tuple[type[_BID], Tuple[bytes]]: ...
|
||||
|
||||
def redis_shard_hash(self) -> int: ...
|
||||
|
||||
|
||||
_UID = TypeVar("_UID", bound=UniqueID)
|
||||
class UniqueID(BaseID):
|
||||
|
||||
def __init__(self, id: bytes) -> None: ...
|
||||
|
||||
@classmethod
|
||||
def nil(cls: type[_UID]) -> _UID: ...
|
||||
|
||||
@classmethod
|
||||
def from_random(cls: type[_UID]) -> _UID: ...
|
||||
|
||||
|
||||
_TID = TypeVar("_TID", bound=TaskID)
|
||||
class TaskID(BaseID):
|
||||
|
||||
def __init__(self, id: bytes) -> None: ...
|
||||
|
||||
def actor_id(self) -> ActorID: ...
|
||||
|
||||
def job_id(self) -> JobID: ...
|
||||
|
||||
@classmethod
|
||||
def nil(cls: type[_TID]) -> _TID: ...
|
||||
|
||||
@classmethod
|
||||
def for_fake_task(cls: type[_TID], job_id: JobID) -> _TID: ...
|
||||
|
||||
@classmethod
|
||||
def for_driver_task(cls: type[_TID], job_id: JobID) -> _TID: ...
|
||||
|
||||
@classmethod
|
||||
def for_actor_creation_task(cls: type[_TID], actor_id: ActorID) -> _TID: ...
|
||||
|
||||
@classmethod
|
||||
def for_actor_task(cls: type[_TID], job_id: JobID, parent_task_id: TaskID,
|
||||
parent_task_counter: int, actor_id: ActorID) -> _TID: ...
|
||||
|
||||
@classmethod
|
||||
def for_normal_task(cls: type[_TID], job_id: JobID, parent_task_id: TaskID, parent_task_counter: int) -> _TID: ...
|
||||
|
||||
|
||||
class NodeID(UniqueID): ...
|
||||
|
||||
_JID = TypeVar("_JID", bound=JobID)
|
||||
class JobID(BaseID):
|
||||
|
||||
def __init__(self, id: bytes) -> None: ...
|
||||
|
||||
@classmethod
|
||||
def from_int(cls: type[_JID], value: int) -> _JID: ...
|
||||
|
||||
@classmethod
|
||||
def nil(cls: type[_JID]) -> _JID: ...
|
||||
|
||||
def int(self) -> int: ...
|
||||
|
||||
|
||||
class WorkerID(UniqueID): ...
|
||||
|
||||
_AID = TypeVar("_AID", bound=ActorID)
|
||||
class ActorID(BaseID):
|
||||
|
||||
def __init__(self, id: bytes) -> None: ...
|
||||
|
||||
@classmethod
|
||||
def of(cls: type[_AID], job_id: JobID, parent_task_id: TaskID, parent_task_counter: int) -> _AID: ...
|
||||
|
||||
@classmethod
|
||||
def nil(cls: type[_AID]) -> _AID: ...
|
||||
|
||||
@classmethod
|
||||
def from_random(cls: type[_AID]) -> _AID: ...
|
||||
|
||||
def _set_id(self, id: bytes) -> None: ...
|
||||
|
||||
@property
|
||||
def job_id(self) -> JobID: ...
|
||||
|
||||
|
||||
class FunctionID(UniqueID): ...
|
||||
class ActorClassID(UniqueID): ...
|
||||
class ClusterID(UniqueID): ...
|
||||
|
||||
|
||||
_PGID = TypeVar("_PGID", bound=PlacementGroupID)
|
||||
class PlacementGroupID(BaseID):
|
||||
|
||||
def __init__(self, id: bytes) -> None: ...
|
||||
|
||||
@classmethod
|
||||
def from_random(cls: type[_PGID]) -> _PGID: ...
|
||||
|
||||
@classmethod
|
||||
def of(cls: type[_PGID], job_id: JobID) -> _PGID: ...
|
||||
|
||||
@classmethod
|
||||
def nil(cls: type[_PGID]) -> _PGID: ...
|
||||
Reference in New Issue
Block a user