import asyncio import contextvars import logging import os import random import time from asyncio import sleep from asyncio.events import AbstractEventLoop from collections import defaultdict from collections.abc import Mapping from dataclasses import dataclass from enum import Enum, auto from typing import ( Any, Callable, DefaultDict, Dict, Iterable, Optional, Set, Tuple, Union, ) import ray from ray._common.utils import get_or_create_event_loop from ray.serve._private.constants import ( DEFAULT_LATENCY_BUCKET_MS, RAY_SERVE_COMPACT_LONG_POLL_METRIC_TAGS, SERVE_LOGGER_NAME, ) from ray.serve.generated.serve_pb2 import ( DeploymentTargetInfo, EndpointInfo as EndpointInfoProto, EndpointSet, LongPollRequest, LongPollResult, UpdatedObject as UpdatedObjectProto, ) from ray.util import metrics logger = logging.getLogger(SERVE_LOGGER_NAME) # Each LongPollClient will send requests to LongPollHost to poll changes # as blocking awaitable. This doesn't scale if we have many client instances # that will slow down, or even block controller actor's event loop if near # its max_concurrency limit. Therefore we timeout a polling request after # a few seconds and let each client retry on their end. # We randomly select a timeout within this range to avoid a "thundering herd" # when there are many clients subscribing at the same time. LISTEN_FOR_CHANGE_REQUEST_TIMEOUT_S = ( float(os.environ.get("LISTEN_FOR_CHANGE_REQUEST_TIMEOUT_S_LOWER_BOUND", "30")), float(os.environ.get("LISTEN_FOR_CHANGE_REQUEST_TIMEOUT_S_UPPER_BOUND", "60")), ) class LongPollNamespace(Enum): def __repr__(self): return f"{self.__class__.__name__}.{self.name}" DEPLOYMENT_TARGETS = auto() ROUTE_TABLE = auto() GLOBAL_LOGGING_CONFIG = auto() DEPLOYMENT_CONFIG = auto() TARGET_GROUPS = auto() FALLBACK_TARGETS = auto() @dataclass class UpdatedObject: object_snapshot: Any # The identifier for the object's version. There is not sequential relation # among different object's snapshot_ids. snapshot_id: int # Timestamp (in seconds since epoch) when notify_changed was called. # Used by clients to measure end-to-end propagation latency. notify_timestamp: float # Type signature for the update state callbacks. E.g. # async def update_state(updated_object: Any): # do_something(updated_object) UpdateStateCallable = Callable[[Any], None] KeyType = Union[str, LongPollNamespace, Tuple[LongPollNamespace, str]] def _get_metric_namespace_tag(key: KeyType) -> str: """Extract the namespace string from a long poll key for metric tags. When RAY_SERVE_COMPACT_LONG_POLL_METRIC_TAGS is enabled, returns only the LongPollNamespace enum name (e.g., "DEPLOYMENT_CONFIG") for tuple keys like (LongPollNamespace.DEPLOYMENT_CONFIG, "deployment_name"), avoiding high-cardinality metric labels that would otherwise include per-deployment identifiers. When disabled (default), returns str(key) preserving deployment-level metric granularity. """ if not RAY_SERVE_COMPACT_LONG_POLL_METRIC_TAGS: return str(key) if isinstance(key, tuple): return key[0].name if isinstance(key[0], LongPollNamespace) else str(key[0]) elif isinstance(key, LongPollNamespace): return key.name else: return str(key) class LongPollState(Enum): TIME_OUT = auto() class LongPollClient: """The asynchronous long polling client. Args: host_actor: handle to actor embedding LongPollHost. key_listeners: a dictionary mapping keys to callbacks to be called on state update for the corresponding keys. call_in_event_loop: an asyncio event loop to post the callback into. client_id: identifier reported back to the host if this client disables itself. """ def __init__( self, host_actor: Any, key_listeners: Dict[KeyType, UpdateStateCallable], call_in_event_loop: AbstractEventLoop, client_id: str, ) -> None: # We used to allow this to be optional, but due to Ray Client issue # we now enforce all long poll client to post callback to event loop # See https://github.com/ray-project/ray/issues/20971 assert call_in_event_loop is not None self.host_actor = host_actor self.key_listeners = key_listeners self.event_loop = call_in_event_loop self.client_id = client_id # The initial snapshot id for each key is < 0, # but real snapshot keys in the long poll host are always >= 0, # so this will always trigger an initial update. self.snapshot_ids: Dict[KeyType, int] = dict.fromkeys( self.key_listeners.keys(), -1 ) self.is_running = True # Metric to track end-to-end latency from controller to client self.long_poll_latency_histogram = metrics.Histogram( "serve_long_poll_latency_ms", description=( "The time in milliseconds for updates to propagate from " "controller to clients." ), boundaries=DEFAULT_LATENCY_BUCKET_MS, tag_keys=("namespace",), ) # NOTE(edoakes): we schedule the initial _poll_next call with an empty context # so that Ray will not recursively cancel the underlying `listen_for_change` # task. See: https://github.com/ray-project/ray/issues/52476. self.event_loop.call_soon_threadsafe( self._poll_next, context=contextvars.Context() ) def stop(self) -> None: """Stop the long poll client after the next RPC returns.""" self.is_running = False def add_key_listeners( self, key_listeners: Dict[KeyType, UpdateStateCallable] ) -> None: """Add more key listeners to the client. The new listeners will only be included in the *next* long poll request; the current request will continue with the existing listeners. If a key is already in the client, the new listener will replace the old one, but the snapshot ID will be preserved, so the new listener will only be called on the *next* update to that key. """ # We need to run the underlying method in the same event loop that runs # the long poll loop, because we need to mutate the mapping of snapshot IDs, # which also needs to be serialized by the long poll's RPC to the # Serve Controller. If those happened concurrently in different threads, # we could get a `RuntimeError: dictionary changed size during iteration`. # See https://github.com/ray-project/ray/pull/52793 for more details. self.event_loop.call_soon_threadsafe(self._add_key_listeners, key_listeners) def _add_key_listeners( self, key_listeners: Dict[KeyType, UpdateStateCallable] ) -> None: """Inner method that actually adds the key listeners, to be called via call_soon_threadsafe for thread safety. """ # Only initialize snapshot ids for *new* keys. self.snapshot_ids.update( {key: -1 for key in key_listeners.keys() if key not in self.key_listeners} ) self.key_listeners.update(key_listeners) def _on_callback_completed(self, trigger_at: int): """Called after a single callback is completed. When the total number of callback completed equals to trigger_at, _poll_next() will be called. This is designed to make sure we only _poll_next() after all the state callbacks completed. This is a way to serialize the callback invocations between object versions. """ self._callbacks_processed_count += 1 if self._callbacks_processed_count == trigger_at: self._poll_next() def _poll_next(self): """Poll the update. The callback is expected to scheduler another _poll_next call. """ if not self.is_running: return self._callbacks_processed_count = 0 self._current_ref = self.host_actor.listen_for_change.remote(self.snapshot_ids) self._current_ref._on_completed(lambda update: self._process_update(update)) def _schedule_to_event_loop(self, callback): # Schedule the next iteration only if the loop is running. # The event loop might not be running if users used a cached # version across loops. if not self.is_running: return if self.event_loop.is_running(): self.event_loop.call_soon_threadsafe(callback) else: reason = "Bound asyncio event loop is not running; controller updates cannot be delivered." logger.error( f"LongPollClient {self.client_id!r} has been disabled: {reason} " f"Keep the loop running for the lifetime of this process." ) self.is_running = False # Fire-and-forget notify so the controller logs this client as disabled. try: self.host_actor.notify_long_poll_client_disabled.remote( self.client_id, reason ) except Exception: logger.exception( "Failed to notify host that LongPollClient " f"{self.client_id!r} disabled itself." ) def _process_update(self, updates: Dict[str, UpdatedObject]): if isinstance(updates, (ray.exceptions.RayActorError)): # This can happen during shutdown where the controller is # intentionally killed, the client should just gracefully # exit. logger.debug("LongPollClient failed to connect to host. Shutting down.") self.is_running = False return if isinstance(updates, ConnectionError): logger.warning("LongPollClient connection failed, shutting down.") self.is_running = False return if isinstance(updates, (ray.exceptions.RayTaskError)): # Some error happened in the controller. It could be a bug or # some undesired state. logger.error("LongPollHost errored\n" + updates.traceback_str) # We must call this in event loop so it works in Ray Client. # See https://github.com/ray-project/ray/issues/20971 self._schedule_to_event_loop(self._poll_next) return if updates == LongPollState.TIME_OUT: logger.debug("LongPollClient polling timed out. Retrying.") self._schedule_to_event_loop(self._poll_next) return logger.debug( f"LongPollClient {self} received updates for keys: {list(updates.keys())}.", extra={"log_to_stderr": False}, ) if not updates: # no updates, no callbacks to run, just poll again self._schedule_to_event_loop(self._poll_next) # Record latency metrics for received updates. # Skip observations that exceed twice the maximum long poll timeout — # these are catch-up updates received by a client that was offline or # missed several polling cycles, and including them would distort the # metric (e.g. a 10-minute spike just because a new replica connected). max_valid_latency_ms = LISTEN_FOR_CHANGE_REQUEST_TIMEOUT_S[1] * 2 * 1000 receive_time = time.time() for key, update in updates.items(): latency_ms = (receive_time - update.notify_timestamp) * 1000 if latency_ms <= max_valid_latency_ms: self.long_poll_latency_histogram.observe( latency_ms, tags={"namespace": _get_metric_namespace_tag(key)}, ) else: logger.debug( f"Skipping long poll latency observation of {latency_ms:.0f}ms " f"for key {key} (exceeds threshold {max_valid_latency_ms:.0f}ms)." ) self.snapshot_ids[key] = update.snapshot_id callback = self.key_listeners[key] # Bind the parameters because closures are late-binding. # https://docs.python-guide.org/writing/gotchas/#late-binding-closures # noqa: E501 def chained(callback=callback, arg=update.object_snapshot): callback(arg) self._on_callback_completed(trigger_at=len(updates)) self._schedule_to_event_loop(chained) class LongPollHost: """The server side object that manages long pulling requests. The desired use case is to embed this in an Ray actor. Client will be expected to call actor.listen_for_change.remote(...). On the host side, you can call host.notify_changed({key: object}) to update the state and potentially notify whoever is polling for these values. Internally, we use snapshot_ids for each object to identify client with outdated object and immediately return the result. If the client has the up-to-date version, then the listen_for_change call will only return when the object is updated. """ def __init__( self, listen_for_change_request_timeout_s: Tuple[ int, int ] = LISTEN_FOR_CHANGE_REQUEST_TIMEOUT_S, ): # Map object_key -> int self.snapshot_ids: Dict[KeyType, int] = {} # Map object_key -> object self.object_snapshots: Dict[KeyType, Any] = {} # Map object_key -> set(asyncio.Event waiting for updates) self.notifier_events: DefaultDict[KeyType, Set[asyncio.Event]] = defaultdict( set ) # Map object_key -> timestamp when notify_changed was called # Used to track latency for propagating updates to clients self._notify_timestamps: Dict[KeyType, float] = {} # Aggregate count of pending clients per namespace tag. Needed because # multiple keys can map to the same low-cardinality namespace tag, so # we must track the total rather than setting per-key counts. self._pending_clients_by_namespace: DefaultDict[str, int] = defaultdict(int) self._listen_for_change_request_timeout_s = listen_for_change_request_timeout_s self.transmission_counter = metrics.Counter( "serve_long_poll_host_transmission_counter", description="The number of times the long poll host transmits data.", tag_keys=("namespace_or_state",), ) self.pending_clients_gauge = metrics.Gauge( "serve_long_poll_pending_clients", description=("The number of clients waiting for updates per namespace."), tag_keys=("namespace",), ) def _get_num_notifier_events(self, key: Optional[KeyType] = None): """Used for testing.""" if key is not None: return len(self.notifier_events[key]) else: return sum(len(events) for events in self.notifier_events.values()) def _get_pending_clients_by_namespace(self, namespace_tag: str) -> int: """Used for testing. Returns the aggregate pending client count.""" return self._pending_clients_by_namespace.get(namespace_tag, 0) def notify_client_disabled(self, client_id: str, reason: str) -> None: """Fire-and-forget hook for clients that are shutting themselves down. LongPollClient calls this before flipping ``is_running`` to False when it cannot keep delivering updates (e.g. its bound asyncio loop is no longer running). """ logger.error( f"LongPollClient {client_id!r} disabled itself and will no longer " f"receive controller updates. Reason: {reason}" ) def _count_send( self, timeout_or_data: Union[LongPollState, Dict[KeyType, UpdatedObject]] ): """Helper method that tracks the data sent by listen_for_change. Records number of times long poll host sends data in the ray_serve_long_poll_host_send_counter metric. """ if isinstance(timeout_or_data, LongPollState): # The only LongPollState is TIME_OUT– the long poll # connection has timed out. self.transmission_counter.inc( value=1, tags={"namespace_or_state": "TIMEOUT"} ) else: data = timeout_or_data for key in data.keys(): self.transmission_counter.inc( value=1, tags={"namespace_or_state": _get_metric_namespace_tag(key)}, ) async def listen_for_change( self, keys_to_snapshot_ids: Dict[KeyType, int], ) -> Union[LongPollState, Dict[KeyType, UpdatedObject]]: """Listen for changed objects. This method will return a dictionary of updated objects. It returns immediately if any of the snapshot_ids are outdated, otherwise it will block until there's an update. """ # If there are no keys to listen for, # just wait for a short time to provide backpressure, # then return an empty update. if not keys_to_snapshot_ids: await sleep(1) updated_objects = {} self._count_send(updated_objects) return updated_objects # If there are any keys with outdated snapshot ids, # return their updated values immediately. updated_objects = {} for key, client_snapshot_id in keys_to_snapshot_ids.items(): try: existing_id = self.snapshot_ids[key] except KeyError: # The caller may ask for keys that we don't know about (yet), # just ignore them. # This can happen when, for example, # a deployment handle is manually created for an app # that hasn't been deployed yet (by bypassing the safety checks). continue if existing_id != client_snapshot_id: updated_objects[key] = UpdatedObject( self.object_snapshots[key], existing_id, self._notify_timestamps[key], ) if len(updated_objects) > 0: self._count_send(updated_objects) return updated_objects # Otherwise, register asyncio events to be waited. async_task_to_events = {} async_task_to_watched_keys = {} for key in keys_to_snapshot_ids.keys(): # Create a new asyncio event for this key. event = asyncio.Event() # Make sure future caller of notify_changed will unblock this # asyncio Event. self.notifier_events[key].add(event) # Update aggregate pending clients gauge for this namespace namespace_tag = _get_metric_namespace_tag(key) self._pending_clients_by_namespace[namespace_tag] += 1 self.pending_clients_gauge.set( self._pending_clients_by_namespace[namespace_tag], tags={"namespace": namespace_tag}, ) task = get_or_create_event_loop().create_task(event.wait()) async_task_to_events[task] = event async_task_to_watched_keys[task] = key done, not_done = await asyncio.wait( async_task_to_watched_keys.keys(), return_when=asyncio.FIRST_COMPLETED, timeout=random.uniform(*self._listen_for_change_request_timeout_s), ) # Collect per-namespace decrements, flush the gauge once per # unique tag after the loop — a single timed-out poll over many # keys can otherwise do N redundant metric writes for the same # namespace. affected_namespaces = set() for task in not_done: task.cancel() event = async_task_to_events[task] key = async_task_to_watched_keys[task] # .get() avoids resurrecting a defaultdict entry for a key # evicted while we were parked. events_set = self.notifier_events.get(key) if events_set is None: continue try: events_set.remove(event) except KeyError: # FIRST_COMPLETED: a sibling wake may have popped this # event via notify_changed. continue if not events_set: self.notifier_events.pop(key, None) namespace_tag = _get_metric_namespace_tag(key) self._pending_clients_by_namespace[namespace_tag] -= 1 affected_namespaces.add(namespace_tag) for namespace_tag in affected_namespaces: self.pending_clients_gauge.set( self._pending_clients_by_namespace[namespace_tag], tags={"namespace": namespace_tag}, ) if len(done) == 0: self._count_send(LongPollState.TIME_OUT) return LongPollState.TIME_OUT else: updated_objects = {} for task in done: updated_object_key = async_task_to_watched_keys[task] # Evicted via remove_keys while parked; skip. if updated_object_key not in self.snapshot_ids: continue updated_objects[updated_object_key] = UpdatedObject( self.object_snapshots[updated_object_key], self.snapshot_ids[updated_object_key], self._notify_timestamps[updated_object_key], ) self._count_send(updated_objects) return updated_objects async def listen_for_change_java( self, keys_to_snapshot_ids_bytes: bytes, ) -> bytes: """Listen for changed objects. only call by java proxy/router now. Args: keys_to_snapshot_ids_bytes: the protobuf bytes of keys_to_snapshot_ids (Dict[str, int]). Returns: The serialized protobuf bytes of the update payload. """ request_proto = LongPollRequest.FromString(keys_to_snapshot_ids_bytes) keys_to_snapshot_ids = { self._parse_xlang_key(xlang_key): snapshot_id for xlang_key, snapshot_id in request_proto.keys_to_snapshot_ids.items() } result = await self.listen_for_change(keys_to_snapshot_ids) # Java long-poll protocol currently doesn't encode LongPollState. # Convert timeout to an empty update payload so the Java client can retry. if result is LongPollState.TIME_OUT: result = {} return self._listen_result_to_proto_bytes(result) def _parse_poll_namespace(self, name: str): if name == LongPollNamespace.ROUTE_TABLE.name: return LongPollNamespace.ROUTE_TABLE elif name == LongPollNamespace.DEPLOYMENT_TARGETS.name: return LongPollNamespace.DEPLOYMENT_TARGETS else: return name def _parse_xlang_key(self, xlang_key: str) -> KeyType: if xlang_key is None: raise ValueError("func _parse_xlang_key: xlang_key is None") if xlang_key.startswith("(") and xlang_key.endswith(")"): fields = xlang_key[1:-1].split(",") if len(fields) == 2: enum_field = self._parse_poll_namespace(fields[0].strip()) if isinstance(enum_field, LongPollNamespace): return enum_field, fields[1].strip() else: return self._parse_poll_namespace(xlang_key) raise ValueError("can not parse key type from xlang_key {}".format(xlang_key)) def _build_xlang_key(self, key: KeyType) -> str: if isinstance(key, tuple): return "(" + key[0].name + "," + key[1] + ")" elif isinstance(key, LongPollNamespace): return key.name else: return key def _object_snapshot_to_proto_bytes( self, key: KeyType, object_snapshot: Any ) -> bytes: if key == LongPollNamespace.ROUTE_TABLE: # object_snapshot is Dict[DeploymentID, EndpointInfo] # NOTE(zcin): the endpoint dictionary broadcasted to Java # HTTP proxies should use string as key because Java does # not yet support 2.x or applications xlang_endpoints = { str(endpoint_tag): EndpointInfoProto(route=endpoint_info.route) for endpoint_tag, endpoint_info in object_snapshot.items() } return EndpointSet(endpoints=xlang_endpoints).SerializeToString() elif isinstance(key, tuple) and key[0] == LongPollNamespace.DEPLOYMENT_TARGETS: # object_snapshot.running_replicas is List[RunningReplicaInfo] actor_name_list = [ replica_info.replica_id.to_full_id_str() for replica_info in object_snapshot.running_replicas ] return DeploymentTargetInfo( replica_names=actor_name_list, is_available=object_snapshot.is_available, ).SerializeToString() else: return str.encode(str(object_snapshot)) def _listen_result_to_proto_bytes( self, keys_to_updated_objects: Dict[KeyType, UpdatedObject] ) -> bytes: xlang_keys_to_updated_objects = { self._build_xlang_key(key): UpdatedObjectProto( snapshot_id=updated_object.snapshot_id, object_snapshot=self._object_snapshot_to_proto_bytes( key, updated_object.object_snapshot ), ) for key, updated_object in keys_to_updated_objects.items() } data = { "updated_objects": xlang_keys_to_updated_objects, } proto = LongPollResult(**data) return proto.SerializeToString() def notify_changed(self, updates: Mapping[KeyType, Any]) -> None: """ Update the current snapshot of some objects and notify any long poll clients. """ notify_time = time.time() for object_key, updated_object in updates.items(): try: self.snapshot_ids[object_key] += 1 except KeyError: # Initial snapshot id must be >= 0, so that the long poll client # can send a negative initial snapshot id to get a fast update. # They should also be randomized; see # https://github.com/ray-project/ray/pull/45881#discussion_r1645243485 self.snapshot_ids[object_key] = random.randint(0, 1_000_000) self.object_snapshots[object_key] = updated_object # Record timestamp for latency tracking self._notify_timestamps[object_key] = notify_time logger.debug(f"LongPollHost: Notify change for key {object_key}.") events_to_notify = self.notifier_events.pop(object_key, set()) if events_to_notify: # Decrement aggregate count by the number of events popped namespace_tag = _get_metric_namespace_tag(object_key) self._pending_clients_by_namespace[namespace_tag] -= len( events_to_notify ) self.pending_clients_gauge.set( self._pending_clients_by_namespace[namespace_tag], tags={"namespace": namespace_tag}, ) for event in events_to_notify: event.set() def remove_keys(self, keys: Iterable[KeyType]) -> None: """Evict per-key state and wake any parked listeners. Do NOT evict a key in the same sync call as ``notify_changed`` on it — the waiter only runs after the call returns, by which point ``listen_for_change``'s done-branch guard would drop the payload. """ affected_namespaces = set() for key in keys: self.snapshot_ids.pop(key, None) self.object_snapshots.pop(key, None) self._notify_timestamps.pop(key, None) events_to_notify = self.notifier_events.pop(key, set()) if events_to_notify: # Decrement before waking: the listen_for_change timeout # cleanup would otherwise skip the decrement on the now- # missing notifier_events entry. namespace_tag = _get_metric_namespace_tag(key) self._pending_clients_by_namespace[namespace_tag] -= len( events_to_notify ) affected_namespaces.add(namespace_tag) for event in events_to_notify: event.set() for namespace_tag in affected_namespaces: self.pending_clients_gauge.set( self._pending_clients_by_namespace[namespace_tag], tags={"namespace": namespace_tag}, )