perf(api): memoize ProviderManager.get_configurations within request/task scope#35088
Closed
YgorLeal wants to merge 4 commits into
Closed
perf(api): memoize ProviderManager.get_configurations within request/task scope#35088YgorLeal wants to merge 4 commits into
YgorLeal wants to merge 4 commits into
Conversation
#26412) - Simplify plugin_data decorator signature by removing the unused optional view parameter, eliminating the ambiguous union return type that caused pyright to treat the decorator as untyped at its 15 call sites. - Convert flask_restx error handler registrations in external_api.py from decorator form to explicit calls, avoiding untyped decorators from a library without stubs. - Remove the reportUntypedFunctionDecorator "hint" override so the rule runs at strict-mode default severity.
Contributor
Pyrefly Diffbase → PR--- /tmp/pyrefly_base.txt 2026-04-13 16:33:49.391009273 +0000
+++ /tmp/pyrefly_pr.txt 2026-04-13 16:33:39.084980534 +0000
@@ -5186,21 +5186,21 @@
ERROR Object of class `ModuleType` has no attribute `trace_manager_queue` [missing-attribute]
--> tests/unit_tests/core/telemetry/test_facade.py:47:5
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:51:20
-ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:64:24
+ --> tests/unit_tests/core/test_provider_manager.py:60:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
--> tests/unit_tests/core/test_provider_manager.py:73:24
-ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:113:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:124:24
+ --> tests/unit_tests/core/test_provider_manager.py:82:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:160:20
+ --> tests/unit_tests/core/test_provider_manager.py:122:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:171:24
+ --> tests/unit_tests/core/test_provider_manager.py:133:24
+ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
+ --> tests/unit_tests/core/test_provider_manager.py:169:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
--> tests/unit_tests/core/test_provider_manager.py:180:24
+ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
+ --> tests/unit_tests/core/test_provider_manager.py:189:24
ERROR `dict[str, str]` is not assignable to TypedDict key `data` with type `BaseNodeData` [bad-typed-dict-key]
--> tests/unit_tests/core/test_trigger_debug_event_selectors.py:56:46
ERROR Object of class `BlobChunkMessage` has no attribute `text`
|
…task scope (#27299) Incremental step for #27299 — assembling `ProviderConfigurations` runs six independent DB queries plus a provider-factory load, costing ~1s per call in production. A single workflow step (e.g. a Retrieval node inside an Iteration) invokes this path many times through independently-created `ProviderManager` instances, so per-instance caching is not enough. - `core/provider_manager.py`: add a module-level `RecyclableContextVar[dict[str, ProviderConfigurations]]` and memoize the result in `get_configurations` for the current request/task scope, keyed by `tenant_id`. The cache is shared across `ProviderManager` instances because callers routinely spawn fresh ones via `create_plugin_provider_manager`. - `extensions/ext_celery.py`: call `RecyclableContextVar.increment_thread_recycles()` inside `FlaskTask.__call__`, matching the Flask `before_request` hook in `app_factory.py` so per-task caches reset cleanly between Celery tasks on recycled worker threads. - `tests/unit_tests/core/test_provider_manager.py`: cover the memoization (hit within a scope, per-tenant isolation, sharing across `ProviderManager` instances, reset on new scope) behind an autouse fixture that bumps the recycle counter for each test. No cross-request invalidation is introduced — the cache lives only for the current request/task, so mutations performed through subsequent requests are always observed. This deliberately scopes the change: it targets the reported hot path (iterative retrieval/LLM nodes) with no staleness risk. ## Verification - uv run --group dev basedpyright -> 0 errors, 0 warnings (project-wide). - pytest tests/unit_tests/core/test_provider_manager.py and related -> 40 passed. - ruff check / ruff format --check -> clean.
49bebfd to
f95146b
Compare
Contributor
Pyrefly Diffbase → PR--- /tmp/pyrefly_base.txt 2026-04-13 16:38:11.978047347 +0000
+++ /tmp/pyrefly_pr.txt 2026-04-13 16:38:01.913744147 +0000
@@ -5186,21 +5186,21 @@
ERROR Object of class `ModuleType` has no attribute `trace_manager_queue` [missing-attribute]
--> tests/unit_tests/core/telemetry/test_facade.py:47:5
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:51:20
+ --> tests/unit_tests/core/test_provider_manager.py:59:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:64:24
+ --> tests/unit_tests/core/test_provider_manager.py:72:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:73:24
+ --> tests/unit_tests/core/test_provider_manager.py:81:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:113:20
+ --> tests/unit_tests/core/test_provider_manager.py:121:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:124:24
+ --> tests/unit_tests/core/test_provider_manager.py:132:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:160:20
+ --> tests/unit_tests/core/test_provider_manager.py:168:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:171:24
+ --> tests/unit_tests/core/test_provider_manager.py:179:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:180:24
+ --> tests/unit_tests/core/test_provider_manager.py:188:24
ERROR `dict[str, str]` is not assignable to TypedDict key `data` with type `BaseNodeData` [bad-typed-dict-key]
--> tests/unit_tests/core/test_trigger_debug_event_selectors.py:56:46
ERROR Object of class `BlobChunkMessage` has no attribute `text`
|
Contributor
Pyrefly Diffbase → PR--- /tmp/pyrefly_base.txt 2026-04-13 18:24:39.738754681 +0000
+++ /tmp/pyrefly_pr.txt 2026-04-13 18:24:29.751581717 +0000
@@ -5186,21 +5186,21 @@
ERROR Object of class `ModuleType` has no attribute `trace_manager_queue` [missing-attribute]
--> tests/unit_tests/core/telemetry/test_facade.py:47:5
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:51:20
+ --> tests/unit_tests/core/test_provider_manager.py:59:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:64:24
+ --> tests/unit_tests/core/test_provider_manager.py:72:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:73:24
+ --> tests/unit_tests/core/test_provider_manager.py:81:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:113:20
+ --> tests/unit_tests/core/test_provider_manager.py:121:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:124:24
+ --> tests/unit_tests/core/test_provider_manager.py:132:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:160:20
+ --> tests/unit_tests/core/test_provider_manager.py:168:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:171:24
+ --> tests/unit_tests/core/test_provider_manager.py:179:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:180:24
+ --> tests/unit_tests/core/test_provider_manager.py:188:24
ERROR `dict[str, str]` is not assignable to TypedDict key `data` with type `BaseNodeData` [bad-typed-dict-key]
--> tests/unit_tests/core/test_trigger_debug_event_selectors.py:56:46
ERROR Object of class `BlobChunkMessage` has no attribute `text`
|
Contributor
Pyrefly Diffbase → PR--- /tmp/pyrefly_base.txt 2026-04-16 04:00:03.850218982 +0000
+++ /tmp/pyrefly_pr.txt 2026-04-16 03:59:53.592942080 +0000
@@ -5166,21 +5166,21 @@
ERROR Object of class `ModuleType` has no attribute `trace_manager_queue` [missing-attribute]
--> tests/unit_tests/core/telemetry/test_facade.py:47:5
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:51:20
+ --> tests/unit_tests/core/test_provider_manager.py:59:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:64:24
+ --> tests/unit_tests/core/test_provider_manager.py:72:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:73:24
+ --> tests/unit_tests/core/test_provider_manager.py:81:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:113:20
+ --> tests/unit_tests/core/test_provider_manager.py:121:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:124:24
+ --> tests/unit_tests/core/test_provider_manager.py:132:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.ProviderModelSetting.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:160:20
+ --> tests/unit_tests/core/test_provider_manager.py:168:20
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:171:24
+ --> tests/unit_tests/core/test_provider_manager.py:179:24
ERROR Argument `Literal['llm']` is not assignable to parameter `model_type` with type `ModelType | SQLCoreOperations[ModelType]` in function `models.provider.LoadBalancingModelConfig.__init__` [bad-argument-type]
- --> tests/unit_tests/core/test_provider_manager.py:180:24
+ --> tests/unit_tests/core/test_provider_manager.py:188:24
ERROR `dict[str, str]` is not assignable to TypedDict key `data` with type `BaseNodeData` [bad-typed-dict-key]
--> tests/unit_tests/core/test_trigger_debug_event_selectors.py:56:46
ERROR Object of class `BlobChunkMessage` has no attribute `text`
|
|
Thanks! This saves me a lot of time. |
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.
Summary
Incremental step for #27299 — assembling
ProviderConfigurationsruns six independent DB queries plus a provider-factory load, costing ~1s per call in production. A single workflow step (e.g. a Retrieval node inside an Iteration) invokes this path many times through independently-createdProviderManagerinstances, so per-instance caching is not enough.This PR memoizes
ProviderManager.get_configurationswithin the current request/task scope, keyed bytenant_id. For a workflow with N retrieval/LLM calls that share the same tenant in a single run, the call path collapses fromN * ~1sof DB work to1 * ~1s+(N-1) * dict-lookup.Reported impact (from #27299 and related comments): within an Iteration node's Retrieve step, each RAG call triggers
ProviderManager.get_configurations, which issues multiple SQL queries and consistently takes ~1s per call — causing multi-second stalls per iteration.Why request/task-scoped (and not cross-request)
A previous attempt (#29305) tried to share a
db.sessionacross helpers insideget_configurations. It was rejected for good reasons (db.sessionis Flask-SQLAlchemy scoped — breaks in Celery/CLI contexts; long-lived session holds transactions open across dozens of queries). This PR takes a different angle: don't touch the queries themselves, but avoid re-running them inside the same request/task.ProviderConfigurations(which would leak decrypted credentials across the trust boundary).core/entities/provider_configuration.py,events/event_handlers/update_provider_when_message_created.py,commands/system.py,core/app/llm/quota.py, etc.Changes
core/provider_manager.py: add a module-levelRecyclableContextVar[dict[str, ProviderConfigurations]]and memoize the result inget_configurationsfor the current scope, keyed bytenant_id. The cache is shared acrossProviderManagerinstances because callers routinely spawn fresh ones viacreate_plugin_provider_manager(seecore/app/llm/model_access.py,services/workflow_service.py, etc.).extensions/ext_celery.py: callRecyclableContextVar.increment_thread_recycles()insideFlaskTask.__call__, matching the Flaskbefore_requesthook inapp_factory.py. This ensures per-task caches reset cleanly between Celery tasks that land on recycled worker threads. Strictly speaking this also benefits any existingRecyclableContextVarused from Celery paths (e.g.plugin_tool_providers).tests/unit_tests/core/test_provider_manager.py: cover the memoization behaviors behind an autouse fixture that bumps the recycle counter for each test.Verification
uv run --group dev basedpyright→ 0 errors, 0 warnings (project-wide).uv run --group dev pytest tests/unit_tests/core/test_provider_manager.py tests/unit_tests/extensions/otel/test_celery_sqlcommenter.py tests/unit_tests/core/test_model_manager.py→ 40 passed.uv run --group dev ruff check/ruff format --check→ clean.Test plan
tenant_idin one scopetenant_ids are cached independentlyProviderManagerinstances in the same scope (the common caller pattern)increment_thread_recycles) starts emptyprovider_manager/celery_sqlcommenterunit tests still pass