mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-25 14:06:18 +00:00
656db1223d
* fix(runtime): drain provider close across cancellation * test(runtime): cover redis provider teardown cancellation
103 lines
3.8 KiB
Python
103 lines
3.8 KiB
Python
"""Cache factory. Mirrors make_stream_bridge: config -> env fallback -> memory."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import hashlib
|
|
import logging
|
|
import os
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
from deerflow.config.app_config import AppConfig
|
|
from deerflow.runtime.checkpoint_cache.base import CACHE_FORMAT_VERSION, CheckpointHistoryCache
|
|
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
|
from deerflow.utils.file_io import await_drained
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_ENV_REDIS_URL = "DEER_FLOW_CHECKPOINT_CACHE_REDIS_URL"
|
|
|
|
|
|
def _resolve_redis_url(config: Any) -> str:
|
|
return config.redis_url or os.getenv(_ENV_REDIS_URL) or os.getenv("REDIS_URL") or "redis://localhost:6379/0"
|
|
|
|
|
|
def _stable_postgres_identity(postgres_url: str) -> str:
|
|
"""Credential-free database identity: host/port/database.
|
|
|
|
Hashing the raw URL would change the cache namespace on every credential
|
|
rotation (cold cache + orphaned keys until TTL) even though the database
|
|
— and thus every cached checkpoint history — is unchanged. Unparseable
|
|
URLs fall back to the raw string (still stable per deployment).
|
|
"""
|
|
if not postgres_url:
|
|
return ""
|
|
try:
|
|
from sqlalchemy.engine.url import make_url
|
|
|
|
parsed = make_url(postgres_url)
|
|
except Exception: # noqa: BLE001 - identity must never fail config load
|
|
return postgres_url
|
|
return f"{parsed.host or 'localhost'}:{parsed.port or 5432}/{parsed.database or ''}"
|
|
|
|
|
|
def checkpoint_cache_db_hash(db_config: Any) -> str:
|
|
"""Deployment-identity hash so two deployments sharing one Redis never collide."""
|
|
backend = getattr(db_config, "backend", "memory")
|
|
if backend == "postgres":
|
|
identity = f"postgres:{_stable_postgres_identity(getattr(db_config, 'postgres_url', ''))}:{getattr(db_config, 'postgres_schema', '')}"
|
|
elif backend == "sqlite":
|
|
identity = f"sqlite:{getattr(db_config, 'checkpointer_sqlite_path', '')}"
|
|
else:
|
|
identity = "memory"
|
|
return hashlib.sha256(identity.encode()).hexdigest()[:12]
|
|
|
|
|
|
def checkpoint_cache_key_prefix(app_config: AppConfig) -> str:
|
|
cache_config = app_config.database.checkpoint_cache
|
|
if cache_config.key_prefix:
|
|
return cache_config.key_prefix
|
|
return f"ckpt-hist:v{CACHE_FORMAT_VERSION}:{checkpoint_cache_db_hash(app_config.database)}"
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def make_checkpoint_cache(
|
|
app_config: AppConfig | None = None,
|
|
*,
|
|
serde: Any,
|
|
) -> AsyncIterator[CheckpointHistoryCache]:
|
|
"""Yield a history cache for the caller's lifetime.
|
|
|
|
``max_entries == 0`` disables the cache uniformly (both types) via a
|
|
disabled memory backend, so the wrapper never needs a None check.
|
|
"""
|
|
config = app_config.database.checkpoint_cache if app_config is not None else None
|
|
|
|
if config is None or config.type == "memory" or config.max_entries == 0:
|
|
max_entries = config.max_entries if config is not None else 128
|
|
cache = MemoryCheckpointHistoryCache(max_entries=max_entries)
|
|
logger.info("Checkpoint history cache initialised: memory (max_entries=%d)", max_entries)
|
|
try:
|
|
yield cache
|
|
finally:
|
|
await await_drained(cache.aclose())
|
|
return
|
|
|
|
if config.type == "redis":
|
|
from deerflow.runtime.checkpoint_cache.redis import RedisCheckpointHistoryCache
|
|
|
|
cache = RedisCheckpointHistoryCache(
|
|
_resolve_redis_url(config),
|
|
serde=serde,
|
|
ttl_seconds=config.ttl_seconds,
|
|
)
|
|
logger.info("Checkpoint history cache initialised: redis (ttl_seconds=%d)", config.ttl_seconds)
|
|
try:
|
|
yield cache
|
|
finally:
|
|
await await_drained(cache.aclose())
|
|
return
|
|
|
|
raise ValueError(f"Unknown checkpoint cache type: {config.type!r}")
|