fix: improve model concurrency limiting

This commit is contained in:
Administrator committed 2026-06-27 18:57:01 +08:00
1 parent aca089e8ef
commit a046d19ce1
7 files changed
+1266 -1160

No files matched your search

+28 -15
View File
@@ -224,15 +224,24 @@ function isRunNotFoundError(message: string): boolean {
* 「already has an active run」。匹配 HTTP 409、英文兜底文案,或后端已中文化 * 「already has an active run」。匹配 HTTP 409、英文兜底文案,或后端已中文化
* 的提示,统一映射成一句友好的中文。 * 的提示,统一映射成一句友好的中文。
*/ */
function isActiveRunConflictError(message: string): boolean { function isActiveRunConflictError(message: string): boolean {
return ( return (
/\b409\b/.test(message) || /\b409\b/.test(message) ||
/already has an active run/i.test(message) || /already has an active run/i.test(message) ||
/还在生成回答/.test(message) /还在生成回答/.test(message)
); );
} }
function buildCurrentUserLabel(): string { function isModelConcurrencyLimitError(message: string): boolean {
return (
/\b429\b/.test(message) ||
/并发已达到上限/.test(message) ||
/使用人数较多/.test(message) ||
/too many requests/i.test(message)
);
}
function buildCurrentUserLabel(): string {
const auth = getStoredAuth(); const auth = getStoredAuth();
if (!auth) { if (!auth) {
return "未登录"; return "未登录";
@@ -570,12 +579,16 @@ export function useThreadStream({
} }
// 上一轮还没结束又发了新问题(后端 409):给一句明确的中文提示,而不是 // 上一轮还没结束又发了新问题(后端 409):给一句明确的中文提示,而不是
// 把原始的英文/HTTP 错误抛给用户。 // 把原始的英文/HTTP 错误抛给用户。
if (isActiveRunConflictError(message)) { if (isActiveRunConflictError(message)) {
toast.info("当前对话还在生成回答,请等待本轮回复结束,或先点「停止」再发送新的问题。"); toast.info("当前对话还在生成回答,请等待本轮回复结束,或先点「停止」再发送新的问题。");
return; return;
} }
toast.error(message); if (isModelConcurrencyLimitError(message)) {
}, toast.info("当前使用人数较多,请稍等一下再试。");
return;
}
toast.error(message);
},
onFinish(state) { onFinish(state) {
listeners.current.onFinish?.(state.values); listeners.current.onFinish?.(state.values);
void queryClient.invalidateQueries({ queryKey: ["threads", "search"] }); void queryClient.invalidateQueries({ queryKey: ["threads", "search"] });
@@ -85,7 +85,7 @@ export default function ModelConcurrencySettingsPage({
className={ className={
embedded embedded
? "flex w-full flex-col gap-6" ? "flex w-full flex-col gap-6"
: "mx-auto flex w-full max-w-3xl flex-col gap-6 px-6 py-10" : "mx-auto flex h-full min-h-0 w-full max-w-3xl flex-col gap-6 overflow-y-auto px-6 py-10"
} }
> >
<div> <div>
@@ -93,7 +93,7 @@ export default function ModelConcurrencySettingsPage({
模型并发限制 模型并发限制
</h1> </h1>
<p className="text-muted-foreground mt-1 text-sm"> <p className="text-muted-foreground mt-1 text-sm">
按模型限制同时进行中的问答数。达到上限后,新请求会收到“当前并发已达到上限,请稍后再试”的提示。 按模型限制同时进行中的问答数。达到上限后,新请求会收到“当前使用人数较多,请稍等一下再试”的提示。
</p> </p>
</div> </div>
File diff suppressed because it is too large. Load diff
@@ -1,46 +1,57 @@
"""LangGraph-compatible runtime — runs, streaming, and lifecycle management. """LangGraph-compatible runtime — runs, streaming, and lifecycle management.
Re-exports the public API of :mod:`~deerflow.runtime.runs` and Re-exports the public API of :mod:`~deerflow.runtime.runs` and
:mod:`~deerflow.runtime.stream_bridge` so that consumers can import :mod:`~deerflow.runtime.stream_bridge` so that consumers can import
directly from ``deerflow.runtime``. 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 ConflictError, DisconnectMode, RunContext, RunManager, RunRecord, RunStatus, UnsupportedStrategyError, run_agent from .runs import (
from .serialization import serialize, serialize_channel_values, serialize_lc_object, serialize_messages_tuple ConflictError,
from .store import get_store, make_store, reset_store, store_context DisconnectMode,
from .stream_bridge import END_SENTINEL, HEARTBEAT_SENTINEL, MemoryStreamBridge, StreamBridge, StreamEvent, make_stream_bridge ModelConcurrencyLimitError,
RunContext,
__all__ = [ RunManager,
# checkpointer RunRecord,
"checkpointer_context", RunStatus,
"get_checkpointer", UnsupportedStrategyError,
"make_checkpointer", run_agent,
"reset_checkpointer", )
# runs from .serialization import serialize, serialize_channel_values, serialize_lc_object, serialize_messages_tuple
"ConflictError", from .store import get_store, make_store, reset_store, store_context
"DisconnectMode", from .stream_bridge import END_SENTINEL, HEARTBEAT_SENTINEL, MemoryStreamBridge, StreamBridge, StreamEvent, make_stream_bridge
"RunContext",
"RunManager", __all__ = [
"RunRecord", # checkpointer
"RunStatus", "checkpointer_context",
"UnsupportedStrategyError", "get_checkpointer",
"run_agent", "make_checkpointer",
# serialization "reset_checkpointer",
"serialize", # runs
"serialize_channel_values", "ConflictError",
"serialize_lc_object", "DisconnectMode",
"serialize_messages_tuple", "ModelConcurrencyLimitError",
# store "RunContext",
"get_store", "RunManager",
"make_store", "RunRecord",
"reset_store", "RunStatus",
"store_context", "UnsupportedStrategyError",
# stream_bridge "run_agent",
"END_SENTINEL", # serialization
"HEARTBEAT_SENTINEL", "serialize",
"MemoryStreamBridge", "serialize_channel_values",
"StreamBridge", "serialize_lc_object",
"StreamEvent", "serialize_messages_tuple",
"make_stream_bridge", # store
] "get_store",
"make_store",
"reset_store",
"store_context",
# stream_bridge
"END_SENTINEL",
"HEARTBEAT_SENTINEL",
"MemoryStreamBridge",
"StreamBridge",
"StreamEvent",
"make_stream_bridge",
]
@@ -1,16 +1,17 @@
"""Run lifecycle management for LangGraph Platform API compatibility.""" """Run lifecycle management for LangGraph Platform API compatibility."""
from .manager import ConflictError, RunManager, RunRecord, UnsupportedStrategyError from .manager import ConflictError, ModelConcurrencyLimitError, RunManager, RunRecord, UnsupportedStrategyError
from .schemas import DisconnectMode, RunStatus from .schemas import DisconnectMode, RunStatus
from .worker import RunContext, run_agent from .worker import RunContext, run_agent
__all__ = [ __all__ = [
"ConflictError", "ConflictError",
"DisconnectMode", "DisconnectMode",
"RunContext", "ModelConcurrencyLimitError",
"RunManager", "RunContext",
"RunRecord", "RunManager",
"RunStatus", "RunRecord",
"UnsupportedStrategyError", "RunStatus",
"run_agent", "UnsupportedStrategyError",
] "run_agent",
]
@@ -1,310 +1,326 @@
"""In-memory run registry with optional persistent RunStore backing.""" """In-memory run registry with optional persistent RunStore backing."""
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import logging import logging
import uuid import uuid
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from deerflow.utils.time import now_iso as _now_iso from deerflow.utils.time import now_iso as _now_iso
from .schemas import DisconnectMode, RunStatus from .schemas import DisconnectMode, RunStatus
if TYPE_CHECKING: if TYPE_CHECKING:
from deerflow.runtime.runs.store.base import RunStore from deerflow.runtime.runs.store.base import RunStore
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@dataclass @dataclass
class RunRecord: class RunRecord:
"""Mutable record for a single run.""" """Mutable record for a single run."""
run_id: str run_id: str
thread_id: str thread_id: str
assistant_id: str | None assistant_id: str | None
status: RunStatus status: RunStatus
on_disconnect: DisconnectMode on_disconnect: DisconnectMode
multitask_strategy: str = "reject" multitask_strategy: str = "reject"
metadata: dict = field(default_factory=dict) metadata: dict = field(default_factory=dict)
kwargs: dict = field(default_factory=dict) kwargs: dict = field(default_factory=dict)
created_at: str = "" created_at: str = ""
updated_at: str = "" updated_at: str = ""
task: asyncio.Task | None = field(default=None, repr=False) task: asyncio.Task | None = field(default=None, repr=False)
abort_event: asyncio.Event = field(default_factory=asyncio.Event, repr=False) abort_event: asyncio.Event = field(default_factory=asyncio.Event, repr=False)
abort_action: str = "interrupt" abort_action: str = "interrupt"
error: str | None = None error: str | None = None
class RunManager: class RunManager:
"""In-memory run registry with optional persistent RunStore backing. """In-memory run registry with optional persistent RunStore backing.
All mutations are protected by an asyncio lock. When a ``store`` is All mutations are protected by an asyncio lock. When a ``store`` is
provided, serializable metadata is also persisted to the store so provided, serializable metadata is also persisted to the store so
that run history survives process restarts. that run history survives process restarts.
""" """
def __init__(self, store: RunStore | None = None) -> None: def __init__(self, store: RunStore | None = None) -> None:
self._runs: dict[str, RunRecord] = {} self._runs: dict[str, RunRecord] = {}
self._lock = asyncio.Lock() self._lock = asyncio.Lock()
self._store = store self._store = store
async def _persist_to_store(self, record: RunRecord) -> None: async def _persist_to_store(self, record: RunRecord) -> None:
"""Best-effort persist run record to backing store.""" """Best-effort persist run record to backing store."""
if self._store is None: if self._store is None:
return return
try: try:
await self._store.put( await self._store.put(
record.run_id, record.run_id,
thread_id=record.thread_id, thread_id=record.thread_id,
assistant_id=record.assistant_id, assistant_id=record.assistant_id,
status=record.status.value, status=record.status.value,
multitask_strategy=record.multitask_strategy, multitask_strategy=record.multitask_strategy,
metadata=record.metadata or {}, metadata=record.metadata or {},
kwargs=record.kwargs or {}, kwargs=record.kwargs or {},
created_at=record.created_at, created_at=record.created_at,
) )
except Exception: except Exception:
logger.warning("Failed to persist run %s to store", record.run_id, exc_info=True) logger.warning("Failed to persist run %s to store", record.run_id, exc_info=True)
async def update_run_completion(self, run_id: str, **kwargs) -> None: async def update_run_completion(self, run_id: str, **kwargs) -> None:
"""Persist token usage and completion data to the backing store.""" """Persist token usage and completion data to the backing store."""
if self._store is not None: if self._store is not None:
try: try:
await self._store.update_run_completion(run_id, **kwargs) await self._store.update_run_completion(run_id, **kwargs)
except Exception: except Exception:
logger.warning("Failed to persist run completion for %s", run_id, exc_info=True) logger.warning("Failed to persist run completion for %s", run_id, exc_info=True)
async def create( async def create(
self, self,
thread_id: str, thread_id: str,
assistant_id: str | None = None, assistant_id: str | None = None,
*, *,
on_disconnect: DisconnectMode = DisconnectMode.cancel, on_disconnect: DisconnectMode = DisconnectMode.cancel,
metadata: dict | None = None, metadata: dict | None = None,
kwargs: dict | None = None, kwargs: dict | None = None,
multitask_strategy: str = "reject", multitask_strategy: str = "reject",
) -> RunRecord: ) -> RunRecord:
"""Create a new pending run and register it.""" """Create a new pending run and register it."""
run_id = str(uuid.uuid4()) run_id = str(uuid.uuid4())
now = _now_iso() now = _now_iso()
record = RunRecord( record = RunRecord(
run_id=run_id, run_id=run_id,
thread_id=thread_id, thread_id=thread_id,
assistant_id=assistant_id, assistant_id=assistant_id,
status=RunStatus.pending, status=RunStatus.pending,
on_disconnect=on_disconnect, on_disconnect=on_disconnect,
multitask_strategy=multitask_strategy, multitask_strategy=multitask_strategy,
metadata=metadata or {}, metadata=metadata or {},
kwargs=kwargs or {}, kwargs=kwargs or {},
created_at=now, created_at=now,
updated_at=now, updated_at=now,
) )
async with self._lock: async with self._lock:
self._runs[run_id] = record self._runs[run_id] = record
await self._persist_to_store(record) await self._persist_to_store(record)
logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id) logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id)
return record return record
def get(self, run_id: str) -> RunRecord | None: def get(self, run_id: str) -> RunRecord | None:
"""Return a run record by ID, or ``None``.""" """Return a run record by ID, or ``None``."""
return self._runs.get(run_id) return self._runs.get(run_id)
def _record_from_store(self, data: dict) -> RunRecord: def _record_from_store(self, data: dict) -> RunRecord:
"""Rebuild a :class:`RunRecord` from a persisted store row. """Rebuild a :class:`RunRecord` from a persisted store row.
The reconstructed record has no live ``task``/``abort_event`` wired to a The reconstructed record has no live ``task``/``abort_event`` wired to a
running coroutine — it represents *history*, not an in-flight run. running coroutine — it represents *history*, not an in-flight run.
""" """
try: try:
status = RunStatus(data.get("status") or RunStatus.success.value) status = RunStatus(data.get("status") or RunStatus.success.value)
except ValueError: except ValueError:
status = RunStatus.success status = RunStatus.success
return RunRecord( return RunRecord(
run_id=data["run_id"], run_id=data["run_id"],
thread_id=data["thread_id"], thread_id=data["thread_id"],
assistant_id=data.get("assistant_id"), assistant_id=data.get("assistant_id"),
status=status, status=status,
# Disconnect mode is irrelevant for a store-only record (no task to # Disconnect mode is irrelevant for a store-only record (no task to
# cancel); ``continue_`` keeps sse_consumer's finally block a no-op. # cancel); ``continue_`` keeps sse_consumer's finally block a no-op.
on_disconnect=DisconnectMode.continue_, on_disconnect=DisconnectMode.continue_,
multitask_strategy=data.get("multitask_strategy", "reject"), multitask_strategy=data.get("multitask_strategy", "reject"),
metadata=data.get("metadata") or {}, metadata=data.get("metadata") or {},
kwargs=data.get("kwargs") or {}, kwargs=data.get("kwargs") or {},
created_at=data.get("created_at") or "", created_at=data.get("created_at") or "",
updated_at=data.get("updated_at") or "", updated_at=data.get("updated_at") or "",
error=data.get("error"), error=data.get("error"),
) )
async def get_persisted(self, run_id: str) -> RunRecord | None: async def get_persisted(self, run_id: str) -> RunRecord | None:
"""Load a run from the backing store (process-restart fallback). """Load a run from the backing store (process-restart fallback).
The in-memory registry is wiped on every Gateway restart, but the The in-memory registry is wiped on every Gateway restart, but the
``RunStore`` row survives. When :meth:`get` misses, callers fall back ``RunStore`` row survives. When :meth:`get` misses, callers fall back
here so a stale ``joinStream``/``GET run`` can report the run's final here so a stale ``joinStream``/``GET run`` can report the run's final
status instead of 404ing. Returns ``None`` when there is no store, no status instead of 404ing. Returns ``None`` when there is no store, no
matching row, or the lookup fails (e.g. no user context). matching row, or the lookup fails (e.g. no user context).
""" """
if self._store is None: if self._store is None:
return None return None
try: try:
data = await self._store.get(run_id) data = await self._store.get(run_id)
except Exception: except Exception:
logger.warning("Failed to load run %s from store", run_id, exc_info=True) logger.warning("Failed to load run %s from store", run_id, exc_info=True)
return None return None
if not data: if not data:
return None return None
return self._record_from_store(data) return self._record_from_store(data)
async def list_by_thread(self, thread_id: str) -> list[RunRecord]: async def list_by_thread(self, thread_id: str) -> list[RunRecord]:
"""Return all runs for a given thread, newest first.""" """Return all runs for a given thread, newest first."""
async with self._lock: async with self._lock:
# Dict insertion order matches creation order, so reversing it gives # Dict insertion order matches creation order, so reversing it gives
# us deterministic newest-first results even when timestamps tie. # us deterministic newest-first results even when timestamps tie.
return [r for r in self._runs.values() if r.thread_id == thread_id] return [r for r in self._runs.values() if r.thread_id == thread_id]
async def set_status(self, run_id: str, status: RunStatus, *, error: str | None = None) -> None: async def set_status(self, run_id: str, status: RunStatus, *, error: str | None = None) -> None:
"""Transition a run to a new status.""" """Transition a run to a new status."""
async with self._lock: async with self._lock:
record = self._runs.get(run_id) record = self._runs.get(run_id)
if record is None: if record is None:
logger.warning("set_status called for unknown run %s", run_id) logger.warning("set_status called for unknown run %s", run_id)
return return
record.status = status record.status = status
record.updated_at = _now_iso() record.updated_at = _now_iso()
if error is not None: if error is not None:
record.error = error record.error = error
if self._store is not None: if self._store is not None:
try: try:
await self._store.update_status(run_id, status.value, error=error) await self._store.update_status(run_id, status.value, error=error)
except Exception: except Exception:
logger.warning("Failed to persist status update for run %s", run_id, exc_info=True) logger.warning("Failed to persist status update for run %s", run_id, exc_info=True)
logger.info("Run %s -> %s", run_id, status.value) logger.info("Run %s -> %s", run_id, status.value)
async def cancel(self, run_id: str, *, action: str = "interrupt") -> bool: async def cancel(self, run_id: str, *, action: str = "interrupt") -> bool:
"""Request cancellation of a run. """Request cancellation of a run.
Args: Args:
run_id: The run ID to cancel. run_id: The run ID to cancel.
action: "interrupt" keeps checkpoint, "rollback" reverts to pre-run state. action: "interrupt" keeps checkpoint, "rollback" reverts to pre-run state.
Sets the abort event with the action reason and cancels the asyncio task. Sets the abort event with the action reason and cancels the asyncio task.
Returns ``True`` if the run was in-flight and cancellation was initiated. Returns ``True`` if the run was in-flight and cancellation was initiated.
""" """
async with self._lock: async with self._lock:
record = self._runs.get(run_id) record = self._runs.get(run_id)
if record is None: if record is None:
return False return False
if record.status not in (RunStatus.pending, RunStatus.running): if record.status not in (RunStatus.pending, RunStatus.running):
return False return False
record.abort_action = action record.abort_action = action
record.abort_event.set() record.abort_event.set()
if record.task is not None and not record.task.done(): if record.task is not None and not record.task.done():
record.task.cancel() record.task.cancel()
record.status = RunStatus.interrupted record.status = RunStatus.interrupted
record.updated_at = _now_iso() record.updated_at = _now_iso()
logger.info("Run %s cancelled (action=%s)", run_id, action) logger.info("Run %s cancelled (action=%s)", run_id, action)
return True return True
async def create_or_reject( async def create_or_reject(
self, self,
thread_id: str, thread_id: str,
assistant_id: str | None = None, assistant_id: str | None = None,
*, *,
on_disconnect: DisconnectMode = DisconnectMode.cancel, on_disconnect: DisconnectMode = DisconnectMode.cancel,
metadata: dict | None = None, metadata: dict | None = None,
kwargs: dict | None = None, kwargs: dict | None = None,
multitask_strategy: str = "reject", multitask_strategy: str = "reject",
) -> RunRecord: model_concurrency_limit: tuple[str, int] | None = None,
"""Atomically check for inflight runs and create a new one. ) -> RunRecord:
"""Atomically check for inflight runs and create a new one.
For ``reject`` strategy, raises ``ConflictError`` if thread
already has a pending/running run. For ``interrupt``/``rollback``, For ``reject`` strategy, raises ``ConflictError`` if thread
cancels inflight runs before creating. already has a pending/running run. For ``interrupt``/``rollback``,
cancels inflight runs before creating.
This method holds the lock across both the check and the insert,
eliminating the TOCTOU race in separate ``has_inflight`` + ``create``. This method holds the lock across both the check and the insert,
""" eliminating the TOCTOU race in separate ``has_inflight`` + ``create``.
run_id = str(uuid.uuid4()) """
now = _now_iso() run_id = str(uuid.uuid4())
now = _now_iso()
_supported_strategies = ("reject", "interrupt", "rollback")
_supported_strategies = ("reject", "interrupt", "rollback")
async with self._lock:
if multitask_strategy not in _supported_strategies: async with self._lock:
raise UnsupportedStrategyError(f"Multitask strategy '{multitask_strategy}' is not yet supported. Supported strategies: {', '.join(_supported_strategies)}") if multitask_strategy not in _supported_strategies:
raise UnsupportedStrategyError(f"Multitask strategy '{multitask_strategy}' is not yet supported. Supported strategies: {', '.join(_supported_strategies)}")
inflight = [r for r in self._runs.values() if r.thread_id == thread_id and r.status in (RunStatus.pending, RunStatus.running)]
inflight = [r for r in self._runs.values() if r.thread_id == thread_id and r.status in (RunStatus.pending, RunStatus.running)]
if multitask_strategy == "reject" and inflight:
raise ConflictError(f"Thread {thread_id} already has an active run") if multitask_strategy == "reject" and inflight:
raise ConflictError(f"Thread {thread_id} already has an active run")
if multitask_strategy in ("interrupt", "rollback") and inflight:
for r in inflight: if multitask_strategy in ("interrupt", "rollback") and inflight:
r.abort_action = multitask_strategy for r in inflight:
r.abort_event.set() r.abort_action = multitask_strategy
if r.task is not None and not r.task.done(): r.abort_event.set()
r.task.cancel() if r.task is not None and not r.task.done():
r.status = RunStatus.interrupted r.task.cancel()
r.updated_at = now r.status = RunStatus.interrupted
logger.info( r.updated_at = now
"Cancelled %d inflight run(s) on thread %s (strategy=%s)", logger.info(
len(inflight), "Cancelled %d inflight run(s) on thread %s (strategy=%s)",
thread_id, len(inflight),
multitask_strategy, thread_id,
) multitask_strategy,
)
record = RunRecord(
run_id=run_id, if model_concurrency_limit is not None:
thread_id=thread_id, model_name, limit = model_concurrency_limit
assistant_id=assistant_id, active_for_model = sum(1 for r in self._runs.values() if r.status in (RunStatus.pending, RunStatus.running) and str((r.metadata or {}).get("model_name") or "").strip() == model_name)
status=RunStatus.pending, if active_for_model >= max(1, int(limit)):
on_disconnect=on_disconnect, raise ModelConcurrencyLimitError(model_name, limit)
multitask_strategy=multitask_strategy,
metadata=metadata or {}, record = RunRecord(
kwargs=kwargs or {}, run_id=run_id,
created_at=now, thread_id=thread_id,
updated_at=now, assistant_id=assistant_id,
) status=RunStatus.pending,
self._runs[run_id] = record on_disconnect=on_disconnect,
multitask_strategy=multitask_strategy,
await self._persist_to_store(record) metadata=metadata or {},
logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id) kwargs=kwargs or {},
return record created_at=now,
updated_at=now,
async def has_inflight(self, thread_id: str) -> bool: )
"""Return ``True`` if *thread_id* has a pending or running run.""" self._runs[run_id] = record
async with self._lock:
return any(r.thread_id == thread_id and r.status in (RunStatus.pending, RunStatus.running) for r in self._runs.values()) await self._persist_to_store(record)
logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id)
async def list_active(self) -> list[RunRecord]: return record
"""Return every in-flight run (status pending/running), oldest first.
async def has_inflight(self, thread_id: str) -> bool:
Powers the admin concurrency monitor: each record represents one live """Return ``True`` if *thread_id* has a pending or running run."""
call into the model runtime (interactive chat, agent chat, roundtable async with self._lock:
seat, …), so ``len(list_active())`` is the current concurrent-call count return any(r.thread_id == thread_id and r.status in (RunStatus.pending, RunStatus.running) for r in self._runs.values())
and each record can be cancelled to relieve model pressure. Completed
records linger in ``_runs`` until :meth:`cleanup`, so we filter by status async def list_active(self) -> list[RunRecord]:
rather than returning the whole registry. """Return every in-flight run (status pending/running), oldest first.
"""
async with self._lock: Powers the admin concurrency monitor: each record represents one live
# Dict insertion order == creation order, so this is oldest-first. call into the model runtime (interactive chat, agent chat, roundtable
return [r for r in self._runs.values() if r.status in (RunStatus.pending, RunStatus.running)] seat, …), so ``len(list_active())`` is the current concurrent-call count
and each record can be cancelled to relieve model pressure. Completed
async def cleanup(self, run_id: str, *, delay: float = 300) -> None: records linger in ``_runs`` until :meth:`cleanup`, so we filter by status
"""Remove a run record after an optional delay.""" rather than returning the whole registry.
if delay > 0: """
await asyncio.sleep(delay) async with self._lock:
async with self._lock: # Dict insertion order == creation order, so this is oldest-first.
self._runs.pop(run_id, None) return [r for r in self._runs.values() if r.status in (RunStatus.pending, RunStatus.running)]
logger.debug("Run record %s cleaned up", run_id)
async def cleanup(self, run_id: str, *, delay: float = 300) -> None:
"""Remove a run record after an optional delay."""
class ConflictError(Exception): if delay > 0:
"""Raised when multitask_strategy=reject and thread has inflight runs.""" await asyncio.sleep(delay)
async with self._lock:
self._runs.pop(run_id, None)
class UnsupportedStrategyError(Exception): logger.debug("Run record %s cleaned up", run_id)
"""Raised when a multitask_strategy value is not yet implemented."""
class ConflictError(Exception):
"""Raised when multitask_strategy=reject and thread has inflight runs."""
class ModelConcurrencyLimitError(Exception):
"""Raised when a model-specific in-flight run limit has been reached."""
def __init__(self, model_name: str, limit: int) -> None:
self.model_name = model_name
self.limit = limit
super().__init__(f"Model {model_name} has reached its concurrency limit")
class UnsupportedStrategyError(Exception):
"""Raised when a multitask_strategy value is not yet implemented."""
@@ -7,15 +7,14 @@ from pathlib import Path
from types import SimpleNamespace from types import SimpleNamespace
import pytest import pytest
from fastapi import HTTPException
from app.gateway.routers import system_settings as ss_router from app.gateway.routers import system_settings as ss_router
from app.gateway.services import ( from app.gateway.services import (
_enforce_model_concurrency_limit, _get_model_concurrency_limit,
_resolve_effective_run_model_name, _resolve_effective_run_model_name,
) )
from deerflow.config import system_settings as ss from deerflow.config import system_settings as ss
from deerflow.runtime.runs.manager import RunManager from deerflow.runtime.runs.manager import ModelConcurrencyLimitError, RunManager
@pytest.fixture @pytest.fixture
@@ -77,9 +76,7 @@ def test_resolve_effective_run_model_name_prefers_requested_model(
) -> None: ) -> None:
app_config = SimpleNamespace( app_config = SimpleNamespace(
models=[SimpleNamespace(name="default-model")], models=[SimpleNamespace(name="default-model")],
get_model_config=lambda name: {"name": name} get_model_config=lambda name: {"name": name} if name in {"default-model", "glm5", "agent-model"} else None,
if name in {"default-model", "glm5", "agent-model"}
else None,
) )
monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config) monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config)
monkeypatch.setattr( monkeypatch.setattr(
@@ -100,9 +97,7 @@ def test_resolve_effective_run_model_name_uses_agent_default_when_request_missin
) -> None: ) -> None:
app_config = SimpleNamespace( app_config = SimpleNamespace(
models=[SimpleNamespace(name="default-model")], models=[SimpleNamespace(name="default-model")],
get_model_config=lambda name: {"name": name} get_model_config=lambda name: {"name": name} if name in {"default-model", "agent-model"} else None,
if name in {"default-model", "agent-model"}
else None,
) )
monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config) monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config)
monkeypatch.setattr( monkeypatch.setattr(
@@ -118,8 +113,50 @@ def test_resolve_effective_run_model_name_uses_agent_default_when_request_missin
assert _resolve_effective_run_model_name(body) == "agent-model" assert _resolve_effective_run_model_name(body) == "agent-model"
def test_resolve_effective_run_model_name_uses_context_agent_id(
monkeypatch: pytest.MonkeyPatch,
) -> None:
app_config = SimpleNamespace(
models=[SimpleNamespace(name="default-model")],
get_model_config=lambda name: {"name": name} if name in {"default-model", "agent-model"} else None,
)
monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config)
monkeypatch.setattr(
"app.gateway.services.load_agent_config",
lambda agent_id: SimpleNamespace(model="agent-model") if agent_id == "custom-agent" else SimpleNamespace(model="default-model"),
)
body = SimpleNamespace(
assistant_id="lead_agent",
context={"agent_id": "custom-agent"},
config=None,
)
assert _resolve_effective_run_model_name(body) == "agent-model"
def test_resolve_effective_run_model_name_context_agent_overrides_assistant_id(
monkeypatch: pytest.MonkeyPatch,
) -> None:
app_config = SimpleNamespace(
models=[SimpleNamespace(name="default-model")],
get_model_config=lambda name: {"name": name} if name in {"default-model", "context-model", "assistant-model"} else None,
)
monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config)
monkeypatch.setattr(
"app.gateway.services.load_agent_config",
lambda agent_id: SimpleNamespace(model="context-model" if agent_id == "context-agent" else "assistant-model"),
)
body = SimpleNamespace(
assistant_id="assistant-agent",
context={"agent_id": "context-agent"},
config=None,
)
assert _resolve_effective_run_model_name(body) == "context-model"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_enforce_model_concurrency_limit_rejects_when_limit_reached( async def test_run_manager_rejects_model_limit_atomically(
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
) -> None: ) -> None:
monkeypatch.setattr( monkeypatch.setattr(
@@ -132,18 +169,21 @@ async def test_enforce_model_concurrency_limit_rejects_when_limit_reached(
) )
), ),
) )
assert _get_model_concurrency_limit("glm5") == 2
run_mgr = RunManager() run_mgr = RunManager()
await run_mgr.create("t1", metadata={"model_name": "glm5"}) await run_mgr.create_or_reject("t1", metadata={"model_name": "glm5"})
await run_mgr.create("t2", metadata={"model_name": "glm5"}) await run_mgr.create_or_reject("t2", metadata={"model_name": "glm5"})
with pytest.raises(HTTPException) as exc: with pytest.raises(ModelConcurrencyLimitError):
await _enforce_model_concurrency_limit(run_mgr, "glm5") await run_mgr.create_or_reject(
assert exc.value.status_code == 429 "t3",
assert "并发已达到上限" in str(exc.value.detail) metadata={"model_name": "glm5"},
model_concurrency_limit=("glm5", 2),
)
@pytest.mark.asyncio def test_get_model_concurrency_limit_skips_when_disabled(
async def test_enforce_model_concurrency_limit_skips_when_disabled(
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
) -> None: ) -> None:
monkeypatch.setattr( monkeypatch.setattr(
@@ -156,6 +196,4 @@ async def test_enforce_model_concurrency_limit_skips_when_disabled(
) )
), ),
) )
run_mgr = RunManager() assert _get_model_concurrency_limit("glm5") is None
await run_mgr.create("t1", metadata={"model_name": "glm5"})
await _enforce_model_concurrency_limit(run_mgr, "glm5")