Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions sdks/python/src/opik/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,11 @@
from .api_objects.trace import Trace
from .configurator.configure import configure
from .decorator.tracker import flush_tracker, track
from .message_processing.data_loss import (
FailedMessageInfo,
FailureReason,
FlushResult,
)
from .evaluation import (
evaluate,
evaluate_experiment,
Expand Down Expand Up @@ -68,6 +73,9 @@
"ExperimentItemReferences",
"track",
"flush_tracker",
"FlushResult",
"FailedMessageInfo",
"FailureReason",
"Opik",
"get_global_client",
"set_global_client",
Expand Down
56 changes: 42 additions & 14 deletions sdks/python/src/opik/api_objects/connection_resources.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -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,
Expand All @@ -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,
)

Expand Down Expand Up @@ -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
)

Expand Down Expand Up @@ -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
Expand All @@ -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)
Comment on lines 365 to 367

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Concurrent flushes race on shared bundle

Two racing release() calls can both pass the entry.refcount > 1 check and then call shared_bundle.flush(timeout) concurrently, so the same shared resource gets drained twice and FlushResult can become inconsistent — should we mark the bundle as flush-in-progress under self._lock?

Severity

Want Baz to fix this for you? Activate Fixer

Other fix methods

Fix in Cursor

Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/api_objects/connection_resources.py around lines 365-375 inside
ConnectionResourceManager.release (the shared pre-flush block where `entry.refcount > 1`
leads to `shared_bundle.flush(timeout)` outside the lock), prevent two racing releases
from both flushing the same shared bundle. Refactor the preflush decision so that, while
holding self._lock, you atomically mark the bundle/entry as “flush in progress”
(e.g., add an `entry.preflush_in_progress` flag and optionally store a flush-result
placeholder) and only the first caller may perform `shared_bundle.flush(timeout)`; any
concurrent callers should not call flush again. If you need consistent FlushReporter
outcomes for all callers, have the non-owner callers either wait for a completion Event
set by the owner (without starting a second flush) or return a clearly defined value
(e.g., flushed=None) while the owner drains and records the authoritative result.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Skipping — this shared-bundle pre-flush block predates this PR (the change here only captures its return value). A concurrent double streamer.flush() is benign: the drain is idempotent (the second flush finds an already-drained queue), and each client builds its own FlushResult from its own marker, so there's no cross-client inconsistency. Adding lock-held flush-in-progress state is the extra concurrency machinery this MVP-scoped PR deliberately avoids.

🤖 Reply posted via /address-github-pr-comments

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, that makes sense — I’ll remember that this pre-existing shared-bundle flush path is intentionally tolerated here and save it to memory once the PR is merged.

Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
77 changes: 70 additions & 7 deletions sdks/python/src/opik/api_objects/opik_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@
url_helpers,
)
from ..message_processing import (
data_loss,
messages,
)
from ..message_processing.batching import sequence_splitter
Expand Down Expand Up @@ -216,6 +217,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(
Expand Down Expand Up @@ -1878,9 +1881,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.

Comment on lines 1936 to 1940

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

end() doc hides drop mode

end() says it submits all pending messages, which is inaccurate for the public flush=False path that deliberately drops queued data — should we qualify this so only flush=True drains and persists the queue?

Severity

Want Baz to fix this for you? Activate Fixer

Other fix methods

Fix in Cursor

Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/api_objects/opik_client.py around lines 1884-1896, in the
`Opik.end()` docstring, the first summary sentence currently says it “submit[s] all
pending messages”, which contradicts the `flush: bool = True` behavior and the
described fire-and-forget teardown when `flush=False`. Update that opening sentence to
explicitly state that pending queued messages (including file uploads) are
drained/submitted only when `flush=True`, and when `flush=False` the method does not
wait and may drop anything still in flight. Keep the rest of the docstring consistent
with the existing Args/Returns sections describing `flush=False` as dropping pending
data.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Commit d8ae4a0 addressed this comment by updating Opik.end()’s docstring to say that pending messages are submitted only when flush=True, and that flush=False drops anything still queued. The surrounding return/behavior docs were also aligned with that distinction.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in d8ae4a0end()'s docstring now states pending messages are submitted only when flush=True, and dropped when flush=False.

🤖 Reply posted via /address-github-pr-comments

Connection resources are shared and ref-counted across clients with a
matching configuration: this releases the current client's reference.
Expand Down Expand Up @@ -1908,28 +1915,84 @@ 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
Comment on lines +1982 to +1995

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Idempotent end overwrites success

end() coerces Lease.release() returning None into bool(flushed) and rebuilds _last_flush_result, so a second end() call overwrites the real flush outcome with FlushResult(flushed=False, ...) even though no new shutdown work ran — should we short-circuit when flushed is None and keep the existing last_flush_result?

Severity

Want Baz to fix this for you? Activate Fixer

Other fix methods

Fix in Cursor

Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/api_objects/opik_client.py around lines 1884-1938 in the `end()`
method, fix idempotency: `self._lease.release(..., flush=flush, ...)` can return `None`
when already released, but the current code still calls
`_flush_reporter.build_result(marker, flushed=bool(flushed))`, overwriting
`_last_flush_result` with a “flushed=False” result on subsequent `end()` calls.
Refactor so that after `flushed = self._lease.release(...)`, if `flush` is True and
`flushed is None`, you short-circuit and return the existing `self._last_flush_result`
(or `None` if it was never set) without rebuilding. Keep the existing behavior for the
first successful release (when `flushed` is not `None`) and still return `None` when
`flush` is False.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Commit d8ae4a0 addressed this comment by short-circuiting end() when self._lease.release(...) returns None. On repeated end() calls, it now returns the existing self._last_flush_result instead of rebuilding a new FlushResult(flushed=False, ...), preserving the original flush outcome.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in d8ae4a0end() now short-circuits when release() returns None (already released), keeping the prior last_flush_result instead of overwriting it with a spurious flushed=False result. Added idempotency test coverage.

🤖 Reply posted via /address-github-pr-comments


def flush(self, timeout: Optional[int] = None) -> bool:
"""
Flush the streamer to ensure all messages are sent.

Covers delivery of trace/span/feedback messages; attachment/file uploads

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 nit | Docs

attachment/file uploads are not reflected in the outcome

This is only true for the data-loss detail (failures / dropped_*). The flushed field — and therefore success and this method's return value — does reflect uploads: streamer.flush() returns upload_flushed and self._all_done(), so a stuck/failed upload makes flushed=False and this returns False.

Consider tightening to something like: "upload failures are not counted in the dropped-message detail, though an incomplete upload still makes the flush report as not fully flushed."

🤖 Review posted via /review-github-pr

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Commit 86e3b46 addressed this comment by documenting that flush covers attachment/file uploads while distinguishing the flush outcome from its data-loss detail.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed — you're right, flushed does reflect uploads (streamer.flush() returns upload_flushed and _all_done()). Reworded the docstring: attachment/file upload failures aren't counted in the data-loss detail (dropped_*/failures), but an incomplete upload still makes the flush report as not fully flushed, so the returned bool does reflect uploads.

are not reflected in the outcome. 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.
Comment on lines 2012 to +2014

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FlushResult hides upload failures

flush() now promises True if all messages were delivered ... with no data loss, but it only reflects Streamer.flush() while FileUploadManager.flush() still returns True after failed uploads, so Opik.flush() and last_flush_result.success can report success even when attachment uploads were lost — should we scope the public contract to trace/span delivery only, or include upload failures in the flush outcome?

Severity

Want Baz to fix this for you? Activate Fixer

Other fix methods

Fix in Cursor

Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/api_objects/opik_client.py around lines 1940-1975, in the `flush()`
method (and the `last_flush_result` property that it feeds), fix the breaking contract
mismatch: today `flush()` computes `last_flush_result` solely from
`self._streamer.flush(timeout)`, but file/attachment uploads can fail independently
(FileUploadManager.flush() can return success even with per-upload failures). Update
`flush()` so the public `last_flush_result.success` and docstring reflect the true
outcome by either (preferred) aggregating streamer message delivery success plus
file-upload failure/partial-loss into the boolean passed to
`_flush_reporter.build_result`, or (alternative) narrowing the doc/result wording to
“trace/span delivery” only and making `last_flush_result` name/semantics match that
narrower scope. Ensure the same semantics are consistent with `end()` (lines ~1884-1939)
and that the reported success includes (or explicitly excludes) attachment upload
failures.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Commit d8ae4a0 addressed this comment by narrowing flush()/last_flush_result to trace/span/feedback delivery and explicitly stating that attachment/file uploads are not reflected in the outcome. It also made end() return the same flush outcome path, so the public semantics are now consistent with that narrower contract.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in d8ae4a0flush()'s docstring now scopes the contract to trace/span/feedback delivery and notes attachment/file uploads aren't reflected in the outcome. (Upload-failure accounting is deferred to a follow-up per this PR's MVP scope.)

🤖 Reply posted via /address-github-pr-comments

"""
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
Comment thread
alexkuzmik marked this conversation as resolved.

@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 __internal_api__drain_to_processors__(
self, timeout: Optional[float] = None
Expand Down
Loading
Loading