[NA] [SDK] fix: evict per-trace state in LangChain OpikTracer after root run finishes - #7517
[NA] [SDK] fix: evict per-trace state in LangChain OpikTracer after root run finishes#7517AnastasisB wants to merge 2 commits into
Conversation
github-actions
Bot
added
python
458669b to
8fcfe43
Compare
![baz-reviewer[bot]](https://avatars.githubusercontent.com/in/933528?s=60&v=4){kind=small}
| def _bookkeeping_sizes(tracer: OpikTracer) -> Dict[str, int]: | ||
| return { | ||
| "span_data_map": len(tracer._span_data_map), | ||
| "created_traces_data_map": len(tracer._created_traces_data_map), | ||
| "externally_created_traces_ids": len(tracer._externally_created_traces_ids), | ||
| "skipped_langgraph_root_run_ids": len(tracer._skipped_langgraph_root_run_ids), | ||
| "langgraph_parent_span_ids": len(tracer._langgraph_parent_span_ids), |
There was a problem hiding this comment.
_bookkeeping_sizes() asserts on five private tracer containers, so a harmless storage refactor breaks the tests even when flush(), created_traces(), and get_current_span_data_for_run() still behave correctly — should we replace _EMPTY_BOOKKEEPING/_bookkeeping_sizes with assertions on fake_backend.trace_trees, tracer.created_traces(), and get_current_span_data_for_run(), as .agents/skills/python-sdk/testing.md suggests?
Want Baz to fix this for you? Activate Fixer You can also update your AI coding guidelines based on this comment by apply pr to [branch name]
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/tests/library_integration/langchain/test_opik_tracer.py around lines
571-589, remove the `_bookkeeping_sizes()` helper and `_EMPTY_BOOKKEEPING` constant
entirely — they assert against private tracer containers (`_span_data_map`,
`_created_traces_data_map`, `_externally_created_traces_ids`,
`_skipped_langgraph_root_run_ids`, `_langgraph_parent_span_ids`) and freeze the internal
bookkeeping layout. Refactor all tests that use them to validate the same behaviors
(completion, chained runs, error paths, concurrency, externally-created trace ID
cleanup) using only public signals per `.agents/skills/python-sdk/testing.md`: -
`fake_backend.trace_trees` contents to confirm spans were emitted -
`tracer.created_traces()` length to confirm trace registration -
`get_current_span_data_for_run()` to assert completed runs have no current span data
after `tracer.flush()`, while in-flight runs remain queryable
| assert { | ||
| name: size for name, size in _bookkeeping_sizes(tracer).items() if size > 0 | ||
| } == in_flight_entries |
There was a problem hiding this comment.
Per-run cleanup not verified
This assertion only checks the map sizes before and after the second trace, so it can miss a regression that evicts the original in_flight_run_id and leaves different same-sized state behind — should we assert that in_flight_run_id is still present in _span_data_map / _created_traces_data_map after the second trace completes?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/tests/library_integration/langchain/test_opik_tracer.py around lines
635-658, in the
`test_opik_tracer__concurrent_in_flight_trace__untouched_by_other_trace_finishing`
logic, replace the final assertion that only compares bookkeeping map sizes (lines
655-657) with assertions that the specific `in_flight_run_id` is still present in the
tracer’s internal state after the second trace finishes. Refactor the test to capture
the pre-flush keys (or explicitly assert membership) for `_span_data_map` and
`_created_traces_data_map` tied to `in_flight_run_id`, and then after `tracer.flush()`
assert those keys are unchanged/present while unrelated trace state was evicted. Keep
the existing size check only as a secondary assertion if useful.
| run_ids = { | ||
| run_id | ||
| for run_id, span_data in self._span_data_map.items() | ||
| if span_data.trace_id == trace_id | ||
| } | ||
| # Child runs map to the same TraceData object as the root run. | ||
| run_ids.update( | ||
| run_id | ||
| for run_id, trace_data in self._created_traces_data_map.items() | ||
| if trace_data.id == trace_id | ||
| ) | ||
| for run_id in run_ids: |
There was a problem hiding this comment.
_evict_finished_trace_state() evicts all entries for a trace_id when the first root run finishes, so a second root run sharing the same external trace_id loses its _span_data_map entries before _process_end_span fires and its final span never emits — should we scope eviction to the finishing root run's own descendants (and only drop shared trace state when a refcount reaches zero), and guard the whole eviction pass with a threading.RLock to prevent concurrent callbacks from racing on the shared dicts?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/integrations/langchain/opik_tracer.py around lines 371-387,
refactor `_evict_finished_trace_state()` to address two related issues: 1. **Correctness
(shared trace_id):** Currently, eviction collects every entry in `_span_data_map` where
`span_data.trace_id == trace_id` and removes them all, which is unsafe when multiple
root runs share the same externally supplied `opik_trace_id`. Refactor so eviction is
scoped to the finishing root run and its own descendants only. For shared/external trace
IDs (i.e., when `_resolve_trace_id()` returns a trace not owned by this root), introduce
a reference count per `trace_id` and only evict shared state (including
`_externally_created_traces_ids`) when the refcount reaches zero. This ensures later
root runs' `_process_end_span` callbacks still find their entries in `_span_data_map`.
2. **Thread safety:** The current code iterates `_span_data_map.items()` and
`_created_traces_data_map.items()` while mutating those same dicts with no
synchronization, which can raise `RuntimeError: dictionary changed size during
iteration` under concurrent callbacks. Introduce a tracer-instance `threading.RLock` and
acquire it for the entire eviction sequence — from `_resolve_trace_id()` through
computing run IDs and popping entries from both maps. Apply the same lock to any other
logic in this file that scans or mutates these bookkeeping maps (e.g.,
`_finalize_trace`, end-span handlers).
…rnal id registration against discard
| with self._external_trace_registration_lock: | ||
| self._create_root_trace_and_span_impl( | ||
| run_id=run_id, | ||
| run_dict=run_dict, | ||
| allow_duplicating_root_span=allow_duplicating_root_span, | ||
| ) |
There was a problem hiding this comment.
_create_root_trace_and_span() and _attach_span_to_local_or_distributed_trace() hold self._external_trace_registration_lock for their entire _..._impl(...) execution, so one slow backend call (__internal_api__trace__/__internal_api__span__, _opik_context_storage.*) blocks all concurrent traces on the shared tracer — according to the PR description this tracer is reused across requests, so should we narrow the lock to just the _externally_created_traces_ids / _span_data_map mutations that need atomicity with _evict_finished_trace_state's check-then-discard?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/integrations/langchain/opik_tracer.py, both
`_create_root_trace_and_span()` (lines ~487-506) and
`_attach_span_to_local_or_distributed_trace()` (lines ~619-634) wrap their entire impl
in `self._external_trace_registration_lock`, serializing unrelated traces and creating a
throughput bottleneck. Refactor by narrowing the lock scope in both methods: 1. Move all
"compute/setup" work (context storage updates via `self._opik_context_storage.*`,
backend calls via `__internal_api__trace__`/`__internal_api__span__`,
`_track_root_run(...)`) outside the lock. 2. Acquire
`self._external_trace_registration_lock` only for the minimal critical section that
mutates or checks shared bookkeeping: adding to `_externally_created_traces_ids` and
writing `_span_data_map`/trace-id associations that `_evict_finished_trace_state`
checks. 3. In `_create_root_trace_and_span_impl`, add to
`_externally_created_traces_ids` only after `_save_span_trace_data_to_local_maps(...)`
has populated the maps for that trace/span, then wrap just that set-add in the lock. 4.
Preserve `_evict_finished_trace_state()`'s lock-held check-then-discard atomicity —
only reduce contention by moving non-bookkeeping operations outside the lock.
| def test_opik_tracer__get_current_span_data_for_run__cleared_once_trace_completes( | ||
| fake_backend, | ||
| ): | ||
| tracer = OpikTracer() | ||
| root_run_id, child_run_id = uuid4(), uuid4() | ||
| tracer.on_chain_start({"name": "root"}, {"input": "x"}, run_id=root_run_id) | ||
| tracer.on_chain_start( | ||
| {"name": "child"}, | ||
| {"input": "y"}, | ||
| run_id=child_run_id, | ||
| parent_run_id=root_run_id, | ||
| ) |
There was a problem hiding this comment.
Missed backend emission regression
get_current_span_data_for_run() only checks that the trace is cleared after completion, so _process_end_span() could stop emitting the final root/child spans to fake_backend and this test would still pass — should we assert the expected fake_backend and __internal_api__span__ calls too?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/tests/library_integration/langchain/test_opik_tracer.py around lines
708-730, in
`test_opik_tracer__get_current_span_data_for_run__cleared_once_trace_completes`,
strengthen the assertions beyond `get_current_span_data_for_run(...) is None` by also
verifying that the end processing emitted the final spans. After `tracer.flush()` (and
before/after the eviction assertions), assert that `fake_backend` recorded span
submissions for the child and root runs (and that the calls include the expected
identifiers/relationships such as `run_id` and the correct parent span linkage), or
alternatively assert that `__internal_api__span__` was called with the expected
parameters for those runs. Refactor the test to capture backend call history (whatever
the fake_backend exposes) and add precise expectations so this would fail if
`_process_end_span()` stops emitting the final spans.
| def test_opik_tracer__shared_external_trace__finishing_root_run_leaves_other_in_flight( | ||
| fake_backend, | ||
| ): | ||
| external_trace_data = trace_module.TraceData(name="external-trace") | ||
| context_storage.set_trace_data(external_trace_data) | ||
| tracer = OpikTracer() | ||
| first_run_id, second_run_id = uuid4(), uuid4() | ||
| try: | ||
| tracer.on_chain_start({"name": "first"}, {"input": "x"}, run_id=first_run_id) | ||
| tracer.on_chain_start({"name": "second"}, {"input": "y"}, run_id=second_run_id) | ||
| tracer.on_chain_end({"output": "x"}, run_id=first_run_id) | ||
|
|
There was a problem hiding this comment.
Cleanup tested without output check
These external-trace tests only check _externally_created_traces_ids and the span-data maps, so a change could clear the bookkeeping while suppressing the final span call and still pass — should we also assert the backend span emission/count the PR description calls out?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/tests/library_integration/langchain/test_opik_tracer.py around lines 733-759
in `test_opik_tracer__shared_external_trace__finishing_root_run_leaves_other_in_flight`
(and similarly check the race test just below), the assertions only validate internal
bookkeeping (`_externally_created_traces_ids` and `get_current_span_data_for_run`) and
do not verify the user-visible outcome that both root runs still emit backend spans.
Refactor the test to also assert the fake backend recorded span/emission for both
run_ids (or the correct span count/name) before/after each `flush`, using the
capabilities of the `fake_backend` fixture. Make the test fail if span emission is
suppressed even when bookkeeping eviction behaves correctly, while keeping the existing
in-flight vs finished membership checks.
| second_thread = threading.Thread(target=start_second_root) | ||
| second_thread.start() | ||
| assert registration_entered.wait(timeout=5) | ||
|
|
||
| # finish the first root while the second is mid-registration; with the | ||
| # lock this blocks until the registration completes, so release it | ||
| # shortly after | ||
| threading.Timer(0.3, resume_registration.set).start() | ||
| tracer.on_chain_end({"output": "x"}, run_id=first_run_id) | ||
| second_thread.join(timeout=5) | ||
| assert not second_thread.is_alive() |
There was a problem hiding this comment.
second_thread can stay blocked in paused_save() if registration_entered.wait(timeout=5) times out before threading.Timer(0.3) fires resume_registration.set(), leaving a live non-daemon thread that can hang the suite — should we move the resume_registration.set() and second_thread.join() into a finally block, and replace the timer with an event-driven handoff so the test doesn't depend on wall-clock scheduling?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/tests/library_integration/langchain/test_opik_tracer.py around lines 800-810
inside
`test_opik_tracer__external_id_registration_races_eviction__id_not_discarded_mid_registration`:
1. Remove the `threading.Timer(0.3)` and its implicit 300ms delay entirely. 2. Refactor
to use additional `threading.Event`(s) to deterministically coordinate the handoff —
e.g., run `tracer.on_chain_end` for the first run in a separate thread and signal an
event when eviction/persist for the first run reaches the point where it would block on
the registration lock; once that event is observed, set `resume_registration`
immediately, then join the first thread. 3. Move the release of `resume_registration`
and the `second_thread.join(timeout=...)` into a `finally` block so they execute on
every exit path (including timeouts and assertion failures), ensuring `second_thread` is
always unblocked before `_save_span_trace_data_to_local_maps` is restored. 4. Keep any
timeout only as a safety guard, not as a timing mechanism.
|
Thank you for working on this @AnastasisB ! Looks like a parallel fix has already been merged into main that fixes this problem so I'm going to close this PR as a duplicate. |
2 participants
Details
Fixes an unbounded memory leak in the LangChain integration.
OpikTraceronly ever adds to its six per-instance bookkeeping structures (_span_data_map,_created_traces_data_map,_created_traces,_externally_created_traces_ids,_skipped_langgraph_root_run_ids,_langgraph_parent_span_ids), so a tracer built once and reused across requests (the documented pattern) retains every prompt/completion payload for the process lifetime. Full analysis, production numbers and an offline repro are in the linked issue.The fix releases a trace's state right after its root run's final end/error callback: a
finallyhook in_process_end_spanand_process_end_span_with_errorfor parentless runs, calling the new_evict_finished_trace_state. That is the earliest safe point: langchain calls_persist_runbefore that callback, and the end-span processing still needs the root run's entries to emit the final span._persist_runis untouched.BaseTracerattaches every child to its parent'schild_runs), so it removes exactly that run's entries and never scans the shared maps. Other in-flight root runs on the same tracer, including ones attached to the same externally created (distributed) trace, are untouched. The external trace id is only discarded once no remaining span references it; a narrow lock shared by the id registration sites and that check-then-discard makes the two atomic, so a concurrent root run cannot have the id discarded from under it. Child span starts stay lock-free.finally, so the error path cannot leak either._persist_runwas tried first: it passes the chain-based tests but silently drops root span emission for direct model invocations (measured 10 invocations -> 10 traces, 0 spans). A regression test now asserts spans reach the backend, so that placement fails the suite._created_tracesis deliberately NOT evicted:created_traces()is documented to accumulate, existing tests and docs rely on it, and Trace objects are small handles without payloads. A long-lived tracer still accumulates one small handle per trace; whether to bound that is your design call, happy to help in a follow-up if you want it bounded.Before/after with the repro from the issue, after 50 invocations:
Change checklist
Issues
AI-WATERMARK
AI-WATERMARK: yes
Testing
Eight regression tests added to
tests/library_integration/langchain/test_opik_tracer.py, following the existingfake_backendconventions: all payload-carrying state evicted after completion while traces still reach the backend with their spans andcreated_traces()accumulates, chained-run child state evicted, eviction on the error path, a concurrent in-flight trace untouched, a root run finishing on a shared external trace leaving the other in-flight root run's state intact, a deterministic registration/eviction race test (a second root run paused between registering the external id and saving its span while the first root run finishes; verified to fail when the lock is neutralized), the publicget_current_span_data_for_runaccessor returning None once the trace completes instead of stale data, an external trace id discarded.Run locally (Python 3.10.12, langchain-core 1.4.9, offline):
pytest tests/library_integration/langchain/test_opik_tracer.py: 89 passedtest_langchain.py,test_langgraph.py,test_langchain_thread_id.py,test_message_converters.py): 32 passed, 2 failed; the same 2 fail identically on pristinemainin this environment (no OPENAI_API_KEY, unrelated to this change)Documentation
No public API or parameter changes,
created_traces()behavior unchanged. One observable difference:get_current_span_data_for_runreturns None once the run's trace completed, where it previously kept returning the retained span data; that is the eviction working and is pinned by a regression test.