291 lines
10 KiB
Python
291 lines
10 KiB
Python
"""Bounded retention cleanup for administrator audit and analytics data.
|
|
|
|
This deliberately removes only data written for operational inspection:
|
|
provider/tool call metrics, compact terminal run summaries, and precomputed
|
|
leaderboard snapshots. It does not touch user threads or LangGraph
|
|
checkpoints, which are user conversation history and require a separate,
|
|
explicit retention policy.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from collections.abc import Awaitable, Callable
|
|
from datetime import UTC, datetime, timedelta
|
|
from time import monotonic
|
|
from typing import Any
|
|
|
|
from sqlalchemy import and_, delete, or_, select
|
|
|
|
from app.gateway.admin_query_window import ADMIN_AUDIT_RETENTION_DAYS
|
|
from deerflow.persistence.admin_stats.model import AdminLeaderboardDailyStatRow
|
|
from deerflow.persistence.types import BEIJING_TZ
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_DEFAULT_BATCH_SIZE = 100
|
|
_MAX_BATCH_SIZE = 500
|
|
_MAX_SNAPSHOT_BATCH_SIZE = 200
|
|
_DEFAULT_MAX_BATCHES_PER_STORE = 4
|
|
_MAX_BATCHES_PER_STORE = 16
|
|
_DEFAULT_INTER_BATCH_PAUSE_SECONDS = 0.05
|
|
_DEFAULT_MAX_SWEEP_SECONDS = 8.0
|
|
|
|
|
|
def _bounded_int(value: int, *, default: int, maximum: int) -> int:
|
|
"""Return a safe positive integer for background maintenance work."""
|
|
try:
|
|
return max(1, min(int(value), maximum))
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
def _bounded_seconds(value: float, *, default: float, maximum: float) -> float:
|
|
"""Return a non-negative, finite-enough scheduler duration."""
|
|
try:
|
|
return max(0.0, min(float(value), maximum))
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
async def _purge_store_batch(
|
|
store: Any,
|
|
cutoff: datetime,
|
|
*,
|
|
batch_size: int,
|
|
method_name: str = "purge_before",
|
|
) -> int:
|
|
"""Delete one committed, primary-key-bounded batch from an audit store."""
|
|
if store is None:
|
|
return 0
|
|
purge = getattr(store, method_name, None)
|
|
if purge is None:
|
|
return 0
|
|
return int(await purge(cutoff, batch_size=batch_size) or 0)
|
|
|
|
|
|
async def _purge_snapshot_batch(
|
|
session_factory,
|
|
cutoff_date: str,
|
|
*,
|
|
batch_size: int,
|
|
) -> int:
|
|
"""Delete one short transaction of old leaderboard rows.
|
|
|
|
Snapshot rows use a composite primary key. Selecting those keys first is
|
|
deliberately more conservative than deleting whole historical dates: a
|
|
single date can have many configuration hashes, and deleting all of them
|
|
in one transaction could otherwise hold a write lock for too long.
|
|
"""
|
|
if session_factory is None:
|
|
return 0
|
|
size = min(batch_size, _MAX_SNAPSHOT_BATCH_SIZE)
|
|
async with session_factory() as session:
|
|
keys = list(
|
|
(
|
|
await session.execute(
|
|
select(
|
|
AdminLeaderboardDailyStatRow.stat_date,
|
|
AdminLeaderboardDailyStatRow.settings_hash,
|
|
)
|
|
.where(AdminLeaderboardDailyStatRow.stat_date < cutoff_date)
|
|
.order_by(
|
|
AdminLeaderboardDailyStatRow.stat_date.asc(),
|
|
AdminLeaderboardDailyStatRow.settings_hash.asc(),
|
|
)
|
|
.limit(size)
|
|
)
|
|
).all()
|
|
)
|
|
if not keys:
|
|
return 0
|
|
result = await session.execute(
|
|
delete(AdminLeaderboardDailyStatRow).where(
|
|
or_(
|
|
*(
|
|
and_(
|
|
AdminLeaderboardDailyStatRow.stat_date == stat_date,
|
|
AdminLeaderboardDailyStatRow.settings_hash == settings_hash,
|
|
)
|
|
for stat_date, settings_hash in keys
|
|
)
|
|
)
|
|
)
|
|
)
|
|
await session.commit()
|
|
rowcount = result.rowcount
|
|
return len(keys) if rowcount is None or rowcount < 0 else int(rowcount)
|
|
|
|
|
|
async def sweep_admin_audit_retention(
|
|
*,
|
|
session_factory,
|
|
llm_metrics_store: Any = None,
|
|
tool_metrics_store: Any = None,
|
|
run_store: Any = None,
|
|
retention_days: int = ADMIN_AUDIT_RETENTION_DAYS,
|
|
now: datetime | None = None,
|
|
batch_size: int = _DEFAULT_BATCH_SIZE,
|
|
max_batches_per_store: int = _DEFAULT_MAX_BATCHES_PER_STORE,
|
|
inter_batch_pause_seconds: float = _DEFAULT_INTER_BATCH_PAUSE_SECONDS,
|
|
max_sweep_seconds: float = _DEFAULT_MAX_SWEEP_SECONDS,
|
|
should_stop: Callable[[], bool] | None = None,
|
|
on_progress: Callable[[dict[str, Any]], Awaitable[None] | None] | None = None,
|
|
) -> dict[str, int]:
|
|
"""Run one safe retention pass and return deletion counts.
|
|
|
|
Every delete is a small independent transaction. The sources are visited
|
|
round-robin, and the scheduler checks its wall-clock budget between
|
|
committed batches while cooperatively pausing between them. A historic
|
|
backlog is therefore drained over later sweeps rather than monopolising
|
|
the database or keeping a long write lock.
|
|
"""
|
|
days = max(1, int(retention_days))
|
|
cutoff = (now or datetime.now(UTC)) - timedelta(days=days)
|
|
cutoff_date = cutoff.astimezone(BEIJING_TZ).date().isoformat()
|
|
safe_batch_size = _bounded_int(
|
|
batch_size,
|
|
default=_DEFAULT_BATCH_SIZE,
|
|
maximum=_MAX_BATCH_SIZE,
|
|
)
|
|
safe_max_batches = _bounded_int(
|
|
max_batches_per_store,
|
|
default=_DEFAULT_MAX_BATCHES_PER_STORE,
|
|
maximum=_MAX_BATCHES_PER_STORE,
|
|
)
|
|
safe_pause_seconds = _bounded_seconds(
|
|
inter_batch_pause_seconds,
|
|
default=_DEFAULT_INTER_BATCH_PAUSE_SECONDS,
|
|
maximum=1.0,
|
|
)
|
|
safe_max_sweep_seconds = _bounded_seconds(
|
|
max_sweep_seconds,
|
|
default=_DEFAULT_MAX_SWEEP_SECONDS,
|
|
maximum=60.0,
|
|
)
|
|
if safe_max_sweep_seconds <= 0:
|
|
safe_max_sweep_seconds = _DEFAULT_MAX_SWEEP_SECONDS
|
|
|
|
async def purge_llm_metrics() -> int:
|
|
return await _purge_store_batch(
|
|
llm_metrics_store,
|
|
cutoff,
|
|
batch_size=safe_batch_size,
|
|
)
|
|
|
|
async def purge_tool_metrics() -> int:
|
|
return await _purge_store_batch(
|
|
tool_metrics_store,
|
|
cutoff,
|
|
batch_size=safe_batch_size,
|
|
)
|
|
|
|
async def purge_run_summaries() -> int:
|
|
return await _purge_store_batch(
|
|
run_store,
|
|
cutoff,
|
|
batch_size=safe_batch_size,
|
|
method_name="purge_terminal_before",
|
|
)
|
|
|
|
async def purge_snapshots() -> int:
|
|
return await _purge_snapshot_batch(
|
|
session_factory,
|
|
cutoff_date,
|
|
batch_size=safe_batch_size,
|
|
)
|
|
|
|
sources: tuple[tuple[str, Callable[[], Awaitable[int]], int], ...] = (
|
|
("llm_metrics", purge_llm_metrics, safe_batch_size),
|
|
("tool_metrics", purge_tool_metrics, safe_batch_size),
|
|
("run_summaries", purge_run_summaries, safe_batch_size),
|
|
("leaderboard_snapshots", purge_snapshots, min(safe_batch_size, _MAX_SNAPSHOT_BATCH_SIZE)),
|
|
)
|
|
result = {name: 0 for name, _purge, _full_batch_size in sources}
|
|
active = {name: True for name in result}
|
|
deadline = monotonic() + safe_max_sweep_seconds
|
|
exhausted_budget = False
|
|
|
|
async def report(*, status: str, round_number: int, current_source: str | None) -> None:
|
|
if on_progress is None:
|
|
return
|
|
payload: dict[str, Any] = {
|
|
"status": status,
|
|
"round": round_number,
|
|
"max_rounds": safe_max_batches,
|
|
"current_source": current_source,
|
|
"deleted": dict(result),
|
|
}
|
|
try:
|
|
pending = on_progress(payload)
|
|
if pending is not None:
|
|
await pending
|
|
except Exception:
|
|
logger.warning("Unable to publish admin audit retention progress", exc_info=True)
|
|
|
|
stopped = False
|
|
for _round in range(safe_max_batches):
|
|
any_deletes = False
|
|
for name, purge, full_batch_size in sources:
|
|
if should_stop is not None and should_stop():
|
|
stopped = True
|
|
break
|
|
if not active[name]:
|
|
continue
|
|
if monotonic() >= deadline:
|
|
exhausted_budget = True
|
|
break
|
|
await report(status="running", round_number=_round + 1, current_source=name)
|
|
try:
|
|
count = await purge()
|
|
except Exception:
|
|
# A busy/temporarily unavailable table must not keep this
|
|
# maintenance sweep alive, nor prevent other audit tables
|
|
# from receiving their small cleanup turn. The next scheduled
|
|
# sweep will retry it.
|
|
logger.warning(
|
|
"admin audit retention skipped %s this sweep; will retry later",
|
|
name,
|
|
exc_info=True,
|
|
)
|
|
active[name] = False
|
|
continue
|
|
result[name] += count
|
|
await report(status="running", round_number=_round + 1, current_source=name)
|
|
any_deletes = any_deletes or count > 0
|
|
if count < full_batch_size:
|
|
active[name] = False
|
|
if count and safe_pause_seconds:
|
|
# The transaction above has committed. Yielding here gives
|
|
# interactive reads/writes a chance to acquire the connection
|
|
# before this low-priority maintenance work continues.
|
|
await asyncio.sleep(safe_pause_seconds)
|
|
if stopped or exhausted_budget or not any(active.values()) or not any_deletes:
|
|
break
|
|
|
|
if stopped or (should_stop is not None and should_stop()):
|
|
await report(status="paused", round_number=min(_round + 1, safe_max_batches), current_source=None)
|
|
|
|
more_work_likely = any(active.values())
|
|
if any(result.values()) or exhausted_budget or more_work_likely:
|
|
logger.info(
|
|
"admin audit retention sweep: llm_metrics=%d tool_metrics=%d run_summaries=%d "
|
|
"leaderboard_snapshots=%d cutoff=%s batch_size=%d max_batches=%d "
|
|
"time_budget_seconds=%.2f budget_exhausted=%s more_work_likely=%s",
|
|
result["llm_metrics"],
|
|
result["tool_metrics"],
|
|
result["run_summaries"],
|
|
result["leaderboard_snapshots"],
|
|
cutoff.isoformat(),
|
|
safe_batch_size,
|
|
safe_max_batches,
|
|
safe_max_sweep_seconds,
|
|
exhausted_budget,
|
|
more_work_likely,
|
|
)
|
|
return result
|
|
|
|
|
|
__all__ = ["sweep_admin_audit_retention"]
|