Add Redis-backed concurrency gate

This commit is contained in:
吉鑫升 committed 2026-06-30 11:25:44 +08:00
1 parent b0db8cc3d8
commit b4206d6421
12 files changed
+654 -46

No files matched your search

@@ -853,8 +853,9 @@ async def _concurrency_sampler_loop(app: FastAPI) -> None:
logger.info("Concurrency sampler disabled by CONCURRENCY_SAMPLE_ENABLED") logger.info("Concurrency sampler disabled by CONCURRENCY_SAMPLE_ENABLED")
return return
store = getattr(app.state, "concurrency_sample_store", None) store = getattr(app.state, "concurrency_sample_store", None)
run_mgr = getattr(app.state, "run_manager", None) run_mgr = getattr(app.state, "run_manager", None)
concurrency_gate = getattr(app.state, "concurrency_gate", None)
if store is None or run_mgr is None: if store is None or run_mgr is None:
logger.info("Concurrency sampler not started (no store / run manager — memory backend?)") logger.info("Concurrency sampler not started (no store / run manager — memory backend?)")
return return
@@ -875,11 +876,16 @@ async def _concurrency_sampler_loop(app: FastAPI) -> None:
try: try:
# Runtime toggle (admin can turn sampling off without a restart). # Runtime toggle (admin can turn sampling off without a restart).
# mtime-cached read, so this is cheap every tick. # mtime-cached read, so this is cheap every tick.
from deerflow.config.system_settings import get_concurrency_monitor_settings from deerflow.config.system_settings import get_concurrency_monitor_settings
if get_concurrency_monitor_settings().enabled: if get_concurrency_monitor_settings().enabled:
active = await run_mgr.list_active() active_count = None
await store.record(len(active)) if concurrency_gate is not None:
active_count = await concurrency_gate.count_active()
if active_count is None:
active = await run_mgr.list_active()
active_count = len(active)
await store.record(active_count)
ticks += 1 ticks += 1
if ticks % purge_every_ticks == 0: if ticks % purge_every_ticks == 0:
from datetime import UTC, datetime, timedelta from datetime import UTC, datetime, timedelta
@@ -367,8 +367,11 @@ async def langgraph_runtime(app: FastAPI) -> AsyncGenerator[None, None]:
run_events_config = getattr(config, "run_events", None) run_events_config = getattr(config, "run_events", None)
app.state.run_event_store = make_run_event_store(run_events_config) app.state.run_event_store = make_run_event_store(run_events_config)
# RunManager with store backing for persistence # RunManager with store backing for persistence
app.state.run_manager = RunManager(store=app.state.run_store) app.state.run_manager = RunManager(store=app.state.run_store)
from deerflow.runtime import make_redis_concurrency_gate
app.state.concurrency_gate = make_redis_concurrency_gate(config)
try: try:
yield yield
@@ -382,7 +385,10 @@ async def langgraph_runtime(app: FastAPI) -> AsyncGenerator[None, None]:
await queue.stop() await queue.stop()
set_auto_ingest_queue(None) set_auto_ingest_queue(None)
await close_engine() await close_engine()
gate = getattr(app.state, "concurrency_gate", None)
if gate is not None:
await gate.aclose()
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -408,7 +414,12 @@ get_run_manager: Callable[[Request], RunManager] = _require("run_manager", "Run
get_checkpointer: Callable[[Request], Checkpointer] = _require("checkpointer", "Checkpointer") get_checkpointer: Callable[[Request], Checkpointer] = _require("checkpointer", "Checkpointer")
get_run_event_store: Callable[[Request], RunEventStore] = _require("run_event_store", "Run event store") get_run_event_store: Callable[[Request], RunEventStore] = _require("run_event_store", "Run event store")
get_feedback_repo: Callable[[Request], FeedbackRepository] = _require("feedback_repo", "Feedback") get_feedback_repo: Callable[[Request], FeedbackRepository] = _require("feedback_repo", "Feedback")
get_run_store: Callable[[Request], RunStore] = _require("run_store", "Run store") get_run_store: Callable[[Request], RunStore] = _require("run_store", "Run store")
def get_concurrency_gate(request: Request):
"""Return the optional cross-worker concurrency gate."""
return getattr(request.app.state, "concurrency_gate", None)
def get_store(request: Request): def get_store(request: Request):
@@ -22,7 +22,7 @@ from fastapi import APIRouter, Query, Request
from pydantic import BaseModel from pydantic import BaseModel
from sqlalchemy import select from sqlalchemy import select
from app.gateway.deps import get_concurrency_sample_store, get_run_manager from app.gateway.deps import get_concurrency_gate, get_concurrency_sample_store, get_run_manager
from app.gateway.routers.admin_users import _require_admin from app.gateway.routers.admin_users import _require_admin
from deerflow.config.system_settings import load_system_settings, save_system_settings from deerflow.config.system_settings import load_system_settings, save_system_settings
from deerflow.persistence.agents.model import AgentRow from deerflow.persistence.agents.model import AgentRow
@@ -149,6 +149,15 @@ def _elapsed_seconds(created_at: str, now: datetime) -> int:
return max(0, int(delta)) return max(0, int(delta))
def _elapsed_seconds_from_epoch(created_at: float, now: datetime) -> int:
try:
if created_at <= 0:
return 0
return max(0, int(now.timestamp() - float(created_at)))
except Exception:
return 0
def _kind_for(metadata: dict | None) -> str: def _kind_for(metadata: dict | None) -> str:
"""Coarse human label for the run's origin, from its thread metadata.""" """Coarse human label for the run's origin, from its thread metadata."""
meta = metadata or {} meta = metadata or {}
@@ -218,14 +227,46 @@ async def list_active_runs(request: Request) -> ActiveRunsResponse:
"""List every in-flight run with who/what is driving it (admin only).""" """List every in-flight run with who/what is driving it (admin only)."""
await _require_admin(request) await _require_admin(request)
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
records = await run_mgr.list_active() gate = get_concurrency_gate(request)
leases = await gate.list_active() if gate is not None else None
records = [] if leases is not None else await run_mgr.list_active()
now = datetime.now(UTC) now = datetime.now(UTC)
maps = await _enrich(records) enrich_source = leases if leases is not None else records
maps = await _enrich(enrich_source)
thread_meta = maps["thread_meta"] thread_meta = maps["thread_meta"]
user_emails = maps["user_emails"] user_emails = maps["user_emails"]
agent_names = maps["agent_names"] agent_names = maps["agent_names"]
items: list[ActiveRunItem] = [] items: list[ActiveRunItem] = []
if leases is not None:
for r in leases:
meta_row = thread_meta.get(r.thread_id)
user_id = r.user_id or (meta_row.user_id if meta_row else "") or ""
metadata = (meta_row.metadata_json if meta_row else None) or {}
assistant_id = r.assistant_id or ""
if assistant_id in _DEFAULT_ASSISTANT_IDS:
agent_name = "榛樿鍔╂墜"
else:
agent_name = agent_names.get(assistant_id) or assistant_id
items.append(
ActiveRunItem(
run_id=r.run_id,
thread_id=r.thread_id,
status="running",
created_at=datetime.fromtimestamp(r.created_at, tz=UTC).isoformat() if r.created_at else "",
elapsed_seconds=_elapsed_seconds_from_epoch(r.created_at, now),
user_id=user_id,
user_name=user_emails.get(user_id) or user_id or "鏈煡鐢ㄦ埛",
agent_id=assistant_id,
agent_name=agent_name,
model_name=r.model_name,
thread_title=(meta_row.display_name if meta_row else "") or "",
kind=_kind_for(metadata),
multitask_strategy="reject",
)
)
return ActiveRunsResponse(count=len(items), server_time=now.isoformat(), runs=items)
for r in records: for r in records:
meta_row = thread_meta.get(r.thread_id) meta_row = thread_meta.get(r.thread_id)
user_id = (meta_row.user_id if meta_row else "") or "" user_id = (meta_row.user_id if meta_row else "") or ""
@@ -12,13 +12,14 @@ import json
import logging import logging
import re import re
import time import time
import uuid
from collections.abc import Mapping from collections.abc import Mapping
from typing import Any from typing import Any
from fastapi import HTTPException, Request from fastapi import HTTPException, Request
from langchain_core.messages import HumanMessage from langchain_core.messages import HumanMessage
from app.gateway.deps import get_run_context, get_run_manager, get_stream_bridge from app.gateway.deps import get_concurrency_gate, get_run_context, get_run_manager, get_stream_bridge
from app.gateway.utils import sanitize_log_param from app.gateway.utils import sanitize_log_param
from deerflow.config import get_app_config from deerflow.config import get_app_config
from deerflow.config.agents_config import load_agent_config, validate_agent_id from deerflow.config.agents_config import load_agent_config, validate_agent_id
@@ -41,6 +42,45 @@ from deerflow.runtime.user_context import get_effective_user_id
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
async def _run_agent_with_concurrency_lease(
*,
concurrency_gate: Any | None,
lease_acquired: bool,
lease_metadata: dict[str, str],
heartbeat_interval_seconds: int,
run_agent_kwargs: dict[str, Any],
) -> None:
"""Run the agent while refreshing and finally releasing an optional lease."""
heartbeat_task: asyncio.Task | None = None
async def _heartbeat_loop() -> None:
if concurrency_gate is None or not lease_acquired:
return
interval = max(5, int(heartbeat_interval_seconds or 20))
while True:
await asyncio.sleep(interval)
await concurrency_gate.heartbeat(**lease_metadata)
if concurrency_gate is not None and lease_acquired:
heartbeat_task = asyncio.create_task(
_heartbeat_loop(),
name=f"concurrency-heartbeat:{lease_metadata.get('run_id')}",
)
try:
await run_agent(**run_agent_kwargs)
finally:
if heartbeat_task is not None:
heartbeat_task.cancel()
try:
await heartbeat_task
except asyncio.CancelledError:
pass
if concurrency_gate is not None and lease_acquired:
await concurrency_gate.release(**lease_metadata)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# SSE formatting # SSE formatting
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -737,6 +777,7 @@ async def start_run(
bridge = get_stream_bridge(request) bridge = get_stream_bridge(request)
run_mgr = get_run_manager(request) run_mgr = get_run_manager(request)
run_ctx = get_run_context(request) run_ctx = get_run_context(request)
concurrency_gate = get_concurrency_gate(request)
disconnect = DisconnectMode.cancel if body.on_disconnect == "cancel" else DisconnectMode.continue_ disconnect = DisconnectMode.cancel if body.on_disconnect == "cancel" else DisconnectMode.continue_
@@ -754,11 +795,46 @@ async def start_run(
effective_user_id = get_effective_user_id() effective_user_id = get_effective_user_id()
model_concurrency_limit = _get_model_concurrency_limit(effective_model_name) model_concurrency_limit = _get_model_concurrency_limit(effective_model_name)
user_model_concurrency_limit = _get_user_model_concurrency_limit(effective_model_name) user_model_concurrency_limit = _get_user_model_concurrency_limit(effective_model_name)
run_id = str(uuid.uuid4())
run_metadata = dict(body.metadata or {}) run_metadata = dict(body.metadata or {})
run_metadata["model_name"] = effective_model_name run_metadata["model_name"] = effective_model_name
run_metadata["user_id"] = effective_user_id run_metadata["user_id"] = effective_user_id
redis_lease_acquired = False
redis_gate_degraded = True
lease_metadata = {
"run_id": run_id,
"thread_id": thread_id,
"model_name": effective_model_name,
"user_id": effective_user_id,
}
if concurrency_gate is not None:
acquire = await concurrency_gate.try_acquire(
run_id=run_id,
thread_id=thread_id,
model_name=effective_model_name,
user_id=effective_user_id,
assistant_id=body.assistant_id or "",
model_limit=model_concurrency_limit,
user_model_limit=user_model_concurrency_limit,
reject_thread=body.multitask_strategy == "reject",
)
redis_lease_acquired = acquire.acquired
redis_gate_degraded = acquire.degraded
if acquire.rejected:
if acquire.reason == "thread":
raise HTTPException(
status_code=409,
detail="Current thread already has an active run. Please wait or stop it before sending another message.",
)
if acquire.reason == "user_model":
raise HTTPException(
status_code=429,
detail="This account has reached the concurrency limit for the selected model. Please try again later.",
)
raise HTTPException(status_code=429, detail="The selected model is busy. Please try again later.")
try: try:
record = await run_mgr.create_or_reject( record = await run_mgr.create_or_reject(
thread_id, thread_id,
@@ -767,24 +843,35 @@ async def start_run(
metadata=run_metadata, metadata=run_metadata,
kwargs={"input": body.input, "config": body.config}, kwargs={"input": body.input, "config": body.config},
multitask_strategy=body.multitask_strategy, multitask_strategy=body.multitask_strategy,
model_concurrency_limit=(effective_model_name, model_concurrency_limit) if model_concurrency_limit is not None else None, model_concurrency_limit=(effective_model_name, model_concurrency_limit) if redis_gate_degraded and model_concurrency_limit is not None else None,
user_model_concurrency_limit=(effective_model_name, effective_user_id, user_model_concurrency_limit) if user_model_concurrency_limit is not None else None, user_model_concurrency_limit=(effective_model_name, effective_user_id, user_model_concurrency_limit) if redis_gate_degraded and user_model_concurrency_limit is not None else None,
run_id=run_id,
) )
except ModelConcurrencyLimitError as exc: except ModelConcurrencyLimitError as exc:
if concurrency_gate is not None and redis_lease_acquired:
await concurrency_gate.release(**lease_metadata)
if getattr(exc, "scope", "") == "user_model": if getattr(exc, "scope", "") == "user_model":
detail = "当前账号对该模型的并发调用已达到上限,请等待已有回答结束后再试。" detail = "This account has reached the concurrency limit for the selected model. Please try again later."
else: else:
detail = "当前使用人数较多,请稍等一下再试。" detail = "The selected model is busy. Please try again later."
raise HTTPException(status_code=429, detail=detail) from exc raise HTTPException(status_code=429, detail=detail) from exc
except ConflictError as exc: except ConflictError as exc:
if concurrency_gate is not None and redis_lease_acquired:
await concurrency_gate.release(**lease_metadata)
# 同一会话上一轮回答还没结束,又提交了新问题 → 拒绝。给用户中文提示, # 同一会话上一轮回答还没结束,又提交了新问题 → 拒绝。给用户中文提示,
# 而不是把内部英文「Thread ... already has an active run」直接抛到前端。 # 而不是把内部英文「Thread ... already has an active run」直接抛到前端。
raise HTTPException( raise HTTPException(
status_code=409, status_code=409,
detail="当前对话还在生成回答,请等待本轮回复结束,或先点「停止」再发送新的问题。", detail="Current thread already has an active run. Please wait or stop it before sending another message.",
) from exc ) from exc
except UnsupportedStrategyError as exc: except UnsupportedStrategyError as exc:
if concurrency_gate is not None and redis_lease_acquired:
await concurrency_gate.release(**lease_metadata)
raise HTTPException(status_code=501, detail=str(exc)) from exc raise HTTPException(status_code=501, detail=str(exc)) from exc
except Exception:
if concurrency_gate is not None and redis_lease_acquired:
await concurrency_gate.release(**lease_metadata)
raise
t_post_create = time.perf_counter() t_post_create = time.perf_counter()
# Upsert thread metadata so the thread appears in /threads/search, # Upsert thread metadata so the thread appears in /threads/search,
@@ -822,7 +909,12 @@ async def start_run(
logger.warning("Failed to upsert thread_meta for %s (non-fatal)", sanitize_log_param(thread_id)) logger.warning("Failed to upsert thread_meta for %s (non-fatal)", sanitize_log_param(thread_id))
t_post_thread_upsert = time.perf_counter() t_post_thread_upsert = time.perf_counter()
agent_factory = resolve_agent_factory(body.assistant_id) try:
agent_factory = resolve_agent_factory(body.assistant_id)
except Exception:
if concurrency_gate is not None and redis_lease_acquired:
await concurrency_gate.release(**lease_metadata)
raise
graph_input = normalize_input(body.input) graph_input = normalize_input(body.input)
# Admin-configurable prompt prefix: when the caller opted in via # Admin-configurable prompt prefix: when the caller opted in via
@@ -860,24 +952,38 @@ async def start_run(
stream_modes = normalize_stream_modes(body.stream_mode) stream_modes = normalize_stream_modes(body.stream_mode)
t_pre_task = time.perf_counter() t_pre_task = time.perf_counter()
task = asyncio.create_task( heartbeat_interval_seconds = 20
run_agent( if concurrency_gate is not None:
bridge, heartbeat_interval_seconds = int(
run_mgr, getattr(getattr(concurrency_gate, "config", None), "heartbeat_interval_seconds", 20) or 20
record,
ctx=run_ctx,
agent_factory=agent_factory,
graph_input=graph_input,
config=config,
stream_modes=stream_modes,
stream_subgraphs=body.stream_subgraphs,
interrupt_before=body.interrupt_before,
interrupt_after=body.interrupt_after,
# 透传 SDK ``Command``:resume / update / goto。Pregel 会据此从最近
# checkpoint 处恢复,而不是把 graph_input 当新一轮输入跑。
command=getattr(body, "command", None),
) )
) try:
task = asyncio.create_task(
_run_agent_with_concurrency_lease(
concurrency_gate=concurrency_gate,
lease_acquired=redis_lease_acquired,
lease_metadata=lease_metadata,
heartbeat_interval_seconds=heartbeat_interval_seconds,
run_agent_kwargs={
"bridge": bridge,
"run_manager": run_mgr,
"record": record,
"ctx": run_ctx,
"agent_factory": agent_factory,
"graph_input": graph_input,
"config": config,
"stream_modes": stream_modes,
"stream_subgraphs": body.stream_subgraphs,
"interrupt_before": body.interrupt_before,
"interrupt_after": body.interrupt_after,
"command": getattr(body, "command", None),
},
)
)
except Exception:
if concurrency_gate is not None and redis_lease_acquired:
await concurrency_gate.release(**lease_metadata)
raise
record.task = task record.task = task
# Title sync is handled by worker.py's finally block which reads the # Title sync is handled by worker.py's finally block which reads the
@@ -9,6 +9,26 @@ log_level: info
gateway: gateway:
port: 8001 port: 8001
# Optional Redis-backed global concurrency coordination.
# Leave disabled to keep the existing in-process behavior.
# When enabled, Redis operations are short-timeout and fail-open by default:
# if Redis is slow or unavailable, user Q&A continues with local-only limits.
concurrency:
redis:
enabled: false
url: redis://127.0.0.1:6379/0
password: null
namespace: deerflow:concurrency
socket_connect_timeout_ms: 80
socket_timeout_ms: 80
acquire_timeout_ms: 120
operation_timeout_ms: 80
circuit_breaker_failures: 3
circuit_breaker_cooldown_seconds: 30
lease_ttl_seconds: 90
heartbeat_interval_seconds: 20
fail_open: true
models: models:
- name: deepseek-chat - name: deepseek-chat
display_name: DeepSeek Chat display_name: DeepSeek Chat
@@ -11,8 +11,9 @@ from pydantic import BaseModel, ConfigDict, Field
from deerflow.config.acp_config import ACPAgentConfig, load_acp_config_from_dict from deerflow.config.acp_config import ACPAgentConfig, load_acp_config_from_dict
from deerflow.config.agents_api_config import AgentsApiConfig, load_agents_api_config_from_dict from deerflow.config.agents_api_config import AgentsApiConfig, load_agents_api_config_from_dict
from deerflow.config.browser_context_config import BrowserContextConfig from deerflow.config.browser_context_config import BrowserContextConfig
from deerflow.config.checkpointer_config import CheckpointerConfig, load_checkpointer_config_from_dict from deerflow.config.checkpointer_config import CheckpointerConfig, load_checkpointer_config_from_dict
from deerflow.config.database_config import DatabaseConfig from deerflow.config.concurrency_config import ConcurrencyConfig
from deerflow.config.database_config import DatabaseConfig
from deerflow.config.extensions_config import ExtensionsConfig from deerflow.config.extensions_config import ExtensionsConfig
from deerflow.config.gateway_config import GatewayConfig from deerflow.config.gateway_config import GatewayConfig
from deerflow.config.guardrails_config import GuardrailsConfig, load_guardrails_config_from_dict from deerflow.config.guardrails_config import GuardrailsConfig, load_guardrails_config_from_dict
@@ -137,8 +138,9 @@ class AppConfig(BaseModel):
gateway: GatewayConfig = Field(default_factory=GatewayConfig, description="Gateway FastAPI listening port (used by intra-process loopback)") gateway: GatewayConfig = Field(default_factory=GatewayConfig, description="Gateway FastAPI listening port (used by intra-process loopback)")
run_events: RunEventsConfig = Field(default_factory=RunEventsConfig, description="Run event storage configuration") run_events: RunEventsConfig = Field(default_factory=RunEventsConfig, description="Run event storage configuration")
checkpointer: CheckpointerConfig | None = Field(default=None, description="Checkpointer configuration") checkpointer: CheckpointerConfig | None = Field(default=None, description="Checkpointer configuration")
stream_bridge: StreamBridgeConfig | None = Field(default=None, description="Stream bridge configuration") stream_bridge: StreamBridgeConfig | None = Field(default=None, description="Stream bridge configuration")
browser_context: BrowserContextConfig = Field(default_factory=BrowserContextConfig, description="Optional browser page capture integration") concurrency: ConcurrencyConfig = Field(default_factory=ConcurrencyConfig, description="Optional cross-worker concurrency coordination")
browser_context: BrowserContextConfig = Field(default_factory=BrowserContextConfig, description="Optional browser page capture integration")
knowledge: KnowledgeConfig = Field(default_factory=KnowledgeConfig, description="Product knowledge base (obsidian-wiki compatible vault) configuration") knowledge: KnowledgeConfig = Field(default_factory=KnowledgeConfig, description="Product knowledge base (obsidian-wiki compatible vault) configuration")
knowledge_ingest: KnowledgeIngestConfig = Field(default_factory=KnowledgeIngestConfig, description="External knowledge-ingest endpoint configuration") knowledge_ingest: KnowledgeIngestConfig = Field(default_factory=KnowledgeIngestConfig, description="External knowledge-ingest endpoint configuration")
task_deeplink: TaskDeeplinkConfig = Field(default_factory=TaskDeeplinkConfig, description="任务深链(rwfx)目的树闸门与跳转地址配置(前端经 /api/public/config/task-deeplink 读取)") task_deeplink: TaskDeeplinkConfig = Field(default_factory=TaskDeeplinkConfig, description="任务深链(rwfx)目的树闸门与跳转地址配置(前端经 /api/public/config/task-deeplink 读取)")
@@ -0,0 +1,35 @@
"""Runtime concurrency coordination configuration."""
from __future__ import annotations
from pydantic import BaseModel, Field
class RedisConcurrencyConfig(BaseModel):
"""Optional Redis-backed cross-worker concurrency coordination."""
enabled: bool = Field(default=False, description="Enable Redis-backed global run concurrency coordination.")
url: str | None = Field(default=None, description="Redis URL, for example redis://127.0.0.1:6379/0.")
password: str | None = Field(
default=None,
description="Optional Redis password. You may also include credentials directly in url.",
)
namespace: str = Field(default="deerflow:concurrency", description="Redis key prefix for concurrency leases.")
socket_connect_timeout_ms: int = Field(default=80, ge=1, le=10_000)
socket_timeout_ms: int = Field(default=80, ge=1, le=10_000)
acquire_timeout_ms: int = Field(default=120, ge=1, le=10_000)
operation_timeout_ms: int = Field(default=80, ge=1, le=10_000)
circuit_breaker_failures: int = Field(default=3, ge=1, le=100)
circuit_breaker_cooldown_seconds: int = Field(default=30, ge=1, le=3600)
lease_ttl_seconds: int = Field(default=90, ge=10, le=3600)
heartbeat_interval_seconds: int = Field(default=20, ge=5, le=600)
fail_open: bool = Field(
default=True,
description="Allow runs to proceed with local-only limits when Redis is unavailable or slow.",
)
class ConcurrencyConfig(BaseModel):
"""Top-level concurrency coordination config."""
redis: RedisConcurrencyConfig = Field(default_factory=RedisConcurrencyConfig)
@@ -7,14 +7,19 @@ directly from ``deerflow.runtime``.
from .checkpointer import checkpointer_context, get_checkpointer, make_checkpointer, reset_checkpointer from .checkpointer import checkpointer_context, get_checkpointer, make_checkpointer, reset_checkpointer
from .runs import ( from .runs import (
AcquireDecision,
AcquireResult,
ActiveLease,
ConflictError, ConflictError,
DisconnectMode, DisconnectMode,
ModelConcurrencyLimitError, ModelConcurrencyLimitError,
RedisConcurrencyGate,
RunContext, RunContext,
RunManager, RunManager,
RunRecord, RunRecord,
RunStatus, RunStatus,
UnsupportedStrategyError, UnsupportedStrategyError,
make_redis_concurrency_gate,
run_agent, run_agent,
) )
from .serialization import serialize, serialize_channel_values, serialize_lc_object, serialize_messages_tuple from .serialization import serialize, serialize_channel_values, serialize_lc_object, serialize_messages_tuple
@@ -28,14 +33,19 @@ __all__ = [
"make_checkpointer", "make_checkpointer",
"reset_checkpointer", "reset_checkpointer",
# runs # runs
"AcquireDecision",
"AcquireResult",
"ActiveLease",
"ConflictError", "ConflictError",
"DisconnectMode", "DisconnectMode",
"ModelConcurrencyLimitError", "ModelConcurrencyLimitError",
"RedisConcurrencyGate",
"RunContext", "RunContext",
"RunManager", "RunManager",
"RunRecord", "RunRecord",
"RunStatus", "RunStatus",
"UnsupportedStrategyError", "UnsupportedStrategyError",
"make_redis_concurrency_gate",
"run_agent", "run_agent",
# serialization # serialization
"serialize", "serialize",
@@ -2,16 +2,22 @@
from .manager import ConflictError, ModelConcurrencyLimitError, RunManager, RunRecord, UnsupportedStrategyError from .manager import ConflictError, ModelConcurrencyLimitError, RunManager, RunRecord, UnsupportedStrategyError
from .schemas import DisconnectMode, RunStatus from .schemas import DisconnectMode, RunStatus
from .concurrency_gate import AcquireDecision, AcquireResult, ActiveLease, RedisConcurrencyGate, make_redis_concurrency_gate
from .worker import RunContext, run_agent from .worker import RunContext, run_agent
__all__ = [ __all__ = [
"AcquireDecision",
"AcquireResult",
"ActiveLease",
"ConflictError", "ConflictError",
"DisconnectMode", "DisconnectMode",
"ModelConcurrencyLimitError", "ModelConcurrencyLimitError",
"RedisConcurrencyGate",
"RunContext", "RunContext",
"RunManager", "RunManager",
"RunRecord", "RunRecord",
"RunStatus", "RunStatus",
"UnsupportedStrategyError", "UnsupportedStrategyError",
"make_redis_concurrency_gate",
"run_agent", "run_agent",
] ]
@@ -0,0 +1,369 @@
"""Optional cross-worker run concurrency gates.
The Redis gate is deliberately fail-open by default: Redis coordinates global
limits when it is fast and healthy, but it must not block the chat path when it
is slow or unavailable.
"""
from __future__ import annotations
import asyncio
import logging
import time
from dataclasses import dataclass
from enum import Enum
from typing import Any
from deerflow.config.concurrency_config import RedisConcurrencyConfig
logger = logging.getLogger(__name__)
class AcquireDecision(str, Enum):
ACQUIRED = "acquired"
REJECTED = "rejected"
DEGRADED = "degraded"
@dataclass(frozen=True)
class AcquireResult:
decision: AcquireDecision
reason: str = ""
source: str = "redis"
@property
def acquired(self) -> bool:
return self.decision == AcquireDecision.ACQUIRED
@property
def rejected(self) -> bool:
return self.decision == AcquireDecision.REJECTED
@property
def degraded(self) -> bool:
return self.decision == AcquireDecision.DEGRADED
@dataclass(frozen=True)
class ActiveLease:
run_id: str
thread_id: str = ""
model_name: str = ""
user_id: str = ""
assistant_id: str = ""
created_at: float = 0.0
expires_at: float = 0.0
class CircuitBreaker:
"""Tiny in-process circuit breaker for Redis failures."""
def __init__(self, *, failure_threshold: int, cooldown_seconds: int) -> None:
self._failure_threshold = max(1, int(failure_threshold))
self._cooldown_seconds = max(1, int(cooldown_seconds))
self._failures = 0
self._opened_until = 0.0
def allow(self) -> bool:
return time.monotonic() >= self._opened_until
def record_success(self) -> None:
self._failures = 0
self._opened_until = 0.0
def record_failure(self) -> None:
self._failures += 1
if self._failures >= self._failure_threshold:
self._opened_until = time.monotonic() + self._cooldown_seconds
@property
def open(self) -> bool:
return not self.allow()
class RedisConcurrencyGate:
"""Redis lease-based global concurrency limiter."""
_ACQUIRE_SCRIPT = """
local now = tonumber(ARGV[1])
local expires_at = tonumber(ARGV[2])
local run_id = ARGV[3]
local model_limit = tonumber(ARGV[4])
local user_model_limit = tonumber(ARGV[5])
local reject_thread = tonumber(ARGV[6])
local payload = ARGV[7]
local ttl_ms = tonumber(ARGV[8])
for i = 1, 4 do
redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now)
end
if reject_thread == 1 and redis.call('ZCARD', KEYS[4]) > 0 then
return {'rejected', 'thread'}
end
if model_limit > 0 and redis.call('ZCARD', KEYS[2]) >= model_limit then
return {'rejected', 'model'}
end
if user_model_limit > 0 and redis.call('ZCARD', KEYS[3]) >= user_model_limit then
return {'rejected', 'user_model'}
end
redis.call('ZADD', KEYS[1], expires_at, run_id)
redis.call('ZADD', KEYS[2], expires_at, run_id)
redis.call('ZADD', KEYS[3], expires_at, run_id)
redis.call('ZADD', KEYS[4], expires_at, run_id)
redis.call('SET', KEYS[5], payload, 'PX', ttl_ms)
return {'acquired', ''}
"""
def __init__(self, config: RedisConcurrencyConfig) -> None:
self.config = config
self._client: Any | None = None
self._breaker = CircuitBreaker(
failure_threshold=config.circuit_breaker_failures,
cooldown_seconds=config.circuit_breaker_cooldown_seconds,
)
@property
def enabled(self) -> bool:
return bool(self.config.enabled and self.config.url)
def health(self) -> dict[str, Any]:
return {
"enabled": self.enabled,
"source": "redis",
"circuit_open": self._breaker.open,
"fail_open": self.config.fail_open,
}
async def aclose(self) -> None:
client = self._client
self._client = None
if client is not None:
try:
await client.aclose()
except Exception:
logger.debug("Failed to close Redis concurrency client", exc_info=True)
def _ns(self) -> str:
return self.config.namespace.strip().rstrip(":") or "deerflow:concurrency"
def _keys(self, *, run_id: str, model_name: str, user_id: str, thread_id: str) -> list[str]:
ns = self._ns()
return [
f"{ns}:all",
f"{ns}:model:{model_name}",
f"{ns}:user_model:{model_name}:{user_id}",
f"{ns}:thread:{thread_id}",
f"{ns}:run:{run_id}",
]
async def _redis(self) -> Any:
if self._client is not None:
return self._client
try:
from redis.asyncio import Redis
except Exception as exc: # pragma: no cover - exercised when dependency is absent
raise RuntimeError("redis package is not installed") from exc
timeout = max(0.001, self.config.socket_timeout_ms / 1000)
connect_timeout = max(0.001, self.config.socket_connect_timeout_ms / 1000)
self._client = Redis.from_url(
self.config.url,
password=self.config.password,
socket_timeout=timeout,
socket_connect_timeout=connect_timeout,
decode_responses=True,
)
return self._client
async def _with_timeout(self, coro, timeout_ms: int) -> Any:
return await asyncio.wait_for(coro, timeout=max(0.001, timeout_ms / 1000))
def _degraded(self, reason: str) -> AcquireResult:
return AcquireResult(AcquireDecision.DEGRADED, reason=reason)
def _on_failure(self, action: str, exc: BaseException) -> None:
self._breaker.record_failure()
logger.warning(
"Redis concurrency %s failed; fail_open=%s circuit_open=%s: %s",
action,
self.config.fail_open,
self._breaker.open,
exc,
)
async def try_acquire(
self,
*,
run_id: str,
thread_id: str,
model_name: str,
user_id: str,
assistant_id: str = "",
model_limit: int | None = None,
user_model_limit: int | None = None,
reject_thread: bool = True,
) -> AcquireResult:
if not self.enabled:
return self._degraded("disabled")
if not self._breaker.allow():
return self._degraded("circuit_open")
now = time.time()
expires_at = now + self.config.lease_ttl_seconds
payload = {
"run_id": run_id,
"thread_id": thread_id,
"model_name": model_name,
"user_id": user_id,
"assistant_id": assistant_id,
"created_at": now,
"expires_at": expires_at,
}
import json
keys = self._keys(run_id=run_id, model_name=model_name, user_id=user_id, thread_id=thread_id)
ttl_ms = int(self.config.lease_ttl_seconds * 1000)
try:
client = await self._redis()
raw = await self._with_timeout(
client.eval(
self._ACQUIRE_SCRIPT,
len(keys),
*keys,
str(now),
str(expires_at),
run_id,
str(max(0, int(model_limit or 0))),
str(max(0, int(user_model_limit or 0))),
"1" if reject_thread else "0",
json.dumps(payload, ensure_ascii=False),
str(ttl_ms),
),
self.config.acquire_timeout_ms,
)
self._breaker.record_success()
except Exception as exc:
self._on_failure("acquire", exc)
if self.config.fail_open:
return self._degraded(type(exc).__name__)
raise
status = raw[0] if isinstance(raw, list) and raw else ""
reason = raw[1] if isinstance(raw, list) and len(raw) > 1 else ""
if status == "acquired":
return AcquireResult(AcquireDecision.ACQUIRED)
return AcquireResult(AcquireDecision.REJECTED, reason=str(reason or "limit"))
async def heartbeat(
self,
*,
run_id: str,
thread_id: str,
model_name: str,
user_id: str,
) -> None:
if not self.enabled or not self._breaker.allow():
return
now = time.time()
expires_at = now + self.config.lease_ttl_seconds
keys = self._keys(run_id=run_id, model_name=model_name, user_id=user_id, thread_id=thread_id)
try:
client = await self._redis()
pipe = client.pipeline(transaction=False)
for key in keys[:4]:
pipe.zadd(key, {run_id: expires_at})
pipe.expire(keys[4], self.config.lease_ttl_seconds)
await self._with_timeout(pipe.execute(), self.config.operation_timeout_ms)
self._breaker.record_success()
except Exception as exc:
self._on_failure("heartbeat", exc)
async def release(
self,
*,
run_id: str,
thread_id: str,
model_name: str,
user_id: str,
) -> None:
if not self.enabled or not self._breaker.allow():
return
keys = self._keys(run_id=run_id, model_name=model_name, user_id=user_id, thread_id=thread_id)
try:
client = await self._redis()
pipe = client.pipeline(transaction=False)
for key in keys[:4]:
pipe.zrem(key, run_id)
pipe.delete(keys[4])
await self._with_timeout(pipe.execute(), self.config.operation_timeout_ms)
self._breaker.record_success()
except Exception as exc:
self._on_failure("release", exc)
async def count_active(self) -> int | None:
if not self.enabled or not self._breaker.allow():
return None
key = f"{self._ns()}:all"
now = time.time()
try:
client = await self._redis()
pipe = client.pipeline(transaction=False)
pipe.zremrangebyscore(key, "-inf", now)
pipe.zcard(key)
result = await self._with_timeout(pipe.execute(), self.config.operation_timeout_ms)
self._breaker.record_success()
return int(result[-1] or 0)
except Exception as exc:
self._on_failure("count", exc)
return None
async def list_active(self, *, limit: int = 1000) -> list[ActiveLease] | None:
if not self.enabled or not self._breaker.allow():
return None
import json
all_key = f"{self._ns()}:all"
now = time.time()
try:
client = await self._redis()
await self._with_timeout(client.zremrangebyscore(all_key, "-inf", now), self.config.operation_timeout_ms)
run_ids = await self._with_timeout(client.zrange(all_key, 0, max(0, limit - 1)), self.config.operation_timeout_ms)
if not run_ids:
self._breaker.record_success()
return []
keys = [f"{self._ns()}:run:{run_id}" for run_id in run_ids]
payloads = await self._with_timeout(client.mget(keys), self.config.operation_timeout_ms)
leases: list[ActiveLease] = []
for run_id, payload in zip(run_ids, payloads, strict=False):
if not payload:
continue
try:
data = json.loads(payload)
except Exception:
continue
leases.append(
ActiveLease(
run_id=str(data.get("run_id") or run_id),
thread_id=str(data.get("thread_id") or ""),
model_name=str(data.get("model_name") or ""),
user_id=str(data.get("user_id") or ""),
assistant_id=str(data.get("assistant_id") or ""),
created_at=float(data.get("created_at") or 0),
expires_at=float(data.get("expires_at") or 0),
)
)
self._breaker.record_success()
return leases
except Exception as exc:
self._on_failure("list", exc)
return None
def make_redis_concurrency_gate(config: Any) -> RedisConcurrencyGate | None:
concurrency = getattr(config, "concurrency", None)
redis_config = getattr(concurrency, "redis", None)
if redis_config is None or not getattr(redis_config, "enabled", False) or not getattr(redis_config, "url", None):
return None
return RedisConcurrencyGate(redis_config)
@@ -219,6 +219,7 @@ class RunManager:
multitask_strategy: str = "reject", multitask_strategy: str = "reject",
model_concurrency_limit: tuple[str, int] | None = None, model_concurrency_limit: tuple[str, int] | None = None,
user_model_concurrency_limit: tuple[str, str, int] | None = None, user_model_concurrency_limit: tuple[str, str, int] | None = None,
run_id: str | None = None,
) -> RunRecord: ) -> RunRecord:
"""Atomically check for inflight runs and create a new one. """Atomically check for inflight runs and create a new one.
@@ -229,7 +230,7 @@ class RunManager:
This method holds the lock across both the check and the insert, This method holds the lock across both the check and the insert,
eliminating the TOCTOU race in separate ``has_inflight`` + ``create``. eliminating the TOCTOU race in separate ``has_inflight`` + ``create``.
""" """
run_id = str(uuid.uuid4()) run_id = run_id or str(uuid.uuid4())
now = _now_iso() now = _now_iso()
_supported_strategies = ("reject", "interrupt", "rollback") _supported_strategies = ("reject", "interrupt", "rollback")
@@ -23,9 +23,10 @@ dependencies = [
"markdownify>=1.2.2", "markdownify>=1.2.2",
"markitdown[all,xlsx]>=0.0.1a2", "markitdown[all,xlsx]>=0.0.1a2",
"pydantic>=2.12.5", "pydantic>=2.12.5",
"pyyaml>=6.0.3", "pyyaml>=6.0.3",
"readabilipy>=0.3.0", "readabilipy>=0.3.0",
"tavily-python>=0.7.17", "redis>=5.0.0",
"tavily-python>=0.7.17",
"firecrawl-py>=1.15.0", "firecrawl-py>=1.15.0",
"tiktoken>=0.8.0", "tiktoken>=0.8.0",
"ddgs>=9.10.0", "ddgs>=9.10.0",