chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,191 @@
|
||||
import asyncio
|
||||
import collections
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
||||
import ray.dashboard.optional_utils as dashboard_optional_utils
|
||||
import ray.dashboard.utils as dashboard_utils
|
||||
from ray._private import ray_constants
|
||||
from ray.core.generated.events_base_event_pb2 import RayEvent
|
||||
from ray.dashboard.consts import (
|
||||
DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX,
|
||||
GCS_RPC_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
routes = dashboard_optional_utils.DashboardHeadRouteTable
|
||||
# Max number of events to cache in memory.
|
||||
MAX_EVENTS_TO_CACHE = 1000
|
||||
|
||||
|
||||
class PlatformEventsHead(dashboard_utils.DashboardHeadModule):
|
||||
"""
|
||||
Dashboard Head Module for managing and serving infrastructure platform events.
|
||||
|
||||
This module acts as a generic, platform-agnostic REST controller and cache
|
||||
for events coming from the underlying hosting environment (e.g., Kubernetes).
|
||||
It delegates platform-specific event watching to dynamic providers and exposes
|
||||
a single unified API endpoint (/api/v0/platform_events) to query the events.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def is_enabled(cls) -> bool:
|
||||
is_enabled = os.environ.get(
|
||||
"RAY_DASHBOARD_INGEST_PLATFORM_EVENTS", "False"
|
||||
).lower() in ("true", "1")
|
||||
if not is_enabled:
|
||||
logger.info(
|
||||
"Skipping PlatformEventsHead because RAY_DASHBOARD_INGEST_PLATFORM_EVENTS is not enabled."
|
||||
)
|
||||
return is_enabled
|
||||
|
||||
def __init__(self, config: dashboard_utils.DashboardHeadModuleConfig):
|
||||
super().__init__(config)
|
||||
self._events = collections.OrderedDict()
|
||||
self._provider = None
|
||||
self._loop = None
|
||||
|
||||
async def run(self):
|
||||
await super().run()
|
||||
self._loop = asyncio.get_running_loop()
|
||||
|
||||
if "PLATFORM_EVENT" in ray_constants.RAY_ENABLE_PYTHON_RAY_EVENT_TYPES:
|
||||
try:
|
||||
from ray._raylet import EventRecorder
|
||||
|
||||
head_node_id_hex = await dashboard_utils.get_head_node_id(
|
||||
self.gcs_client, timeout=GCS_RPC_TIMEOUT_SECONDS
|
||||
)
|
||||
|
||||
if head_node_id_hex:
|
||||
key = f"{DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX}{head_node_id_hex}"
|
||||
value = await self.gcs_client.async_internal_kv_get(
|
||||
key.encode(),
|
||||
namespace=ray_constants.KV_NAMESPACE_DASHBOARD,
|
||||
timeout=GCS_RPC_TIMEOUT_SECONDS,
|
||||
)
|
||||
if value:
|
||||
ip, _, grpc_port = json.loads(value.decode())
|
||||
EventRecorder.initialize(
|
||||
aggregator_port=int(grpc_port),
|
||||
node_ip=ip,
|
||||
node_id_hex=head_node_id_hex,
|
||||
max_buffer_size=10000,
|
||||
metric_source="platform_events",
|
||||
)
|
||||
logger.info("Initialized EventRecorder for PlatformEventsHead.")
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Failed to initialize EventRecorder in PlatformEventsHead: {e}"
|
||||
)
|
||||
|
||||
# Detect infrastructure platform.
|
||||
if "KUBERNETES_SERVICE_HOST" in os.environ:
|
||||
try:
|
||||
from ray.dashboard.modules.platform_events.providers.k8s_provider import (
|
||||
KubernetesEventProvider,
|
||||
)
|
||||
|
||||
self._provider = KubernetesEventProvider(self._process_event_callback)
|
||||
logger.info("Initialized Kubernetes platform event provider.")
|
||||
await self._provider.run()
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error running Kubernetes platform event provider: {e}"
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"No supported platform environment detected (e.g. KUBERNETES_SERVICE_HOST not found). "
|
||||
"Platform events will be disabled."
|
||||
)
|
||||
|
||||
def _process_event_callback(self, ray_event: RayEvent):
|
||||
"""Thread-safe entry point that dispatches event caching to the main asyncio loop."""
|
||||
if self._loop and self._loop.is_running():
|
||||
self._loop.call_soon_threadsafe(self._process_event_internal, ray_event)
|
||||
else:
|
||||
# Shutdown path
|
||||
self._process_event_internal(ray_event)
|
||||
|
||||
def _process_event_internal(self, ray_event: RayEvent):
|
||||
"""Internal callback running in the main asyncio loop to cache events."""
|
||||
uid = ray_event.event_id.decode()
|
||||
|
||||
if uid in self._events:
|
||||
# Refresh with the latest delivery
|
||||
self._events[uid] = ray_event
|
||||
self._events.move_to_end(uid)
|
||||
else:
|
||||
# Add new event and enforce max size
|
||||
if len(self._events) >= MAX_EVENTS_TO_CACHE:
|
||||
self._events.popitem(last=False)
|
||||
self._events[uid] = ray_event
|
||||
|
||||
if "PLATFORM_EVENT" in ray_constants.RAY_ENABLE_PYTHON_RAY_EVENT_TYPES:
|
||||
try:
|
||||
from ray._common.observability.platform_events import (
|
||||
PlatformEventBuilder,
|
||||
)
|
||||
from ray._raylet import EventRecorder
|
||||
|
||||
builder = PlatformEventBuilder(
|
||||
event_uid=uid,
|
||||
platform=ray_event.platform_event.source.platform,
|
||||
object_kind=ray_event.platform_event.object_kind,
|
||||
object_name=ray_event.platform_event.object_name,
|
||||
reason=ray_event.platform_event.reason,
|
||||
message=ray_event.message,
|
||||
severity=ray_event.severity,
|
||||
component=ray_event.platform_event.source.component,
|
||||
source_metadata=dict(ray_event.platform_event.source.metadata),
|
||||
custom_fields=dict(ray_event.platform_event.custom_fields),
|
||||
)
|
||||
cython_event = builder.build(
|
||||
event_id=uid.encode(),
|
||||
timestamp_ns=(
|
||||
ray_event.timestamp.seconds * 1_000_000_000
|
||||
+ ray_event.timestamp.nanos
|
||||
),
|
||||
)
|
||||
EventRecorder.emit(cython_event)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to emit platform event via EventRecorder: {e}")
|
||||
|
||||
@routes.get("/api/v0/platform_events")
|
||||
@dashboard_optional_utils.aiohttp_cache
|
||||
async def get_platform_events(self, req):
|
||||
"""Return recently observed platform events as a JSON array."""
|
||||
# Sort by timestamp descending (newest first)
|
||||
sorted_events = sorted(
|
||||
self._events.values(),
|
||||
key=lambda e: e.timestamp.seconds + e.timestamp.nanos / 1e9,
|
||||
reverse=True,
|
||||
)
|
||||
payload = [
|
||||
dashboard_utils.message_to_dict(
|
||||
e, always_print_fields_with_no_presence=True
|
||||
)
|
||||
for e in sorted_events
|
||||
]
|
||||
return dashboard_optional_utils.rest_response(
|
||||
status_code=dashboard_utils.HTTPStatusCode.OK,
|
||||
message="Retrieved platform events",
|
||||
events=payload,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def is_minimal_module():
|
||||
return False
|
||||
|
||||
async def cleanup(self):
|
||||
if self._provider:
|
||||
await self._provider.cleanup()
|
||||
if "PLATFORM_EVENT" in ray_constants.RAY_ENABLE_PYTHON_RAY_EVENT_TYPES:
|
||||
try:
|
||||
from ray._raylet import EventRecorder
|
||||
|
||||
EventRecorder.shutdown()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to shutdown EventRecorder: {e}")
|
||||
logger.info("Platform events watcher stopped.")
|
||||
@@ -0,0 +1,37 @@
|
||||
import abc
|
||||
from typing import Callable
|
||||
|
||||
from ray.core.generated.events_base_event_pb2 import RayEvent
|
||||
|
||||
|
||||
class PlatformEventProvider(abc.ABC):
|
||||
"""
|
||||
Base interface for all platform-specific event providers (e.g. Kubernetes).
|
||||
"""
|
||||
|
||||
def __init__(self, callback: Callable[[RayEvent], None]):
|
||||
"""
|
||||
Initialize the provider.
|
||||
Implementations of run() may invoke the callback from background worker
|
||||
threads, so consumers are responsible for passing a callback that is
|
||||
fully thread-safe.
|
||||
|
||||
Args:
|
||||
callback: A callback function to deliver new or updated RayEvent
|
||||
protobuf objects to the central manager.
|
||||
"""
|
||||
self._callback = callback
|
||||
|
||||
@abc.abstractmethod
|
||||
async def run(self) -> None:
|
||||
"""
|
||||
Start the provider. This method should run the event watch/poll loop.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def cleanup(self) -> None:
|
||||
"""
|
||||
Stop the provider and release all resources gracefully.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,314 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from google.protobuf.timestamp_pb2 import Timestamp
|
||||
|
||||
from ray.core.generated.events_base_event_pb2 import RayEvent
|
||||
from ray.core.generated.platform_event_pb2 import PlatformEvent, Source
|
||||
from ray.dashboard.modules.platform_events.providers.base import PlatformEventProvider
|
||||
|
||||
try:
|
||||
from kubernetes import client, config as k8s_config, watch
|
||||
from kubernetes.client.rest import ApiException
|
||||
except ImportError:
|
||||
client = None
|
||||
k8s_config = None
|
||||
watch = None
|
||||
ApiException = None
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class KubernetesEventProvider(PlatformEventProvider):
|
||||
"""
|
||||
Kubernetes-specific implementation of PlatformEventProvider.
|
||||
|
||||
Watches the Kubernetes Event API for events related to the active RayCluster
|
||||
and its parent owners (such as RayJob or RayService). It parses raw K8s event
|
||||
objects, maps them to standardized, platform-agnostic RayEvent protobuf
|
||||
objects, and delivers them to the manager's callback. The callback may be
|
||||
invoked from any number of watcher threads, so consumers must pass a
|
||||
thread-safe callback.
|
||||
"""
|
||||
|
||||
def __init__(self, callback: Callable[[RayEvent], None]):
|
||||
super().__init__(callback)
|
||||
self._cluster_name: Optional[str] = None
|
||||
self._ray_job_name: Optional[str] = None
|
||||
self._ray_service_name: Optional[str] = None
|
||||
self._namespace: Optional[str] = None
|
||||
|
||||
# Dict mapping "kind/name" target IDs to their last seen Kubernetes resource versions
|
||||
self._last_resource_versions: Dict[str, Optional[str]] = {}
|
||||
|
||||
# The Kubernetes CoreV1Api client, populated in `_init_k8s_client`
|
||||
self._k8s_v1_api: Optional[Any] = None
|
||||
|
||||
# Lifecycle and thread synchronization controls
|
||||
self._stop_event: threading.Event = threading.Event()
|
||||
self._async_stop_event: asyncio.Event = asyncio.Event()
|
||||
|
||||
# Dict mapping "kind/name" target IDs to active kubernetes.watch.Watch instances
|
||||
self._active_watches: Dict[str, Any] = {}
|
||||
self._watches_lock: threading.Lock = threading.Lock()
|
||||
self._threads: List[threading.Thread] = []
|
||||
self._cleaned_up: bool = False
|
||||
|
||||
async def _init_k8s_client(self) -> bool:
|
||||
if not k8s_config:
|
||||
logger.info(
|
||||
"Kubernetes dependencies not found. Kubernetes events will be disabled."
|
||||
)
|
||||
return False
|
||||
|
||||
# RAY_CLUSTER_NAME is set by KubeRay Operator on every Ray cluster.
|
||||
cluster_name = os.environ.get("RAY_CLUSTER_NAME")
|
||||
if not cluster_name:
|
||||
logger.warning(
|
||||
"RAY_CLUSTER_NAME not found. Kubernetes events will not be available."
|
||||
)
|
||||
return False
|
||||
|
||||
# RAY_CLUSTER_NAMESPACE is set by KubeRay Operator on every Ray cluster pod.
|
||||
cluster_namespace = os.environ.get("RAY_CLUSTER_NAMESPACE")
|
||||
if not cluster_namespace:
|
||||
logger.warning(
|
||||
"RAY_CLUSTER_NAMESPACE not found. Kubernetes events will not be available."
|
||||
)
|
||||
return False
|
||||
|
||||
self._namespace = cluster_namespace
|
||||
|
||||
try:
|
||||
# Try in-cluster config first
|
||||
k8s_config.load_incluster_config()
|
||||
except k8s_config.ConfigException:
|
||||
try:
|
||||
# Fallback to kubeconfig (local dev)
|
||||
k8s_config.load_kube_config()
|
||||
except k8s_config.ConfigException:
|
||||
logger.warning(
|
||||
"Kubernetes client not installed or configuration failed. "
|
||||
"Kubernetes events will not be available."
|
||||
)
|
||||
return False
|
||||
|
||||
self._k8s_v1_api = client.CoreV1Api()
|
||||
logger.info("Kubernetes client initialized.")
|
||||
|
||||
self._cluster_name = cluster_name
|
||||
logger.info(
|
||||
f"Discovered RayCluster name: {self._cluster_name} in namespace: {self._namespace}"
|
||||
)
|
||||
try:
|
||||
custom_api = client.CustomObjectsApi()
|
||||
loop = asyncio.get_running_loop()
|
||||
cluster_obj = await loop.run_in_executor(
|
||||
None,
|
||||
lambda: custom_api.get_namespaced_custom_object(
|
||||
group="ray.io",
|
||||
version="v1",
|
||||
namespace=self._namespace,
|
||||
plural="rayclusters",
|
||||
name=self._cluster_name,
|
||||
_request_timeout=30,
|
||||
),
|
||||
)
|
||||
c_metadata = cluster_obj.get("metadata", {})
|
||||
owner_refs = c_metadata.get("ownerReferences", [])
|
||||
for ref in owner_refs:
|
||||
if ref.get("kind") == "RayJob":
|
||||
self._ray_job_name = ref.get("name")
|
||||
break
|
||||
elif ref.get("kind") == "RayService":
|
||||
self._ray_service_name = ref.get("name")
|
||||
break
|
||||
except Exception as ce:
|
||||
logger.info(f"Could not read RayCluster owners for discovery: {ce}")
|
||||
if self._ray_job_name:
|
||||
logger.info(f"Discovered parent RayJob: {self._ray_job_name}")
|
||||
if self._ray_service_name:
|
||||
logger.info(f"Discovered parent RayService: {self._ray_service_name}")
|
||||
|
||||
return True
|
||||
|
||||
async def run(self) -> None:
|
||||
if await self._init_k8s_client():
|
||||
# Targets to watch
|
||||
targets = [("RayCluster", self._cluster_name)]
|
||||
if self._ray_job_name:
|
||||
targets.append(("RayJob", self._ray_job_name))
|
||||
if self._ray_service_name:
|
||||
targets.append(("RayService", self._ray_service_name))
|
||||
|
||||
# Pre-initialize keys to avoid race conditions during dict resizing in threads
|
||||
for kind, name in targets:
|
||||
target_id = f"{kind}/{name}"
|
||||
self._last_resource_versions[target_id] = None
|
||||
|
||||
try:
|
||||
# Start a dedicated, named OS thread for each target to ensure strict execution guarantees
|
||||
for kind, name in targets:
|
||||
t = threading.Thread(
|
||||
target=self._run_k8s_watch,
|
||||
args=(kind, name),
|
||||
name=f"platform_events_watch_{kind}_{name}",
|
||||
daemon=True,
|
||||
)
|
||||
t.start()
|
||||
self._threads.append(t)
|
||||
|
||||
# Block the run coroutine until cleanup is called
|
||||
await self._async_stop_event.wait()
|
||||
finally:
|
||||
await self.cleanup()
|
||||
|
||||
def _run_k8s_watch(self, kind: str, name: str) -> None:
|
||||
"""Blocking method to run in a separate thread."""
|
||||
if not name:
|
||||
logger.error(
|
||||
f"Attempted to watch events for kind {kind} without a specific name. Aborting."
|
||||
)
|
||||
return
|
||||
|
||||
target_id = f"{kind}/{name}"
|
||||
logger.info(f"Starting Kubernetes event watcher thread for {target_id}")
|
||||
if not self._k8s_v1_api or not watch:
|
||||
logger.warning(
|
||||
f"Halting Kubernetes event watcher for {target_id}: "
|
||||
"Kubernetes client or watch library is not initialized."
|
||||
)
|
||||
return
|
||||
|
||||
w = watch.Watch()
|
||||
with self._watches_lock:
|
||||
self._active_watches[target_id] = w
|
||||
try:
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
stream = w.stream(
|
||||
self._k8s_v1_api.list_namespaced_event,
|
||||
namespace=self._namespace,
|
||||
field_selector=(
|
||||
f"involvedObject.kind={kind},involvedObject.name={name}"
|
||||
),
|
||||
resource_version=self._last_resource_versions.get(target_id),
|
||||
timeout_seconds=60,
|
||||
_request_timeout=70,
|
||||
)
|
||||
|
||||
for event in stream:
|
||||
k8s_event_obj = event["object"]
|
||||
rv = k8s_event_obj.metadata.resource_version
|
||||
self._last_resource_versions[target_id] = rv
|
||||
self._process_k8s_event(k8s_event_obj)
|
||||
|
||||
except ApiException as e:
|
||||
if e.status == 410:
|
||||
logger.warning(
|
||||
f"Resource version expired for {target_id}, resetting watch"
|
||||
)
|
||||
self._last_resource_versions[target_id] = None
|
||||
else:
|
||||
logger.error(
|
||||
f"Kubernetes API error watching events for {target_id}: {e}"
|
||||
)
|
||||
self._stop_event.wait(5)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error watching Kubernetes events for {target_id}: {e}"
|
||||
)
|
||||
self._stop_event.wait(5)
|
||||
finally:
|
||||
with self._watches_lock:
|
||||
self._active_watches.pop(target_id, None)
|
||||
|
||||
def _process_k8s_event(self, event_obj: Any) -> None:
|
||||
"""Parse a K8s event into a RayEvent proto and deliver it via the
|
||||
consumer callback. Invoked from a watcher worker thread."""
|
||||
involved_object = event_obj.involved_object
|
||||
uid = event_obj.metadata.uid
|
||||
|
||||
event_timestamp = event_obj.last_timestamp or event_obj.first_timestamp
|
||||
timestamp = Timestamp()
|
||||
if event_timestamp:
|
||||
timestamp.FromDatetime(event_timestamp)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Kubernetes event {uid} involvedObject={involved_object.kind}/{involved_object.name} "
|
||||
"is missing both last_timestamp and first_timestamp. Leaving event timestamp unpopulated."
|
||||
)
|
||||
|
||||
k8s_event_type = event_obj.type or "Normal"
|
||||
severity = (
|
||||
RayEvent.Severity.WARNING
|
||||
if k8s_event_type == "Warning"
|
||||
else RayEvent.Severity.INFO
|
||||
)
|
||||
|
||||
source_metadata = {"namespace": event_obj.metadata.namespace or self._namespace}
|
||||
if self._cluster_name:
|
||||
source_metadata["ray_cluster_name"] = self._cluster_name
|
||||
|
||||
source_proto = Source(
|
||||
platform=Source.Platform.KUBERNETES,
|
||||
component=event_obj.source.component if event_obj.source else "",
|
||||
metadata=source_metadata,
|
||||
)
|
||||
|
||||
custom_fields = {}
|
||||
if event_obj.count and event_obj.count > 1:
|
||||
custom_fields["count"] = str(event_obj.count)
|
||||
|
||||
platform_event = PlatformEvent(
|
||||
source=source_proto,
|
||||
object_kind=involved_object.kind,
|
||||
object_name=involved_object.name,
|
||||
message=event_obj.message or "",
|
||||
reason=event_obj.reason or "",
|
||||
custom_fields=custom_fields,
|
||||
)
|
||||
|
||||
ray_event = RayEvent(
|
||||
event_id=uid.encode(),
|
||||
source_type=RayEvent.SourceType.CLUSTER_LIFECYCLE,
|
||||
event_type=RayEvent.EventType.PLATFORM_EVENT,
|
||||
timestamp=timestamp,
|
||||
severity=severity,
|
||||
message=event_obj.message or "",
|
||||
platform_event=platform_event,
|
||||
)
|
||||
|
||||
# Deliver the parsed RayEvent to the central callback
|
||||
self._callback(ray_event)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
if self._cleaned_up:
|
||||
return
|
||||
self._cleaned_up = True
|
||||
|
||||
self._stop_event.set()
|
||||
self._async_stop_event.set()
|
||||
|
||||
with self._watches_lock:
|
||||
watches = list(self._active_watches.values())
|
||||
for w in watches:
|
||||
try:
|
||||
w.stop()
|
||||
except Exception as e:
|
||||
logger.warning(f"Error stopping watch: {e}")
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def join_all_threads():
|
||||
for t in self._threads:
|
||||
try:
|
||||
t.join(timeout=1.0)
|
||||
except Exception as e:
|
||||
logger.warning(f"Error joining thread {t.name}: {e}")
|
||||
|
||||
await loop.run_in_executor(None, join_all_threads)
|
||||
logger.info("Kubernetes events watcher stopped.")
|
||||
@@ -0,0 +1,191 @@
|
||||
"""Unit tests for KubernetesEventProvider.
|
||||
|
||||
Tests verify that raw Kubernetes events are correctly parsed, translated,
|
||||
and mapped to standardized RayEvent protobuf structures.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.core.generated.events_base_event_pb2 import RayEvent
|
||||
from ray.core.generated.platform_event_pb2 import Source
|
||||
from ray.dashboard.modules.platform_events.providers.k8s_provider import (
|
||||
KubernetesEventProvider,
|
||||
)
|
||||
|
||||
|
||||
def _make_k8s_event(
|
||||
uid: str = "uid-abc-123",
|
||||
kind: str = "RayCluster",
|
||||
name: str = "my-cluster",
|
||||
message: str = "Pod started",
|
||||
reason: str = "Started",
|
||||
event_type: str = "Normal",
|
||||
namespace: str = "default",
|
||||
count: int = 1,
|
||||
last_timestamp: datetime = None,
|
||||
first_timestamp: datetime = None,
|
||||
component: str = "kubelet",
|
||||
) -> MagicMock:
|
||||
evt = MagicMock()
|
||||
evt.metadata.uid = uid
|
||||
evt.metadata.namespace = namespace
|
||||
evt.metadata.resource_version = "12345"
|
||||
evt.involved_object.kind = kind
|
||||
evt.involved_object.name = name
|
||||
evt.message = message
|
||||
evt.reason = reason
|
||||
evt.type = event_type
|
||||
evt.count = count
|
||||
evt.last_timestamp = last_timestamp
|
||||
evt.first_timestamp = first_timestamp
|
||||
evt.source.component = component
|
||||
return evt
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_k8s_client_missing_cluster_name():
|
||||
provider = KubernetesEventProvider(lambda e: None)
|
||||
result = await provider._init_k8s_client()
|
||||
assert result is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_k8s_client_missing_cluster_namespace(monkeypatch):
|
||||
monkeypatch.setenv("RAY_CLUSTER_NAME", "my-cluster")
|
||||
monkeypatch.delenv("RAY_CLUSTER_NAMESPACE", raising=False)
|
||||
|
||||
provider = KubernetesEventProvider(lambda e: None)
|
||||
result = await provider._init_k8s_client()
|
||||
assert result is False
|
||||
|
||||
|
||||
def test_proto_event_type_and_source_type():
|
||||
delivered_event = None
|
||||
|
||||
def callback(event: RayEvent):
|
||||
nonlocal delivered_event
|
||||
delivered_event = event
|
||||
|
||||
provider = KubernetesEventProvider(callback)
|
||||
provider._cluster_name = "my-cluster"
|
||||
|
||||
evt = _make_k8s_event()
|
||||
provider._process_k8s_event(evt)
|
||||
|
||||
assert delivered_event is not None
|
||||
assert delivered_event.event_type == RayEvent.EventType.PLATFORM_EVENT
|
||||
assert delivered_event.source_type == RayEvent.SourceType.CLUSTER_LIFECYCLE
|
||||
|
||||
|
||||
def test_platform_source_is_kubernetes():
|
||||
delivered_event = None
|
||||
|
||||
def callback(event: RayEvent):
|
||||
nonlocal delivered_event
|
||||
delivered_event = event
|
||||
|
||||
provider = KubernetesEventProvider(callback)
|
||||
provider._cluster_name = "prod-cluster"
|
||||
|
||||
evt = _make_k8s_event(component="kubelet", namespace="my-namespace")
|
||||
provider._process_k8s_event(evt)
|
||||
|
||||
assert delivered_event is not None
|
||||
source = delivered_event.platform_event.source
|
||||
assert source.platform == Source.Platform.KUBERNETES
|
||||
assert source.component == "kubelet"
|
||||
assert source.metadata["namespace"] == "my-namespace"
|
||||
assert source.metadata["ray_cluster_name"] == "prod-cluster"
|
||||
|
||||
|
||||
def test_platform_event_object_fields():
|
||||
delivered_event = None
|
||||
|
||||
def callback(event: RayEvent):
|
||||
nonlocal delivered_event
|
||||
delivered_event = event
|
||||
|
||||
provider = KubernetesEventProvider(callback)
|
||||
evt = _make_k8s_event(
|
||||
uid="uid-abc-123",
|
||||
kind="RayJob",
|
||||
name="my-job",
|
||||
message="Job failed",
|
||||
reason="BackoffLimitExceeded",
|
||||
)
|
||||
provider._process_k8s_event(evt)
|
||||
|
||||
assert delivered_event is not None
|
||||
assert delivered_event.event_id == b"uid-abc-123"
|
||||
|
||||
pe = delivered_event.platform_event
|
||||
assert pe.object_kind == "RayJob"
|
||||
assert pe.object_name == "my-job"
|
||||
assert pe.message == "Job failed"
|
||||
assert pe.reason == "BackoffLimitExceeded"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"event_type,expected_severity",
|
||||
[
|
||||
("Warning", RayEvent.Severity.WARNING),
|
||||
("Normal", RayEvent.Severity.INFO),
|
||||
(None, RayEvent.Severity.INFO),
|
||||
],
|
||||
)
|
||||
def test_severity_mapping(event_type, expected_severity):
|
||||
delivered_event = None
|
||||
|
||||
def callback(event: RayEvent):
|
||||
nonlocal delivered_event
|
||||
delivered_event = event
|
||||
|
||||
provider = KubernetesEventProvider(callback)
|
||||
evt = _make_k8s_event(event_type=event_type)
|
||||
provider._process_k8s_event(evt)
|
||||
|
||||
assert delivered_event is not None
|
||||
assert delivered_event.severity == expected_severity
|
||||
|
||||
|
||||
def test_timestamp_uses_last_timestamp():
|
||||
delivered_event = None
|
||||
|
||||
def callback(event: RayEvent):
|
||||
nonlocal delivered_event
|
||||
delivered_event = event
|
||||
|
||||
provider = KubernetesEventProvider(callback)
|
||||
ts = datetime(2025, 6, 1, 12, 0, 0, tzinfo=timezone.utc)
|
||||
evt = _make_k8s_event(
|
||||
last_timestamp=ts, first_timestamp=datetime(2025, 1, 1, tzinfo=timezone.utc)
|
||||
)
|
||||
provider._process_k8s_event(evt)
|
||||
|
||||
assert delivered_event is not None
|
||||
assert delivered_event.timestamp.seconds == int(ts.timestamp())
|
||||
|
||||
|
||||
def test_timestamp_falls_back_to_first_timestamp():
|
||||
delivered_event = None
|
||||
|
||||
def callback(event: MagicMock):
|
||||
nonlocal delivered_event
|
||||
delivered_event = event
|
||||
|
||||
provider = KubernetesEventProvider(callback)
|
||||
ts = datetime(2025, 3, 15, 8, 30, 0, tzinfo=timezone.utc)
|
||||
evt = _make_k8s_event(last_timestamp=None, first_timestamp=ts)
|
||||
provider._process_k8s_event(evt)
|
||||
|
||||
assert delivered_event is not None
|
||||
assert delivered_event.timestamp.seconds == int(ts.timestamp())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,166 @@
|
||||
"""Unit tests for PlatformEventsHead module.
|
||||
|
||||
Tests verify the core caching, deduplication, eviction, and REST API
|
||||
logic inside the central PlatformEventsHead controller.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from ray._private import ray_constants
|
||||
from ray.core.generated.events_base_event_pb2 import RayEvent
|
||||
from ray.core.generated.platform_event_pb2 import PlatformEvent, Source
|
||||
from ray.dashboard.modules.platform_events.platform_event_head import (
|
||||
MAX_EVENTS_TO_CACHE,
|
||||
PlatformEventsHead,
|
||||
)
|
||||
from ray.dashboard.utils import DashboardHeadModuleConfig
|
||||
|
||||
|
||||
def _make_config() -> DashboardHeadModuleConfig:
|
||||
return DashboardHeadModuleConfig(
|
||||
minimal=False,
|
||||
cluster_id_hex="deadbeef",
|
||||
session_name="test_session",
|
||||
gcs_address="127.0.0.1:6379",
|
||||
log_dir="/tmp",
|
||||
temp_dir="/tmp",
|
||||
session_dir="/tmp",
|
||||
ip="127.0.0.1",
|
||||
http_host="127.0.0.1",
|
||||
http_port=8265,
|
||||
)
|
||||
|
||||
|
||||
def _make_head(**kwargs) -> PlatformEventsHead:
|
||||
head = PlatformEventsHead(_make_config())
|
||||
for k, v in kwargs.items():
|
||||
setattr(head, k, v)
|
||||
return head
|
||||
|
||||
|
||||
def test_head_event_caching_and_update():
|
||||
head = _make_head()
|
||||
|
||||
# Deliver a brand new RayEvent
|
||||
ray_event = RayEvent(
|
||||
event_id=b"test_uid",
|
||||
message="some message",
|
||||
)
|
||||
ray_event.platform_event.message = "some message"
|
||||
ray_event.platform_event.custom_fields["count"] = "1"
|
||||
|
||||
head._process_event_callback(ray_event)
|
||||
assert len(head._events) == 1
|
||||
|
||||
cached_event = next(iter(head._events.values()))
|
||||
assert cached_event.message == "some message"
|
||||
assert cached_event.platform_event.custom_fields["count"] == "1"
|
||||
|
||||
# Deliver an updated RayEvent with same event_id
|
||||
updated_ray_event = RayEvent(
|
||||
event_id=b"test_uid",
|
||||
message="updated message",
|
||||
)
|
||||
updated_ray_event.platform_event.message = "updated message"
|
||||
updated_ray_event.platform_event.custom_fields["count"] = "5"
|
||||
|
||||
head._process_event_callback(updated_ray_event)
|
||||
assert len(head._events) == 1
|
||||
|
||||
cached_event_updated = next(iter(head._events.values()))
|
||||
assert cached_event_updated.message == "updated message"
|
||||
assert cached_event_updated.platform_event.custom_fields["count"] == "5"
|
||||
|
||||
|
||||
def test_dedup_suppresses_repeated_uid():
|
||||
head = _make_head()
|
||||
ray_event = RayEvent(event_id=b"uid-repeated")
|
||||
|
||||
head._process_event_callback(ray_event)
|
||||
head._process_event_callback(ray_event)
|
||||
|
||||
assert len(head._events) == 1
|
||||
|
||||
|
||||
def test_dedup_uid_eviction_keeps_recent():
|
||||
"""Verify that cache size does not exceed MAX_EVENTS_TO_CACHE and keeps recent items."""
|
||||
head = _make_head()
|
||||
total = MAX_EVENTS_TO_CACHE + 10
|
||||
for i in range(total):
|
||||
ray_event = RayEvent(event_id=f"uid-{i}".encode())
|
||||
head._process_event_callback(ray_event)
|
||||
|
||||
assert len(head._events) == MAX_EVENTS_TO_CACHE
|
||||
# Oldest UIDs should be gone
|
||||
assert "uid-0" not in head._events
|
||||
# Newest should still be there
|
||||
assert f"uid-{total - 1}" in head._events
|
||||
|
||||
|
||||
def test_head_event_emission_via_event_recorder(monkeypatch):
|
||||
head = _make_head()
|
||||
ray_event = RayEvent(
|
||||
event_id=b"test_uid_emission",
|
||||
message="test emission message",
|
||||
severity=RayEvent.Severity.WARNING,
|
||||
)
|
||||
ray_event.platform_event.message = "test emission message"
|
||||
ray_event.platform_event.source.platform = Source.Platform.KUBERNETES
|
||||
ray_event.platform_event.source.component = "kubelet"
|
||||
ray_event.platform_event.source.metadata["namespace"] = "default"
|
||||
ray_event.platform_event.object_kind = "Pod"
|
||||
ray_event.platform_event.object_name = "test-pod"
|
||||
ray_event.platform_event.reason = "OOMKilled"
|
||||
ray_event.platform_event.custom_fields["count"] = "3"
|
||||
monkeypatch.setattr(
|
||||
ray_constants,
|
||||
"RAY_ENABLE_PYTHON_RAY_EVENT_TYPES",
|
||||
frozenset({"PLATFORM_EVENT"}),
|
||||
)
|
||||
captured = {}
|
||||
|
||||
def fake_ray_event_ctor(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return MagicMock(name="cython_ray_event")
|
||||
|
||||
mock_recorder = MagicMock()
|
||||
with (
|
||||
patch(
|
||||
"ray._common.observability.internal_event.RayEvent",
|
||||
side_effect=fake_ray_event_ctor,
|
||||
),
|
||||
patch.dict(
|
||||
sys.modules, {"ray._raylet": MagicMock(EventRecorder=mock_recorder)}
|
||||
),
|
||||
):
|
||||
head._process_event_callback(ray_event)
|
||||
|
||||
mock_recorder.emit.assert_called_once()
|
||||
assert captured["entity_id"] == "test_uid_emission"
|
||||
assert captured["event_id"] == b"test_uid_emission"
|
||||
assert captured["source_type"] == RayEvent.SourceType.CLUSTER_LIFECYCLE
|
||||
assert captured["event_type"] == RayEvent.EventType.PLATFORM_EVENT
|
||||
assert captured["severity"] == RayEvent.Severity.WARNING
|
||||
assert captured["nested_event_field_number"] == (
|
||||
RayEvent.PLATFORM_EVENT_FIELD_NUMBER
|
||||
)
|
||||
|
||||
nested = PlatformEvent()
|
||||
nested.ParseFromString(captured["serialized_data"])
|
||||
assert nested.source.platform == Source.Platform.KUBERNETES
|
||||
assert nested.source.component == "kubelet"
|
||||
assert nested.source.metadata["namespace"] == "default"
|
||||
assert nested.object_kind == "Pod"
|
||||
assert nested.object_name == "test-pod"
|
||||
assert nested.reason == "OOMKilled"
|
||||
assert nested.message == "test emission message"
|
||||
assert nested.custom_fields["count"] == "3"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Reference in New Issue
Block a user