evaluator.py

"""Deployment strategy evaluator — orchestrates per-deployment FSM evaluation (BEP-1049).

Loads policies and non-terminated routes in bulk, dispatches each deployment to
the appropriate strategy FSM, and aggregates route mutations.  The applier is
responsible for applying the aggregated route changes and updating sub_step in DB.
"""

from __future__ import annotations

import logging
from collections import defaultdict
from collections.abc import Sequence
from uuid import UUID

from ai.backend.common.data.model_deployment.types import DeploymentStrategy
from ai.backend.common.exception import BackendAIError
from ai.backend.logging import BraceStyleAdapter
from ai.backend.manager.data.deployment.types import (
    DeploymentInfo,
    DeploymentPolicyData,
    RouteInfo,
    RouteStatus,
)
from ai.backend.manager.errors.deployment import (
    InvalidDeploymentStrategy,
    InvalidDeploymentStrategySpec,
)
from ai.backend.manager.repositories.base import (
    BatchQuerier,
    NoPagination,
    QueryCondition,
    combine_conditions_or,
)
from ai.backend.manager.repositories.deployment.options import (
    DeploymentPolicyConditions,
    RouteConditions,
)
from ai.backend.manager.repositories.deployment.repository import DeploymentRepository
from ai.backend.manager.sokovan.deployment.recorder import DeploymentRecorderContext

from .types import (
    AbstractDeploymentStrategy,
    DeploymentStrategyRegistry,
    EvaluationErrorData,
    RouteChanges,
    StrategyEvaluationSummary,
)

log = BraceStyleAdapter(logging.getLogger(__name__))


class DeploymentStrategyEvaluator:
    """Evaluates DEPLOYING deployments and produces sub_step assignments + route mutations."""

    def __init__(
        self,
        deployment_repo: DeploymentRepository,
        strategy_registry: DeploymentStrategyRegistry,
    ) -> None:
        self._deployment_repo = deployment_repo
        self._strategy_registry = strategy_registry

    async def evaluate(
        self,
        deployments: Sequence[DeploymentInfo],
    ) -> StrategyEvaluationSummary:
        """Evaluate all DEPLOYING deployments in a single cycle.

        Steps:
            1. Bulk-load policies and non-terminated routes.
            2. Per-deployment: dispatch to strategy FSM.
            3. Aggregate route changes and sub_step assignments.
        """
        result = StrategyEvaluationSummary()

        if not deployments:
            return result

        endpoint_ids = {deployment.id for deployment in deployments}

        # ── 1. Bulk-load policies and routes ──
        policy_search = await self._deployment_repo.search_deployment_policies(
            BatchQuerier(
                pagination=NoPagination(),
                conditions=[DeploymentPolicyConditions.by_endpoint_ids(endpoint_ids)],
            )
        )
        policy_map = {policy.endpoint: policy for policy in policy_search.items}
        # Fetch non-terminated routes + terminated routes belonging to a
        # deploying revision.  The FSM needs terminated new-revision routes
        # to count accumulated failures for rollback detection, but old
        # terminated routes are irrelevant and would bloat the result set.
        deploying_revision_ids = {
            deployment.deploying_revision_id
            for deployment in deployments
            if deployment.deploying_revision_id is not None
        }
        route_conditions: list[QueryCondition] = [
            RouteConditions.by_endpoint_ids(endpoint_ids),
        ]
        if deploying_revision_ids:
            route_conditions.append(
                combine_conditions_or([
                    RouteConditions.exclude_statuses([RouteStatus.TERMINATED]),
                    RouteConditions.by_revision_ids(deploying_revision_ids),
                ])
            )
        else:
            route_conditions.append(
                RouteConditions.exclude_statuses([RouteStatus.TERMINATED]),
            )
        route_search = await self._deployment_repo.search_routes(
            BatchQuerier(
                pagination=NoPagination(),
                conditions=route_conditions,
            )
        )
        route_map: defaultdict[UUID, list[RouteInfo]] = defaultdict(list)
        for route in route_search.items:
            route_map[route.endpoint_id].append(route)

        # ── 2. Per-deployment evaluation ──
        for deployment in deployments:
            policy = policy_map.get(deployment.id)
            if policy is None:
                log.warning("deployment {}: no policy found — skipping", deployment.id)
                result.errors.append(
                    EvaluationErrorData(deployment=deployment, reason="No deployment policy found")
                )
                continue

            routes: list[RouteInfo] = list(route_map.get(deployment.id, []))

            try:
                strategy = self._create_strategy(policy.strategy, policy)
                cycle_result = strategy.evaluate_cycle(deployment, routes, policy.strategy_spec)
            except BackendAIError as e:
                log.warning("deployment {}: evaluation error — {}", deployment.id, e)
                result.errors.append(EvaluationErrorData(deployment=deployment, reason=str(e)))
                continue
            except Exception:
                log.exception("deployment {}: unexpected evaluation error", deployment.id)
                result.errors.append(
                    EvaluationErrorData(deployment=deployment, reason="Unexpected evaluation error")
                )
                continue

            # ── 3. Aggregate route changes and record sub-steps ──
            changes = cycle_result.route_changes
            result.route_changes.rollout_specs.extend(changes.rollout_specs)
            result.route_changes.drain_route_ids.extend(changes.drain_route_ids)
            self._record_route_changes(deployment, changes)

            # Classify into assignments
            result.assignments[deployment.id] = cycle_result.sub_step

        return result

    def _record_route_changes(self, deployment: DeploymentInfo, changes: RouteChanges) -> None:
        """Record rollout/drain operations as sub-steps for observability."""
        if not changes.rollout_specs and not changes.drain_route_ids:
            return
        pool = DeploymentRecorderContext.current_pool()
        recorder = pool.recorder(deployment.id)
        with recorder.phase("route_mutations"):
            if changes.rollout_specs:
                with recorder.step(
                    "rollout",
                    success_detail=f"{len(changes.rollout_specs)} new route(s)",
                ):
                    pass
            if changes.drain_route_ids:
                with recorder.step(
                    "drain",
                    success_detail=f"{len(changes.drain_route_ids)} route(s)",
                ):
                    pass

    def _create_strategy(
        self,
        strategy: DeploymentStrategy,
        policy: DeploymentPolicyData,
    ) -> AbstractDeploymentStrategy:
        """Create a strategy instance for the given deployment policy."""
        entry = self._strategy_registry.get(strategy)
        if entry is None:
            raise InvalidDeploymentStrategy(
                extra_msg=f"Unsupported deployment strategy: {strategy}"
            )
        spec = policy.strategy_spec
        if not isinstance(spec, entry.spec_type):
            raise InvalidDeploymentStrategySpec(
                extra_msg=(
                    f"Expected {entry.spec_type.__name__} for {strategy.name} strategy,"
                    f" got {type(spec).__name__}"
                ),
            )
        return entry.strategy_cls()