wehub-resource-sync
1004 lines
39 KiB
Cython
1004 lines
39 KiB
Cython
# cython: profile=False
|
|
# distutils: language = c++
|
|
# cython: embedsignature = True
|
|
# cython: language_level = 3
|
|
# cython: c_string_encoding = default
|
|
|
|
"""
|
|
Binding of C++ ray::gcs::GcsClient.
|
|
"""
|
|
|
|
# This file is a .pxi which is included in _raylet.pyx. This means any already-imported
|
|
# symbols can directly be used here. This is not ideal, but we can't easily split this
|
|
# out to a separate translation unit because we need to access the singleton thread.
|
|
#
|
|
# We need to best-effort import everything we need.
|
|
#
|
|
# For how async API are implemented, see src/ray/common/python_callbacks.h
|
|
from asyncio import Future
|
|
from ray._common.utils import get_or_create_event_loop
|
|
from typing import Dict, List, Sequence, Tuple
|
|
from libcpp.utility cimport move
|
|
import concurrent.futures
|
|
from ray.core.generated.gcs_service_pb2 import GetAllResourceUsageReply
|
|
from ray.includes.common cimport (
|
|
CGcsClient,
|
|
CGetAllResourceUsageReply,
|
|
ConnectOnSingletonIoContext,
|
|
MultiItemPyCallback,
|
|
OptionalItemPyCallback,
|
|
StatusPyCallback,
|
|
CGetClusterStatusReply,
|
|
CStatusOr,
|
|
CGcsNodeState,
|
|
CNodeSelector,
|
|
CGcsNodeInfo,
|
|
CAddEventsRequest,
|
|
CAddEventsReply,
|
|
CRayStatus,
|
|
)
|
|
from ray.includes.optional cimport optional, make_optional
|
|
from ray.core.generated import gcs_pb2, autoscaler_pb2
|
|
from ray.core.generated.gcs_service_pb2 import GetAllNodeInfoRequest
|
|
from cython.operator import dereference, postincrement
|
|
cimport cpython
|
|
|
|
|
|
cdef class InnerGcsClient:
|
|
cdef:
|
|
shared_ptr[CGcsClient] inner
|
|
|
|
# Creates and connects a standalone GcsClient.
|
|
# cluster_id is in hex, if any.
|
|
# TODO(ryw): we can also reuse the CoreWorker's GcsClient to save resources.
|
|
@staticmethod
|
|
def standalone(gcs_address: str,
|
|
cluster_id: Optional[str],
|
|
timeout_ms: int) -> "InnerGcsClient":
|
|
cdef GcsClientOptions gcs_options = None
|
|
if cluster_id:
|
|
gcs_options = GcsClientOptions.create(
|
|
gcs_address, cluster_id, allow_cluster_id_nil=False,
|
|
fetch_cluster_id_if_nil=False)
|
|
else:
|
|
gcs_options = GcsClientOptions.create(
|
|
gcs_address, None, allow_cluster_id_nil=True,
|
|
fetch_cluster_id_if_nil=True)
|
|
cdef CGcsClientOptions* native = gcs_options.native()
|
|
cdef shared_ptr[CGcsClient] inner = make_shared[CGcsClient](
|
|
dereference(native))
|
|
cdef int64_t c_timeout_ms = timeout_ms
|
|
|
|
with nogil:
|
|
check_status_timeout_as_rpc_error(
|
|
ConnectOnSingletonIoContext(dereference(inner), c_timeout_ms))
|
|
|
|
gcs_client = InnerGcsClient()
|
|
gcs_client.inner = inner
|
|
return gcs_client
|
|
|
|
@property
|
|
def address(self) -> str:
|
|
cdef c_pair[c_string, int] pair = self.inner.get().GetGcsServerAddress()
|
|
host = pair.first.decode("utf-8")
|
|
port = pair.second
|
|
return build_address(host, port)
|
|
|
|
@property
|
|
def cluster_id(self) -> ray.ClusterID:
|
|
cdef CClusterID cluster_id = self.inner.get().GetClusterId()
|
|
return ray.ClusterID.from_binary(cluster_id.Binary())
|
|
|
|
#############################################################
|
|
# Internal KV sync methods
|
|
#############################################################
|
|
def internal_kv_get(
|
|
self, c_string key, namespace=None, timeout=None
|
|
) -> Optional[bytes]:
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
optional[c_string] opt_value = c_string()
|
|
CRayStatus status
|
|
with nogil:
|
|
status = self.inner.get().InternalKV().Get(
|
|
ns, key, timeout_ms, opt_value.value())
|
|
return raise_or_return(
|
|
convert_optional_str_none_for_not_found(status, move(opt_value)))
|
|
|
|
def internal_kv_multi_get(
|
|
self, keys: List[bytes], namespace=None, timeout=None
|
|
) -> Dict[bytes, bytes]:
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
c_vector[c_string] c_keys = [key for key in keys]
|
|
optional[unordered_map[c_string, c_string]] opt_values = \
|
|
unordered_map[c_string, c_string]()
|
|
CRayStatus status
|
|
with nogil:
|
|
status = self.inner.get().InternalKV().MultiGet(
|
|
ns, c_keys, timeout_ms, opt_values.value())
|
|
return raise_or_return(convert_optional_multi_get(status, move(opt_values)))
|
|
|
|
def internal_kv_put(self, c_string key, c_string value, c_bool overwrite=False,
|
|
namespace=None, timeout=None) -> int:
|
|
"""
|
|
Returns 1 if the key is newly added, 0 if the key is overwritten.
|
|
"""
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
optional[c_bool] opt_added = 0
|
|
CRayStatus status
|
|
with nogil:
|
|
status = self.inner.get().InternalKV().Put(
|
|
ns, key, value, overwrite, timeout_ms, opt_added.value())
|
|
added = raise_or_return(convert_optional_bool(status, move(opt_added)))
|
|
return 1 if added else 0
|
|
|
|
def internal_kv_del(self, c_string key, c_bool del_by_prefix,
|
|
namespace=None, timeout=None) -> int:
|
|
"""
|
|
Returns number of keys deleted.
|
|
"""
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
optional[int] opt_num_deleted = 0
|
|
CRayStatus status
|
|
with nogil:
|
|
status = self.inner.get().InternalKV().Del(
|
|
ns, key, del_by_prefix, timeout_ms, opt_num_deleted.value())
|
|
return raise_or_return(convert_optional_int(status, move(opt_num_deleted)))
|
|
|
|
def internal_kv_keys(
|
|
self, c_string prefix, namespace=None, timeout=None
|
|
) -> List[bytes]:
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
optional[c_vector[c_string]] opt_keys = c_vector[c_string]()
|
|
CRayStatus status
|
|
with nogil:
|
|
status = self.inner.get().InternalKV().Keys(
|
|
ns, prefix, timeout_ms, opt_keys.value())
|
|
return raise_or_return(convert_optional_vector_str(status, move(opt_keys)))
|
|
|
|
def internal_kv_exists(self, c_string key, namespace=None, timeout=None) -> bool:
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
optional[c_bool] opt_exists = 0
|
|
CRayStatus status
|
|
with nogil:
|
|
status = self.inner.get().InternalKV().Exists(
|
|
ns, key, timeout_ms, opt_exists.value())
|
|
return raise_or_return(convert_optional_bool(status, move(opt_exists)))
|
|
|
|
#############################################################
|
|
# Internal KV async methods
|
|
#############################################################
|
|
|
|
def async_internal_kv_get(
|
|
self, c_string key, namespace=None, timeout=None
|
|
) -> Future[Optional[bytes]]:
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
fut = incremented_fut()
|
|
with nogil:
|
|
self.inner.get().InternalKV().AsyncInternalKVGet(
|
|
ns, key, timeout_ms,
|
|
OptionalItemPyCallback[c_string](
|
|
&convert_optional_str_none_for_not_found,
|
|
assign_and_decrement_fut,
|
|
fut))
|
|
return asyncio.wrap_future(fut)
|
|
|
|
def async_internal_kv_multi_get(
|
|
self, keys: List[bytes], namespace=None, timeout=None
|
|
) -> Future[Dict[bytes, bytes]]:
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
c_vector[c_string] c_keys = [key for key in keys]
|
|
fut = incremented_fut()
|
|
with nogil:
|
|
self.inner.get().InternalKV().AsyncInternalKVMultiGet(
|
|
ns, c_keys, timeout_ms,
|
|
OptionalItemPyCallback[unordered_map[c_string, c_string]](
|
|
&convert_optional_multi_get,
|
|
assign_and_decrement_fut,
|
|
fut))
|
|
return asyncio.wrap_future(fut)
|
|
|
|
def async_internal_kv_put(
|
|
self, c_string key, c_string value, c_bool overwrite=False, namespace=None,
|
|
timeout=None
|
|
) -> Future[bool]:
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
fut = incremented_fut()
|
|
with nogil:
|
|
self.inner.get().InternalKV().AsyncInternalKVPut(
|
|
ns, key, value, overwrite, timeout_ms,
|
|
OptionalItemPyCallback[c_bool](
|
|
&convert_optional_bool,
|
|
assign_and_decrement_fut,
|
|
fut))
|
|
return asyncio.wrap_future(fut)
|
|
|
|
def async_internal_kv_del(self, c_string key, c_bool del_by_prefix,
|
|
namespace=None, timeout=None) -> Future[int]:
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
fut = incremented_fut()
|
|
with nogil:
|
|
self.inner.get().InternalKV().AsyncInternalKVDel(
|
|
ns, key, del_by_prefix, timeout_ms,
|
|
OptionalItemPyCallback[int](
|
|
&convert_optional_int,
|
|
assign_and_decrement_fut,
|
|
fut))
|
|
return asyncio.wrap_future(fut)
|
|
|
|
def async_internal_kv_keys(self, c_string prefix, namespace=None, timeout=None
|
|
) -> Future[List[bytes]]:
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
fut = incremented_fut()
|
|
with nogil:
|
|
self.inner.get().InternalKV().AsyncInternalKVKeys(
|
|
ns, prefix, timeout_ms,
|
|
OptionalItemPyCallback[c_vector[c_string]](
|
|
&convert_optional_vector_str,
|
|
assign_and_decrement_fut,
|
|
fut))
|
|
return asyncio.wrap_future(fut)
|
|
|
|
def async_internal_kv_exists(self, c_string key, namespace=None, timeout=None
|
|
) -> Future[bool]:
|
|
cdef:
|
|
c_string ns = namespace or b""
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
fut = incremented_fut()
|
|
with nogil:
|
|
self.inner.get().InternalKV().AsyncInternalKVExists(
|
|
ns, key, timeout_ms,
|
|
OptionalItemPyCallback[c_bool](
|
|
&convert_optional_bool,
|
|
assign_and_decrement_fut,
|
|
fut))
|
|
return asyncio.wrap_future(fut)
|
|
|
|
#############################################################
|
|
# NodeInfo methods
|
|
#############################################################
|
|
def check_alive(
|
|
self, node_ids: List[NodeID], timeout: Optional[int | float] = None
|
|
) -> List[bool]:
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
c_vector[CNodeID] c_node_ids;
|
|
c_vector[c_bool] results
|
|
CRayStatus status
|
|
c_node_ids.reserve(len(node_ids));
|
|
for node_id in node_ids:
|
|
c_node_ids.push_back((<NodeID>node_id).native())
|
|
with nogil:
|
|
status = self.inner.get().Nodes().CheckAlive(
|
|
c_node_ids, timeout_ms, results)
|
|
return raise_or_return(convert_multi_bool(status, move(results)))
|
|
|
|
def async_check_alive(
|
|
self, node_ids: List[NodeID], timeout: Optional[int | float] = None
|
|
) -> Future[List[bool]]:
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
c_vector[CNodeID] c_node_ids;
|
|
fut = incremented_fut()
|
|
c_node_ids.reserve(len(node_ids));
|
|
for node_id in node_ids:
|
|
c_node_ids.push_back((<NodeID>node_id).native())
|
|
with nogil:
|
|
self.inner.get().Nodes().AsyncCheckAlive(
|
|
c_node_ids, timeout_ms,
|
|
MultiItemPyCallback[c_bool](
|
|
&convert_multi_bool,
|
|
assign_and_decrement_fut,
|
|
fut))
|
|
return asyncio.wrap_future(fut)
|
|
|
|
def drain_nodes(
|
|
self, node_ids: Sequence[bytes], timeout: Optional[int | float] = None
|
|
) -> List[bytes]:
|
|
"""returns a list of node_ids that are successfully drained."""
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
c_vector[CNodeID] c_node_ids
|
|
c_vector[c_string] results
|
|
CRayStatus status
|
|
c_node_ids.reserve(len(node_ids));
|
|
for node_id in node_ids:
|
|
c_node_ids.push_back(<CNodeID>CUniqueID.FromBinary(node_id))
|
|
with nogil:
|
|
status = self.inner.get().Nodes().DrainNodes(
|
|
c_node_ids, timeout_ms, results)
|
|
return raise_or_return(convert_multi_str(status, move(results)))
|
|
|
|
def get_all_node_info(
|
|
self, timeout: Optional[int | float] = None,
|
|
node_selectors: Optional[List[GetAllNodeInfoRequest.NodeSelector]] = None,
|
|
state_filter: Optional[int] = None,
|
|
) -> Dict[NodeID, gcs_pb2.GcsNodeInfo]:
|
|
"""Get all node info with optional filters.
|
|
|
|
Args:
|
|
timeout: Timeout in seconds
|
|
node_selectors: Optional list of GetAllNodeInfoRequest.NodeSelector
|
|
proto objects to filter nodes.
|
|
state_filter: Optional int representing the GcsNodeState enum value
|
|
to filter by node state.
|
|
|
|
Returns:
|
|
Dictionary mapping NodeID to GcsNodeInfo
|
|
"""
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
c_vector[CGcsNodeInfo] reply
|
|
CRayStatus status
|
|
optional[CStatusOr[c_vector[CGcsNodeInfo]]] status_or
|
|
optional[CGcsNodeState] c_state_filter = nullopt
|
|
c_vector[CNodeSelector] c_node_selectors
|
|
CNodeSelector c_node_selector
|
|
|
|
# Convert state_filter int to CGcsNodeState
|
|
if state_filter is not None:
|
|
c_state_filter.emplace(<CGcsNodeState>state_filter)
|
|
|
|
# Convert Python NodeSelector protos to C++ CNodeSelector
|
|
if node_selectors is not None:
|
|
for py_selector in node_selectors:
|
|
c_node_selector = CNodeSelector()
|
|
c_node_selector.ParseFromString(py_selector.SerializeToString())
|
|
c_node_selectors.push_back(c_node_selector)
|
|
|
|
with nogil:
|
|
status_or = self.inner.get().Nodes().GetAllNoCache(timeout_ms, c_state_filter, c_node_selectors)
|
|
status = status_or.value().status()
|
|
if status_or.value().ok():
|
|
reply = move(status_or.value().value())
|
|
return raise_or_return(convert_get_all_node_info(status, move(reply)))
|
|
|
|
def async_get_all_node_info(
|
|
self, timeout: Optional[int | float] = None,
|
|
node_selectors: Optional[List[GetAllNodeInfoRequest.NodeSelector]] = None,
|
|
state_filter: Optional[int] = None,
|
|
limit: Optional[int] = None,
|
|
) -> Future[Tuple[Dict[NodeID, gcs_pb2.GcsNodeInfo], int]]:
|
|
"""Asynchronously get all node info with filter metrics
|
|
|
|
Args:
|
|
timeout: Timeout in seconds
|
|
node_selectors: Optional list of GetAllNodeInfoRequest.NodeSelector
|
|
proto objects to filter nodes.
|
|
state_filter: Optional int representing the GcsNodeState enum value
|
|
to filter by node state.
|
|
limit: Maximum number of nodes to return
|
|
|
|
Returns:
|
|
A tuple of (node_infos, num_filtered) where:
|
|
- node_infos: Dict[NodeID, GcsNodeInfo] mapping node IDs to their info
|
|
- num_filtered: int, the number of nodes filtered out by the query
|
|
"""
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
optional[CGcsNodeState] c_state_filter = nullopt
|
|
c_vector[CNodeSelector] c_node_selectors
|
|
CNodeSelector c_node_selector
|
|
optional[int64_t] c_limit = nullopt
|
|
fut = incremented_fut()
|
|
|
|
# Convert state_filter int to CGcsNodeState
|
|
if state_filter is not None:
|
|
c_state_filter.emplace(<CGcsNodeState>state_filter)
|
|
|
|
# Convert Python NodeSelector protos to C++ CNodeSelector
|
|
if node_selectors is not None:
|
|
for py_selector in node_selectors:
|
|
c_node_selector = CNodeSelector()
|
|
c_node_selector.ParseFromString(py_selector.SerializeToString())
|
|
c_node_selectors.push_back(c_node_selector)
|
|
|
|
# Set limit if provided
|
|
if limit is not None:
|
|
c_limit = <int64_t>limit
|
|
|
|
with nogil:
|
|
self.inner.get().Nodes().AsyncGetAll(
|
|
OptionalItemPyCallback[c_pair[c_vector[CGcsNodeInfo], int64_t]](
|
|
convert_get_all_node_info_results,
|
|
assign_and_decrement_fut,
|
|
fut),
|
|
timeout_ms,
|
|
c_state_filter,
|
|
c_node_selectors,
|
|
c_limit)
|
|
return asyncio.wrap_future(fut)
|
|
|
|
#############################################################
|
|
# NodeResources methods
|
|
#############################################################
|
|
def get_all_resource_usage(
|
|
self, timeout: Optional[int | float] = None
|
|
) -> GetAllResourceUsageReply:
|
|
cdef int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
cdef CGetAllResourceUsageReply c_reply
|
|
cdef c_string serialized_reply
|
|
with nogil:
|
|
check_status_timeout_as_rpc_error(
|
|
self.inner.get()
|
|
.NodeResources()
|
|
.GetAllResourceUsage(timeout_ms, c_reply)
|
|
)
|
|
serialized_reply = c_reply.SerializeAsString()
|
|
ret = GetAllResourceUsageReply()
|
|
ret.ParseFromString(serialized_reply)
|
|
return ret
|
|
|
|
#############################################################
|
|
# Actor methods
|
|
#############################################################
|
|
|
|
def async_get_all_actor_info(
|
|
self,
|
|
actor_id: Optional[ActorID] = None,
|
|
job_id: Optional[JobID] = None,
|
|
actor_state_name: Optional[str] = None,
|
|
timeout: Optional[int | float] = None
|
|
) -> Future[Dict[ActorID, gcs_pb2.ActorTableData]]:
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
optional[CActorID] c_actor_id
|
|
optional[CJobID] c_job_id
|
|
optional[c_string] c_actor_state_name
|
|
fut = incremented_fut()
|
|
if actor_id is not None:
|
|
c_actor_id = (<ActorID>actor_id).native()
|
|
if job_id is not None:
|
|
c_job_id = (<JobID>job_id).native()
|
|
if actor_state_name is not None:
|
|
c_actor_state_name = <c_string>actor_state_name.encode()
|
|
|
|
with nogil:
|
|
self.inner.get().Actors().AsyncGetAllByFilter(
|
|
c_actor_id, c_job_id, c_actor_state_name,
|
|
MultiItemPyCallback[CActorTableData](
|
|
&convert_get_all_actor_info,
|
|
assign_and_decrement_fut,
|
|
fut),
|
|
timeout_ms)
|
|
return asyncio.wrap_future(fut)
|
|
|
|
#############################################################
|
|
# Job methods
|
|
#############################################################
|
|
|
|
def get_all_job_info(
|
|
self, *, job_or_submission_id: Optional[str] = None,
|
|
skip_submission_job_info_field: bool = False,
|
|
skip_is_running_tasks_field: bool = False,
|
|
timeout: Optional[int | float] = None
|
|
) -> Dict[JobID, gcs_pb2.JobTableData]:
|
|
cdef c_string c_job_or_submission_id
|
|
cdef optional[c_string] c_optional_job_or_submission_id = nullopt
|
|
cdef int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
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
|
|
cdef CRayStatus status
|
|
cdef c_vector[CJobTableData] reply
|
|
if job_or_submission_id:
|
|
c_job_or_submission_id = job_or_submission_id
|
|
c_optional_job_or_submission_id = \
|
|
make_optional[c_string](c_job_or_submission_id)
|
|
with nogil:
|
|
status = self.inner.get().Jobs().GetAll(
|
|
c_optional_job_or_submission_id, c_skip_submission_job_info_field,
|
|
c_skip_is_running_tasks_field, reply, timeout_ms)
|
|
return raise_or_return((convert_get_all_job_info(status, move(reply))))
|
|
|
|
def async_get_all_job_info(
|
|
self, *, job_or_submission_id: Optional[str] = None,
|
|
skip_submission_job_info_field: bool = False,
|
|
skip_is_running_tasks_field: bool = False,
|
|
timeout: Optional[int | float] = None
|
|
) -> Future[Dict[JobID, gcs_pb2.JobTableData]]:
|
|
cdef:
|
|
c_string c_job_or_submission_id
|
|
optional[c_string] c_optional_job_or_submission_id = nullopt
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
c_bool c_skip_submission_job_info_field = skip_submission_job_info_field
|
|
c_bool c_skip_is_running_tasks_field = skip_is_running_tasks_field
|
|
fut = incremented_fut()
|
|
if job_or_submission_id:
|
|
c_job_or_submission_id = job_or_submission_id
|
|
c_optional_job_or_submission_id = \
|
|
make_optional[c_string](c_job_or_submission_id)
|
|
with nogil:
|
|
self.inner.get().Jobs().AsyncGetAll(
|
|
c_optional_job_or_submission_id,
|
|
c_skip_submission_job_info_field,
|
|
c_skip_is_running_tasks_field,
|
|
MultiItemPyCallback[CJobTableData](
|
|
&convert_get_all_job_info,
|
|
assign_and_decrement_fut,
|
|
fut),
|
|
timeout_ms)
|
|
return asyncio.wrap_future(fut)
|
|
|
|
#############################################################
|
|
# Runtime Env methods
|
|
#############################################################
|
|
def pin_runtime_env_uri(self, str uri, int expiration_s, timeout=None):
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
c_string c_uri = uri.encode()
|
|
with nogil:
|
|
check_status_timeout_as_rpc_error(
|
|
self.inner.get()
|
|
.RuntimeEnvs()
|
|
.PinRuntimeEnvUri(c_uri, expiration_s, timeout_ms)
|
|
)
|
|
|
|
#############################################################
|
|
# Autoscaler methods
|
|
#############################################################
|
|
def request_cluster_resource_constraint(
|
|
self,
|
|
bundles: c_vector[unordered_map[c_string, cython.double]],
|
|
label_selectors: c_vector[unordered_map[c_string, c_string]],
|
|
count_array: c_vector[int64_t],
|
|
timeout_s=None):
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout_s) if timeout_s else -1
|
|
with nogil:
|
|
check_status_timeout_as_rpc_error(
|
|
self.inner.get()
|
|
.Autoscaler()
|
|
.RequestClusterResourceConstraint(timeout_ms, bundles, label_selectors, count_array)
|
|
)
|
|
|
|
def get_cluster_resource_state(
|
|
self,
|
|
timeout_s=None):
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout_s) if timeout_s else -1
|
|
c_string serialized_reply
|
|
with nogil:
|
|
check_status_timeout_as_rpc_error(
|
|
self.inner.get()
|
|
.Autoscaler()
|
|
.GetClusterResourceState(timeout_ms, serialized_reply)
|
|
)
|
|
|
|
return serialized_reply
|
|
|
|
def get_cluster_status(
|
|
self,
|
|
timeout_s=None):
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout_s) if timeout_s else -1
|
|
c_string serialized_reply
|
|
with nogil:
|
|
check_status_timeout_as_rpc_error(
|
|
self.inner.get()
|
|
.Autoscaler()
|
|
.GetClusterStatus(timeout_ms, serialized_reply)
|
|
)
|
|
|
|
return serialized_reply
|
|
|
|
def async_get_cluster_status(
|
|
self,
|
|
timeout_s=None
|
|
) -> Future[autoscaler_pb2.GetClusterStatusReply]:
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout_s) if timeout_s else -1
|
|
fut = incremented_fut()
|
|
with nogil:
|
|
self.inner.get().Autoscaler().AsyncGetClusterStatus(
|
|
timeout_ms,
|
|
OptionalItemPyCallback[CGetClusterStatusReply](
|
|
&convert_get_cluster_status_reply,
|
|
assign_and_decrement_fut,
|
|
fut))
|
|
return asyncio.wrap_future(fut)
|
|
|
|
def report_autoscaling_state(
|
|
self,
|
|
serialzied_state: c_string,
|
|
timeout_s=None
|
|
):
|
|
"""Report autoscaling state to GCS"""
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout_s) if timeout_s else -1
|
|
with nogil:
|
|
check_status_timeout_as_rpc_error(
|
|
self.inner.get()
|
|
.Autoscaler()
|
|
.ReportAutoscalingState(timeout_ms, serialzied_state)
|
|
)
|
|
|
|
def drain_node(
|
|
self,
|
|
node_id: c_string,
|
|
reason: int32_t,
|
|
reason_message: c_string,
|
|
deadline_timestamp_ms: int64_t):
|
|
"""Send a DrainNode request to GCS to gracefully terminate a node.
|
|
|
|
Used by the `ray drain-node` CLI command and by autoscaler v2's
|
|
ray_stopper for idle and preemption-based node termination.
|
|
|
|
Args:
|
|
node_id: Binary node ID of the target node.
|
|
reason: A `DrainNodeReason` enum value. `IDLE_TERMINATION`
|
|
requests are rejectable by the raylet; `PREEMPTION`
|
|
requests are non-rejectable.
|
|
reason_message: Human-readable explanation, used for
|
|
observability.
|
|
deadline_timestamp_ms: Timestamp (ms) when the node will be
|
|
force-killed. Used as a hint so workloads can drain
|
|
before the deadline.
|
|
|
|
Returns:
|
|
Tuple of (is_accepted, rejection_reason_message). When
|
|
`is_accepted` is False, `rejection_reason_message` describes
|
|
why the raylet rejected the request.
|
|
"""
|
|
cdef:
|
|
int64_t timeout_ms = -1
|
|
c_bool is_accepted = False
|
|
c_string rejection_reason_message
|
|
with nogil:
|
|
check_status_timeout_as_rpc_error(self.inner.get().Autoscaler().DrainNode(
|
|
node_id, reason, reason_message,
|
|
deadline_timestamp_ms, timeout_ms, is_accepted,
|
|
rejection_reason_message))
|
|
|
|
return (is_accepted, rejection_reason_message.decode())
|
|
|
|
def resize_raylet_resource_instances(
|
|
self,
|
|
node_id: c_string,
|
|
resources: unordered_map[c_string, cython.double],
|
|
timeout_s=None):
|
|
"""Send the ResizeRayletResourceInstances request to GCS."""
|
|
cdef:
|
|
int64_t timeout_ms = round(1000 * timeout_s) if timeout_s else -1
|
|
unordered_map[c_string, cython.double] total_resources
|
|
with nogil:
|
|
check_status_timeout_as_rpc_error(
|
|
self.inner.get().Autoscaler().ResizeRayletResourceInstances(
|
|
node_id, resources, timeout_ms, total_resources
|
|
)
|
|
)
|
|
return {key.decode(): value for key, value in total_resources}
|
|
|
|
#############################################################
|
|
# Publisher methods
|
|
#############################################################
|
|
|
|
def publish_error(self, key_id: bytes, error_type: str, message: str,
|
|
job_id: Optional[JobID] = None, timeout = None):
|
|
cdef:
|
|
CErrorTableData error_info
|
|
c_string c_key_id = key_id
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
|
|
if job_id is None:
|
|
job_id = ray.JobID.nil()
|
|
error_info.set_job_id(job_id.binary())
|
|
error_info.set_type(error_type)
|
|
error_info.set_error_message(message)
|
|
error_info.set_timestamp(time.time())
|
|
|
|
with nogil:
|
|
self.inner.get().Publisher().PublishError(
|
|
move(c_key_id), move(error_info), timeout_ms)
|
|
|
|
def publish_logs(self, log_json: dict, timeout = None):
|
|
cdef:
|
|
CLogBatch log_batch
|
|
c_string c_key_id
|
|
int64_t timeout_ms = round(1000 * timeout) if timeout else -1
|
|
|
|
job_id = log_json.get("job")
|
|
log_batch.set_ip(log_json.get("ip") if log_json.get("ip") else b"")
|
|
log_batch.set_pid(
|
|
str(log_json.get("pid")).encode() if log_json.get("pid") else b"")
|
|
log_batch.set_job_id(job_id.encode() if job_id else b"")
|
|
log_batch.set_is_error(bool(log_json.get("is_err")))
|
|
for line in log_json.get("lines", []):
|
|
log_batch.add_lines(line)
|
|
actor_name = log_json.get("actor_name")
|
|
log_batch.set_actor_name(actor_name.encode() if actor_name else b"")
|
|
task_name = log_json.get("task_name")
|
|
log_batch.set_task_name(task_name.encode() if task_name else b"")
|
|
|
|
c_key_id = job_id.encode() if job_id else b""
|
|
|
|
with nogil:
|
|
check_status_timeout_as_rpc_error(
|
|
self.inner.get().Publisher().PublishLogs(
|
|
move(c_key_id), move(log_batch), timeout_ms))
|
|
|
|
def async_publish_node_resource_usage(
|
|
self, key_id: str, node_resource_usage_json: str) -> Future[None]:
|
|
cdef:
|
|
c_string c_key_id = key_id
|
|
c_string c_node_resource_usage_json = node_resource_usage_json.encode()
|
|
fut = incremented_fut()
|
|
with nogil:
|
|
self.inner.get().Publisher().AsyncPublishNodeResourceUsage(
|
|
move(c_key_id),
|
|
move(c_node_resource_usage_json),
|
|
StatusPyCallback(convert_status, assign_and_decrement_fut, fut))
|
|
return asyncio.wrap_future(fut)
|
|
|
|
def report_cluster_config(
|
|
self,
|
|
serialized_cluster_config: c_string):
|
|
"""Report cluster config to GCS"""
|
|
cdef:
|
|
int64_t timeout_ms = -1
|
|
with nogil:
|
|
check_status_timeout_as_rpc_error(
|
|
self.inner.get().Autoscaler().ReportClusterConfig(
|
|
timeout_ms, serialized_cluster_config
|
|
)
|
|
)
|
|
|
|
#############################################################
|
|
# TaskInfo methods
|
|
#############################################################
|
|
async def async_add_events(self, serialized_request: bytes, timeout_s=None, executor=None):
|
|
"""Send async AddEvents request to GCS."""
|
|
cdef:
|
|
CAddEventsRequest c_req
|
|
int64_t timeout_ms
|
|
fut = incremented_fut()
|
|
timeout_ms = round(1000 * timeout_s) if timeout_s else -1
|
|
|
|
# Parse the protobuf payload
|
|
cdef c_string payload = serialized_request
|
|
cdef bint parsed = False
|
|
if executor is not None:
|
|
parsed = await get_or_create_event_loop().run_in_executor(
|
|
executor,
|
|
lambda: c_req.ParseFromString(payload),
|
|
)
|
|
else:
|
|
parsed = c_req.ParseFromString(payload)
|
|
|
|
if not parsed:
|
|
# Fail fast on parse error
|
|
assign_and_decrement_fut((None, ValueError("Invalid AddEventsRequest payload")), fut)
|
|
return await asyncio.wrap_future(fut)
|
|
|
|
with nogil:
|
|
self.inner.get().Tasks().AsyncAddEvents(
|
|
move(c_req),
|
|
StatusPyCallback(convert_status, assign_and_decrement_fut, fut),
|
|
timeout_ms)
|
|
return await asyncio.wrap_future(fut)
|
|
|
|
#############################################################
|
|
# Converter functions: C++ types -> Python types, use by both Sync and Async APIs.
|
|
# They have to be defined here as pure functions because a function pointer is passed
|
|
# to C++ for Async APIs.
|
|
#
|
|
# Each function accepts what the C++ callback passes, typically a Status and a value.
|
|
# Returns `Tuple[object, Optional[Exception]]` (we are all gophers now lol).
|
|
# Must not raise exceptions, or it crashes the process.
|
|
#############################################################
|
|
|
|
cdef convert_get_all_node_info(
|
|
CRayStatus status, c_vector[CGcsNodeInfo] c_data) with gil:
|
|
# -> Dict[NodeID, gcs_pb2.GcsNodeInfo]
|
|
# No GIL block for C++ looping && serialization.
|
|
# GIL block for Pyhton deserialization and dict building.
|
|
# Not doing per-object GIL lock because it's expensive.
|
|
cdef c_vector[c_string] serialized_reply
|
|
try:
|
|
check_status_timeout_as_rpc_error(status)
|
|
with nogil:
|
|
serialized_reply.reserve(c_data.size())
|
|
for c_proto in c_data:
|
|
serialized_reply.push_back(c_proto.SerializeAsString())
|
|
node_table_data = {}
|
|
for b in serialized_reply:
|
|
proto = gcs_pb2.GcsNodeInfo()
|
|
proto.ParseFromString(b)
|
|
node_table_data[NodeID.from_binary(proto.node_id)] = proto
|
|
return node_table_data, None
|
|
except Exception as e:
|
|
return None, e
|
|
|
|
cdef convert_get_all_node_info_results(
|
|
CRayStatus status, optional[c_pair[c_vector[CGcsNodeInfo], int64_t]] c_data) with gil:
|
|
# -> Tuple[Dict[NodeID, GcsNodeInfo], int], Exception
|
|
try:
|
|
check_status_timeout_as_rpc_error(status)
|
|
if not c_data.has_value():
|
|
raise ValueError("Reply is empty")
|
|
node_table_data, exc = convert_get_all_node_info(
|
|
status, c_data.value().first)
|
|
if exc is not None:
|
|
return None, exc
|
|
return (node_table_data, c_data.value().second), None
|
|
except Exception as e:
|
|
return None, e
|
|
|
|
cdef convert_get_all_job_info(
|
|
CRayStatus status, c_vector[CJobTableData] c_data) with gil:
|
|
# -> Dict[JobID, gcs_pb2.JobTableData]
|
|
# No GIL block for C++ looping && serialization.
|
|
# GIL block for Pyhton deserialization and dict building.
|
|
# Not doing per-object GIL lock because it's expensive.
|
|
cdef c_vector[c_string] serialized_reply
|
|
try:
|
|
check_status_timeout_as_rpc_error(status)
|
|
with nogil:
|
|
serialized_reply.reserve(c_data.size())
|
|
for c_proto in c_data:
|
|
serialized_reply.push_back(c_proto.SerializeAsString())
|
|
job_table_data = {}
|
|
for b in serialized_reply:
|
|
proto = gcs_pb2.JobTableData()
|
|
proto.ParseFromString(b)
|
|
job_table_data[JobID.from_binary(proto.job_id)] = proto
|
|
return job_table_data, None
|
|
except Exception as e:
|
|
return None, e
|
|
|
|
cdef convert_get_all_actor_info(
|
|
CRayStatus status, c_vector[CActorTableData] c_data) with gil:
|
|
# -> Dict[ActorID, gcs_pb2.ActorTableData]
|
|
cdef c_vector[c_string] serialized_reply
|
|
try:
|
|
check_status_timeout_as_rpc_error(status)
|
|
with nogil:
|
|
serialized_reply.reserve(c_data.size())
|
|
for c_proto in c_data:
|
|
serialized_reply.push_back(c_proto.SerializeAsString())
|
|
actor_table_data = {}
|
|
for b in serialized_reply:
|
|
proto = gcs_pb2.ActorTableData()
|
|
proto.ParseFromString(b)
|
|
actor_table_data[ActorID.from_binary(proto.actor_id)] = proto
|
|
return actor_table_data, None
|
|
except Exception as e:
|
|
return None, e
|
|
|
|
cdef convert_get_cluster_status_reply(
|
|
CRayStatus status, optional[CGetClusterStatusReply] c_data
|
|
) with gil: # -> Tuple[autoscaler_pb2.GetClusterStatusReply, Exception]
|
|
cdef c_string serialized_reply
|
|
try:
|
|
check_status_timeout_as_rpc_error(status)
|
|
assert c_data.has_value()
|
|
with nogil:
|
|
serialized_reply = c_data.value().SerializeAsString()
|
|
proto = autoscaler_pb2.GetClusterStatusReply()
|
|
proto.ParseFromString(serialized_reply)
|
|
return proto, None
|
|
except Exception as e:
|
|
return None, e
|
|
|
|
cdef convert_status(CRayStatus status) with gil:
|
|
# This function is currently only used by `async_kill_actor` to
|
|
# convert RayStatus to an HTTP status code.
|
|
#
|
|
# Returns:
|
|
# Tuple[int, Optional[Exception]]:
|
|
# - int: HTTP status code.
|
|
# (1) 200: Success
|
|
# (2) 404: Actor not found
|
|
# (3) 500: Other errors
|
|
# - Optional[Exception]: Exception raised by RayStatus
|
|
try:
|
|
if status.IsNotFound():
|
|
return 404, None
|
|
check_status_timeout_as_rpc_error(status)
|
|
return 200, None
|
|
except Exception as e:
|
|
return 500, e
|
|
|
|
cdef convert_optional_str_none_for_not_found(
|
|
CRayStatus status, optional[c_string] c_str) with gil:
|
|
# If status is NotFound, return None.
|
|
# If status is OK, return the value.
|
|
# Else, raise exception.
|
|
# -> Optional[bytes]
|
|
try:
|
|
if status.IsNotFound():
|
|
return None, None
|
|
check_status_timeout_as_rpc_error(status)
|
|
assert c_str.has_value()
|
|
return move(c_str.value()), None
|
|
except Exception as e:
|
|
return None, e
|
|
|
|
cdef convert_optional_multi_get(
|
|
CRayStatus status,
|
|
optional[unordered_map[c_string, c_string]] c_map) with gil:
|
|
# -> Dict[str, str]
|
|
cdef unordered_map[c_string, c_string].iterator it
|
|
try:
|
|
check_status_timeout_as_rpc_error(status)
|
|
assert c_map.has_value()
|
|
|
|
result = {}
|
|
it = c_map.value().begin()
|
|
while it != c_map.value().end():
|
|
result[dereference(it).first] = move(dereference(it).second)
|
|
postincrement(it)
|
|
return result, None
|
|
except Exception as e:
|
|
return None, e
|
|
|
|
cdef convert_optional_int(CRayStatus status, optional[int] c_int) with gil:
|
|
# -> int
|
|
try:
|
|
check_status_timeout_as_rpc_error(status)
|
|
assert c_int.has_value()
|
|
return c_int.value(), None
|
|
except Exception as e:
|
|
return None, e
|
|
|
|
cdef convert_optional_vector_str(
|
|
CRayStatus status, optional[c_vector[c_string]] c_vec) with gil:
|
|
# -> List[bytes]
|
|
try:
|
|
check_status_timeout_as_rpc_error(status)
|
|
return convert_multi_str(status, move(c_vec.value()))
|
|
except Exception as e:
|
|
return None, e
|
|
|
|
cdef convert_optional_bool(CRayStatus status, optional[c_bool] b) with gil:
|
|
# -> bool
|
|
try:
|
|
check_status_timeout_as_rpc_error(status)
|
|
assert b.has_value()
|
|
return b.value(), None
|
|
except Exception as e:
|
|
return None, e
|
|
|
|
cdef convert_multi_bool(CRayStatus status, c_vector[c_bool] c_data) with gil:
|
|
# -> List[bool]
|
|
try:
|
|
check_status_timeout_as_rpc_error(status)
|
|
return [b for b in c_data], None
|
|
except Exception as e:
|
|
return None, e
|
|
|
|
cdef convert_multi_str(CRayStatus status, c_vector[c_string] vec) with gil:
|
|
# -> List[bytes]
|
|
cdef c_vector[c_string].iterator it
|
|
try:
|
|
check_status_timeout_as_rpc_error(status)
|
|
|
|
it = vec.begin()
|
|
result = []
|
|
while it != vec.end():
|
|
result.append(move(dereference(it)))
|
|
postincrement(it)
|
|
return result, None
|
|
except Exception as e:
|
|
return None, e
|