wehub-resource-sync
290 lines
11 KiB
Cython
290 lines
11 KiB
Cython
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},
|
|
}
|