-
Notifications
You must be signed in to change notification settings - Fork 167
fix: caller-controlled scope IDs grow Runtime caches without bound #1325
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: master
Are you sure you want to change the base?
Changes from 1 commit
7041a64
a782e4d
136ac7e
8f287ed
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,136 @@ | ||
| """Bounded lifecycle for scope-local Runtime resources.""" | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import asyncio | ||
| from collections import OrderedDict | ||
| from collections.abc import Callable, Iterator | ||
| from contextlib import contextmanager, suppress | ||
| from dataclasses import dataclass, field | ||
|
|
||
| DEFAULT_SCOPE_CACHE_SIZE = 128 | ||
|
|
||
| ScopeEvictor = Callable[[str], None] | ||
| ScopeCacheObserver = Callable[[int, int], None] | ||
|
|
||
|
|
||
| @dataclass(frozen=True, slots=True) | ||
| class ScopeCacheCounts: | ||
| """Low-cardinality snapshot of cached and currently active scopes.""" | ||
|
|
||
| cached: int | ||
| active: int | ||
|
|
||
|
|
||
| @dataclass(slots=True) | ||
| class _ScopeEntry: | ||
| lock: asyncio.Lock = field(default_factory=asyncio.Lock) | ||
| leases: int = 0 | ||
|
|
||
|
|
||
| class ScopeCache: | ||
| """Retain a bounded LRU of inactive scopes without evicting in-flight work.""" | ||
|
|
||
| def __init__( | ||
| self, | ||
| capacity: int, | ||
| *, | ||
| evictor: ScopeEvictor | None = None, | ||
| observer: ScopeCacheObserver | None = None, | ||
| ) -> None: | ||
| if capacity < 1: | ||
| raise ValueError("scope cache capacity must be positive") # noqa: TRY003 | ||
| self.capacity = capacity | ||
| self._evictor = evictor | ||
| self._observer = observer | ||
| self._entries: OrderedDict[str, _ScopeEntry] = OrderedDict() | ||
| self._observe() | ||
|
|
||
| @contextmanager | ||
| def lease(self, scope_id: str, /) -> Iterator[None]: | ||
| """Keep one scope and its serialization lock alive for an operation.""" | ||
|
|
||
| entry = self._entries.get(scope_id) | ||
| if entry is None: | ||
| self._make_room() | ||
| entry = _ScopeEntry() | ||
| self._entries[scope_id] = entry | ||
| entry.leases += 1 | ||
| self._entries.move_to_end(scope_id) | ||
| self._observe() | ||
| try: | ||
| yield | ||
| finally: | ||
| current = self._entries.get(scope_id) | ||
| if current is not entry or entry.leases < 1: | ||
| raise RuntimeError("scope cache lease invariant violated") # noqa: TRY003 | ||
| entry.leases -= 1 | ||
| self._entries.move_to_end(scope_id) | ||
| self._trim() | ||
| self._observe() | ||
|
|
||
| def lock(self, scope_id: str, /) -> asyncio.Lock: | ||
| """Return the serialization lock protected by the caller's current lease.""" | ||
|
|
||
| entry = self._entries.get(scope_id) | ||
| if entry is None or entry.leases < 1: | ||
| raise RuntimeError("scope cache entry is not leased") # noqa: TRY003 | ||
| return entry.lock | ||
|
|
||
| @property | ||
| def counts(self) -> ScopeCacheCounts: | ||
| """Return cached and active scope counts without scope identifiers.""" | ||
|
|
||
| return ScopeCacheCounts( | ||
| cached=len(self._entries), | ||
| active=sum(entry.leases > 0 for entry in self._entries.values()), | ||
| ) | ||
|
thunguo marked this conversation as resolved.
Outdated
|
||
|
|
||
| def clear(self) -> None: | ||
| """Evict all entries after the Runtime has drained its operations.""" | ||
|
|
||
| if any(entry.leases > 0 for entry in self._entries.values()): | ||
| raise RuntimeError("cannot clear active scope cache entries") # noqa: TRY003 | ||
| for scope_id in tuple(self._entries): | ||
| self._evict(scope_id) | ||
| self._observe() | ||
|
|
||
| def _make_room(self) -> None: | ||
| while len(self._entries) >= self.capacity: | ||
| scope_id = self._oldest_inactive_scope() | ||
| if scope_id is None: | ||
| return | ||
| self._evict(scope_id) | ||
|
|
||
| def _trim(self) -> None: | ||
| while len(self._entries) > self.capacity: | ||
| scope_id = self._oldest_inactive_scope() | ||
| if scope_id is None: | ||
| return | ||
| self._evict(scope_id) | ||
|
|
||
| def _oldest_inactive_scope(self) -> str | None: | ||
| return next((scope_id for scope_id, entry in self._entries.items() if entry.leases == 0), None) | ||
|
|
||
| def _evict(self, scope_id: str) -> None: | ||
| entry = self._entries.pop(scope_id) | ||
| if entry.leases > 0: | ||
| raise RuntimeError("cannot evict an active scope cache entry") # noqa: TRY003 | ||
| if self._evictor is not None: | ||
| self._evictor(scope_id) | ||
|
|
||
| def _observe(self) -> None: | ||
| if self._observer is None: | ||
| return | ||
| counts = self.counts | ||
| with suppress(Exception): | ||
| self._observer(counts.cached, counts.active) | ||
|
|
||
|
|
||
| __all__ = [ | ||
| "DEFAULT_SCOPE_CACHE_SIZE", | ||
| "ScopeCache", | ||
| "ScopeCacheCounts", | ||
| "ScopeCacheObserver", | ||
| "ScopeEvictor", | ||
| ] | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -9,7 +9,7 @@ | |
| from datetime import UTC, datetime | ||
| from pathlib import Path | ||
| from time import perf_counter | ||
| from typing import TYPE_CHECKING | ||
| from typing import TYPE_CHECKING, Any | ||
|
|
||
| from powercontext._logging import log_safely | ||
| from powercontext.artifacts import ArtifactRef | ||
|
|
@@ -48,6 +48,12 @@ | |
| from powercontext.builtin.inference.usage import bind_usage_reporter | ||
| from powercontext.builtin.review.generation import GeneratedCandidateResult, ReviewedGenerationService | ||
| from powercontext.builtin.review.service import ReviewService | ||
| from powercontext.builtin.runtime._scope_cache import ( | ||
| DEFAULT_SCOPE_CACHE_SIZE, | ||
| ScopeCache, | ||
| ScopeCacheObserver, | ||
| ScopeEvictor, | ||
| ) | ||
| from powercontext.builtin.runtime.errors import InvalidRuntimeRequestError | ||
| from powercontext.builtin.runtime.models import ( | ||
| ApproveArtifactCandidateRequest, | ||
|
|
@@ -185,7 +191,7 @@ def __init__(self, runtime: BuiltinRuntime, scope_id: str) -> None: | |
| self.scope_id = validate_scope_id(scope_id) | ||
|
|
||
| async def overview(self, *, period: StatisticsPeriod = StatisticsPeriod.THIRTY_DAYS) -> Statistics: | ||
| async with self._runtime._operation(): | ||
| async with self._runtime._scope_operation(self.scope_id): | ||
| return await self._runtime._statistics(self.scope_id).overview(period, self._runtime._clock()) | ||
|
|
||
| async def record_model_usage( | ||
|
|
@@ -196,12 +202,13 @@ async def record_model_usage( | |
| /, | ||
| ) -> None: | ||
| try: | ||
| await self._runtime._statistics(self.scope_id).record( | ||
| purpose, | ||
| operation, | ||
| usage, | ||
| self._runtime._clock().astimezone(UTC).date(), | ||
| ) | ||
| async with self._runtime._scope_operation(self.scope_id): | ||
| await self._runtime._statistics(self.scope_id).record( | ||
| purpose, | ||
| operation, | ||
| usage, | ||
| self._runtime._clock().astimezone(UTC).date(), | ||
| ) | ||
| except Exception as error: | ||
| log_safely( | ||
| logger, | ||
|
|
@@ -219,10 +226,11 @@ async def record_model_usage( | |
|
|
||
| async def record_recall(self, measurement: RecallTokenMeasurement, /) -> None: | ||
| try: | ||
| await self._runtime._statistics(self.scope_id).record_recall( | ||
| measurement, | ||
| self._runtime._clock().astimezone(UTC).date(), | ||
| ) | ||
| async with self._runtime._scope_operation(self.scope_id): | ||
| await self._runtime._statistics(self.scope_id).record_recall( | ||
| measurement, | ||
| self._runtime._clock().astimezone(UTC).date(), | ||
| ) | ||
| except Exception as error: | ||
| log_safely( | ||
| logger, | ||
|
|
@@ -256,6 +264,10 @@ def __init__(self, runtime: BuiltinRuntime, scope_id: str) -> None: | |
| self.scope_id = validate_scope_id(scope_id) | ||
|
|
||
| async def prepare(self, request: PrepareContextRequest, /) -> PreparedContext: | ||
| async with self._runtime._scope_operation(self.scope_id): | ||
| return await self._prepare(request) | ||
|
|
||
| async def _prepare(self, request: PrepareContextRequest, /) -> PreparedContext: | ||
| builder = PreparedContextBuilder() | ||
| async with ( | ||
| self._runtime._context(self.scope_id, embedding_purpose=ModelUsagePurpose.MEMORY_RECALL) as context, | ||
|
|
@@ -903,6 +915,9 @@ def __init__( | |
| provider: PowerContextProvider[BuiltinSources, BuiltinArtifacts, BuiltinTriggers], | ||
| capabilities: RuntimeCapabilities, | ||
| source_window_limit: int = 100, | ||
| scope_cache_size: int = DEFAULT_SCOPE_CACHE_SIZE, | ||
| scope_evictor: ScopeEvictor | None = None, | ||
| scope_cache_observer: ScopeCacheObserver | None = None, | ||
| scope_ids: ScopeIds | None = None, | ||
| review_service: ReviewServiceFactory | None = None, | ||
| generation_service: GenerationServiceFactory | None = None, | ||
|
|
@@ -917,6 +932,8 @@ def __init__( | |
| ) -> None: | ||
| if source_window_limit < 1: | ||
| raise _RuntimeConfigurationError("source_window_limit") | ||
| if scope_cache_size < 1: | ||
| raise _RuntimeConfigurationError("scope_cache_size") | ||
| self._provider = provider | ||
| self._capabilities = capabilities | ||
| self._review_service = review_service | ||
|
|
@@ -930,11 +947,16 @@ def __init__( | |
| self._readiness = RuntimeReadinessChecks() if readiness is None else readiness | ||
| self._clock = _utc_now if clock is None else clock | ||
| self.source_window_limit = source_window_limit | ||
| self._locks: dict[str, asyncio.Lock] = {} | ||
| self._scope_cache = ScopeCache( | ||
| scope_cache_size, | ||
| evictor=scope_evictor, | ||
| observer=scope_cache_observer, | ||
| ) | ||
| self._processor_lock = asyncio.Lock() | ||
| self._close_lock = asyncio.Lock() | ||
| self._lifecycle = asyncio.Condition() | ||
| self._active_operations = 0 | ||
| self._operation_depths: dict[asyncio.Task[Any], int] = {} | ||
| self._closing = False | ||
| self._closed = False | ||
| self._scheduler: AsyncIOScheduler | None = None | ||
|
|
@@ -1063,21 +1085,40 @@ async def close(self) -> None: | |
| unregister_processor(self._scheduler_runtime_key) | ||
| self._scheduler_runtime_key = None | ||
| self._scheduler = None | ||
| self._scope_cache.clear() | ||
| self._closed = True | ||
|
|
||
| @asynccontextmanager | ||
| async def _operation(self) -> AsyncIterator[None]: | ||
| task = asyncio.current_task() | ||
| if task is None: | ||
| raise RuntimeError("Runtime operations require an asyncio Task") # noqa: TRY003 | ||
| async with self._lifecycle: | ||
| if self._closing or self._closed: | ||
| raise _RuntimeStateError("closed") | ||
| self._active_operations += 1 | ||
| depth = self._operation_depths.get(task, 0) | ||
| if depth == 0: | ||
| if self._closing or self._closed: | ||
| raise _RuntimeStateError("closed") | ||
| self._active_operations += 1 | ||
| self._operation_depths[task] = depth + 1 | ||
| try: | ||
| yield | ||
| finally: | ||
| async with self._lifecycle: | ||
| self._active_operations -= 1 | ||
| if self._active_operations == 0: | ||
| self._lifecycle.notify_all() | ||
| depth = self._operation_depths[task] - 1 | ||
| if depth > 0: | ||
| self._operation_depths[task] = depth | ||
| else: | ||
| del self._operation_depths[task] | ||
| self._active_operations -= 1 | ||
| if self._active_operations == 0: | ||
| self._lifecycle.notify_all() | ||
|
|
||
| @asynccontextmanager | ||
| async def _scope_operation(self, scope_id: str) -> AsyncIterator[None]: | ||
| scope = validate_scope_id(scope_id) | ||
| async with self._operation(): | ||
| with self._scope_cache.lease(scope): | ||
| yield | ||
|
|
||
| @asynccontextmanager | ||
| async def _scoped_operation( | ||
|
|
@@ -1088,7 +1129,7 @@ async def _scoped_operation( | |
| embedding_purpose: ModelUsagePurpose | None = None, | ||
| ) -> AsyncIterator[None]: | ||
| scope = validate_scope_id(scope_id) | ||
| async with self._operation(): | ||
| async with self._scope_operation(scope): | ||
| with bind_usage_reporter( | ||
| self.statistics.for_scope(scope).record_model_usage, | ||
| generation_purpose=generation_purpose, | ||
|
|
@@ -1112,7 +1153,7 @@ async def _context( | |
| yield await self._provider.get(validate_scope_id(scope_id)) | ||
|
|
||
| def _lock(self, scope_id: str) -> asyncio.Lock: | ||
| return self._locks.setdefault(validate_scope_id(scope_id), asyncio.Lock()) | ||
| return self._scope_cache.lock(validate_scope_id(scope_id)) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [P0] 删除 这里把 scope lock 存储切到了 |
||
|
|
||
| def _review(self, scope_id: str) -> ReviewService: | ||
| if self._review_service is None: | ||
|
|
||
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.
[P1] 两个新增 Python 文件缺少许可证头
exact Head 的 License Check 已明确报告
invalid: 2,缺失文件是本文件和tests/builtin/runtime/test_scope_cache.py,因此当前强制门禁退出 1。请按仓库现有 Python 文件格式在两个文件顶部补齐 Apache 2.0 许可证头,并重新运行 License Check。