ray-project--ray/python/ray/includes/event_recorder.pxi

"""Cython bindings for RayEventRecorder.

This module provides Python access to the C++ RayEventRecorder for emitting
internal Ray events from Python code (e.g., submission job events).
"""

from ray.includes.event_recorder cimport (
    CRayEventInterface,
    CPythonEventRecorder,
    CreatePythonRayEvent,
)
from ray.includes.common cimport move
from libc.stdint cimport int64_t
from libcpp.memory cimport unique_ptr
from libcpp.vector cimport vector as c_vector
from libcpp.string cimport string as c_string

import logging
import threading

logger = logging.getLogger(__name__)

cdef class RayEvent:
    """Python wrapper holding event data for transfer to C++.

    This class stores event metadata and serialized protobuf data. When added
    to the EventRecorder, it creates the underlying C++ RayEvent object.

    Args:
        source_type: Integer value of RayEvent.SourceType enum.
        event_type: Integer value of RayEvent.EventType enum.
        severity: Integer value of RayEvent.Severity enum.
        entity_id: Unique identifier for the event entity (e.g., submission_id).
        message: Optional message associated with the event.
        session_name: The Ray session name.
        serialized_data: Serialized protobuf bytes of the nested event message.
        nested_event_field_number: The field number in RayEvent proto for the
            nested event message. Use RayEventProto.<FIELD>_FIELD_NUMBER constants.
        event_id: Optional explicit event id bytes. Pass empty bytes to let the
            C++ layer generate a random id (matching the convention used by the
            other RayEventInterface subclasses). A non-empty value is preserved
            as-is — use this to reuse an id from an upstream source
            (e.g., a Kubernetes event uid).
        timestamp_ns: Optional explicit event timestamp in nanoseconds since the
            unix epoch. Pass 0 to let the C++ layer capture the time at
            construction via absl::Now().
    """
    cdef:
        int _source_type
        int _event_type
        int _severity
        str _entity_id
        str _message
        str _session_name
        bytes _serialized_data
        int _nested_event_field_number
        bytes _event_id
        int64_t _timestamp_ns

    def __init__(
        self,
        int source_type,
        int event_type,
        int severity,
        str entity_id,
        str message,
        str session_name,
        bytes serialized_data,
        int nested_event_field_number,
        bytes event_id = b"",
        int64_t timestamp_ns = 0,
    ):
        self._source_type = source_type
        self._event_type = event_type
        self._severity = severity
        self._entity_id = entity_id
        self._message = message
        self._session_name = session_name
        self._serialized_data = serialized_data
        self._nested_event_field_number = nested_event_field_number
        self._event_id = event_id
        self._timestamp_ns = timestamp_ns

    cdef unique_ptr[CRayEventInterface] to_cpp_event(self):
        """Create the underlying C++ event. Ownership is transferred to caller."""
        return CreatePythonRayEvent(
            self._source_type,
            self._event_type,
            self._severity,
            self._entity_id.encode("utf-8"),
            self._message.encode("utf-8"),
            self._session_name.encode("utf-8"),
            self._serialized_data,
            self._nested_event_field_number,
            self._event_id,
            self._timestamp_ns,
        )

    @property
    def entity_id(self):
        return self._entity_id

    @property
    def event_type(self):
        return self._event_type


# module-level Singleton instance, lazily created by EventRecorder.initialize().
_event_recorder_instance = None
# Guards singleton lifecycle and emission against concurrent shutdown/initialize.
_event_recorder_lock = threading.RLock()


cdef class EventRecorder:
    """Per-process singleton for recording Ray events."""
    cdef unique_ptr[CPythonEventRecorder] _recorder

    def __dealloc__(self):
        """Safety-net cleanup. C++ destructor also calls Shutdown()."""
        if self._recorder.get() != NULL:
            self._recorder.get().Shutdown()
        self._recorder.reset()

    @staticmethod
    def initialize(
        int aggregator_port,
        str node_ip,
        str node_id_hex,
        size_t max_buffer_size,
        str metric_source = "python",
    ):
        """Initialize the per-process event recorder.

        Creates the underlying C++ PythonEventRecorder with a background I/O
        thread and gRPC client. No-op if already initialized.

        Args:
            aggregator_port: Port of the event aggregator server (bound on 127.0.0.1).
            node_ip: IP address of the current node.
            node_id_hex: Hex-encoded node ID.
            max_buffer_size: Maximum number of events to buffer.
            metric_source: Label for the "Source" tag on dropped-events metrics
                (default "python").
        """
        global _event_recorder_instance
        cdef EventRecorder rec
        with _event_recorder_lock:
            if _event_recorder_instance is not None:
                return

            rec = EventRecorder()
            rec._recorder.reset(
                new CPythonEventRecorder(
                    aggregator_port,
                    node_ip.encode("utf-8"),
                    node_id_hex.encode("utf-8"),
                    max_buffer_size,
                    metric_source.encode("utf-8"),
                )
            )
            _event_recorder_instance = rec

    @staticmethod
    def instance():
        """Get the per-process EventRecorder singleton.

        Returns:
            The EventRecorder instance if initialized, None otherwise.
        """
        with _event_recorder_lock:
            return _event_recorder_instance

    @staticmethod
    def shutdown():
        """Shutdown the event recorder.

        Stops exporting events, performs a final flush, and releases all
        C++ resources. After this call, emit() and emit_batch() will be
        no-ops until initialize() is called again.
        """
        global _event_recorder_instance
        cdef EventRecorder rec
        with _event_recorder_lock:
            if _event_recorder_instance is None:
                return

            rec = <EventRecorder>_event_recorder_instance
            if rec._recorder.get() != NULL:
                with nogil:
                    rec._recorder.get().Shutdown()
            rec._recorder.reset()
            _event_recorder_instance = None

    @staticmethod
    def emit(RayEvent event):
        """Emit a single event. No-op if not initialized.

        Args:
            event: A RayEvent object (created via InternalEventBuilder.build()).

        Returns:
            True if the event was successfully queued, False otherwise.
        """
        return EventRecorder.emit_batch([event])

    @staticmethod
    def emit_batch(list events):
        """Emit multiple events. No-op if not initialized.

        Args:
            events: List of RayEvent objects to emit.

        Returns:
            True if events were successfully queued, False otherwise.
        """
        cdef EventRecorder rec
        cdef c_vector[unique_ptr[CRayEventInterface]] cpp_events
        cdef RayEvent ev

        if not events:
            return True

        for ev in events:
            cpp_events.push_back(move(ev.to_cpp_event()))

        with _event_recorder_lock:
            if _event_recorder_instance is None:
                logger.debug(
                    "Event recorder not initialized, dropping %d events",
                    len(events),
                )
                return False

            rec = <EventRecorder>_event_recorder_instance

            with nogil:
                rec._recorder.get().AddEvents(move(cpp_events))

            return True