344 lines
14 KiB
Python
344 lines
14 KiB
Python
"""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 cancellation was initiated or the run was already stopped.
|
|
"""
|
|
async with self._lock:
|
|
record = self._runs.get(run_id)
|
|
if record is None:
|
|
return True
|
|
if record.status not in (RunStatus.pending, RunStatus.running):
|
|
return True
|
|
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,
|
|
user_model_concurrency_limit: tuple[str, str, int] | None = None,
|
|
run_id: str | 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 = run_id or 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, scope="model")
|
|
|
|
if user_model_concurrency_limit is not None:
|
|
model_name, user_id, limit = user_model_concurrency_limit
|
|
active_for_user_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
|
|
and str((r.metadata or {}).get("user_id") or "").strip() == user_id
|
|
and str((r.metadata or {}).get("run_scope") or "external_root").strip() == "external_root"
|
|
)
|
|
if active_for_user_model >= max(1, int(limit)):
|
|
raise ModelConcurrencyLimitError(model_name, limit, scope="user_model", user_id=user_id)
|
|
|
|
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, *, scope: str = "model", user_id: str | None = None) -> None:
|
|
self.model_name = model_name
|
|
self.limit = limit
|
|
self.scope = scope
|
|
self.user_id = user_id
|
|
super().__init__(f"Model {model_name} has reached its concurrency limit")
|
|
|
|
|
|
class UnsupportedStrategyError(Exception):
|
|
"""Raised when a multitask_strategy value is not yet implemented."""
|