feat(BA-6929): add DB record retention repository for simple/grouped categories (#12952)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: HyeockJinKim

changes/12952.feature.md

@@ -0,0 +1 @@
+Add a DB record retention repository that purges accumulated log, reconcile-history, login, roles/invitation, and usage records past their configured age boundary.

src/ai/backend/manager/data/retention/types.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import enum
+from dataclasses import dataclass
 
 
 class RetentionCategory(enum.StrEnum):
@@ -17,3 +18,15 @@ class RetentionCategory(enum.StrEnum):
     SESSIONS = "sessions"
     USAGE_RECORDS = "usage_records"
     USAGE_BUCKETS = "usage_buckets"
+
+
+@dataclass(frozen=True)
+class RetentionPurgeResult:
+    """Outcome of purging one category's older-than-threshold rows.
+
+    ``deleted_count`` is the total rows removed across the category's tables,
+    letting the sweep account the result against its per-tick budget.
+    """
+
+    category: RetentionCategory
+    deleted_count: int

src/ai/backend/manager/errors/retention.py

@@ -0,0 +1,33 @@
+"""Retention layer exceptions."""
+
+from __future__ import annotations
+
+from typing import override
+
+from ai.backend.common.exception import (
+    ErrorCode,
+    ErrorDetail,
+    ErrorDomain,
+    ErrorOperation,
+)
+from ai.backend.manager.errors.repository import RepositoryError
+
+
+class RetentionCategoryNotSupportedError(RepositoryError):
+    """Raised when a retention category has no code-side cleanup wired yet.
+
+    The ordered-delete categories (``sessions``, ``deployments``,
+    ``usage_buckets``) are implemented separately; requesting one here fails
+    loudly instead of silently deleting nothing.
+    """
+
+    error_type = "https://api.backend.ai/probs/retention-category-not-supported"
+    error_title = "Retention category is not supported."
+
+    @override
+    def error_code(self) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.BACKENDAI,
+            operation=ErrorOperation.GENERIC,
+            error_detail=ErrorDetail.INTERNAL_ERROR,
+        )

src/ai/backend/manager/repositories/repositories.py

@@ -64,6 +64,7 @@
 from ai.backend.manager.repositories.resource_usage_history.repositories import (
     ResourceUsageHistoryRepositories,
 )
+from ai.backend.manager.repositories.retention.repositories import RetentionRepositories
 from ai.backend.manager.repositories.role_preset.repositories import RolePresetRepositories
 from ai.backend.manager.repositories.runtime_variant.repositories import RuntimeVariantRepositories
 from ai.backend.manager.repositories.runtime_variant_preset.repositories import (
@@ -142,6 +143,7 @@ class Repositories:
     events: EventsRepositories
     storage_namespace: StorageNamespaceRepositories
     audit_log: AuditLogRepositories
+    retention: RetentionRepositories
 
     @classmethod
     def create(cls, args: RepositoryArgs) -> Self:
@@ -198,6 +200,7 @@ def create(cls, args: RepositoryArgs) -> Self:
         events_repositories = EventsRepositories.create(args)
         storage_namespace_repositories = StorageNamespaceRepositories.create(args)
         audit_log_repositories = AuditLogRepositories.create(args)
+        retention_repositories = RetentionRepositories.create(args)
 
         return cls(
             agent=agent_repositories,
@@ -251,4 +254,5 @@ def create(cls, args: RepositoryArgs) -> Self:
             events=events_repositories,
             storage_namespace=storage_namespace_repositories,
             audit_log=audit_log_repositories,
+            retention=retention_repositories,
         )

src/ai/backend/manager/repositories/retention/__init__.py

@@ -0,0 +1,189 @@
+"""Database source for retention cleanup.
+
+Holds the ``category -> tables`` catalog (kept inside the repository, never a
+module global) and drains each table with chunk-based delete-and-advance via
+the shared ``batch_purge``.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy.sql.elements import ColumnElement
+
+from ai.backend.logging.utils import BraceStyleAdapter
+from ai.backend.manager.data.auth.login_session_types import LoginSessionStatus
+from ai.backend.manager.data.permission.status import RoleStatus
+from ai.backend.manager.data.retention.types import RetentionCategory, RetentionPurgeResult
+from ai.backend.manager.data.role_invitation.types import RoleInvitationState
+from ai.backend.manager.data.vfolder.types import VFolderInvitationState
+from ai.backend.manager.errors.retention import RetentionCategoryNotSupportedError
+from ai.backend.manager.models.audit_log.row import AuditLogRow
+from ai.backend.manager.models.base import Base
+from ai.backend.manager.models.error_logs import ErrorLogRow
+from ai.backend.manager.models.event_log.row import EventLogRow
+from ai.backend.manager.models.login_session.row import LoginHistoryRow, LoginSessionRow
+from ai.backend.manager.models.rbac_models.role import RoleRow
+from ai.backend.manager.models.replica_group_history.row import ReplicaGroupHistoryRow
+from ai.backend.manager.models.resource_usage_history.row import KernelUsageRecordRow
+from ai.backend.manager.models.role_invitation.row import RoleInvitationRow
+from ai.backend.manager.models.scheduling_history.row import (
+    DeploymentHistoryRow,
+    KernelSchedulingHistoryRow,
+    RouteHistoryRow,
+    SessionSchedulingHistoryRow,
+)
+from ai.backend.manager.models.vfolder.row import VFolderInvitationRow
+from ai.backend.manager.repositories.base import BatchPurger, BatchPurgerSpec
+from ai.backend.manager.repositories.ops import DBOpsProvider
+from ai.backend.manager.repositories.retention.purgers import TimestampBoundaryPurgerSpec
+
+log = BraceStyleAdapter(logging.getLogger(__spec__.name))
+
+# Invitation states that are terminal (no further update expected), so their
+# proxy timestamp is a safe grace boundary. Only PENDING is non-terminal.
+_TERMINAL_ROLE_INVITATION_STATES = (
+    RoleInvitationState.ACCEPTED,
+    RoleInvitationState.REJECTED,
+    RoleInvitationState.CANCELED,
+)
+_TERMINAL_VFOLDER_INVITATION_STATES = (
+    VFolderInvitationState.ACCEPTED,
+    VFolderInvitationState.REJECTED,
+    VFolderInvitationState.CANCELED,
+)
+
+
+@dataclass(frozen=True)
+class _BoundaryTable:
+    """One table's fixed cleanup definition; only ``threshold`` varies per sweep.
+
+    The boundary column is chosen by table nature: append-only logs use
+    ``created_at``, in-place-merged history uses ``updated_at``, and lifecycle
+    records use their terminal timestamp plus a terminal-status
+    ``extra_conditions`` filter.
+    """
+
+    row_class: type[Base]
+    boundary: Any
+    extra_conditions: Sequence[ColumnElement[bool]] = field(default_factory=tuple)
+
+
+class RetentionDBSource:
+    _ops: DBOpsProvider
+    _catalog: Mapping[RetentionCategory, Sequence[_BoundaryTable]]
+
+    def __init__(self, ops_provider: DBOpsProvider) -> None:
+        self._ops = ops_provider
+        self._catalog = self._build_catalog()
+
+    @staticmethod
+    def _build_catalog() -> Mapping[RetentionCategory, Sequence[_BoundaryTable]]:
+        """Build the fixed ``category -> tables`` catalog once (column refs are
+        constant; only the threshold is applied per sweep).
+
+        Categories with a bespoke ordered delete (sessions, deployments,
+        usage_buckets) are implemented separately and intentionally absent.
+        """
+        return {
+            RetentionCategory.LOGS: (
+                _BoundaryTable(EventLogRow, EventLogRow.created_at),
+                _BoundaryTable(AuditLogRow, AuditLogRow.created_at),
+                # error_logs purges purely on the boundary — is_read/is_cleared
+                # flags are intentionally ignored (all rows past boundary go).
+                _BoundaryTable(ErrorLogRow, ErrorLogRow.created_at),
+            ),
+            # updated_at (not created_at): attempts++ merges touch updated_at,
+            # so a recently-retried row keeps an old created_at but survives.
+            RetentionCategory.RECONCILE_HISTORY: (
+                _BoundaryTable(SessionSchedulingHistoryRow, SessionSchedulingHistoryRow.updated_at),
+                _BoundaryTable(KernelSchedulingHistoryRow, KernelSchedulingHistoryRow.updated_at),
+                _BoundaryTable(DeploymentHistoryRow, DeploymentHistoryRow.updated_at),
+                _BoundaryTable(RouteHistoryRow, RouteHistoryRow.updated_at),
+                _BoundaryTable(ReplicaGroupHistoryRow, ReplicaGroupHistoryRow.updated_at),
+            ),
+            RetentionCategory.LOGIN: (
+                _BoundaryTable(LoginHistoryRow, LoginHistoryRow.created_at),
+                _BoundaryTable(
+                    LoginSessionRow,
+                    LoginSessionRow.invalidated_at,
+                    (
+                        LoginSessionRow.status.in_((
+                            LoginSessionStatus.INVALIDATED,
+                            LoginSessionStatus.REVOKED,
+                        )),
+                    ),
+                ),
+            ),
+            RetentionCategory.ROLES_INVITATIONS: (
+                _BoundaryTable(
+                    RoleRow,
+                    RoleRow.deleted_at,
+                    (RoleRow.status == RoleStatus.DELETED,),
+                ),
+                _BoundaryTable(
+                    RoleInvitationRow,
+                    RoleInvitationRow.updated_at,
+                    (RoleInvitationRow.state.in_(_TERMINAL_ROLE_INVITATION_STATES),),
+                ),
+                _BoundaryTable(
+                    VFolderInvitationRow,
+                    VFolderInvitationRow.modified_at,
+                    (VFolderInvitationRow.state.in_(_TERMINAL_VFOLDER_INVITATION_STATES),),
+                ),
+            ),
+            RetentionCategory.USAGE_RECORDS: (
+                _BoundaryTable(KernelUsageRecordRow, KernelUsageRecordRow.period_end),
+            ),
+        }
+
+    def _purger_specs(
+        self,
+        category: RetentionCategory,
+        threshold: datetime,
+    ) -> list[BatchPurgerSpec[Any]]:
+        """Look up the category's tables and bind ``threshold`` into each spec."""
+        tables = self._catalog.get(category)
+        if tables is None:
+            raise RetentionCategoryNotSupportedError(
+                f"Retention category '{category.value}' has no simple/grouped "
+                "cleanup wired in this repository."
+            )
+        return [
+            TimestampBoundaryPurgerSpec(
+                t.row_class, t.boundary, threshold, extra_conditions=t.extra_conditions
+            )
+            for t in tables
+        ]
+
+    async def purge_older_than(
+        self,
+        category: RetentionCategory,
+        threshold: datetime,
+        batch_size: int,
+    ) -> RetentionPurgeResult:
+        """Delete every row of ``category`` older than ``threshold``.
+
+        Each of the category's tables is drained in its own transaction via
+        ``batch_purge``, which deletes in ``batch_size`` chunks (delete-and-
+        advance) so a large backlog never becomes a single huge DELETE.
+        """
+        specs = self._purger_specs(category, threshold)
+        total_deleted = 0
+
+        for spec in specs:
+            async with self._ops.write_ops() as w:
+                result = await w.batch_purge(BatchPurger(spec=spec, batch_size=batch_size))
+                total_deleted += result.deleted_count
+
+        log.debug(
+            "retention purge category={} threshold={} deleted={}",
+            category.value,
+            threshold,
+            total_deleted,
+        )
+        return RetentionPurgeResult(category=category, deleted_count=total_deleted)

src/ai/backend/manager/repositories/retention/db_source/__init__.py

@@ -0,0 +1,48 @@
+"""Batch-purge specs for retention cleanup.
+
+A single :class:`TimestampBoundaryPurgerSpec` covers every retention target:
+common-column groups (``logs`` on ``created_at``, ``reconcile_history`` on
+``updated_at``) reuse one instance per table, simple categories reuse it
+directly, and lifecycle categories add a terminal-state filter through
+``extra_conditions``. Keeping the boundary logic in one spec is the shared
+"mixin" the design calls for — the ``category -> specs`` mapping lives in the
+DB source, not here.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, override
+
+import sqlalchemy as sa
+from sqlalchemy.sql.elements import ColumnElement
+
+from ai.backend.manager.models.base import Base
+from ai.backend.manager.repositories.base.purger import BatchPurgerSpec
+
+
+@dataclass
+class TimestampBoundaryPurgerSpec[TRow: Base](BatchPurgerSpec[TRow]):
+    """Selects rows older than ``threshold`` on a single boundary timestamp.
+
+    Rows whose ``boundary`` column is NULL are never selected, so a lifecycle
+    record still lacking its terminal timestamp is preserved. ``extra_conditions``
+    appends terminal-status / lifecycle filters (e.g. ``status == DELETED``).
+    """
+
+    row_class: type[TRow]
+    # A mapped timestamp column expression. Any-typed because targets range over
+    # declaratively-mapped attributes and imperatively-mapped ones (error_logs),
+    # and over both non-null and nullable lifecycle columns.
+    boundary: Any
+    threshold: datetime
+    extra_conditions: Sequence[ColumnElement[bool]] = field(default_factory=tuple)
+
+    @override
+    def build_subquery(self) -> sa.sql.Select[tuple[TRow]]:
+        stmt = sa.select(self.row_class).where(self.boundary < self.threshold)
+        for condition in self.extra_conditions:
+            stmt = stmt.where(condition)
+        return stmt

src/ai/backend/manager/repositories/retention/db_source/db_source.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+from typing import Self
+
+from ai.backend.manager.repositories.retention.repository import RetentionRepository
+from ai.backend.manager.repositories.types import RepositoryArgs
+
+
+@dataclass
+class RetentionRepositories:
+    repository: RetentionRepository
+
+    @classmethod
+    def create(cls, args: RepositoryArgs) -> Self:
+        return cls(
+            repository=RetentionRepository(args.ops_provider),
+        )

src/ai/backend/manager/repositories/retention/purgers.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+
+from ai.backend.logging import BraceStyleAdapter
+from ai.backend.manager.data.retention.types import RetentionCategory, RetentionPurgeResult
+from ai.backend.manager.repositories.ops import DBOpsProvider
+from ai.backend.manager.repositories.retention.db_source.db_source import RetentionDBSource
+
+log = BraceStyleAdapter(logging.getLogger(__spec__.name))
+
+
+class RetentionRepository:
+    """Deletes accumulated DB records past their category's age boundary.
+
+    A single repository owns the ``category -> tables`` mapping and reuses the
+    ``BatchPurger`` framework for chunk-based (delete-and-advance) deletes. The
+    leader sweep (separate task) computes ``threshold = now - retention_period``
+    per policy and calls :meth:`purge_older_than`.
+    """
+
+    _db_source: RetentionDBSource
+
+    def __init__(self, ops_provider: DBOpsProvider) -> None:
+        self._db_source = RetentionDBSource(ops_provider)
+
+    async def purge_older_than(
+        self,
+        category: RetentionCategory,
+        threshold: datetime,
+        batch_size: int,
+    ) -> RetentionPurgeResult:
+        """Purge rows of ``category`` older than ``threshold``, chunked per table."""
+        return await self._db_source.purge_older_than(category, threshold, batch_size)