Skip to content

perf(google_adk): stop redoing request-side attribute work on every streaming chunk - #3591

Open
KCNyu wants to merge 1 commit into
Arize-ai:mainfrom
KCNyu:perf/google-adk-per-chunk-attrs
Open

perf(google_adk): stop redoing request-side attribute work on every streaming chunk#3591
KCNyu wants to merge 1 commit into
Arize-ai:mainfrom
KCNyu:perf/google-adk-per-chunk-attrs

Conversation

@KCNyu

@KCNyu KCNyu commented Aug 20, 2026

Copy link
Copy Markdown

Description

Resolves #3590

ADK opens the call_llm span outside its streaming loop and calls trace_call_llm from inside it
(google/adk/flows/llm_flows/base_llm_flow.py:1596 / :1668 / :1669 on main), so
_TraceCallLlm.__call__ runs once per streamed chunk. llm_request is bound outside the loop and
does not change, so every chunk re-serialized the whole request, re-derived every tool declaration
and re-emitted the full message history, then wrote the same values onto the same span. Only the
last pass survived.

What changed

  1. Return early when the span is not recording. This is placed after the span-kind set_attribute
    so OpenInferenceSpan's important-attribute bookkeeping still runs. ADK's own trace_call_llm
    already does the same thing (src/google/adk/telemetry/tracing.py:376), as does
    OpenInferenceSpan.set_attribute, so a sampled-out span was previously costing more with this
    instrumentation than without it.

  2. Derive the request-side attributes once per span. _TraceCallLlm keeps a small per-instance
    record of which request has already been written for a given span id, and only records the span
    once the whole block has completed, so a pass that raises part way through is retried on the next
    chunk instead of being suppressed. Response-side attributes still update for every chunk.

    Three notes on that bookkeeping:

    • The span's own attributes cannot be the marker. Checking whether INPUT_VALUE is already set
      looks like the obvious approach, but span attributes are bounded
      (SpanLimits.max_attributes, 128 by default) and evict oldest-first, so a request with a long
      enough message history overflows the limit during the first chunk and drops INPUT_VALUE,
      which would turn the guard off for exactly the requests that cost the most to re-derive.
      test_guard_survives_the_span_attribute_limit covers this.
    • The recorded value is the request's identity rather than a plain flag. ADK binds one
      llm_request per call_llm span today, but keying on the request means a span that ever
      serves a second one still gets its attributes.
    • The record is bounded at 1024 entries, least-recently-used first. Losing an entry costs one
      redundant re-derivation and can never leave an attribute unwritten, and every access is a
      single dict operation, so concurrent callers can at worst duplicate work.

Benchmark

Generated tools and requests only, no model and no network. main vs. this branch in the same
process: 7 tools x 27 params, 300 chunks, min of 5 repeats, Linux x86_64, CPython 3.12.3.

google-adk tools before after per chunk
2.7.1 7 x 27 Annotated[..., Field(...)] 0.802 s 0.026 s 2.673 -> 0.088 ms
1.2.1 (repo test pin) 7 x 27 plain type hints 3.130 s 0.033 s 10.435 -> 0.109 ms

The second row uses plainly annotated parameters because google-adk 1.2.1 cannot parse
Annotated[..., Field(...)] signatures at all; those are the more expensive kind.

The underlying cost is FunctionTool._get_declaration(), which on google-adk < 2.6.0 rebuilds the
declaration from the function signature on every call, so the total scales as chunks x tools. On
1.39.0, timeit number=20 repeat=7 min:

annotation style params schema bytes per call per param
Annotated + Field 5 / 15 / 27 815 / 2250 / 3990 B 1.962 / 5.486 / 9.651 ms ~357-392 us
Annotated + Field + Literal 5 / 15 / 27 733 / 2018 / 3590 B 2.216 / 6.299 / 11.311 ms ~419-443 us
plain type hints 5 / 15 / 27 296 / 691 / 1171 B 0.317 / 0.796 / 1.377 ms ~51-63 us

Cost tracks the number of parameters carrying Annotated[..., pydantic.Field(...)], not the size of
the resulting schema: the 3990 B row is faster than the 3590 B one.

Tests

tests/test_trace_call_llm_per_chunk.py, run through OITracer so the spans are the
OpenInferenceSpan proxies the instrumentor installs:

  • request-side attributes derived once across five chunks on one span, asserted with a FunctionTool
    subclass that counts declaration builds, with the span still carrying INPUT_VALUE,
    LLM_MODEL_NAME and the tool schema
  • response-side attributes still written for every chunk
  • each new span gets its own request attributes
  • a second, different request on the same span is still recorded
  • a pass that raises part way through is retried on the next chunk
  • nothing derived at all for a non-recording span
  • the guard still holds when the request overflows the span attribute limit
  • build count stays at 1 for 1, 5 and 20 chunks

Eight of the ten fail on main, on both test-google_adk (adk 1.2.1) and test-google_adk-latest
(adk 2.7.1). The two that pass either way are the per-chunk response test and the single-chunk case.
ruff and mypy are clean.

#3451 changes the span.set_status(StatusCode.OK) line just above this diff, happy to rebase on
whichever lands first.

Checklist:

  • Follows OpenInference configuration to hide sensitive info
  • Spans properly inherit from context attributes
  • Properly respects suppress tracing context

…treaming chunk

ADK creates the `call_llm` span outside its streaming loop and calls `trace_call_llm`
from inside it, so `_TraceCallLlm.__call__` runs once per chunk against the same span
while `llm_request` never changes. Every chunk re-serialized the whole `LlmRequest`,
walked `tools_dict` and re-derived every tool declaration, then overwrote the same
attributes with the same values -- all of it discarded except the last pass.

Return early when the span is not recording, and derive the request-side attributes
once per span. Response-side attributes still update for every chunk.
@KCNyu
KCNyu requested a review from a team as a code owner August 20, 2026 13:52
@dosubot dosubot Bot added the size:M This PR changes 30-99 lines, ignoring generated files. label Aug 20, 2026
@github-actions

github-actions Bot commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

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

@KCNyu

KCNyu commented Aug 20, 2026

Copy link
Copy Markdown
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 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size:M This PR changes 30-99 lines, ignoring generated files.

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

google-adk: request-side span attributes are re-derived on every streamed chunk

1 participant