Skip to content

fix: tracer leak and cost tracking guards - #15084

Merged
RogerHYang merged 7 commits into
Arize-ai:mainfrom
DivyaNarahari97:fix/tracer-leak-and-cost-tracking-guards
Aug 8, 2026
Merged

fix: tracer leak and cost tracking guards#15084
RogerHYang merged 7 commits into
Arize-ai:mainfrom
DivyaNarahari97:fix/tracer-leak-and-cost-tracking-guards

Conversation

@DivyaNarahari97

Copy link
Copy Markdown
Contributor

Three independent bugs found while reviewing recently-changed code. Each is a separate commit and can be reviewed on its own.

1. fix(cost-tracking): non-numeric token counts crash the threshold calculator

ThresholdBasedTokenCostCalculator.calculate_cost compared the prompt token count against its threshold without a type check. OTLP preserves whatever type the client sent, so a token count arriving as a string or list raised:

TypeError: '>' not supported between instances of 'str' and 'float'

Every span-ingestion path catches and logs per span, so the span silently lost its entire cost record rather than getting a wrong one. The exception is Tracer.get_db_traces, which is called outside a try in tracers.py.

This path only became reachable in #14329, which populated the manifest with 77 threshold customizations and taught the facilitator to sync them into token_prices.customization — covering the flagship models (claude-sonnet-4*, gemini-2.5-pro, gpt-5.*) on all four of their token prices. Note the output price tiers on the prompt count, so a span with a string prompt count and a valid integer completion count still reaches the calculator through phase 2 of the details calculator.

Non-numeric values now bill at the base rate, matching how get_aggregated_tokens and phase 1 of calculate_details already read the same attributes.

Adds the first test coverage for this calculator — there was no test_token_cost_calculator.py. 5 of the new cases fail against the previous implementation.

2. fix(tracing): a TracerProvider is retained per agent turn and per work item

Tracer builds a fresh TracerProvider per request, and the SDK's shutdown_on_exit default registers atexit.register(provider.shutdown):

opentelemetry/sdk/trace/__init__.py:1316    shutdown_on_exit: bool = True,
opentelemetry/sdk/trace/__init__.py:1347        self._atexit_handler = atexit.register(self.shutdown)
opentelemetry/sdk/trace/__init__.py:1479            atexit.unregister(self._atexit_handler)

That stores a bound method, so the registry holds a strong reference and the provider is never collected — and neither is anything it reaches: its SimpleSpanProcessor, the _BufferedSpanExporter, and every captured span. Those spans carry full chat message histories by design, since span limits are raised to max_span_attributes=100_000 so nothing is evicted.

Only shutdown() unregisters the handler, and the experiment runner never called it — TaskWorkItem.execute and EvalWorkItem.execute each build a tracer per work item with no teardown on any path. An experiment over 1,000 examples retained ~1,000 providers plus every message history they captured, for the life of the server process, growing with each run.

shutdown_on_exit is meant for one process-wide provider — which is what server/telemetry.py has, and it correctly keeps the default. I checked every other TracerProvider( construction in the repo; this is the only per-request one.

Both work items now shut down in a finally, and TaskWorkItem's tracer moves below the template-error branch that returns without tracing anything, so the finally spans its whole lifetime.

The regression test takes a weakref, drops the tracer, and forces a collection — on the previous implementation the provider survives.

3. fix(agents): trace persistence in finally masks errors and skips teardown

Three finally blocks flushed the tracer, persisted what it captured, and shut the provider down, with nothing guarded. Any DB hiccup in _ensure_project_exists / _persist_db_traces_and_emit_event, or a cost-calculation failure inside get_db_traces (bug 1), had two consequences:

The exception replaced the one already in flight. In the summarize endpoint that is unambiguous — except SummarizationError: raise HTTPException(502, ...) is immediately followed by this finally, so a telemetry-side failure discarded the intended 502 and its detail and answered an opaque 500. On the success path it turned a completed summarization into a 500 for a bookkeeping problem.

And it skipped shutdown(), retaining the provider (bug 2) and stranding the BatchSpanProcessor worker thread whenever remote export is on.

The three copies collapse into one guarded helper. Every other place that persists spans already guards this way (bulk_inserter, span_cost_calculator); walking every finally in src/phoenix for unguarded awaited DB work returns only these three sites, so this was an omission rather than house style.

Verification

Check Result
ruff check All checks passed
ruff format --check 6 files already formatted
mypy Success: no issues found in 1037 source files
Affected unit suites 711 passed, 49 skipped
New tests vs. previous code 7 of 22 fail as expected

Affected suites: test_experiment_runner.py, test_tracers.py, server/cost_tracking/, server/api/routers/test_agents.py, server/agents/.

Notes for reviewers

  • No behavior change for valid input in bug 1 — a numeric 0 still falls through to the base rate exactly as the previous falsy check did.
  • Bug 2's fix means providers are no longer shut down at process exit. That is safe here: the experiment-runner factory builds tracers with enable_remote_export=False, so there is no BatchSpanProcessor to flush, and the agent router paths shut down explicitly.
  • Bugs 2 and 3 are related — 3 is one of the ways shutdown() was being skipped — but they are independent defects and either fix stands alone.

DivyaNarahari97 and others added 3 commits August 4, 2026 20:24
ThresholdBasedTokenCostCalculator compared the prompt token count against
its threshold without checking the type. OTLP preserves whatever type the
client sent, so a token count that arrives as a string or a list raised
`TypeError: '>' not supported between instances of 'str' and 'float'`.

Every ingestion path catches and logs per span, so the span silently lost
its entire cost record rather than getting a wrong one. The one unguarded
path is `Tracer.get_db_traces`, which is called outside a try in
`tracers.py` and would fail the whole batch.

The path only became reachable in Arize-ai#14329, which populated the manifest
with 77 threshold customizations and taught the facilitator to sync them
into `token_prices.customization` — covering the flagship models
(claude-sonnet-4*, gemini-2.5-pro, gpt-5.*) on all four of their token
prices. The `output` price tiers on the prompt count too, so a span with
a string prompt count and a valid integer completion count reaches the
calculator through phase 2 of the details calculator.

Anything not numeric now bills at the base rate, matching how
`get_aggregated_tokens` and phase 1 of `calculate_details` already read
the same attributes. `bool` is excluded because it is an `int` subclass
and never a token count.

Adds the first test coverage for this calculator: 5 of the new cases fail
against the previous implementation.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… item

`Tracer` builds a fresh `TracerProvider` per request, and the SDK's
`shutdown_on_exit` default registers `atexit.register(provider.shutdown)`.
That stores a bound method, so the registry holds a strong reference and
the provider is never collected — and neither is anything it reaches: its
SimpleSpanProcessor, the `_BufferedSpanExporter`, and every captured span.
Those spans carry full chat message histories by design, since the span
limits are raised to `max_span_attributes=100_000` so nothing is evicted.

Only `shutdown()` unregisters the handler, and the experiment runner never
called it: `TaskWorkItem.execute` and `EvalWorkItem.execute` each build a
tracer per work item with no teardown on any path. An experiment over
1,000 examples retained ~1,000 providers plus every message history they
captured, for the life of the server process, growing with each run.

`shutdown_on_exit` is meant for one process-wide provider — which is what
`server/telemetry.py` has, and it correctly keeps the default. These are
request-scoped, so opt out and let each owner tear its own provider down.

Both work items now shut down in a `finally`. `TaskWorkItem`'s tracer moves
below the template-error branch that returns without tracing anything, so
the `finally` spans its whole lifetime.

The regression test takes a weakref, drops the tracer, and forces a
collection: on the previous implementation the provider survives.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… teardown

Three `finally` blocks flushed the tracer, persisted what it captured, and
shut the provider down — with nothing guarded. Any DB hiccup in
`_ensure_project_exists` or `_persist_db_traces_and_emit_event`, or a cost
calculation failure inside `get_db_traces`, had two consequences.

The exception replaced the one already in flight. In the summarize endpoint
that is unambiguous: `except SummarizationError: raise HTTPException(502)`
is immediately followed by this `finally`, so a telemetry-side failure
discarded the intended 502 and its detail and answered an opaque 500. On
the success path it turned a completed summarization into a 500 for a
bookkeeping problem.

And it skipped `shutdown()`, retaining the provider and stranding the
BatchSpanProcessor's worker thread whenever remote export is on.

The three copies collapse into one helper that logs persistence failures
and shuts down in its own `finally`. Recording telemetry is bookkeeping:
a turn whose traces cannot be written still produced its answer.

Every other place that persists spans already guards this way
(`bulk_inserter`, `span_cost_calculator`); walking every `finally` in
src/phoenix for unguarded awaited DB work returns only these three sites.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@DivyaNarahari97
DivyaNarahari97 requested a review from a team as a code owner August 5, 2026 03:32
@github-project-automation github-project-automation Bot moved this to 📘 Todo in phoenix Aug 5, 2026
@dosubot dosubot Bot added the size:L This PR changes 100-499 lines, ignoring generated files. label Aug 5, 2026
@github-actions github-actions Bot added the triage issues that need triage label Aug 5, 2026
@github-actions

github-actions Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

All contributors have signed the CLA ✍️ ✅
Posted by the CLA Assistant Lite bot.

@DivyaNarahari97

Copy link
Copy Markdown
Contributor Author

I have read the CLA Document and I hereby sign the CLA

github-actions Bot added a commit that referenced this pull request Aug 5, 2026
@RogerHYang RogerHYang changed the title Fix/tracer leak and cost tracking guards fix: tracer leak and cost tracking guards Aug 5, 2026
@RogerHYang RogerHYang self-assigned this Aug 5, 2026
@RogerHYang RogerHYang removed the triage issues that need triage label Aug 5, 2026
`Tracer` had six owners and each hand-rolled its lifecycle differently:
three in the agents router (flush, persist, shut down), two in the
experiment runner (shut down only), and one in the playground
subscription that never tore down at all. Nothing enforced the contract,
so whether a tracer was released depended on which statement in a
`finally` happened to raise first.

Give `Tracer` `__enter__`/`__exit__` and `__aenter__`/`__aexit__`, and
scope every owner with `with` or `async with`. Releasing shuts the
provider down and then clears the buffer, and both are needed:
`shutdown()` runs the processors' shutdown hooks, but it reaches
`_BufferedSpanExporter.shutdown()`, which returns without touching
`_finished_spans`, so every captured span and the message histories on it
stayed reachable.

The order matters, and it is shutdown first. With a remote exporter the
shutdown drains the batch queue, bounded as described below, and
`SimpleSpanProcessor` keeps filling the buffer throughout, having no
shutdown check in SDK 1.43.0. Clearing first left behind whatever ended
during the drain — a window as wide as the collector is slow. Releasing
still does not seal the tracer: only the batch processor refuses spans
afterwards, so a span ending later lands in the buffer unread, and
nothing enforces a bound on that beyond the tracer going out of scope
with its block.

A failure while tearing down is logged rather than raised. Every owner
releases from a block that may already be unwinding an agent error or a
cancellation, and a broken exporter must not stand in for it.

`__enter__` rejects a second entry, which makes the class docstring's
"use a separate tracer for distinct operations" a contract rather than a
recommendation. Reuse fails silently otherwise: the first release stops
the batch processor but not `SimpleSpanProcessor`, so spans recorded
afterwards reach the local buffer while remote export refuses them, and
the two disagree with nothing raised. The flag is never cleared, so reuse
after the block is rejected as well as re-entry inside it. No owner does
either — the guard is here to keep it that way.

The async release also owns the remote drain, on a worker thread. That
drain is a full OTLP round-trip, and inline it stalled every other
request on the worker — `test_async_block_does_not_hold_the_event_loop`
pins that with a 0.3 s exporter.

`BatchProcessor.shutdown` joins its worker for up to 30 seconds, which
waits out an export already in flight. Past that it abandons what is still
queued and calls the exporter's `shutdown()`; for the OTLP HTTP exporter
that closes the session and sets the flag its retry loop polls, so a
retrying export gives up at its next checkpoint.

The previous code called `force_flush()` first. That is unbounded at the
SDK layer — `BatchProcessor.force_flush` ignores the `timeout_millis` it
accepts — but the stock OTLP HTTP exporter bounds each export at 10
seconds and does not retry a read timeout, so against a hung collector the
old path was already bounded in practice. This narrows the window rather
than converting a hang into a bound, and a custom exporter with no
deadline of its own is the case where the difference is real.

`__aexit__` is shielded. The usual way a turn ends early is a client
disconnecting, which cancels the scope the release runs in, and an
unshielded `await` raises before the release happens — leaving the tracer
not torn down at all, which is worse than the stall it was avoiding. The
shield covers anyio scope cancellation, which is what Starlette delivers.
A native `asyncio.Task.cancel` arriving while the release is queued for a
thread-pool token still skips it; no current owner combines a remote
exporter with native cancellation. Without a remote exporter there is
nothing to drain, so that case skips the thread rather than paying for a
hop it cannot use.

Two properties now hold structurally rather than by convention: the drain
necessarily runs after everything in the block, so a hung collector
cannot delay the local write, and there is no flush statement left to
order by hand or to delete by accident.

One gap stays open, and it is not new. The two streaming agents endpoints
build their tracer in the endpoint frame but enter the block inside the
streaming generator, so an `AgentError` raised in between leaves the
tracer unreleased. The provider and every captured span are still
collected, which is what `shutdown_on_exit=False` buys, but with remote
export enabled the batch processor's daemon thread survives — one per
failed request. The same gap is on `main`. Closing it needs the ownership
transfer that `contextlib.AsyncExitStack.pop_all` exists for, plus a
behaviour change, since the error response would then wait for the drain.
It is left to its own change.

The agents router keeps only the database write, which is shielded and
bounded for the same reason `_persist_run` in the experiment runner is:
unshielded, a disconnect raises `CancelledError` before the database is
reached — a `BaseException`, so `except Exception` never sees it and
nothing is logged. The turn produced its answer and its traces were
dropped silently. The bound stops a burst of disconnects against a
stalled database from piling up handlers. It applies on every path, so a
legitimately slow write that exceeds it is logged and abandoned too.

With the drain moved, that helper only persists, so it is renamed
`_persist_agent_traces` and loses its `ingest_traces` flag — the
parameter made the whole call a no-op when false, which is a decision the
three call sites can state themselves.
@dosubot dosubot Bot added size:XXL This PR changes 1000+ lines, ignoring generated files. and removed size:L This PR changes 100-499 lines, ignoring generated files. labels Aug 6, 2026
The lifecycle docstrings had grown to 23 and 24 lines against methods of
four and five. Most of the excess was justification aimed at a reader who
had just watched the change happen, not at one arriving cold.

Drop the paragraph in `__aexit__` explaining `BatchProcessor.shutdown`'s
internal algorithm — the worker join, the queue abandonment, the exporter
shutdown, the retry-loop flag. It describes another package's code, will
drift on the next SDK bump, and has been rewritten and found wrong more
than once. What the caller needs from it is one clause: bounded near 30
seconds, drops what it could not send.

Replace the counterfactual constructions in `__exit__` — "clearing first
would leave behind", "neither call is sufficient alone" — with what the
code does and why. Same for the class docstring, which still called a
separate tracer per operation "recommended" after `__enter__` began
rejecting a second entry, and for the `shutdown_on_exit` comment, which
enumerated the current owners and so needed editing whenever one is added.

No behaviour change. 22 lines out of two docstrings.
The comments and docstrings on this branch were written against the change
they accompanied rather than against the code a later reader will find.
Rewrite them so each one stands on its own.

- Drop history: "previously tore its tracer down", "used to raise
  TypeError", "the manifest now ships", and a reference to a force_flush
  call this branch deleted.
- Drop narrative: arguments for why a guard was added, counterfactuals
  about what a missing `async with` would hide, and an anthropomorphism.
- Drop pointers that defer their justification elsewhere. The bound on
  _TRACE_PERSIST_TIMEOUT_SECONDS cited the experiment runner, which sets
  the same bound with a bare literal and no rationale of its own; two
  "shields/holds the invariant the same way" asides sat after sentences
  that already stated the rationale in full.
- Fix imprecision: "Both phases" named nothing, and the atexit test
  described its flag backwards.
- Unwind stacked em-dash clauses into one claim per sentence.

Mechanism and invariants are kept, including the OTLP attribute-type
rationale in ThresholdBasedTokenCostCalculator and the SDK version that
scopes the SimpleSpanProcessor claim in Tracer.__exit__. Cross-references
that name a real consistency constraint are kept, such as
get_aggregated_tokens reading the same attribute with the same fallback.
Pre-existing comments that the branch only reindented are untouched, as
are exception and assertion message strings.
_persist_agent_traces opened two sessions: one inside
_ensure_project_exists and one inside _persist_db_traces_and_emit_event,
with nothing but in-memory model building between them. On SQLite every
session — reads included — serializes on a process-wide asyncio.Lock that
waits indefinitely, so the split queued for that lock twice per turn. On
Postgres it cost a second pool checkout, the pool_pre_ping SELECT 1 that
each checkout carries, and a second commit.

Open one session for the whole write. _ensure_project_exists takes an
AsyncSession and derives its dialect from the bind, matching
_upsert_project_sessions in this module. _persist_db_traces_and_emit_event
had no other caller and is inlined at its one call site; the insert event
is still emitted after the transaction commits, and the empty-traces short
circuit it provided survives as a guard on the persist call.

Resolve the project with a read before writing, as span ingestion already
does. INSERT ... ON CONFLICT DO NOTHING returns no row once the project
exists, so the previous shape paid an unconditional no-op write and a
SELECT on every turn. Reading first makes the steady-state path a single
statement; the conflict insert and a re-select run only when the project is
missing or a concurrent caller wins the insert.

TestPersistDbTracesAndEmitEvent targeted the inlined wrapper, so it now
drives _persist_agent_traces end to end and keeps its assertion that the
event carries the right project ids. TestEnsureProjectExists covers the
creation path, which is what exercises RETURNING on the conflict insert.
The concurrent-insert branch stays uncovered: provoking it needs two
simultaneous transactions, which the SQLite lock rules out.

@RogerHYang RogerHYang left a comment

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.

Thank you for your contribution!

@github-project-automation github-project-automation Bot moved this from 📘 Todo to 👍 Approved in phoenix Aug 8, 2026
@RogerHYang
RogerHYang merged commit 70c1225 into Arize-ai:main Aug 8, 2026
49 checks passed
@github-project-automation github-project-automation Bot moved this from 👍 Approved to ✅ Done in phoenix Aug 8, 2026
axiomofjoy added a commit that referenced this pull request Aug 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size:XXL This PR changes 1000+ lines, ignoring generated files.

Projects

Status: ✅ Done

Development

Successfully merging this pull request may close these issues.

2 participants