fix: tracer leak and cost tracking guards - #15084
Merged
RogerHYang merged 7 commits intoAug 8, 2026
Merged
Conversation
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>
Contributor
|
All contributors have signed the CLA ✍️ ✅ |
Contributor
Author
|
I have read the CLA Document and I hereby sign the CLA |
`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.
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
approved these changes
Aug 8, 2026
RogerHYang
left a comment
Contributor
There was a problem hiding this comment.
Thank you for your contribution!
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 calculatorThresholdBasedTokenCostCalculator.calculate_costcompared 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: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 intracers.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 theoutputprice 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_tokensand phase 1 ofcalculate_detailsalready 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): aTracerProvideris retained per agent turn and per work itemTracerbuilds a freshTracerProviderper request, and the SDK'sshutdown_on_exitdefault registersatexit.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 span limits are raised tomax_span_attributes=100_000so nothing is evicted.Only
shutdown()unregisters the handler, and the experiment runner never called it —TaskWorkItem.executeandEvalWorkItem.executeeach 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_exitis meant for one process-wide provider — which is whatserver/telemetry.pyhas, and it correctly keeps the default. I checked every otherTracerProvider(construction in the repo; this is the only per-request one.Both work items now shut down in a
finally, andTaskWorkItem's tracer moves below the template-error branch that returns without tracing anything, so thefinallyspans 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 infinallymasks errors and skips teardownThree
finallyblocks 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 insideget_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 thisfinally, 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 theBatchSpanProcessorworker 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 everyfinallyinsrc/phoenixfor unguarded awaited DB work returns only these three sites, so this was an omission rather than house style.Verification
ruff checkruff format --checkmypyAffected suites:
test_experiment_runner.py,test_tracers.py,server/cost_tracking/,server/api/routers/test_agents.py,server/agents/.Notes for reviewers
0still falls through to the base rate exactly as the previous falsy check did.enable_remote_export=False, so there is noBatchSpanProcessorto flush, and the agent router paths shut down explicitly.shutdown()was being skipped — but they are independent defects and either fix stands alone.