perf(google_adk): stop redoing request-side attribute work on every streaming chunk - #3591
Open
KCNyu wants to merge 1 commit into
Open
perf(google_adk): stop redoing request-side attribute work on every streaming chunk#3591KCNyu wants to merge 1 commit into
KCNyu wants to merge 1 commit into
Conversation
…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.
Contributor
|
CLA Assistant Lite bot All contributors have signed the CLA ✍️ ✅ |
Author
|
I have read the CLA Document and I hereby sign the CLA |
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.
Description
Resolves #3590
ADK opens the
call_llmspan outside its streaming loop and callstrace_call_llmfrom inside it(
google/adk/flows/llm_flows/base_llm_flow.py:1596/:1668/:1669onmain), so_TraceCallLlm.__call__runs once per streamed chunk.llm_requestis bound outside the loop anddoes 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
Return early when the span is not recording. This is placed after the span-kind
set_attributeso
OpenInferenceSpan's important-attribute bookkeeping still runs. ADK's owntrace_call_llmalready does the same thing (
src/google/adk/telemetry/tracing.py:376), as doesOpenInferenceSpan.set_attribute, so a sampled-out span was previously costing more with thisinstrumentation than without it.
Derive the request-side attributes once per span.
_TraceCallLlmkeeps a small per-instancerecord 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:
INPUT_VALUEis already setlooks like the obvious approach, but span attributes are bounded
(
SpanLimits.max_attributes, 128 by default) and evict oldest-first, so a request with a longenough 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_limitcovers this.llm_requestpercall_llmspan today, but keying on the request means a span that everserves a second one still gets its attributes.
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.
mainvs. this branch in the sameprocess: 7 tools x 27 params, 300 chunks, min of 5 repeats, Linux x86_64, CPython 3.12.3.
Annotated[..., Field(...)]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 thedeclaration from the function signature on every call, so the total scales as chunks x tools. On
1.39.0,
timeitnumber=20 repeat=7 min:Annotated+FieldAnnotated+Field+LiteralCost tracks the number of parameters carrying
Annotated[..., pydantic.Field(...)], not the size ofthe resulting schema: the 3990 B row is faster than the 3590 B one.
Tests
tests/test_trace_call_llm_per_chunk.py, run throughOITracerso the spans are theOpenInferenceSpanproxies the instrumentor installs:FunctionToolsubclass that counts declaration builds, with the span still carrying
INPUT_VALUE,LLM_MODEL_NAMEand the tool schemaEight of the ten fail on
main, on bothtest-google_adk(adk 1.2.1) andtest-google_adk-latest(adk 2.7.1). The two that pass either way are the per-chunk response test and the single-chunk case.
ruffandmypyare clean.#3451 changes the
span.set_status(StatusCode.OK)line just above this diff, happy to rebase onwhichever lands first.
Checklist: