Skip to content

perf(api): memoize ProviderManager.get_configurations within request/task scope#35088

Closed
YgorLeal wants to merge 4 commits into
langgenius:mainfrom
YgorLeal:perf/provider-manager-scope-cache
Closed

perf(api): memoize ProviderManager.get_configurations within request/task scope#35088
YgorLeal wants to merge 4 commits into
langgenius:mainfrom
YgorLeal:perf/provider-manager-scope-cache

Conversation

@YgorLeal
Copy link
Copy Markdown
Contributor

Summary

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.

This PR memoizes ProviderManager.get_configurations within the current request/task scope, keyed by tenant_id. For a workflow with N retrieval/LLM calls that share the same tenant in a single run, the call path collapses from N * ~1s of DB work to 1 * ~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.session across helpers inside get_configurations. It was rejected for good reasons (db.session is 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.

  • No Redis serialization of ProviderConfigurations (which would leak decrypted credentials across the trust boundary).
  • No SQLAlchemy event listeners or invalidation at the ~30 provider write sites spread across core/entities/provider_configuration.py, events/event_handlers/update_provider_when_message_created.py, commands/system.py, core/app/llm/quota.py, etc.
  • Cache dies at the end of the request/task — zero staleness risk for quota, credential, or load-balancing updates observed across requests.

Changes

  • core/provider_manager.py: add a module-level RecyclableContextVar[dict[str, ProviderConfigurations]] and memoize the result in get_configurations for the current scope, keyed by tenant_id. The cache is shared across ProviderManager instances because callers routinely spawn fresh ones via create_plugin_provider_manager (see core/app/llm/model_access.py, services/workflow_service.py, etc.).
  • extensions/ext_celery.py: call RecyclableContextVar.increment_thread_recycles() inside FlaskTask.__call__, matching the Flask before_request hook in app_factory.py. This ensures per-task caches reset cleanly between Celery tasks that land on recycled worker threads. Strictly speaking this also benefits any existing RecyclableContextVar used 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

  • Memoization returns the same object for repeated calls with the same tenant_id in one scope
  • Different tenant_ids are cached independently
  • Cache is shared across distinct ProviderManager instances in the same scope (the common caller pattern)
  • A new scope (increment_thread_recycles) starts empty
  • No cross-request leakage (cache is request/task-scoped only)
  • Existing provider_manager / celery_sqlcommenter unit tests still pass

#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.
@dosubot dosubot Bot added the size:M This PR changes 30-99 lines, ignoring generated files. label Apr 13, 2026
@github-actions
Copy link
Copy Markdown
Contributor

Pyrefly Diff

base → 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.
@YgorLeal YgorLeal force-pushed the perf/provider-manager-scope-cache branch from 49bebfd to f95146b Compare April 13, 2026 16:36
@github-actions
Copy link
Copy Markdown
Contributor

Pyrefly Diff

base → 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`

@github-actions
Copy link
Copy Markdown
Contributor

Pyrefly Diff

base → 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`

@github-actions
Copy link
Copy Markdown
Contributor

Pyrefly Diff

base → 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`

@youaodu
Copy link
Copy Markdown

youaodu commented Apr 21, 2026

Thanks! This saves me a lot of time.

@YgorLeal YgorLeal closed this by deleting the head repository May 19, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants