Add Redis-backed concurrency gate
This commit is contained in:
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 读取)")
|
||||||
|
|||||||
+35
@@ -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",
|
||||||
]
|
]
|
||||||
+369
@@ -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",
|
||||||
|
|||||||
Reference in new issue
Block a user