wsk/deerflow-code
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: improve model concurrency limiting

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

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