diff --git a/sdks/python/src/opik/__init__.py b/sdks/python/src/opik/__init__.py index 36974dc3b4e..7e609a82a8a 100644 --- a/sdks/python/src/opik/__init__.py +++ b/sdks/python/src/opik/__init__.py @@ -24,6 +24,12 @@ from .api_objects.trace import Trace from .configurator.configure import configure from .decorator.tracker import flush_tracker, track +from .message_processing.data_loss import ( + ErrorsReport, + FailedMessageInfo, + FailureReason, + FlushResult, +) from .evaluation import ( evaluate, evaluate_experiment, @@ -68,6 +74,10 @@ "ExperimentItemReferences", "track", "flush_tracker", + "FlushResult", + "FailedMessageInfo", + "FailureReason", + "ErrorsReport", "Opik", "get_global_client", "set_global_client", diff --git a/sdks/python/src/opik/api_objects/connection_resources.py b/sdks/python/src/opik/api_objects/connection_resources.py index 63364628520..919ebd925ca 100644 --- a/sdks/python/src/opik/api_objects/connection_resources.py +++ b/sdks/python/src/opik/api_objects/connection_resources.py @@ -30,6 +30,8 @@ from ..file_upload import upload_manager from ..healthcheck import connection_monitor, connection_probe from ..message_processing import ( + data_loss, + flush_reporter, message_queue, permissions, streamer, @@ -61,6 +63,7 @@ def __init__( file_upload_manager: upload_manager.FileUploadManager, replay_manager: replay_manager.ReplayManager, streamer: streamer.Streamer, + data_loss_tracker: data_loss.DataLossTracker, flush_timeout: Optional[int], ) -> None: self.httpx_client = httpx_client @@ -69,14 +72,18 @@ def __init__( self.file_upload_manager = file_upload_manager self.replay_manager = replay_manager self.streamer = streamer + self.flush_reporter = flush_reporter.FlushReporter( + streamer=streamer, + data_loss_tracker=data_loss_tracker, + ) self.flush_timeout = flush_timeout - def close(self, timeout: Optional[int], *, flush: bool) -> None: + def close(self, timeout: Optional[int], *, flush: bool) -> bool: # Drain/stop the streamer (consumer threads, replay, batch preprocessor); # on flush=True it also flushes pending file uploads. # Closing the streamer also stops and joins the replay manager (its own # daemon thread), so there is no separate replay teardown to do here. - self.streamer.close(timeout, flush=flush) + flushed = self.streamer.close(timeout, flush=flush) # Stop the upload worker pool too, so eviction doesn't leave its threads # running. wait=flush mirrors the streamer: block for in-flight uploads # on a durable close, return immediately on fire-and-forget teardown. @@ -90,19 +97,23 @@ def close(self, timeout: Optional[int], *, flush: bool) -> None: # daemon threads to finish in-flight requests, so closing the pool # here would race them — leave it for GC / process-exit close_all. self.httpx_client.close() + return flushed - def flush(self, timeout: Optional[int]) -> None: + def flush(self, timeout: Optional[int]) -> bool: """Drain the shared message queue without tearing the bundle down. Used when a handle releases with ``flush=True`` while other handles still share the bundle: the queued data is persisted now, but the transport stays alive for the remaining handles. + + Returns whether the queue drained fully within ``timeout``. """ - self.streamer.flush(timeout) + return self.streamer.flush(timeout) def _create_replay_manager( - config: opik_config.OpikConfig, httpx_client: httpx.Client + config: opik_config.OpikConfig, + httpx_client: httpx.Client, ) -> replay_manager.ReplayManager: probe = connection_probe.ConnectionProbe( base_url=config.url_override, @@ -156,6 +167,8 @@ def create_connection_resources( worker_count=config.file_upload_background_workers, ) + data_loss_tracker = data_loss.DataLossTracker() + fallback_replay = _create_replay_manager(config, httpx_client_) message_processor = message_processors_chain.create_message_processors_chain( @@ -166,6 +179,7 @@ def create_connection_resources( retry_interval_seconds=config.unauthorized_message_type_retry_interval, max_retry_count=config.unauthorized_message_type_max_retry_count, ), + data_loss_tracker=data_loss_tracker, ) streamer_ = streamer_constructors.construct_online_streamer( file_uploader=file_uploader, @@ -186,6 +200,7 @@ def create_connection_resources( file_upload_manager=file_uploader, replay_manager=fallback_replay, streamer=streamer_, + data_loss_tracker=data_loss_tracker, flush_timeout=config.default_flush_timeout, ) @@ -236,12 +251,16 @@ def __init__( def release( self, timeout: Optional[int], *, flush: bool = True, close_on_zero: bool - ) -> None: + ) -> Optional[bool]: + """Release this handle's reference. Returns the authoritative flush + outcome when this release performed the drain (an explicit + ``flush=True`` release), ``None`` otherwise (already released, or a GC + finalizer that does no network I/O).""" with self._once_lock: if self._released: - return + return None self._released = True - self._manager.release( + return self._manager.release( self._key, timeout, flush=flush, close_on_zero=close_on_zero ) @@ -326,7 +345,7 @@ def release( *, flush: bool = True, close_on_zero: bool, - ) -> None: + ) -> Optional[bool]: # Durability under sharing: an explicit ``end(flush=True)`` on a handle # that still shares its bundle must drain the shared queue *before* this # handle gives up its reference. Flushing while our reference is still @@ -337,6 +356,12 @@ def release( # bundle; a sole holder's ``close(flush=True)`` below already drains # durably. A GC finalizer (``close_on_zero=False``) never does network # I/O, so it never pre-flushes. + # Authoritative flush outcome for the caller: set by whichever branch + # below actually drains — the shared pre-flush or the last-reference + # close. Stays None when this release did no draining (GC finalizer, or + # a bundle already released elsewhere), so the caller can tell "not + # confirmed" apart from "confirmed not flushed". + flushed: Optional[bool] = None if flush and close_on_zero: with self._lock: entry = self._entries.get(key) @@ -346,7 +371,7 @@ def release( else None ) if shared_bundle is not None: - shared_bundle.flush(timeout) + flushed = shared_bundle.flush(timeout) # Now drop our reference. Because we only decrement here — after any # pre-flush above has completed — a close can never run while another @@ -355,10 +380,10 @@ def release( with self._lock: entry = self._entries.get(key) if entry is None: - return + return flushed entry.refcount -= 1 if entry.refcount > 0: - return + return flushed if not close_on_zero: # The last reference was dropped by a GC finalizer (see # ``Opik._acquire_shared_resources``). Only the refcount @@ -367,13 +392,16 @@ def release( # never happen inside garbage collection. Leave the bundle # cached so a later same-identity ``acquire`` reuses it, or # ``close_all`` disposes it at process exit. - return + return flushed # Evict before close, under the lock, so a concurrent acquire never # receives a bundle that is being torn down. del self._entries[key] bundle = entry.resources - bundle.close(timeout, flush=flush) + closed_flushed = bundle.close(timeout, flush=flush) + # A non-draining teardown has no flush outcome to report — return None + # (not close()'s bool) so the result stays "no drain happened here". + return closed_flushed if flush else None def close_all(self, *, flush: bool = True) -> None: """Close and evict every cached bundle. Registered as the process diff --git a/sdks/python/src/opik/api_objects/opik_client.py b/sdks/python/src/opik/api_objects/opik_client.py index 640424230d0..295984c1f04 100644 --- a/sdks/python/src/opik/api_objects/opik_client.py +++ b/sdks/python/src/opik/api_objects/opik_client.py @@ -71,6 +71,7 @@ url_helpers, ) from ..message_processing import ( + data_loss, messages, ) from ..message_processing.batching import sequence_splitter @@ -217,6 +218,8 @@ def _bind_resources(self) -> None: self._rest_client = self._resources.rest_client self.__internal_api__message_processor__ = self._resources.message_processor self._streamer = self._resources.streamer + self._flush_reporter = self._resources.flush_reporter + self._last_flush_result: Optional[data_loss.FlushResult] = None def _display_trace_url(self, trace_id: str, project_name: str) -> None: project_url = url_helpers.get_project_url_by_trace_id( @@ -1927,9 +1930,13 @@ def get_experiment_by_id(self, id: str) -> experiment.Experiment: project_name=experiment_public.project_name, ) - def end(self, timeout: Optional[int] = None, *, flush: bool = True) -> None: + def end( + self, timeout: Optional[int] = None, *, flush: bool = True + ) -> Optional[data_loss.FlushResult]: """ - End the Opik session and submit all pending messages. + End the Opik session, releasing this client's connection reference. When + ``flush`` is True (the default), all pending messages are submitted + first; when ``flush`` is False, anything still queued is dropped. Connection resources are shared and ref-counted across clients with a matching configuration: this releases the current client's reference. @@ -1957,28 +1964,102 @@ def end(self, timeout: Optional[int] = None, *, flush: bool = True) -> None: is shared — it may still succeed by riding another live client's resources. Do not rely on either outcome; create a new client instead. + The outcome is also available afterwards via :attr:`last_flush_result`. + Returns: - None + The flush outcome (including any data-loss detail) when ``flush`` is + True; ``None`` when ``flush`` is False (nothing was flushed). """ timeout = timeout if timeout is not None else self._flush_timeout + marker = self._flush_reporter.marker() # Explicit teardown on a user thread, so close on the last reference # (close_on_zero=True). Releasing is idempotent, so the detached GC - # finalizer cannot double-decrement. - self._lease.release(timeout, flush=flush, close_on_zero=True) + # finalizer cannot double-decrement. release() returns the authoritative + # flush outcome computed inside the drain (streamer.flush) — the same + # source flush() uses — rather than the weaker queue_size()==0 proxy, + # which can read empty on the pop-vs-processed race and while file + # uploads are still in flight. + flushed = self._lease.release(timeout, flush=flush, close_on_zero=True) self._finalizer.detach() + if not flush: + return None + if flushed is None: + # No drain ran on this call — e.g. a repeated end() after the client + # was already released. Keep the outcome from the release that did + # the work rather than overwriting it with a spurious not-flushed + # result, so end() is idempotent. + return self._last_flush_result + self._last_flush_result = self._flush_reporter.build_result( + marker, flushed=flushed + ) + return self._last_flush_result def flush(self, timeout: Optional[int] = None) -> bool: """ Flush the streamer to ensure all messages are sent. + Attachment/file upload *failures* are not counted in the data-loss + detail (``dropped_*`` / ``failures``), but an incomplete upload still + makes the flush report as not fully flushed — so ``flushed`` (and hence + the returned bool) does reflect uploads. Never raises and never blocks + beyond ``timeout``: an observability SDK must not disrupt the app it + instruments. Detailed outcome — including any data that was dropped — is + available via :attr:`last_flush_result`. + Args: timeout (Optional[int]): The timeout for flushing the streamer. Once the timeout is reached, the flush method will return regardless of whether all messages have been sent. Returns: - True if all messages have been sent within specified timeout, False otherwise. + True if all messages were delivered within the timeout with no data + loss; False if the timeout was hit or any message was dropped. """ timeout = timeout if timeout is not None else self._flush_timeout - return self._streamer.flush(timeout) + try: + marker = self._flush_reporter.marker() + flushed = self._streamer.flush(timeout) + self._last_flush_result = self._flush_reporter.build_result( + marker, flushed=flushed + ) + return self._last_flush_result.success + except Exception: + # An observability SDK must not disrupt the app it instruments: a + # failure inside flush is reported as "not flushed", never raised. + # Record a failed outcome so last_flush_result reflects this attempt + # rather than keeping a stale prior success. Built directly (not via + # build_result, which may itself be what raised) so it cannot re-raise. + LOGGER.error("Opik flush failed unexpectedly", exc_info=True) + self._last_flush_result = data_loss.FlushResult( + flushed=False, + remaining_queue_size=0, + dropped_messages=0, + dropped_items=0, + failures=(), + ) + return False + + @property + def last_flush_result(self) -> Optional[data_loss.FlushResult]: + """Outcome of the most recent ``flush()``/``end()`` on this client. + + ``None`` until the first flush. + """ + return self._last_flush_result + + def get_errors_report(self) -> data_loss.ErrorsReport: + """Report of messages the background sender terminally dropped. + + Unlike :attr:`last_flush_result`, which is scoped to a single flush, this + reports the sender's retained data-loss history — including drops that + happened before or between flushes. + + The report is **capped**: the total counts are exact, but the per-drop + ``failures`` list keeps only the most recent entries (bounded, drop-oldest) + so it never grows without bound — see :class:`~opik.ErrorsReport`. + + The sender is shared across clients with a matching configuration, so + the report may include drops from sibling clients on the same connection. + """ + return self._flush_reporter.build_errors_report() def __internal_api__drain_to_processors__( self, timeout: Optional[float] = None diff --git a/sdks/python/src/opik/message_processing/data_loss.py b/sdks/python/src/opik/message_processing/data_loss.py new file mode 100644 index 00000000000..5a1a20029e7 --- /dev/null +++ b/sdks/python/src/opik/message_processing/data_loss.py @@ -0,0 +1,178 @@ +"""Tracking of terminally-dropped (never-delivered) messages. + +The background sender never raises into user code and never blocks the app's +critical path. That leaves a gap: when a batch of traces/spans is dropped after +all retries and recovery are exhausted, the caller has no in-band way to learn +about it. :class:`DataLossTracker` records those terminal drops so they can be +surfaced through :class:`FlushResult` when the caller flushes or ends the client. + +Only *terminal* drops are recorded here — messages that will never be sent +again. Transient states that are still expected to be delivered (rate-limit +re-enqueues, connection errors parked for replay) are deliberately excluded. +""" + +import collections +import dataclasses +import enum +import time +from typing import Deque, List, Optional, Tuple + +# Opaque token returned by ``DataLossTracker.marker`` and passed back to +# ``drops_since``: the running (message, item) totals at a point in time. +DropMarker = Tuple[int, int] + + +class FailureReason(str, enum.Enum): + HTTP_CLIENT_ERROR = "http_client_error" + HTTP_SERVER_ERROR = "http_server_error" + UNAUTHORIZED = "unauthorized" + SERIALIZATION = "serialization" + UNKNOWN = "unknown" + + @staticmethod + def from_status_code(status_code: Optional[int]) -> "FailureReason": + if status_code is not None and 400 <= status_code < 500: + return FailureReason.HTTP_CLIENT_ERROR + if status_code is not None and 500 <= status_code < 600: + return FailureReason.HTTP_SERVER_ERROR + return FailureReason.UNKNOWN + + +@dataclasses.dataclass(frozen=True) +class FailedMessageInfo: + """A single terminal drop: one message the SDK gave up on delivering.""" + + message_type: str + reason: FailureReason + item_count: int + status_code: Optional[int] = None + detail: Optional[str] = None + timestamp: float = dataclasses.field(default_factory=time.time) + + +@dataclasses.dataclass(frozen=True) +class FlushResult: + """Outcome of a ``flush()``/``end()`` call. + + Attributes: + flushed: Whether the queue drained within the timeout. + remaining_queue_size: Messages still queued when the call returned. + dropped_messages: Terminally-dropped messages observed during this flush. + dropped_items: Traces/spans lost across those dropped messages. + failures: Details of the drops observed during this flush (best-effort; + bounded by the tracker's capacity). + """ + + flushed: bool + remaining_queue_size: int + dropped_messages: int + dropped_items: int + failures: Tuple[FailedMessageInfo, ...] + + @property + def success(self) -> bool: + """True only if this flush drained the queue with no data loss.""" + return self.flushed and self.dropped_messages == 0 + + +@dataclasses.dataclass(frozen=True) +class ErrorsReport: + """Snapshot of terminal data loss recorded by the background sender. + + Sender-wide and not tied to a single flush. + + .. note:: + The report is **capped**. ``total_dropped_messages`` / + ``total_dropped_items`` are always exact (kept as running totals), but + ``failures`` holds only a bounded, most-recent window of the per-drop + details. Once that limit is reached the oldest details are discarded, so + ``failures`` can contain fewer entries than ``total_dropped_messages``, + and ``first_failure_at`` reflects the oldest *retained* detail, not + necessarily the first-ever drop. + + Attributes: + total_dropped_messages: Messages terminally dropped since the sender + started (exact). + total_dropped_items: Traces/spans lost across those messages (exact). + failures: Most-recent per-drop details, capped (see note); each carries + its own ``timestamp``. + generated_at: Unix time when this report was produced. + """ + + total_dropped_messages: int + total_dropped_items: int + failures: Tuple[FailedMessageInfo, ...] + generated_at: float + + @property + def first_failure_at(self) -> Optional[float]: + """Timestamp of the oldest retained failure (bounded window), or None.""" + return min((failure.timestamp for failure in self.failures), default=None) + + @property + def last_failure_at(self) -> Optional[float]: + """Timestamp of the newest retained failure, or None.""" + return max((failure.timestamp for failure in self.failures), default=None) + + +class DataLossTracker: + """Bounded record of terminally-dropped messages. + + Shared across all :class:`opik.Opik` handles on one connection identity (the + background sender is shared). Per-flush attribution is done with an opaque + monotonic marker: :meth:`marker` before the flush, :meth:`drops_since` + after. + + Lock-free: writes come from the sender's background threads and rely on + ``deque`` being thread-safe. The running counters are plain integers, so + under concurrent drops a count may momentarily lag or round off — that + imprecision is deliberately accepted (a data-loss tally need not be exact to + the message, and a lock would add contention on the hot sending path). + """ + + def __init__(self, max_entries: int = 1000): + self._entries: Deque[FailedMessageInfo] = collections.deque(maxlen=max_entries) + # Running totals kept independently of the bounded ``_entries`` window, + # so counts survive eviction of the oldest details. + self._recorded_count = 0 + self._recorded_items = 0 + + def record(self, failure: FailedMessageInfo) -> None: + self._entries.append(failure) + self._recorded_count += 1 + self._recorded_items += failure.item_count + + def marker(self) -> DropMarker: + """Opaque token marking the current point in the drop history. + + Carries the running (message, item) totals; pass it to + :meth:`drops_since` to get the deltas observed since. + """ + return self._recorded_count, self._recorded_items + + def drops_since( + self, marker: DropMarker + ) -> Tuple[int, int, List[FailedMessageInfo]]: + """Drops recorded since ``marker``. + + Returns ``(message_count, item_count, failures)`` — the counts from the + running totals, and the retained per-drop details (bounded to the most + recent ``max_entries``, oldest evicted once capacity is exceeded). + """ + marker_count, marker_items = marker + count = self._recorded_count - marker_count + items = self._recorded_items - marker_items + window = list(self._entries) + window_size = min(count, len(window)) + failures = window[-window_size:] if window_size > 0 else [] + return count, items, failures + + def total_drops(self) -> Tuple[int, int, List[FailedMessageInfo]]: + """All-time drop totals plus retained details. + + Returns ``(message_count, item_count, failures)``, independent of any + flush boundary — answers "has anything been lost?" across the sender's + lifetime. ``failures`` is bounded to the most recent ``max_entries``, + older details evicted. + """ + return self._recorded_count, self._recorded_items, list(self._entries) diff --git a/sdks/python/src/opik/message_processing/flush_reporter.py b/sdks/python/src/opik/message_processing/flush_reporter.py new file mode 100644 index 00000000000..d655576c591 --- /dev/null +++ b/sdks/python/src/opik/message_processing/flush_reporter.py @@ -0,0 +1,70 @@ +"""Assembles :class:`~opik.message_processing.data_loss.FlushResult` values. + +The read side of data-loss reporting. It holds the two collaborators needed to +describe a flush — the queue (via the streamer) and the :class:`DataLossTracker` +— so that no caller has to gather them itself. One instance is owned per +connection bundle and shared by every client on it. +""" + +import logging +import time +from typing import TYPE_CHECKING + +from . import data_loss + +if TYPE_CHECKING: + from . import streamer as streamer_module + + +LOGGER = logging.getLogger(__name__) + + +class FlushReporter: + def __init__( + self, + streamer: "streamer_module.Streamer", + data_loss_tracker: data_loss.DataLossTracker, + ) -> None: + self._streamer = streamer + self._data_loss_tracker = data_loss_tracker + + def marker(self) -> "data_loss.DropMarker": + """Opaque token identifying the current point in the drop history. + + Take one before a flush; pass it to :meth:`build_result` afterwards to + attribute only the drops observed in between to that flush. + """ + return self._data_loss_tracker.marker() + + def build_result( + self, marker: "data_loss.DropMarker", *, flushed: bool + ) -> "data_loss.FlushResult": + dropped_messages, dropped_items, failures = self._data_loss_tracker.drops_since( + marker + ) + result = data_loss.FlushResult( + flushed=flushed, + remaining_queue_size=self._streamer.queue_size(), + dropped_messages=dropped_messages, + dropped_items=dropped_items, + failures=tuple(failures), + ) + if not result.success: + LOGGER.error( + "Opik flush completed with data loss: %d message(s) / %d item(s) " + "dropped, %d still queued. Inspect Opik.last_flush_result for details.", + result.dropped_messages, + result.dropped_items, + result.remaining_queue_size, + ) + return result + + def build_errors_report(self) -> "data_loss.ErrorsReport": + """Sender-wide data-loss snapshot, independent of any single flush.""" + total_messages, total_items, failures = self._data_loss_tracker.total_drops() + return data_loss.ErrorsReport( + total_dropped_messages=total_messages, + total_dropped_items=total_items, + failures=tuple(failures), + generated_at=time.time(), + ) diff --git a/sdks/python/src/opik/message_processing/messages.py b/sdks/python/src/opik/message_processing/messages.py index 4e8034cd391..0e032c7936c 100644 --- a/sdks/python/src/opik/message_processing/messages.py +++ b/sdks/python/src/opik/message_processing/messages.py @@ -71,6 +71,16 @@ def as_payload_dict(self) -> Dict[str, Any]: def as_db_message_dict(self) -> Dict[str, Any]: return {**self.__dict__} + @property + def item_count(self) -> int: + """Number of data items (traces/spans/...) this message carries. + + Batch messages hold many; a plain message counts as one. Used to report + how much data is lost when a message is dropped. + """ + batch = getattr(self, "batch", None) + return len(batch) if batch is not None else 1 + @dataclasses.dataclass class CreateTraceMessage(BaseMessage): diff --git a/sdks/python/src/opik/message_processing/processors/message_processors_chain.py b/sdks/python/src/opik/message_processing/processors/message_processors_chain.py index dd41b23a773..92c876325d6 100644 --- a/sdks/python/src/opik/message_processing/processors/message_processors_chain.py +++ b/sdks/python/src/opik/message_processing/processors/message_processors_chain.py @@ -8,7 +8,7 @@ message_processors, online_message_processor, ) -from .. import permissions +from .. import data_loss, permissions from ..emulation import local_emulator_message_processor from ..replay import replay_manager @@ -21,6 +21,7 @@ def create_message_processors_chain( file_upload_manager: base_upload_manager.BaseFileUploadManager, fallback_replay_manager: replay_manager.ReplayManager, unauthorized_message_types_registry: permissions.UnauthorizedMessageTypeRegistry, + data_loss_tracker: data_loss.DataLossTracker, ) -> message_processors.ChainedMessageProcessor: """ Creates a chain of message processors by combining an online processor and a @@ -50,6 +51,7 @@ def create_message_processors_chain( file_upload_manager=file_upload_manager, fallback_replay_manager=fallback_replay_manager, unauthorized_message_types_registry=unauthorized_message_types_registry, + data_loss_tracker=data_loss_tracker, ) # is not active by default - will be activated during evaluation local = local_emulator_message_processor.LocalEmulatorMessageProcessor(active=False) diff --git a/sdks/python/src/opik/message_processing/processors/online_message_processor.py b/sdks/python/src/opik/message_processing/processors/online_message_processor.py index f43f3aae8cf..109a7c2a96d 100644 --- a/sdks/python/src/opik/message_processing/processors/online_message_processor.py +++ b/sdks/python/src/opik/message_processing/processors/online_message_processor.py @@ -1,5 +1,5 @@ import logging -from typing import Callable, Dict, Type, Any +from typing import Callable, Dict, Optional, Type, Any import httpx import pydantic @@ -18,7 +18,7 @@ ) from . import assertion_results_processor, message_processors -from .. import encoder_helpers, messages, permissions +from .. import data_loss, encoder_helpers, messages, permissions from ..replay import replay_manager, db_manager @@ -35,6 +35,7 @@ def __init__( file_upload_manager: base_upload_manager.BaseFileUploadManager, fallback_replay_manager: replay_manager.ReplayManager, unauthorized_message_types_registry: permissions.UnauthorizedMessageTypeRegistry, + data_loss_tracker: data_loss.DataLossTracker, batch_memory_limit_mb: int = 50, active: bool = True, ): @@ -44,6 +45,7 @@ def __init__( self._is_active = active self._replay_manager = fallback_replay_manager self._unauthorized_message_types_registry = unauthorized_message_types_registry + self._data_loss_tracker = data_loss_tracker self._assertion_results_processor = ( assertion_results_processor.AssertionResultsMessageProcessor( @@ -93,6 +95,9 @@ def process(self, message: messages.BaseMessage) -> None: "Unauthorized message type: '%s' - ignored from processing.", message.message_type, ) + self._record_data_loss( + message, data_loss.FailureReason.UNAUTHORIZED, status_code=401 + ) return message_type = type(message) @@ -151,6 +156,12 @@ def process(self, message: messages.BaseMessage) -> None: ) # register a message type as unauthorized to avoid re-sending it to the backend self._unauthorized_message_types_registry.add(message.message_type) + self._record_data_loss( + message, + data_loss.FailureReason.UNAUTHORIZED, + status_code=401, + detail=str(exception.body), + ) else: error_tracking_extra = _generate_error_tracking_extra( exception, message @@ -161,6 +172,12 @@ def process(self, message: messages.BaseMessage) -> None: str(exception), extra={"error_tracking_extra": error_tracking_extra}, ) + self._record_data_loss( + message, + data_loss.FailureReason.from_status_code(exception.status_code), + status_code=exception.status_code, + detail=str(exception), + ) except tenacity.RetryError as retry_error: cause = retry_error.last_attempt.exception() error_tracking_extra = _generate_error_tracking_extra(cause, message) @@ -171,6 +188,14 @@ def process(self, message: messages.BaseMessage) -> None: extra={"error_tracking_extra": error_tracking_extra}, ) LOGGER.warning(logging_messages.MAKE_SURE_OPIK_IS_CONFIGURED_CORRECTLY) + self._record_data_loss( + message, + data_loss.FailureReason.from_status_code( + error_tracking_extra.get("status_code") + ), + status_code=error_tracking_extra.get("status_code"), + detail=f"{cause.__class__.__name__} - {cause}", + ) except pydantic.ValidationError as validation_error: error_tracking_extra = _generate_error_tracking_extra( validation_error, message @@ -182,6 +207,11 @@ def process(self, message: messages.BaseMessage) -> None: exc_info=True, extra={"error_tracking_extra": error_tracking_extra}, ) + self._record_data_loss( + message, + data_loss.FailureReason.SERIALIZATION, + detail=str(validation_error), + ) except (httpx.ConnectError, httpx.TimeoutException) as ex: should_unregister_message = False LOGGER.warning( @@ -203,11 +233,31 @@ def process(self, message: messages.BaseMessage) -> None: extra={"error_tracking_extra": error_tracking_extra}, ) LOGGER.warning(logging_messages.MAKE_SURE_OPIK_IS_CONFIGURED_CORRECTLY) + self._record_data_loss( + message, data_loss.FailureReason.UNKNOWN, detail=str(exception) + ) # unregister a message from the reply manager because it is delivered or other error occurred if should_unregister_message: self._replay_manager.unregister_message(message.message_id) # type: ignore + def _record_data_loss( + self, + message: messages.BaseMessage, + reason: data_loss.FailureReason, + status_code: Optional[int] = None, + detail: Optional[str] = None, + ) -> None: + self._data_loss_tracker.record( + data_loss.FailedMessageInfo( + message_type=type(message).__name__, + reason=reason, + item_count=message.item_count, + status_code=status_code, + detail=detail, + ) + ) + def _process_create_span_message( self, message: messages.CreateSpanMessage, diff --git a/sdks/python/src/opik/message_processing/streamer.py b/sdks/python/src/opik/message_processing/streamer.py index e7545d7719f..cbea1ba34bb 100644 --- a/sdks/python/src/opik/message_processing/streamer.py +++ b/sdks/python/src/opik/message_processing/streamer.py @@ -94,6 +94,11 @@ def close(self, timeout: Optional[int] = None, *, flush: bool = True) -> bool: teardowns where pending data can be dropped (e.g. per-test cleanup in e2e tests where assertions already polled the backend during the test body). + + Returns: + Whether all data was flushed to the backend. Authoritative on a + ``flush=True`` close (the result of the internal ``flush(timeout)``); + ``False`` on ``flush=False`` (pending data is deliberately dropped). """ with self._lock: if self._drain: @@ -116,8 +121,9 @@ def close(self, timeout: Optional[int] = None, *, flush: bool = True) -> bool: # actually drain before releasing the caller. Consumers must keep # running while the queue drains, so close them at the very end. self._fallback_replay_manager.join(timeout) - self.flush(timeout) + flushed = self.flush(timeout) self._close_queue_consumers() + return flushed else: # Fire-and-forget: drop pending messages so the stop-signalled # consumers see an empty queue and exit on their own. No joins — @@ -135,8 +141,7 @@ def close(self, timeout: Optional[int] = None, *, flush: bool = True) -> bool: ) self._message_queue.clear() self._close_queue_consumers() - - return self._message_queue.empty() + return False def drain_to_processors(self, timeout: Optional[float] = None) -> bool: """Lightweight drain: ensure every message submitted so far has diff --git a/sdks/python/tests/unit/api_objects/test_connection_resource_manager.py b/sdks/python/tests/unit/api_objects/test_connection_resource_manager.py index 60a3b81abee..4286dea32b6 100644 --- a/sdks/python/tests/unit/api_objects/test_connection_resource_manager.py +++ b/sdks/python/tests/unit/api_objects/test_connection_resource_manager.py @@ -5,6 +5,7 @@ from opik import config as opik_config from opik.api_objects import connection_resources +from opik.message_processing import data_loss class FakeBundle: @@ -351,6 +352,7 @@ def _bundle_with_mock_transport(flush_timeout=None): replay_manager=mock.Mock(), streamer=streamer, flush_timeout=flush_timeout, + data_loss_tracker=data_loss.DataLossTracker(), ) return bundle, streamer, file_upload_manager, httpx_client diff --git a/sdks/python/tests/unit/api_objects/test_opik_client.py b/sdks/python/tests/unit/api_objects/test_opik_client.py index 7d79b511bb8..b3420064a0b 100644 --- a/sdks/python/tests/unit/api_objects/test_opik_client.py +++ b/sdks/python/tests/unit/api_objects/test_opik_client.py @@ -9,7 +9,7 @@ from opik.api_objects.dataset.test_suite import TestSuite from opik.api_objects import prompt as prompt_module from opik.api_objects.prompt import client as prompt_client_module -from opik.message_processing import messages +from opik.message_processing import data_loss, messages from opik.types import BatchFeedbackScoreDict from opik import context_storage from opik.api_objects.trace import trace_data as trace_data_mod @@ -1292,6 +1292,114 @@ def test_search_prompts__returns_text_and_chat_prompts_with_resolved_project_nam assert results[1].project_name == "default-project" +def _flush_result(*, flushed: bool, dropped_messages: int = 0) -> data_loss.FlushResult: + return data_loss.FlushResult( + flushed=flushed, + remaining_queue_size=0, + dropped_messages=dropped_messages, + dropped_items=0, + failures=(), + ) + + +class TestOpikClientFlushResult: + """flush()/end()/last_flush_result behaviors.""" + + @pytest.fixture + def client(self): + # Bypass __init__ so the test doesn't acquire real connection resources + # or spin up background threads just to mock them out; set only what + # flush()/end()/last_flush_result touch. + client_ = opik_client.Opik.__new__(opik_client.Opik) + client_._streamer = MagicMock() + client_._flush_reporter = MagicMock() + client_._lease = MagicMock() + client_._finalizer = MagicMock() + client_._flush_timeout = 5 + client_._last_flush_result = None + return client_ + + def test_flush__no_data_loss__returns_true_and_stores_result(self, client): + result = _flush_result(flushed=True) + client._flush_reporter.build_result.return_value = result + + assert client.flush() is True + assert client.last_flush_result is result + + def test_flush__dropped_messages__returns_false(self, client): + client._flush_reporter.build_result.return_value = _flush_result( + flushed=True, dropped_messages=2 + ) + + assert client.flush() is False + + def test_flush__build_result_raises__never_raises_returns_false(self, client): + client._flush_reporter.build_result.side_effect = RuntimeError("boom") + + assert client.flush() is False + + def test_flush__failure_after_prior_success__result_not_stale(self, client): + client._flush_reporter.build_result.return_value = _flush_result(flushed=True) + assert client.flush() is True + assert client.last_flush_result.success is True + + # A later flush that fails inside build_result must not leave the stale + # success in last_flush_result. + client._flush_reporter.build_result.side_effect = RuntimeError("boom") + assert client.flush() is False + assert client.last_flush_result.flushed is False + assert client.last_flush_result.success is False + + def test_end__flush_true__reports_authoritative_flushed_not_queue_size( + self, client + ): + # release() returns the authoritative drain outcome. end() must use it, + # not the weaker queue_size()==0 proxy — so even with a non-empty queue, + # a True from release yields flushed=True. + client._lease.release.return_value = True + client._streamer.queue_size.return_value = 5 + client._flush_reporter.marker.return_value = 11 + result = _flush_result(flushed=True) + client._flush_reporter.build_result.return_value = result + + returned = client.end() + + client._flush_reporter.build_result.assert_called_once_with(11, flushed=True) + assert returned is result + assert client.last_flush_result is result + + def test_end__release_reports_not_flushed__flushed_false(self, client): + client._lease.release.return_value = False + client._streamer.queue_size.return_value = 0 # queue looks empty... + client._flush_reporter.build_result.return_value = _flush_result(flushed=False) + + client.end() + + # ...but release said not flushed, so flushed=False is reported. + _, kwargs = client._flush_reporter.build_result.call_args + assert kwargs["flushed"] is False + + def test_end__flush_false__returns_none_and_skips_build_result(self, client): + assert client.end(flush=False) is None + client._flush_reporter.build_result.assert_not_called() + + def test_end__called_again_after_release__keeps_prior_result(self, client): + # First end() drains and stores the real outcome. + client._lease.release.return_value = True + first = _flush_result(flushed=True) + client._flush_reporter.build_result.return_value = first + assert client.end() is first + + # Second end(): the client is already released, so release() returns + # None (no drain ran). end() must stay idempotent — keep the prior + # result rather than overwriting it with a spurious not-flushed one. + client._lease.release.return_value = None + client._flush_reporter.build_result.reset_mock() + assert client.end() is first + client._flush_reporter.build_result.assert_not_called() + assert client.last_flush_result is first + + @pytest.mark.parametrize( "exclude,expected", [ diff --git a/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_data_loss.py b/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_data_loss.py new file mode 100644 index 00000000000..96d734c6bbe --- /dev/null +++ b/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_data_loss.py @@ -0,0 +1,200 @@ +"""Unit tests for terminal-drop (data-loss) recording in OpikMessageProcessor. + +Verifies that messages abandoned after a non-recoverable error are recorded in +the DataLossTracker, while transient/retried states (429 with usable headers) +and successful sends are not. +""" + +from unittest import mock + +import pydantic +import pytest +import tenacity + +from opik import exceptions +from opik.message_processing import data_loss, messages, permissions +from opik.message_processing.processors import online_message_processor +from opik.message_processing.replay import replay_manager +from opik.rest_api import core as rest_api_core + + +def _spans_batch_message(item_count: int = 2, message_id: int = 1): + msg = messages.CreateSpansBatchMessage(batch=[]) + msg.batch = [f"span-{index}" for index in range(item_count)] + msg.message_id = message_id + return msg + + +def _recorded( + tracker: data_loss.DataLossTracker, +) -> list: + """All drops the tracker has recorded, read via its public API.""" + # (0, 0) is the zero marker — the running totals before anything was recorded. + _count, _items, failures = tracker.drops_since((0, 0)) + return failures + + +@pytest.fixture +def tracker() -> data_loss.DataLossTracker: + return data_loss.DataLossTracker() + + +@pytest.fixture +def rest_client() -> mock.MagicMock: + return mock.MagicMock() + + +@pytest.fixture +def processor( + rest_client: mock.MagicMock, tracker: data_loss.DataLossTracker +) -> online_message_processor.OpikMessageProcessor: + registry = mock.MagicMock(spec=permissions.UnauthorizedMessageTypeRegistry) + registry.is_authorized.return_value = True + return online_message_processor.OpikMessageProcessor( + rest_client=rest_client, + file_upload_manager=mock.MagicMock(), + fallback_replay_manager=mock.MagicMock(spec=replay_manager.ReplayManager), + unauthorized_message_types_registry=registry, + data_loss_tracker=tracker, + ) + + +def test_process__batch_403__recorded_as_client_error_data_loss( + processor, rest_client, tracker +): + rest_client.spans.create_spans.side_effect = rest_api_core.ApiError( + status_code=403, body="Forbidden" + ) + + processor.process(_spans_batch_message(item_count=2)) + + failures = _recorded(tracker) + assert len(failures) == 1 + assert failures[0].reason == data_loss.FailureReason.HTTP_CLIENT_ERROR + assert failures[0].status_code == 403 + assert failures[0].item_count == 2 + assert failures[0].message_type == "CreateSpansBatchMessage" + + +def test_process__batch_500__recorded_as_server_error_data_loss( + processor, rest_client, tracker +): + rest_client.spans.create_spans.side_effect = rest_api_core.ApiError( + status_code=500, body="oops" + ) + + processor.process(_spans_batch_message()) + + assert _recorded(tracker)[0].reason == data_loss.FailureReason.HTTP_SERVER_ERROR + + +def test_process__unexpected_exception__recorded_as_unknown_data_loss( + processor, rest_client, tracker +): + rest_client.spans.create_spans.side_effect = ValueError("boom") + + processor.process(_spans_batch_message()) + + assert _recorded(tracker)[0].reason == data_loss.FailureReason.UNKNOWN + + +def test_process__429_with_usable_headers__retried_not_recorded( + processor, rest_client, tracker +): + rest_client.spans.create_spans.side_effect = rest_api_core.ApiError( + status_code=429, headers={"retry-after": "1"} + ) + + with mock.patch.object( + online_message_processor.rate_limit, + "parse_rate_limit", + return_value=mock.Mock(retry_after=mock.Mock(return_value=1.0)), + ): + with pytest.raises(exceptions.OpikCloudRequestsRateLimited): + processor.process(_spans_batch_message()) + + assert _recorded(tracker) == [] + + +def test_process__successful_send__nothing_recorded(processor, tracker): + processor.process(_spans_batch_message()) + + assert _recorded(tracker) == [] + + +def test_process__unauthorized_type_skipped__recorded_as_unauthorized( + rest_client, tracker +): + registry = mock.MagicMock(spec=permissions.UnauthorizedMessageTypeRegistry) + registry.is_authorized.return_value = False + processor = online_message_processor.OpikMessageProcessor( + rest_client=rest_client, + file_upload_manager=mock.MagicMock(), + fallback_replay_manager=mock.MagicMock(spec=replay_manager.ReplayManager), + unauthorized_message_types_registry=registry, + data_loss_tracker=tracker, + ) + + processor.process(_spans_batch_message()) + + assert _recorded(tracker)[0].reason == data_loss.FailureReason.UNAUTHORIZED + + +def test_process__batch_401__recorded_as_unauthorized_and_type_registered( + rest_client, tracker +): + registry = mock.MagicMock(spec=permissions.UnauthorizedMessageTypeRegistry) + registry.is_authorized.return_value = True + processor = online_message_processor.OpikMessageProcessor( + rest_client=rest_client, + file_upload_manager=mock.MagicMock(), + fallback_replay_manager=mock.MagicMock(spec=replay_manager.ReplayManager), + unauthorized_message_types_registry=registry, + data_loss_tracker=tracker, + ) + rest_client.spans.create_spans.side_effect = rest_api_core.ApiError( + status_code=401, body="no access" + ) + + processor.process(_spans_batch_message(item_count=2)) + + failure = _recorded(tracker)[0] + assert failure.reason == data_loss.FailureReason.UNAUTHORIZED + assert failure.status_code == 401 + assert failure.item_count == 2 + # 401 also registers the type so it is not re-sent. + registry.add.assert_called_once_with("CreateSpansBatchMessage") + + +def test_process__validation_error__recorded_as_serialization( + processor, rest_client, tracker +): + class _Model(pydantic.BaseModel): + value: int + + try: + _Model(value="not-an-int") + except pydantic.ValidationError as validation_error: + rest_client.spans.create_spans.side_effect = validation_error + + processor.process(_spans_batch_message(item_count=3)) + + failure = _recorded(tracker)[0] + assert failure.reason == data_loss.FailureReason.SERIALIZATION + assert failure.item_count == 3 + + +def test_process__retry_error__recorded_with_cause_status_code( + processor, rest_client, tracker +): + last_attempt = mock.Mock() + last_attempt.exception.return_value = rest_api_core.ApiError( + status_code=500, body="upstream down" + ) + rest_client.spans.create_spans.side_effect = tenacity.RetryError(last_attempt) + + processor.process(_spans_batch_message()) + + failure = _recorded(tracker)[0] + assert failure.reason == data_loss.FailureReason.HTTP_SERVER_ERROR + assert failure.status_code == 500 diff --git a/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_replay.py b/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_replay.py index 7da676caa38..4561c6a41f5 100644 --- a/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_replay.py +++ b/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_replay.py @@ -15,6 +15,7 @@ import pytest import tenacity +from opik.message_processing import data_loss from opik.message_processing import messages from opik.message_processing.processors import online_message_processor from opik.message_processing.processors.online_message_processor import ( @@ -114,6 +115,7 @@ def processor( file_upload_manager=mock_file_uploader, fallback_replay_manager=mock_replay, unauthorized_message_types_registry=mock.Mock(), + data_loss_tracker=data_loss.DataLossTracker(), ) @@ -129,8 +131,8 @@ def test_process__known_message__registers_before_handler( mock_replay.register_message.side_effect = lambda m: call_order.append( "register" ) - mock_rest_client.traces.create_trace.side_effect = ( - lambda **kw: call_order.append("handler") + mock_rest_client.traces.create_trace.side_effect = lambda **kw: ( + call_order.append("handler") ) msg = _create_trace_message() @@ -178,6 +180,7 @@ def test_process__inactive_processor__does_not_register( fallback_replay_manager=mock_replay, active=False, unauthorized_message_types_registry=mock.Mock(), + data_loss_tracker=data_loss.DataLossTracker(), ) msg = _create_trace_message() @@ -208,6 +211,7 @@ def offline_processor( file_upload_manager=mock_file_uploader, fallback_replay_manager=offline_replay, unauthorized_message_types_registry=mock.Mock(), + data_loss_tracker=data_loss.DataLossTracker(), ) def test_process__no_connection__registers_message_as_failed( @@ -284,8 +288,8 @@ def test_process__successful_handler__unregisters_after_handler( ): """unregister_message must be called after the handler completes.""" call_order = [] - mock_rest_client.traces.create_trace.side_effect = ( - lambda **kw: call_order.append("handler") + mock_rest_client.traces.create_trace.side_effect = lambda **kw: ( + call_order.append("handler") ) mock_replay.unregister_message.side_effect = lambda mid: call_order.append( "unregister" @@ -816,6 +820,7 @@ def test_process__attachment_no_connection__registers_as_failed_and_skips_upload file_upload_manager=mock_file_uploader, fallback_replay_manager=offline_replay, unauthorized_message_types_registry=mock.Mock(), + data_loss_tracker=data_loss.DataLossTracker(), ) msg = _create_attachment_message(message_id=80) @@ -883,6 +888,7 @@ def test_process__ignored_type__no_replay_called_on_offline_connection__handler_ file_upload_manager=mock_file_uploader, fallback_replay_manager=offline_replay, unauthorized_message_types_registry=mock.Mock(), + data_loss_tracker=data_loss.DataLossTracker(), ) noop_handler = mock.MagicMock( diff --git a/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_unauthorized_message_type.py b/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_unauthorized_message_type.py index 8b8ef6eeb51..3f983094408 100644 --- a/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_unauthorized_message_type.py +++ b/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_unauthorized_message_type.py @@ -16,6 +16,7 @@ import pytest +from opik.message_processing import data_loss from opik.message_processing import messages from opik.message_processing import permissions from opik.message_processing.processors import online_message_processor @@ -107,6 +108,7 @@ def processor( file_upload_manager=mock_file_uploader, fallback_replay_manager=mock_replay, unauthorized_message_types_registry=mock_registry, + data_loss_tracker=data_loss.DataLossTracker(), ) @@ -142,6 +144,7 @@ def unauthorized_processor( file_upload_manager=mock_file_uploader, fallback_replay_manager=mock_replay, unauthorized_message_types_registry=registry, + data_loss_tracker=data_loss.DataLossTracker(), ) def test_process__unauthorized_message_type__handler_not_called( @@ -203,6 +206,7 @@ def test_process__unauthorized_message_type__is_authorized_called_with_message_t file_upload_manager=mock_file_uploader, fallback_replay_manager=mock_replay, unauthorized_message_types_registry=registry, + data_loss_tracker=data_loss.DataLossTracker(), ) msg = _create_trace_message() @@ -318,6 +322,7 @@ def test_process__after_401_error__same_type_blocked( file_upload_manager=mock_file_uploader, fallback_replay_manager=mock_replay, unauthorized_message_types_registry=real_registry, + data_loss_tracker=data_loss.DataLossTracker(), ) # The first call gets a 401 → type is registered as unauthorized @@ -358,6 +363,7 @@ def test_process__after_401_error__different_type_not_blocked( file_upload_manager=mock_file_uploader, fallback_replay_manager=mock_replay, unauthorized_message_types_registry=real_registry, + data_loss_tracker=data_loss.DataLossTracker(), ) # Cause CreateTraceMessage to be blocked via a 401 @@ -390,6 +396,7 @@ def test_process__after_retry_interval__type_authorized_again( file_upload_manager=mock_file_uploader, fallback_replay_manager=mock_replay, unauthorized_message_types_registry=real_registry, + data_loss_tracker=data_loss.DataLossTracker(), ) # Manually add the type with a timestamp far in the past so the diff --git a/sdks/python/tests/unit/message_processing/test_data_loss.py b/sdks/python/tests/unit/message_processing/test_data_loss.py new file mode 100644 index 00000000000..17739bdcaee --- /dev/null +++ b/sdks/python/tests/unit/message_processing/test_data_loss.py @@ -0,0 +1,215 @@ +from unittest import mock + +from opik.message_processing import data_loss, flush_reporter, messages + + +def _failure( + reason: data_loss.FailureReason = data_loss.FailureReason.HTTP_CLIENT_ERROR, + item_count: int = 1, +) -> data_loss.FailedMessageInfo: + return data_loss.FailedMessageInfo( + message_type="CreateSpansBatchMessage", + reason=reason, + item_count=item_count, + ) + + +class TestFailureReason: + def test_from_status_code__4xx__client_error(self): + assert ( + data_loss.FailureReason.from_status_code(403) + == data_loss.FailureReason.HTTP_CLIENT_ERROR + ) + + def test_from_status_code__5xx__server_error(self): + assert ( + data_loss.FailureReason.from_status_code(503) + == data_loss.FailureReason.HTTP_SERVER_ERROR + ) + + def test_from_status_code__none__unknown(self): + assert ( + data_loss.FailureReason.from_status_code(None) + == data_loss.FailureReason.UNKNOWN + ) + + +class TestFlushResultSuccess: + def test_success__drained_no_drops__true(self): + result = data_loss.FlushResult( + flushed=True, + remaining_queue_size=0, + dropped_messages=0, + dropped_items=0, + failures=(), + ) + assert result.success is True + + def test_success__dropped_messages__false(self): + result = data_loss.FlushResult( + flushed=True, + remaining_queue_size=0, + dropped_messages=1, + dropped_items=5, + failures=(_failure(item_count=5),), + ) + assert result.success is False + + def test_success__not_flushed__false(self): + result = data_loss.FlushResult( + flushed=False, + remaining_queue_size=3, + dropped_messages=0, + dropped_items=0, + failures=(), + ) + assert result.success is False + + +class TestDataLossTracker: + def test_drops_since__records_after_marker__exact_delta(self): + tracker = data_loss.DataLossTracker() + tracker.record(_failure()) + marker = tracker.marker() + tracker.record(_failure(item_count=2)) + tracker.record(_failure(item_count=3)) + + count, items, failures = tracker.drops_since(marker) + assert count == 2 + assert items == 5 + assert len(failures) == 2 + + def test_drops_since__no_new_records__empty(self): + tracker = data_loss.DataLossTracker() + tracker.record(_failure()) + marker = tracker.marker() + + count, items, failures = tracker.drops_since(marker) + assert count == 0 + assert items == 0 + assert failures == [] + + def test_drops_since__details_evicted__counts_still_exact(self): + tracker = data_loss.DataLossTracker(max_entries=2) + marker = tracker.marker() + for _ in range(5): + tracker.record(_failure(item_count=4)) + + count, items, failures = tracker.drops_since(marker) + # Both counts are exact (running totals); only the retained details are + # bounded to capacity. + assert count == 5 + assert items == 20 + assert len(failures) == 2 + + def test_total_drops__exact_counts_and_details(self): + tracker = data_loss.DataLossTracker() + tracker.record(_failure()) + tracker.record(_failure(item_count=2)) + + count, items, failures = tracker.total_drops() + assert count == 2 + assert items == 3 + assert len(failures) == 2 + + def test_total_drops__details_bounded_but_counts_exact(self): + tracker = data_loss.DataLossTracker(max_entries=2) + recorded = [ + data_loss.FailedMessageInfo( + message_type="CreateSpansBatchMessage", + reason=data_loss.FailureReason.HTTP_CLIENT_ERROR, + item_count=1, + detail=str(index), + ) + for index in range(5) + ] + for failure in recorded: + tracker.record(failure) + + count, _items, failures = tracker.total_drops() + assert count == 5 + # Only the two most recent details are retained, in order. + assert [failure.detail for failure in failures] == ["3", "4"] + + +class TestFlushReporter: + def _reporter(self, tracker, *, queue_size=0): + streamer = mock.Mock() + streamer.queue_size.return_value = queue_size + return flush_reporter.FlushReporter(streamer, tracker) + + def test_build_result__drop_after_marker__reported_as_data_loss(self): + tracker = data_loss.DataLossTracker() + reporter = self._reporter(tracker) + marker = reporter.marker() + tracker.record(_failure(item_count=3)) + + result = reporter.build_result(marker, flushed=True) + + assert result.success is False + assert result.dropped_messages == 1 + assert result.dropped_items == 3 + assert result.failures[0].reason == data_loss.FailureReason.HTTP_CLIENT_ERROR + + def test_build_result__no_drops__success(self): + tracker = data_loss.DataLossTracker() + reporter = self._reporter(tracker) + marker = reporter.marker() + + result = reporter.build_result(marker, flushed=True) + + assert result.success is True + + def test_build_errors_report__surfaces_drops_outside_flush_window(self): + # A drop that happened before the flush window is not in the flush + # result, but is still discoverable via the sender-wide report. + tracker = data_loss.DataLossTracker() + reporter = self._reporter(tracker) + tracker.record(_failure()) + + marker = reporter.marker() + result = reporter.build_result(marker, flushed=True) + report = reporter.build_errors_report() + + assert result.dropped_messages == 0 + assert report.total_dropped_messages == 1 + assert len(report.failures) == 1 + + def test_build_errors_report__carries_timestamps(self): + tracker = data_loss.DataLossTracker() + reporter = self._reporter(tracker) + tracker.record( + data_loss.FailedMessageInfo( + message_type="CreateSpansBatchMessage", + reason=data_loss.FailureReason.HTTP_CLIENT_ERROR, + item_count=1, + timestamp=1000.0, + ) + ) + + report = reporter.build_errors_report() + + assert report.total_dropped_messages == 1 + assert report.generated_at > 0 + assert report.first_failure_at == 1000.0 + assert report.last_failure_at == 1000.0 + + def test_build_errors_report__no_drops__empty(self): + tracker = data_loss.DataLossTracker() + reporter = self._reporter(tracker) + + report = reporter.build_errors_report() + + assert report.total_dropped_messages == 0 + assert report.failures == () + assert report.first_failure_at is None + + +class TestMessageItemCount: + def test_item_count__batch_message__batch_length(self): + message = messages.CreateSpansBatchMessage(batch=[]) + message.batch = ["span1", "span2", "span3"] + assert message.item_count == 3 + + def test_item_count__non_batch_message__one(self): + assert messages.BaseMessage().item_count == 1 diff --git a/sdks/python/tests/unit/message_processing/test_uploads_streaming.py b/sdks/python/tests/unit/message_processing/test_uploads_streaming.py index 71438b921c6..c7476c219b4 100644 --- a/sdks/python/tests/unit/message_processing/test_uploads_streaming.py +++ b/sdks/python/tests/unit/message_processing/test_uploads_streaming.py @@ -6,7 +6,7 @@ import pytest from opik.file_upload import upload_manager -from opik.message_processing import messages, streamer, streamer_constructors +from opik.message_processing import data_loss, messages, streamer, streamer_constructors from opik.message_processing.processors import online_message_processor NOT_USED = sentinel.NOT_USED @@ -31,6 +31,7 @@ def upload_noop(**kwargs): file_upload_manager=file_upload_manager, fallback_replay_manager=mock.Mock(), unauthorized_message_types_registry=mock.Mock(), + data_loss_tracker=data_loss.DataLossTracker(), ) streamer_ = streamer_constructors.construct_streamer( message_processor=online,