chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,293 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
from collections import deque
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import grpc
|
||||
from grpc import aio as aiogrpc
|
||||
|
||||
import ray._private.gcs_utils as gcs_utils
|
||||
from ray._common.utils import get_or_create_event_loop
|
||||
from ray.core.generated import (
|
||||
gcs_pb2,
|
||||
gcs_service_pb2,
|
||||
gcs_service_pb2_grpc,
|
||||
pubsub_pb2,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_OBSERVABILITY_PUBSUB_CHANNELS = (
|
||||
pubsub_pb2.RAY_ERROR_INFO_CHANNEL,
|
||||
pubsub_pb2.RAY_LOG_CHANNEL,
|
||||
pubsub_pb2.RAY_NODE_RESOURCE_USAGE_CHANNEL,
|
||||
)
|
||||
|
||||
|
||||
class _SubscriberBase:
|
||||
def __init__(self, worker_id: bytes = None):
|
||||
self._worker_id = worker_id
|
||||
# self._subscriber_id needs to match the binary format of a random
|
||||
# SubscriberID / UniqueID, which is 28 (kUniqueIDSize) random bytes.
|
||||
self._subscriber_id = bytes(bytearray(random.getrandbits(8) for _ in range(28)))
|
||||
self._last_batch_size = 0
|
||||
self._max_processed_sequence_id = 0
|
||||
self._publisher_id = b""
|
||||
|
||||
# Batch size of the result from last poll. Used to indicate whether the
|
||||
# subscriber can keep up.
|
||||
@property
|
||||
def last_batch_size(self):
|
||||
return self._last_batch_size
|
||||
|
||||
def _subscribe_request(self, channel):
|
||||
cmd = pubsub_pb2.Command(channel_type=channel, subscribe_message={})
|
||||
req = gcs_service_pb2.GcsSubscriberCommandBatchRequest(
|
||||
subscriber_id=self._subscriber_id, sender_id=self._worker_id, commands=[cmd]
|
||||
)
|
||||
return req
|
||||
|
||||
def _poll_request(self):
|
||||
return gcs_service_pb2.GcsSubscriberPollRequest(
|
||||
subscriber_id=self._subscriber_id,
|
||||
max_processed_sequence_id=self._max_processed_sequence_id,
|
||||
publisher_id=self._publisher_id,
|
||||
)
|
||||
|
||||
def _unsubscribe_request(self, channels):
|
||||
req = gcs_service_pb2.GcsSubscriberCommandBatchRequest(
|
||||
subscriber_id=self._subscriber_id, sender_id=self._worker_id, commands=[]
|
||||
)
|
||||
for channel in channels:
|
||||
req.commands.append(
|
||||
pubsub_pb2.Command(channel_type=channel, unsubscribe_message={})
|
||||
)
|
||||
return req
|
||||
|
||||
@staticmethod
|
||||
def _should_terminate_polling(e: grpc.RpcError) -> None:
|
||||
# Caller only expects polling to be terminated after deadline exceeded.
|
||||
if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
|
||||
return True
|
||||
# Could be a temporary connection issue. Suppress error.
|
||||
# TODO: reconnect GRPC channel?
|
||||
if e.code() == grpc.StatusCode.UNAVAILABLE:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class _AioSubscriber(_SubscriberBase):
|
||||
"""Async io subscriber to GCS.
|
||||
|
||||
Usage example common to Aio subscribers:
|
||||
subscriber = GcsAioXxxSubscriber(address="...")
|
||||
await subscriber.subscribe()
|
||||
while running:
|
||||
...... = await subscriber.poll()
|
||||
......
|
||||
await subscriber.close()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pubsub_channel_type,
|
||||
worker_id: bytes = None,
|
||||
address: str = None,
|
||||
channel: aiogrpc.Channel = None,
|
||||
):
|
||||
super().__init__(worker_id)
|
||||
|
||||
if address:
|
||||
assert channel is None, "address and channel cannot both be specified"
|
||||
channel = gcs_utils.create_gcs_channel(address, aio=True)
|
||||
else:
|
||||
assert channel is not None, "One of address and channel must be specified"
|
||||
if pubsub_channel_type in _OBSERVABILITY_PUBSUB_CHANNELS:
|
||||
self._stub = gcs_service_pb2_grpc.ObservabilityPubSubServiceStub(channel)
|
||||
else:
|
||||
self._stub = gcs_service_pb2_grpc.ControlPlanePubSubGcsServiceStub(channel)
|
||||
|
||||
# Type of the channel.
|
||||
self._channel = pubsub_channel_type
|
||||
# A queue of received PubMessage.
|
||||
self._queue = deque()
|
||||
# Indicates whether the subscriber has closed.
|
||||
self._close = asyncio.Event()
|
||||
|
||||
async def subscribe(self) -> None:
|
||||
"""Registers a subscription for the subscriber's channel type.
|
||||
|
||||
Before the registration, published messages in the channel will not be
|
||||
saved for the subscriber.
|
||||
"""
|
||||
if self._close.is_set():
|
||||
return
|
||||
req = self._subscribe_request(self._channel)
|
||||
await self._stub.GcsSubscriberCommandBatch(req, timeout=30)
|
||||
|
||||
async def _poll_call(self, req, timeout=None):
|
||||
# Wrap GRPC _AioCall as a coroutine.
|
||||
return await self._stub.GcsSubscriberPoll(req, timeout=timeout)
|
||||
|
||||
async def _poll(self, timeout=None) -> None:
|
||||
while len(self._queue) == 0:
|
||||
req = self._poll_request()
|
||||
poll = get_or_create_event_loop().create_task(
|
||||
self._poll_call(req, timeout=timeout)
|
||||
)
|
||||
close = get_or_create_event_loop().create_task(self._close.wait())
|
||||
done, others = await asyncio.wait(
|
||||
[poll, close], timeout=timeout, return_when=asyncio.FIRST_COMPLETED
|
||||
)
|
||||
# Cancel the other task if needed to prevent memory leak.
|
||||
other_task = others.pop()
|
||||
if not other_task.done():
|
||||
other_task.cancel()
|
||||
if poll not in done or close in done:
|
||||
# Request timed out or subscriber closed.
|
||||
break
|
||||
try:
|
||||
self._last_batch_size = len(poll.result().pub_messages)
|
||||
if poll.result().publisher_id != self._publisher_id:
|
||||
if self._publisher_id != b"":
|
||||
logger.debug(
|
||||
f"replied publisher_id {poll.result().publisher_id} "
|
||||
f"different from {self._publisher_id}, this should "
|
||||
"only happen during gcs failover."
|
||||
)
|
||||
self._publisher_id = poll.result().publisher_id
|
||||
self._max_processed_sequence_id = 0
|
||||
for msg in poll.result().pub_messages:
|
||||
if msg.sequence_id <= self._max_processed_sequence_id:
|
||||
logger.warning(f"Ignoring out of order message {msg}")
|
||||
continue
|
||||
self._max_processed_sequence_id = msg.sequence_id
|
||||
self._queue.append(msg)
|
||||
except grpc.RpcError as e:
|
||||
if self._should_terminate_polling(e):
|
||||
return
|
||||
raise
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Closes the subscriber and its active subscription."""
|
||||
|
||||
# Mark close to terminate inflight polling and prevent future requests.
|
||||
if self._close.is_set():
|
||||
return
|
||||
self._close.set()
|
||||
req = self._unsubscribe_request(channels=[self._channel])
|
||||
try:
|
||||
await self._stub.GcsSubscriberCommandBatch(req, timeout=5)
|
||||
except Exception:
|
||||
pass
|
||||
self._stub = None
|
||||
|
||||
|
||||
class GcsAioResourceUsageSubscriber(_AioSubscriber):
|
||||
def __init__(
|
||||
self,
|
||||
worker_id: bytes = None,
|
||||
address: str = None,
|
||||
channel: grpc.Channel = None,
|
||||
):
|
||||
super().__init__(
|
||||
pubsub_pb2.RAY_NODE_RESOURCE_USAGE_CHANNEL, worker_id, address, channel
|
||||
)
|
||||
|
||||
async def poll(self, timeout: Optional[float] = None) -> Tuple[bytes, str]:
|
||||
"""Polls for new resource usage message.
|
||||
|
||||
Args:
|
||||
timeout: Optional timeout in seconds for the poll request.
|
||||
|
||||
Returns:
|
||||
A tuple of string reporter ID and resource usage json string.
|
||||
"""
|
||||
await self._poll(timeout=timeout)
|
||||
return self._pop_resource_usage(self._queue)
|
||||
|
||||
@staticmethod
|
||||
def _pop_resource_usage(queue):
|
||||
if len(queue) == 0:
|
||||
return None, None
|
||||
msg = queue.popleft()
|
||||
return msg.key_id.decode(), msg.node_resource_usage_message.json
|
||||
|
||||
|
||||
class GcsAioActorSubscriber(_AioSubscriber):
|
||||
def __init__(
|
||||
self,
|
||||
worker_id: bytes = None,
|
||||
address: str = None,
|
||||
channel: grpc.Channel = None,
|
||||
):
|
||||
super().__init__(pubsub_pb2.GCS_ACTOR_CHANNEL, worker_id, address, channel)
|
||||
|
||||
@property
|
||||
def queue_size(self):
|
||||
return len(self._queue)
|
||||
|
||||
async def poll(
|
||||
self, batch_size: int, timeout: Optional[float] = None
|
||||
) -> List[Tuple[bytes, gcs_pb2.ActorTableData]]:
|
||||
"""Polls for new actor message.
|
||||
|
||||
Args:
|
||||
batch_size: Maximum number of messages to return.
|
||||
timeout: Optional timeout in seconds for the poll request.
|
||||
|
||||
Returns:
|
||||
A list of tuples of binary actor ID and actor table data.
|
||||
"""
|
||||
await self._poll(timeout=timeout)
|
||||
return self._pop_actors(self._queue, batch_size=batch_size)
|
||||
|
||||
@staticmethod
|
||||
def _pop_actors(queue, batch_size):
|
||||
if len(queue) == 0:
|
||||
return []
|
||||
popped = 0
|
||||
msgs = []
|
||||
while len(queue) > 0 and popped < batch_size:
|
||||
msg = queue.popleft()
|
||||
msgs.append((msg.key_id, msg.actor_message))
|
||||
popped += 1
|
||||
return msgs
|
||||
|
||||
|
||||
class GcsAioNodeInfoSubscriber(_AioSubscriber):
|
||||
def __init__(
|
||||
self,
|
||||
worker_id: bytes = None,
|
||||
address: str = None,
|
||||
channel: grpc.Channel = None,
|
||||
):
|
||||
super().__init__(pubsub_pb2.GCS_NODE_INFO_CHANNEL, worker_id, address, channel)
|
||||
|
||||
async def poll(
|
||||
self, batch_size: int, timeout: Optional[float] = None
|
||||
) -> List[Tuple[bytes, gcs_pb2.GcsNodeInfo]]:
|
||||
"""Polls for new node info message.
|
||||
|
||||
Args:
|
||||
batch_size: Maximum number of messages to return.
|
||||
timeout: Optional timeout in seconds for the poll request.
|
||||
|
||||
Returns:
|
||||
A list of tuples of (node_id, GcsNodeInfo).
|
||||
"""
|
||||
await self._poll(timeout=timeout)
|
||||
return self._pop_node_infos(self._queue, batch_size=batch_size)
|
||||
|
||||
@staticmethod
|
||||
def _pop_node_infos(queue, batch_size):
|
||||
if len(queue) == 0:
|
||||
return []
|
||||
popped = 0
|
||||
msgs = []
|
||||
while len(queue) > 0 and popped < batch_size:
|
||||
msg = queue.popleft()
|
||||
msgs.append((msg.key_id, msg.node_info_message))
|
||||
popped += 1
|
||||
return msgs
|
||||
Reference in New Issue
Block a user