chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,291 @@
|
||||
import json
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
import ray.dashboard.utils as dashboard_utils
|
||||
from ray._common.utils import get_or_create_event_loop
|
||||
from ray._private.protobuf_compat import message_to_json
|
||||
from ray._raylet import GcsClient
|
||||
from ray.core.generated import (
|
||||
events_base_event_pb2,
|
||||
events_event_aggregator_service_pb2,
|
||||
)
|
||||
from ray.dashboard.modules.aggregator.publisher.configs import (
|
||||
GCS_EXPOSABLE_EVENT_TYPES,
|
||||
HTTP_EXPOSABLE_EVENT_TYPES,
|
||||
PUBLISHER_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PublishStats:
|
||||
"""Data class that represents stats of publishing a batch of events."""
|
||||
|
||||
# Whether the publish was successful
|
||||
is_publish_successful: bool
|
||||
# Number of events published
|
||||
num_events_published: int
|
||||
# Number of events filtered out
|
||||
num_events_filtered_out: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class PublishBatch:
|
||||
"""Data class that represents a batch of events to publish."""
|
||||
|
||||
# The list of events to publish
|
||||
events: list[events_base_event_pb2.RayEvent]
|
||||
# dropped task events metadata
|
||||
task_events_metadata: Optional[
|
||||
events_event_aggregator_service_pb2.TaskEventsMetadata
|
||||
] = None
|
||||
|
||||
|
||||
class PublisherClientInterface(ABC):
|
||||
"""Abstract interface for publishing Ray event batches to external destinations.
|
||||
|
||||
Implementations should handle the actual publishing logic, filtering,
|
||||
and format conversion appropriate for their specific destination type.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._exposable_event_types_list: List[str] = []
|
||||
self._allow_all_event_types: bool = False
|
||||
|
||||
def count_num_events_in_batch(self, batch: PublishBatch) -> int:
|
||||
"""Count the number of events in a given PublishBatch."""
|
||||
return len(batch.events)
|
||||
|
||||
def _can_expose_event(self, event) -> bool:
|
||||
"""
|
||||
Check if an event should be allowed to be published.
|
||||
"""
|
||||
if self._allow_all_event_types:
|
||||
return True
|
||||
if not self._exposable_event_types_list:
|
||||
return False
|
||||
|
||||
event_type_name = events_base_event_pb2.RayEvent.EventType.Name(
|
||||
event.event_type
|
||||
)
|
||||
return event_type_name in self._exposable_event_types_list
|
||||
|
||||
@abstractmethod
|
||||
async def publish(self, batch: PublishBatch) -> PublishStats:
|
||||
"""Publish a batch of events to the destination."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def close(self) -> None:
|
||||
"""Clean up any resources used by this client. Should be called when the publisherClient is no longer required"""
|
||||
pass
|
||||
|
||||
|
||||
class AsyncHttpPublisherClient(PublisherClientInterface):
|
||||
"""Client for publishing ray event batches to an external HTTP service."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
endpoint: str,
|
||||
executor: ThreadPoolExecutor,
|
||||
timeout: float = PUBLISHER_TIMEOUT_SECONDS,
|
||||
preserve_proto_field_name: bool = False,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._endpoint = endpoint
|
||||
self._executor = executor
|
||||
self._timeout = aiohttp.ClientTimeout(total=timeout)
|
||||
self._session = None
|
||||
self._preserve_proto_field_name = preserve_proto_field_name
|
||||
|
||||
if HTTP_EXPOSABLE_EVENT_TYPES.strip().upper() == "ALL":
|
||||
self._allow_all_event_types = True
|
||||
self._exposable_event_types_list = []
|
||||
else:
|
||||
self._exposable_event_types_list = [
|
||||
event_type.strip()
|
||||
for event_type in HTTP_EXPOSABLE_EVENT_TYPES.split(",")
|
||||
if event_type.strip()
|
||||
]
|
||||
|
||||
async def publish(self, batch: PublishBatch) -> PublishStats:
|
||||
events_batch: list[events_base_event_pb2.RayEvent] = batch.events
|
||||
if not events_batch:
|
||||
# Nothing to publish -> success but nothing published
|
||||
return PublishStats(
|
||||
is_publish_successful=True,
|
||||
num_events_published=0,
|
||||
num_events_filtered_out=0,
|
||||
)
|
||||
filtered = [e for e in events_batch if self._can_expose_event(e)]
|
||||
num_filtered_out = len(events_batch) - len(filtered)
|
||||
if not filtered:
|
||||
# All filtered out -> success but nothing published
|
||||
return PublishStats(
|
||||
is_publish_successful=True,
|
||||
num_events_published=0,
|
||||
num_events_filtered_out=num_filtered_out,
|
||||
)
|
||||
|
||||
# Convert protobuf objects to python dictionaries for HTTP POST. Run in executor to avoid blocking the event loop.
|
||||
filtered_json = await get_or_create_event_loop().run_in_executor(
|
||||
self._executor,
|
||||
lambda: [
|
||||
json.loads(
|
||||
message_to_json(
|
||||
e,
|
||||
always_print_fields_with_no_presence=True,
|
||||
preserving_proto_field_name=self._preserve_proto_field_name,
|
||||
)
|
||||
)
|
||||
for e in filtered
|
||||
],
|
||||
)
|
||||
|
||||
try:
|
||||
# Create session on first use (lazy initialization)
|
||||
if not self._session:
|
||||
self._session = aiohttp.ClientSession(timeout=self._timeout)
|
||||
|
||||
return await self._send_http_request(filtered_json, num_filtered_out)
|
||||
except Exception as e:
|
||||
logger.error("Failed to send events to external service. Error: %r", e)
|
||||
return PublishStats(
|
||||
is_publish_successful=False,
|
||||
num_events_published=0,
|
||||
num_events_filtered_out=0,
|
||||
)
|
||||
|
||||
async def _send_http_request(self, json_data, num_filtered_out) -> PublishStats:
|
||||
async with self._session.post(
|
||||
self._endpoint,
|
||||
json=json_data,
|
||||
) as resp:
|
||||
resp.raise_for_status()
|
||||
return PublishStats(
|
||||
is_publish_successful=True,
|
||||
num_events_published=len(json_data),
|
||||
num_events_filtered_out=num_filtered_out,
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Closes the http session if one was created. Should be called when the publisherClient is no longer required"""
|
||||
if self._session:
|
||||
await self._session.close()
|
||||
self._session = None
|
||||
|
||||
def set_session(self, session) -> None:
|
||||
"""Inject an HTTP client session.
|
||||
|
||||
If a session is set explicitly, it will be used and managed by close().
|
||||
"""
|
||||
self._session = session
|
||||
|
||||
|
||||
class AsyncGCSTaskEventsPublisherClient(PublisherClientInterface):
|
||||
"""Client for publishing ray event batches to GCS."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
gcs_client: GcsClient,
|
||||
executor: ThreadPoolExecutor,
|
||||
timeout_s: float = PUBLISHER_TIMEOUT_SECONDS,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._gcs_client = gcs_client
|
||||
self._executor = executor
|
||||
self._timeout_s = timeout_s
|
||||
|
||||
self._exposable_event_types_list = GCS_EXPOSABLE_EVENT_TYPES
|
||||
|
||||
async def publish(
|
||||
self,
|
||||
batch: PublishBatch,
|
||||
) -> PublishStats:
|
||||
events = batch.events
|
||||
task_events_metadata = batch.task_events_metadata
|
||||
has_dropped_task_attempts = (
|
||||
task_events_metadata and task_events_metadata.dropped_task_attempts
|
||||
)
|
||||
if not events and not has_dropped_task_attempts:
|
||||
# Nothing to publish -> success but nothing published
|
||||
return PublishStats(
|
||||
is_publish_successful=True,
|
||||
num_events_published=0,
|
||||
num_events_filtered_out=0,
|
||||
)
|
||||
|
||||
# Filter events based on exposable event types
|
||||
filtered_events = [e for e in events if self._can_expose_event(e)]
|
||||
num_filtered_out = len(events) - len(filtered_events)
|
||||
|
||||
if not filtered_events and not has_dropped_task_attempts:
|
||||
# all events filtered out and no task events metadata -> success but nothing published
|
||||
return PublishStats(
|
||||
is_publish_successful=True,
|
||||
num_events_published=0,
|
||||
num_events_filtered_out=num_filtered_out,
|
||||
)
|
||||
|
||||
try:
|
||||
events_data = self._create_ray_events_data(
|
||||
filtered_events, task_events_metadata
|
||||
)
|
||||
request = events_event_aggregator_service_pb2.AddEventsRequest(
|
||||
events_data=events_data
|
||||
)
|
||||
serialized_request = await get_or_create_event_loop().run_in_executor(
|
||||
self._executor,
|
||||
lambda: request.SerializeToString(),
|
||||
)
|
||||
status_code = await self._gcs_client.async_add_events(
|
||||
serialized_request, self._timeout_s, self._executor
|
||||
)
|
||||
|
||||
if status_code != dashboard_utils.HTTPStatusCode.OK:
|
||||
logger.error(f"GCS AddEvents failed: {status_code}")
|
||||
return PublishStats(
|
||||
is_publish_successful=False,
|
||||
num_events_published=0,
|
||||
num_events_filtered_out=0,
|
||||
)
|
||||
return PublishStats(
|
||||
is_publish_successful=True,
|
||||
num_events_published=len(filtered_events),
|
||||
num_events_filtered_out=num_filtered_out,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to send events to GCS: {e}")
|
||||
return PublishStats(
|
||||
is_publish_successful=False,
|
||||
num_events_published=0,
|
||||
num_events_filtered_out=0,
|
||||
)
|
||||
|
||||
def _create_ray_events_data(
|
||||
self,
|
||||
event_batch: List[events_base_event_pb2.RayEvent],
|
||||
task_events_metadata: Optional[
|
||||
events_event_aggregator_service_pb2.TaskEventsMetadata
|
||||
] = None,
|
||||
) -> events_event_aggregator_service_pb2.RayEventsData:
|
||||
"""
|
||||
Helper method to create RayEventsData from event batch and metadata.
|
||||
"""
|
||||
events_data = events_event_aggregator_service_pb2.RayEventsData()
|
||||
events_data.events.extend(event_batch)
|
||||
|
||||
if task_events_metadata:
|
||||
events_data.task_events_metadata.CopyFrom(task_events_metadata)
|
||||
|
||||
return events_data
|
||||
|
||||
async def close(self) -> None:
|
||||
pass
|
||||
@@ -0,0 +1,54 @@
|
||||
# Environment variables for the aggregator agent publisher component.
|
||||
import os
|
||||
|
||||
from ray._private import ray_constants
|
||||
|
||||
env_var_prefix = "RAY_DASHBOARD_AGGREGATOR_AGENT_PUBLISHER"
|
||||
# Timeout for the publisher to publish events to the destination
|
||||
PUBLISHER_TIMEOUT_SECONDS = ray_constants.env_integer(
|
||||
f"{env_var_prefix}_TIMEOUT_SECONDS", 3
|
||||
)
|
||||
# Maximum number of retries for publishing events to the destination, if less than 0, will retry indefinitely
|
||||
PUBLISHER_MAX_RETRIES = ray_constants.env_integer(f"{env_var_prefix}_MAX_RETRIES", -1)
|
||||
# Initial backoff time for publishing events to the destination
|
||||
PUBLISHER_INITIAL_BACKOFF_SECONDS = ray_constants.env_float(
|
||||
f"{env_var_prefix}_INITIAL_BACKOFF_SECONDS", 0.01
|
||||
)
|
||||
# Maximum backoff time for publishing events to the destination
|
||||
PUBLISHER_MAX_BACKOFF_SECONDS = ray_constants.env_float(
|
||||
f"{env_var_prefix}_MAX_BACKOFF_SECONDS", 5.0
|
||||
)
|
||||
# Jitter ratio for publishing events to the destination
|
||||
PUBLISHER_JITTER_RATIO = ray_constants.env_float(f"{env_var_prefix}_JITTER_RATIO", 0.1)
|
||||
# Maximum sleep time between sending batches of events to the destination, should be greater than 0.0 to avoid busy looping
|
||||
PUBLISHER_MAX_BUFFER_SEND_INTERVAL_SECONDS = ray_constants.env_float(
|
||||
f"{env_var_prefix}_MAX_BUFFER_SEND_INTERVAL_SECONDS", 0.1
|
||||
)
|
||||
|
||||
# HTTP Publisher specific configurations
|
||||
# Comma-separated list of event types that are allowed to be exposed to external HTTP services
|
||||
# Valid values: TASK_DEFINITION_EVENT, TASK_LIFECYCLE_EVENT, ACTOR_TASK_DEFINITION_EVENT, etc.
|
||||
# Set to "ALL" to allow all event types.
|
||||
# The list of all supported event types can be found in src/ray/protobuf/public/events_base_event.proto (EventType enum)
|
||||
# By default TASK_PROFILE_EVENT is not exposed to external services
|
||||
DEFAULT_HTTP_EXPOSABLE_EVENT_TYPES = (
|
||||
"TASK_DEFINITION_EVENT,TASK_LIFECYCLE_EVENT,ACTOR_TASK_DEFINITION_EVENT,"
|
||||
"DRIVER_JOB_DEFINITION_EVENT,DRIVER_JOB_LIFECYCLE_EVENT,"
|
||||
"ACTOR_DEFINITION_EVENT,ACTOR_LIFECYCLE_EVENT,"
|
||||
"NODE_DEFINITION_EVENT,NODE_LIFECYCLE_EVENT,"
|
||||
"PLATFORM_EVENT,"
|
||||
)
|
||||
HTTP_EXPOSABLE_EVENT_TYPES = os.environ.get(
|
||||
"RAY_DASHBOARD_AGGREGATOR_AGENT_EXPOSABLE_EVENT_TYPES",
|
||||
DEFAULT_HTTP_EXPOSABLE_EVENT_TYPES,
|
||||
)
|
||||
|
||||
# GCS Publisher specific configurations
|
||||
# List of event types that are allowed to be exposed to GCS, not overridden by environment variable
|
||||
# as GCS only supports Task event types
|
||||
GCS_EXPOSABLE_EVENT_TYPES = [
|
||||
"TASK_DEFINITION_EVENT",
|
||||
"TASK_LIFECYCLE_EVENT",
|
||||
"TASK_PROFILE_EVENT",
|
||||
"ACTOR_TASK_DEFINITION_EVENT",
|
||||
]
|
||||
@@ -0,0 +1,53 @@
|
||||
from ray._private.telemetry.open_telemetry_metric_recorder import (
|
||||
OpenTelemetryMetricRecorder,
|
||||
)
|
||||
from ray.dashboard.modules.aggregator.constants import (
|
||||
AGGREGATOR_AGENT_METRIC_PREFIX,
|
||||
)
|
||||
|
||||
# OpenTelemetry metrics setup (registered once at import time)
|
||||
metric_recorder = OpenTelemetryMetricRecorder()
|
||||
|
||||
# Counter metrics
|
||||
published_counter_name = f"{AGGREGATOR_AGENT_METRIC_PREFIX}_published_events"
|
||||
metric_recorder.register_counter_metric(
|
||||
published_counter_name,
|
||||
"Total number of events successfully published to the destination.",
|
||||
)
|
||||
|
||||
filtered_counter_name = f"{AGGREGATOR_AGENT_METRIC_PREFIX}_filtered_events"
|
||||
metric_recorder.register_counter_metric(
|
||||
filtered_counter_name,
|
||||
"Total number of events filtered out before publishing to the destination.",
|
||||
)
|
||||
|
||||
failed_counter_name = f"{AGGREGATOR_AGENT_METRIC_PREFIX}_publish_failures"
|
||||
metric_recorder.register_counter_metric(
|
||||
failed_counter_name,
|
||||
"Total number of events that failed to publish after retries.",
|
||||
)
|
||||
|
||||
# Histogram metric
|
||||
publish_latency_hist_name = f"{AGGREGATOR_AGENT_METRIC_PREFIX}_publish_latency_seconds"
|
||||
metric_recorder.register_histogram_metric(
|
||||
publish_latency_hist_name,
|
||||
"Duration of publish calls in seconds.",
|
||||
[0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2, 5],
|
||||
)
|
||||
|
||||
# Gauge metrics
|
||||
consecutive_failures_gauge_name = (
|
||||
f"{AGGREGATOR_AGENT_METRIC_PREFIX}_consecutive_failures_since_last_success"
|
||||
)
|
||||
metric_recorder.register_gauge_metric(
|
||||
consecutive_failures_gauge_name,
|
||||
"Number of consecutive failed publish attempts since the last success.",
|
||||
)
|
||||
|
||||
time_since_last_success_gauge_name = (
|
||||
f"{AGGREGATOR_AGENT_METRIC_PREFIX}_time_since_last_success_seconds"
|
||||
)
|
||||
metric_recorder.register_gauge_metric(
|
||||
time_since_last_success_gauge_name,
|
||||
"Seconds since the last successful publish to the destination.",
|
||||
)
|
||||
@@ -0,0 +1,289 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, Optional
|
||||
|
||||
from ray.dashboard.modules.aggregator.constants import (
|
||||
CONSUMER_TAG_KEY,
|
||||
)
|
||||
from ray.dashboard.modules.aggregator.multi_consumer_event_buffer import (
|
||||
MultiConsumerEventBuffer,
|
||||
)
|
||||
from ray.dashboard.modules.aggregator.publisher.async_publisher_client import (
|
||||
PublishBatch,
|
||||
PublisherClientInterface,
|
||||
)
|
||||
from ray.dashboard.modules.aggregator.publisher.configs import (
|
||||
PUBLISHER_INITIAL_BACKOFF_SECONDS,
|
||||
PUBLISHER_JITTER_RATIO,
|
||||
PUBLISHER_MAX_BACKOFF_SECONDS,
|
||||
PUBLISHER_MAX_BUFFER_SEND_INTERVAL_SECONDS,
|
||||
PUBLISHER_MAX_RETRIES,
|
||||
)
|
||||
from ray.dashboard.modules.aggregator.publisher.metrics import (
|
||||
consecutive_failures_gauge_name,
|
||||
failed_counter_name,
|
||||
filtered_counter_name,
|
||||
metric_recorder,
|
||||
publish_latency_hist_name,
|
||||
published_counter_name,
|
||||
time_since_last_success_gauge_name,
|
||||
)
|
||||
from ray.dashboard.modules.aggregator.task_events_metadata_buffer import (
|
||||
TaskEventsMetadataBuffer,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RayEventPublisherInterface(ABC):
|
||||
"""Abstract interface for publishing Ray event batches to external destinations."""
|
||||
|
||||
@abstractmethod
|
||||
async def run_forever(self) -> None:
|
||||
"""Run the publisher forever until cancellation or process death."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def wait_until_running(self, timeout: Optional[float] = None) -> bool:
|
||||
"""Wait until the publisher has started."""
|
||||
pass
|
||||
|
||||
|
||||
class RayEventPublisher(RayEventPublisherInterface):
|
||||
"""RayEvents publisher that publishes batches of events to a destination by running a worker loop.
|
||||
|
||||
The worker loop continuously pulls batches from the event buffer and publishes them to the destination.
|
||||
"""
|
||||
|
||||
# Cap the exponent to avoid computing unnecessarily large intermediate
|
||||
_MAX_BACKOFF_EXPONENT = 30
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
publish_client: PublisherClientInterface,
|
||||
event_buffer: MultiConsumerEventBuffer,
|
||||
common_metric_tags: Optional[Dict[str, str]] = None,
|
||||
task_metadata_buffer: Optional[TaskEventsMetadataBuffer] = None,
|
||||
max_retries: int = PUBLISHER_MAX_RETRIES,
|
||||
initial_backoff: float = PUBLISHER_INITIAL_BACKOFF_SECONDS,
|
||||
max_backoff: float = PUBLISHER_MAX_BACKOFF_SECONDS,
|
||||
jitter_ratio: float = PUBLISHER_JITTER_RATIO,
|
||||
) -> None:
|
||||
"""Initialize a RayEventsPublisher.
|
||||
|
||||
Args:
|
||||
name: Name identifier for this publisher instance
|
||||
publish_client: Client for publishing events to the destination
|
||||
event_buffer: Buffer for reading batches of events
|
||||
common_metric_tags: Common labels for all prometheus metrics
|
||||
task_metadata_buffer: Buffer for reading a batch of droppedtask metadata
|
||||
max_retries: Maximum number of retries for failed publishes
|
||||
initial_backoff: Initial backoff time between retries in seconds
|
||||
max_backoff: Maximum backoff time between retries in seconds
|
||||
jitter_ratio: Random jitter ratio to add to backoff times
|
||||
"""
|
||||
self._name = name
|
||||
self._common_metric_tags = dict(common_metric_tags or {})
|
||||
self._common_metric_tags[CONSUMER_TAG_KEY] = name
|
||||
self._max_retries = int(max_retries)
|
||||
self._initial_backoff = float(initial_backoff)
|
||||
self._max_backoff = float(max_backoff)
|
||||
self._jitter_ratio = float(jitter_ratio)
|
||||
self._publish_client = publish_client
|
||||
self._event_buffer = event_buffer
|
||||
self._task_metadata_buffer = task_metadata_buffer
|
||||
|
||||
# Event set once the publisher has registered as a consumer and is ready to publish events
|
||||
self._started_event: asyncio.Event = asyncio.Event()
|
||||
|
||||
async def run_forever(self) -> None:
|
||||
"""Run the publisher forever until cancellation or process death.
|
||||
|
||||
Registers as a consumer, starts the worker loop, and handles cleanup on cancellation.
|
||||
"""
|
||||
await self._event_buffer.register_consumer(self._name)
|
||||
|
||||
# Signal that the publisher is ready to publish events
|
||||
self._started_event.set()
|
||||
|
||||
try:
|
||||
logger.info(f"Starting publisher {self._name}")
|
||||
while True:
|
||||
events_batch = await self._event_buffer.wait_for_batch(
|
||||
self._name,
|
||||
PUBLISHER_MAX_BUFFER_SEND_INTERVAL_SECONDS,
|
||||
)
|
||||
publish_batch = PublishBatch(events=events_batch)
|
||||
|
||||
if self._task_metadata_buffer is not None:
|
||||
task_metadata_batch = self._task_metadata_buffer.get()
|
||||
publish_batch.task_events_metadata = task_metadata_batch
|
||||
|
||||
await self._async_publish_with_retries(publish_batch)
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"Publisher {self._name} cancelled, shutting down gracefully")
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Publisher {self._name} encountered error: {e}")
|
||||
raise
|
||||
finally:
|
||||
self._started_event.clear()
|
||||
await self._publish_client.close()
|
||||
|
||||
async def wait_until_running(self, timeout: Optional[float] = None) -> bool:
|
||||
"""Wait until the publisher has started.
|
||||
|
||||
Args:
|
||||
timeout: Maximum time to wait in seconds. If None, waits indefinitely.
|
||||
|
||||
Returns:
|
||||
True if the publisher started before the timeout, False otherwise.
|
||||
If timeout is None, waits indefinitely.
|
||||
"""
|
||||
if timeout is None:
|
||||
await self._started_event.wait()
|
||||
return True
|
||||
try:
|
||||
await asyncio.wait_for(self._started_event.wait(), timeout)
|
||||
return True
|
||||
except asyncio.TimeoutError:
|
||||
return False
|
||||
|
||||
async def _async_publish_with_retries(self, batch) -> None:
|
||||
"""Attempts to publish a batch with retries.
|
||||
|
||||
Will retry failed publishes up to max_retries times with increasing delays.
|
||||
"""
|
||||
num_events_in_batch = self._publish_client.count_num_events_in_batch(batch)
|
||||
failed_attempts_since_last_success = 0
|
||||
while True:
|
||||
start = asyncio.get_running_loop().time()
|
||||
result = await self._publish_client.publish(batch)
|
||||
duration = asyncio.get_running_loop().time() - start
|
||||
|
||||
if result.is_publish_successful:
|
||||
await self._record_success(
|
||||
num_published=int(result.num_events_published),
|
||||
num_filtered=int(result.num_events_filtered_out),
|
||||
duration=float(duration),
|
||||
)
|
||||
failed_attempts_since_last_success = 0
|
||||
return
|
||||
|
||||
# Failed attempt
|
||||
# case 1: if max retries are exhausted mark as failed and break out, retry indefinitely if max_retries is less than 0
|
||||
if (
|
||||
self._max_retries >= 0
|
||||
and failed_attempts_since_last_success >= self._max_retries
|
||||
):
|
||||
await self._record_final_failure(
|
||||
num_failed_events=int(num_events_in_batch),
|
||||
duration=float(duration),
|
||||
)
|
||||
return
|
||||
|
||||
# case 2: max retries not exhausted, increment failed attempts counter and add latency to failure list, retry publishing batch with backoff
|
||||
failed_attempts_since_last_success += 1
|
||||
await self._record_retry_failure(
|
||||
duration=float(duration),
|
||||
failed_attempts=int(failed_attempts_since_last_success),
|
||||
)
|
||||
|
||||
await self._async_sleep_with_backoff(failed_attempts_since_last_success)
|
||||
|
||||
async def _async_sleep_with_backoff(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff and optional jitter.
|
||||
|
||||
Args:
|
||||
attempt: The current attempt number (0-based)
|
||||
"""
|
||||
capped_attempt = min(attempt, self._MAX_BACKOFF_EXPONENT)
|
||||
delay = min(
|
||||
self._max_backoff,
|
||||
self._initial_backoff * (2**capped_attempt),
|
||||
)
|
||||
if self._jitter_ratio > 0:
|
||||
jitter = delay * self._jitter_ratio
|
||||
delay = max(0.0, random.uniform(delay - jitter, delay + jitter))
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
async def _record_success(
|
||||
self, num_published: int, num_filtered: int, duration: float
|
||||
) -> None:
|
||||
"""Update in-memory stats and Prometheus metrics for a successful publish."""
|
||||
if num_published > 0:
|
||||
metric_recorder.set_metric_value(
|
||||
published_counter_name,
|
||||
self._common_metric_tags,
|
||||
int(num_published),
|
||||
)
|
||||
if num_filtered > 0:
|
||||
metric_recorder.set_metric_value(
|
||||
filtered_counter_name, self._common_metric_tags, int(num_filtered)
|
||||
)
|
||||
metric_recorder.set_metric_value(
|
||||
consecutive_failures_gauge_name, self._common_metric_tags, 0
|
||||
)
|
||||
metric_recorder.set_metric_value(
|
||||
time_since_last_success_gauge_name, self._common_metric_tags, 0
|
||||
)
|
||||
metric_recorder.set_metric_value(
|
||||
publish_latency_hist_name,
|
||||
{**self._common_metric_tags, "Outcome": "success"},
|
||||
float(duration),
|
||||
)
|
||||
|
||||
async def _record_retry_failure(
|
||||
self, duration: float, failed_attempts: int
|
||||
) -> None:
|
||||
"""Update Prometheus metrics for a retryable failure attempt."""
|
||||
metric_recorder.set_metric_value(
|
||||
consecutive_failures_gauge_name,
|
||||
self._common_metric_tags,
|
||||
int(failed_attempts),
|
||||
)
|
||||
metric_recorder.set_metric_value(
|
||||
publish_latency_hist_name,
|
||||
{**self._common_metric_tags, "Outcome": "failure"},
|
||||
float(duration),
|
||||
)
|
||||
|
||||
async def _record_final_failure(
|
||||
self, num_failed_events: int, duration: float
|
||||
) -> None:
|
||||
"""Update in-memory stats and Prometheus metrics for a final (non-retryable) failure."""
|
||||
if num_failed_events > 0:
|
||||
metric_recorder.set_metric_value(
|
||||
failed_counter_name,
|
||||
self._common_metric_tags,
|
||||
int(num_failed_events),
|
||||
)
|
||||
metric_recorder.set_metric_value(
|
||||
consecutive_failures_gauge_name, self._common_metric_tags, 0
|
||||
)
|
||||
metric_recorder.set_metric_value(
|
||||
publish_latency_hist_name,
|
||||
{**self._common_metric_tags, "Outcome": "failure"},
|
||||
float(duration),
|
||||
)
|
||||
|
||||
|
||||
class NoopPublisher(RayEventPublisherInterface):
|
||||
"""A no-op publisher that adheres to the minimal interface used by AggregatorAgent.
|
||||
|
||||
Used when a destination is disabled. It runs forever but does nothing.
|
||||
"""
|
||||
|
||||
async def run_forever(self) -> None:
|
||||
"""Run forever doing nothing until cancellation."""
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
logger.info("NoopPublisher cancelled")
|
||||
raise
|
||||
|
||||
async def wait_until_running(self, timeout: Optional[float] = None) -> bool:
|
||||
return True
|
||||
Reference in New Issue
Block a user