wehub-resource-sync
68 lines
2.3 KiB
Cython
68 lines
2.3 KiB
Cython
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)
|