-
Notifications
You must be signed in to change notification settings - Fork 1.1k
⚡️ Speed up function _map_usage
by 172%
#2197
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
⚡️ Speed up function _map_usage
by 172%
#2197
Conversation
Here is an optimized version of the provided program. **Key optimizations with rationale:** - Avoids repeated/expensive `isinstance()` checks by merging logic. - Avoids unnecessary dictionary comprehensions and allocations if unused. - Minimizes `details.get()` calls and reuse of variables. - Uses local variable assignment to reduce attribute lookups. - Minimizes creation of empty dicts and `Usage()` when possible. - Uses tuple membership checks for event classes to condense branching ("flat is better than nested"). - Moves `model_dump().items()` out of dict comprehension when it's not needed. - Handles the common/early-exit (no usage info) case with a constant. **Summary of changes:** - Single type comparison and attribute fetch per code path; avoids multiple checks and data flows. - Uses a static `_EMPTY_USAGE` for "no details" paths—eliminates unnecessary object allocations. - Handles dictionary and token computation as fast/local as possible (no repeated dict-get, minimal fallbacks). - Preserves function signature, behavior, and all comment clarifications. This implementation should provide improved (lower) latency per call, specifically for high-throughput scenarios.
elif msg_type is BetaRawMessageStartEvent: | ||
response_usage = cast(BetaRawMessageStartEvent, message).message.usage | ||
elif msg_type is BetaRawMessageDeltaEvent: | ||
response_usage = cast(BetaRawMessageDeltaEvent, message).usage |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Are the performance improvement metrics in the description still accurate after the latest changes to this PR?
I think this change reduces readability (and consistency with other code) significantly so I'd only want to do it if the performance impact (on real examples, not just the test_usage
test) are significant. And if it is significant, should we do this everywhere we currently use isinstance
(and don't care about subclasses)? Otherwise it seems odd to do it only in this specific place.
Saurabh's comments - The type checks are expensive, also this seems to be used very frequently.
This speeds up the test
models/test_anthropic.py::test_usage
by 25%📄 172% (1.72x) speedup for
_map_usage
inpydantic_ai_slim/pydantic_ai/models/anthropic.py
⏱️ Runtime :
1.12 milliseconds
→413 microseconds
(best of76
runs)📝 Explanation and details
Here is an optimized version of the provided program.
Key optimizations with rationale:
isinstance()
checks by merging logic.details.get()
calls and reuse of variables.Usage()
when possible.model_dump().items()
out of dict comprehension when it's not needed.Summary of changes:
_EMPTY_USAGE
for "no details" paths—eliminates unnecessary object allocations.This implementation should provide improved (lower) latency per call, specifically for high-throughput scenarios.
✅ Correctness verification report:
⚙️ Existing Unit Tests and Runtime
models/test_anthropic.py::test_usage
test_pytest_inlinesnapshotdisable_testsproviderstest_bedrock_py_testsproviderstest_google_gla_py_teststes__replay_test_0.py::test_pydantic_ai_models_anthropic__map_usage
test_pytest_inlinesnapshotdisable_teststest_messages_py_teststest_mcp_py_teststest_deps_py__replay_test_0.py::test_pydantic_ai_models_anthropic__map_usage
🌀 Generated Regression Tests and Runtime
⏪ Replay Tests and Runtime
models/test_anthropic.py::test_usage
test_pytest_inlinesnapshotdisable_testsproviderstest_bedrock_py_testsproviderstest_google_gla_py_teststes__replay_test_0.py::test_pydantic_ai_models_anthropic__map_usage
test_pytest_inlinesnapshotdisable_teststest_messages_py_teststest_mcp_py_teststest_deps_py__replay_test_0.py::test_pydantic_ai_models_anthropic__map_usage
To edit these changes
git checkout codeflash/optimize-_map_usage-md0d6dgh
and push.