Skip to content

Commit 49bebfd

Browse files
committed
perf(api): memoize ProviderManager.get_configurations within request/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.
1 parent 7862b61 commit 49bebfd

3 files changed

Lines changed: 129 additions & 0 deletions

File tree

api/core/provider_manager.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import contextlib
44
from collections import defaultdict
55
from collections.abc import Sequence
6+
from contextvars import ContextVar
67
from json import JSONDecodeError
78
from typing import TYPE_CHECKING, Any
89

@@ -20,6 +21,7 @@
2021
from sqlalchemy.orm import Session
2122

2223
from configs import dify_config
24+
from contexts.wrapper import RecyclableContextVar
2325
from core.entities.model_entities import DefaultModelEntity, DefaultModelProviderEntity
2426
from core.entities.provider_configuration import ProviderConfiguration, ProviderConfigurations, ProviderModelBundle
2527
from core.entities.provider_entities import (
@@ -60,6 +62,23 @@
6062

6163
_credentials_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any])
6264

65+
# Request/task-scoped memoization of `get_configurations`.
66+
#
67+
# Assembling a `ProviderConfigurations` runs six separate DB queries plus a
68+
# provider-factory load, costing ~1s per call in production. A single workflow
69+
# step (e.g. a Retrieval node inside an Iteration) can invoke this path dozens
70+
# of times through independently-created `ProviderManager` instances, so
71+
# per-instance caching is not enough — the cache must live on the execution
72+
# context instead.
73+
#
74+
# `RecyclableContextVar` is Dify's existing pattern for per-request/per-task
75+
# state. It is reset by `RecyclableContextVar.increment_thread_recycles()`,
76+
# which is already called in Flask `before_request` and is called per Celery
77+
# task in `FlaskTask.__call__` (see `extensions/ext_celery.py`).
78+
_provider_configurations_cache: RecyclableContextVar[dict[str, ProviderConfigurations]] = RecyclableContextVar(
79+
ContextVar("provider_configurations_cache")
80+
)
81+
6382

6483
class ProviderManager:
6584
"""
@@ -114,6 +133,16 @@ def get_configurations(self, tenant_id: str) -> ProviderConfigurations:
114133
:param tenant_id:
115134
:return:
116135
"""
136+
try:
137+
scope_cache = _provider_configurations_cache.get()
138+
except LookupError:
139+
scope_cache = {}
140+
_provider_configurations_cache.set(scope_cache)
141+
142+
cached = scope_cache.get(tenant_id)
143+
if cached is not None:
144+
return cached
145+
117146
# Get all provider records of the workspace
118147
provider_name_to_provider_records_dict = self._get_all_providers(tenant_id)
119148

@@ -273,6 +302,8 @@ def get_configurations(self, tenant_id: str) -> ProviderConfigurations:
273302

274303
provider_configurations[str(provider_id_entity)] = provider_configuration
275304

305+
scope_cache[tenant_id] = provider_configurations
306+
276307
# Return the encapsulated object
277308
return provider_configurations
278309

api/extensions/ext_celery.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,11 +75,15 @@ def get_celery_broker_transport_options() -> CelerySentinelTransportDict | dict[
7575
def init_app(app: DifyApp) -> Celery:
7676
class FlaskTask(Task):
7777
def __call__(self, *args: object, **kwargs: object) -> object:
78+
from contexts.wrapper import RecyclableContextVar
7879
from core.logging.context import init_request_context
7980

8081
with app.app_context():
8182
# Initialize logging context for this task (similar to before_request in Flask)
8283
init_request_context()
84+
# Reset RecyclableContextVar-backed per-task caches, matching the
85+
# Flask before_request hook in app_factory.py.
86+
RecyclableContextVar.increment_thread_recycles()
8387
return self.run(*args, **kwargs)
8488

8589
broker_transport_options = get_celery_broker_transport_options()

api/tests/unit_tests/core/test_provider_manager.py

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from contextlib import ExitStack
12
from types import SimpleNamespace
23
from unittest.mock import MagicMock, Mock, PropertyMock, patch
34

@@ -6,12 +7,20 @@
67
from graphon.model_runtime.entities.model_entities import ModelType
78
from pytest_mock import MockerFixture
89

10+
from contexts.wrapper import RecyclableContextVar
911
from core.entities.provider_entities import ModelSettings
1012
from core.provider_manager import ProviderManager
1113
from models.provider import LoadBalancingModelConfig, ProviderModelSetting, TenantDefaultModel
1214
from models.provider_ids import ModelProviderID
1315

1416

17+
@pytest.fixture(autouse=True)
18+
def _reset_provider_configurations_cache() -> None:
19+
# Advance the recycle counter so the per-scope memoization inside
20+
# ProviderManager.get_configurations starts fresh for every test.
21+
RecyclableContextVar.increment_thread_recycles()
22+
23+
1524
def _build_provider_manager(mocker: MockerFixture) -> ProviderManager:
1625
return ProviderManager(model_runtime=mocker.Mock())
1726

@@ -563,6 +572,91 @@ def test_get_all_provider_load_balancing_configs_returns_empty_when_cached_flag_
563572
mock_session_cls.assert_not_called()
564573

565574

575+
def _patch_get_configurations_dependencies(manager: ProviderManager):
576+
return (
577+
patch.object(manager, "_get_all_providers", return_value={}),
578+
patch.object(manager, "_init_trial_provider_records", return_value={}),
579+
patch.object(manager, "_get_all_provider_models", return_value={}),
580+
patch.object(manager, "_get_all_preferred_model_providers", return_value={}),
581+
patch.object(manager, "_get_all_provider_model_settings", return_value={}),
582+
patch.object(manager, "_get_all_provider_load_balancing_configs", return_value={}),
583+
patch.object(manager, "_get_all_provider_model_credentials", return_value={}),
584+
patch("core.provider_manager.ModelProviderFactory"),
585+
)
586+
587+
588+
def test_get_configurations_memoizes_result_within_scope(mocker: MockerFixture) -> None:
589+
RecyclableContextVar.increment_thread_recycles()
590+
manager = _build_provider_manager(mocker)
591+
592+
with ExitStack() as stack:
593+
for cm in _patch_get_configurations_dependencies(manager):
594+
stack.enter_context(cm)
595+
all_providers = mocker.patch.object(manager, "_get_all_providers", return_value={})
596+
597+
first = manager.get_configurations("tenant-id")
598+
second = manager.get_configurations("tenant-id")
599+
600+
assert first is second
601+
all_providers.assert_called_once_with("tenant-id")
602+
603+
604+
def test_get_configurations_memoization_is_per_tenant(mocker: MockerFixture) -> None:
605+
RecyclableContextVar.increment_thread_recycles()
606+
manager = _build_provider_manager(mocker)
607+
608+
with ExitStack() as stack:
609+
for cm in _patch_get_configurations_dependencies(manager):
610+
stack.enter_context(cm)
611+
all_providers = mocker.patch.object(manager, "_get_all_providers", return_value={})
612+
613+
first = manager.get_configurations("tenant-a")
614+
second = manager.get_configurations("tenant-b")
615+
first_again = manager.get_configurations("tenant-a")
616+
617+
assert first is not second
618+
assert first is first_again
619+
assert all_providers.call_count == 2
620+
621+
622+
def test_get_configurations_memoization_shared_across_manager_instances(mocker: MockerFixture) -> None:
623+
RecyclableContextVar.increment_thread_recycles()
624+
manager_one = _build_provider_manager(mocker)
625+
manager_two = _build_provider_manager(mocker)
626+
627+
with ExitStack() as stack:
628+
for cm in _patch_get_configurations_dependencies(manager_one):
629+
stack.enter_context(cm)
630+
for cm in _patch_get_configurations_dependencies(manager_two):
631+
stack.enter_context(cm)
632+
one_calls = mocker.patch.object(manager_one, "_get_all_providers", return_value={})
633+
two_calls = mocker.patch.object(manager_two, "_get_all_providers", return_value={})
634+
635+
first = manager_one.get_configurations("tenant-id")
636+
second = manager_two.get_configurations("tenant-id")
637+
638+
assert first is second
639+
one_calls.assert_called_once_with("tenant-id")
640+
two_calls.assert_not_called()
641+
642+
643+
def test_get_configurations_memoization_reset_on_new_scope(mocker: MockerFixture) -> None:
644+
RecyclableContextVar.increment_thread_recycles()
645+
manager = _build_provider_manager(mocker)
646+
647+
with ExitStack() as stack:
648+
for cm in _patch_get_configurations_dependencies(manager):
649+
stack.enter_context(cm)
650+
all_providers = mocker.patch.object(manager, "_get_all_providers", return_value={})
651+
652+
first = manager.get_configurations("tenant-id")
653+
RecyclableContextVar.increment_thread_recycles()
654+
second = manager.get_configurations("tenant-id")
655+
656+
assert first is not second
657+
assert all_providers.call_count == 2
658+
659+
566660
def test_get_all_provider_load_balancing_configs_populates_cache_and_groups_configs() -> None:
567661
session = Mock()
568662
openai_config = SimpleNamespace(provider_name="openai")

0 commit comments

Comments
 (0)