fix: improve model concurrency limiting
This commit is contained in:
1 parent
aca089e8ef
commit
a046d19ce1
7 files changed
+1266
-1160
No files matched your search
@@ -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",
|
||||||
|
]
|
||||||
+17
-16
@@ -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",
|
||||||
|
]
|
||||||
+326
-310
@@ -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")
|
|
||||||
Reference in new issue
Block a user