deploying.py

"""Handlers for DEPLOYING sub-steps (BEP-1049).

Two DEPLOYING handlers are registered in the coordinator's HandlerRegistry:

- **DeployingProvisioningHandler**: Runs the strategy FSM each cycle to
  create/drain routes and check for completion.
- **DeployingRollingBackHandler**: Clears ``deploying_revision`` and
  transitions directly to READY.

Sub-step flow::

    PROVISIONING ──(need_retry)──▸ PROVISIONING  (route mutations, logged)
         │
         │ (success)
         ▼
       READY  (completed — all routes replaced)

    PROVISIONING ──(timeout)──▸ ROLLING_BACK ──(success)──▸ READY

The evaluator determines sub-step assignments and route mutations;
the applier persists them to DB atomically.  Each handler classifies
deployments into successes (transition forward), need_retry (route mutations
with history logged), and skipped (no change — waiting).
"""

from __future__ import annotations

import asyncio
import dataclasses
import logging
from collections.abc import Coroutine, Sequence
from typing import Any, override
from uuid import UUID

from ai.backend.logging import BraceStyleAdapter
from ai.backend.manager.data.deployment.types import (
    DeploymentLifecycleStatus,
    DeploymentLifecycleSubStep,
    DeploymentStatusTransitions,
)
from ai.backend.manager.data.model_serving.types import EndpointLifecycle
from ai.backend.manager.defs import LockID
from ai.backend.manager.repositories.deployment.repository import DeploymentRepository
from ai.backend.manager.sokovan.deployment.deployment_controller import DeploymentController
from ai.backend.manager.sokovan.deployment.executor import DeploymentExecutor
from ai.backend.manager.sokovan.deployment.route.route_controller import RouteController
from ai.backend.manager.sokovan.deployment.route.types import RouteLifecycleType
from ai.backend.manager.sokovan.deployment.strategy.applier import (
    StrategyResultApplier,
)
from ai.backend.manager.sokovan.deployment.strategy.evaluator import (
    DeploymentStrategyEvaluator,
)
from ai.backend.manager.sokovan.deployment.types import (
    DeploymentExecutionError,
    DeploymentExecutionResult,
    DeploymentLifecycleType,
    DeploymentWithHistory,
)

from .base import DeploymentHandler

log = BraceStyleAdapter(logging.getLogger(__name__))


@dataclasses.dataclass(frozen=True)
class _EndpointRegistrationBatch:
    deployments: list[DeploymentWithHistory]
    coroutines: list[Coroutine[Any, Any, str]]


# ---------------------------------------------------------------------------
# DEPLOYING sub-step handlers
# ---------------------------------------------------------------------------


class DeployingProvisioningHandler(DeploymentHandler):
    """Handler for the DEPLOYING / PROVISIONING sub-step.

    Runs the strategy FSM each cycle to create/drain routes and check
    for completion.  Classification:

    - **Route mutations executed** (create/drain): need_retry — stays in
      PROVISIONING with a new history record for progress tracking.
      Never escalated to give_up (normal progress).
    - **No changes** (routes still warming up): skipped — no history.
    - **Completed** (all old routes replaced): success → READY.
    """

    def __init__(
        self,
        deployment_controller: DeploymentController,
        route_controller: RouteController,
        evaluator: DeploymentStrategyEvaluator,
        applier: StrategyResultApplier,
        deployment_executor: DeploymentExecutor,
        deployment_repo: DeploymentRepository,
    ) -> None:
        self._deployment_controller = deployment_controller
        self._route_controller = route_controller
        self._evaluator = evaluator
        self._applier = applier
        self._deployment_executor = deployment_executor
        self._deployment_repo = deployment_repo

    @classmethod
    @override
    def name(cls) -> str:
        return "deploying-provisioning"

    @property
    @override
    def lock_id(self) -> LockID | None:
        return LockID.LOCKID_DEPLOYMENT_DEPLOYING

    @classmethod
    @override
    def target_statuses(cls) -> list[DeploymentLifecycleStatus]:
        return [
            DeploymentLifecycleStatus(
                lifecycle=EndpointLifecycle.DEPLOYING,
                sub_step=DeploymentLifecycleSubStep.DEPLOYING_PROVISIONING,
            ),
        ]

    @classmethod
    @override
    def status_transitions(cls) -> DeploymentStatusTransitions:
        return DeploymentStatusTransitions(
            success=DeploymentLifecycleStatus(
                lifecycle=EndpointLifecycle.READY,
                sub_step=None,
            ),
            need_retry=DeploymentLifecycleStatus(
                lifecycle=EndpointLifecycle.DEPLOYING,
                sub_step=DeploymentLifecycleSubStep.DEPLOYING_PROVISIONING,
            ),
            expired=DeploymentLifecycleStatus(
                lifecycle=EndpointLifecycle.DEPLOYING,
                sub_step=DeploymentLifecycleSubStep.DEPLOYING_ROLLING_BACK,
            ),
        )

    async def _ensure_endpoints_registered(
        self, deployments: Sequence[DeploymentWithHistory]
    ) -> None:
        """Register endpoints in appproxy for deployments that have no URL yet.

        Deployments that entered DEPLOYING via ActivateRevision skip
        check_pending (which normally registers them), so this method
        ensures they are registered before route provisioning begins.
        """
        unregistered = [
            d
            for d in deployments
            if not d.deployment_info.network.url and d.deployment_info.deploying_revision_id
        ]
        if not unregistered:
            return

        batch = await self._build_registration_batch(unregistered)
        if not batch.coroutines:
            return

        url_updates = await self._execute_registration_batch(batch)
        if url_updates:
            await self._deployment_repo.update_endpoint_urls_bulk(url_updates)

    async def _build_registration_batch(
        self, deployments: Sequence[DeploymentWithHistory]
    ) -> _EndpointRegistrationBatch:
        """Build registration coroutines for deployments with valid proxy targets."""
        scaling_groups = {d.deployment_info.metadata.resource_group for d in deployments}
        scaling_group_targets = await self._deployment_repo.fetch_scaling_group_proxy_targets(
            scaling_groups
        )

        valid_deployments: list[DeploymentWithHistory] = []
        coroutines: list[Coroutine[Any, Any, str]] = []
        for deployment in deployments:
            info = deployment.deployment_info
            target = scaling_group_targets.get(info.metadata.resource_group)
            if not target:
                log.warning(
                    "No proxy target for scaling group {}, skipping endpoint registration for {}",
                    info.metadata.resource_group,
                    info.id,
                )
                continue
            if info.deploying_revision_id is None:
                # May have been cleared between handler start and registration.
                continue
            coroutines.append(
                self._deployment_executor.register_endpoint(
                    info, target, info.deploying_revision_id
                )
            )
            valid_deployments.append(deployment)
        return _EndpointRegistrationBatch(valid_deployments, coroutines)

    @staticmethod
    async def _execute_registration_batch(
        batch: _EndpointRegistrationBatch,
    ) -> dict[UUID, str]:
        """Execute registration batch and collect successful URL updates."""
        results = await asyncio.gather(*batch.coroutines, return_exceptions=True)

        url_updates: dict[UUID, str] = {}
        for deployment, result in zip(batch.deployments, results, strict=True):
            if isinstance(result, BaseException):
                log.error(
                    "Failed to register endpoint for deployment {}: {}",
                    deployment.deployment_info.id,
                    result,
                )
            else:
                url_updates[deployment.deployment_info.id] = result
        return url_updates

    @override
    async def execute(
        self, deployments: Sequence[DeploymentWithHistory]
    ) -> DeploymentExecutionResult:
        # Register endpoints in appproxy for deployments that entered DEPLOYING
        # via ActivateRevision without passing through check_pending.
        await self._ensure_endpoints_registered(deployments)

        deployment_infos = [d.deployment_info for d in deployments]
        deployment_map = {d.deployment_info.id: d for d in deployments}

        summary = await self._evaluator.evaluate(deployment_infos)
        apply_result = await self._applier.apply(summary)

        # Filter out deployments marked for destruction during DEPLOYING.
        destroying_ids = {
            d.deployment_info.id
            for d in deployments
            if d.deployment_info.state.lifecycle
            in (EndpointLifecycle.DESTROYING, EndpointLifecycle.DESTROYED)
        }
        if destroying_ids:
            log.warning(
                "Skipping {} deployments with DESTROYING/DESTROYED lifecycle during DEPLOYING",
                len(destroying_ids),
            )

        successes: list[DeploymentWithHistory] = []
        failures: list[DeploymentExecutionError] = []
        skipped: list[DeploymentWithHistory] = []

        # COMPLETED → success (coordinator transitions to READY)
        for endpoint_id in apply_result.completed_ids:
            if endpoint_id in destroying_ids:
                continue
            deployment = deployment_map.get(endpoint_id)
            if deployment is not None:
                successes.append(deployment)

        # Evaluation errors → failures
        for error_data in summary.errors:
            deployment = deployment_map.get(error_data.deployment.id)
            if deployment is not None:
                failures.append(
                    DeploymentExecutionError(
                        deployment_info=deployment,
                        reason=error_data.reason,
                        error_detail=error_data.reason,
                    )
                )

        # Classify rest: route mutations happened → failures (recoverable, coordinator
        # will classify as need_retry), no changes → skipped (no history).
        completed_or_error_ids = apply_result.completed_ids | {
            e.deployment.id for e in summary.errors
        }
        has_route_mutations = bool(apply_result.routes_created or apply_result.routes_drained)
        for deployment in deployments:
            endpoint_id = deployment.deployment_info.id
            if endpoint_id in completed_or_error_ids or endpoint_id in destroying_ids:
                continue
            if has_route_mutations:
                failures.append(
                    DeploymentExecutionError(
                        deployment_info=deployment,
                        reason="Route mutations in progress",
                        error_detail="Waiting for route provisioning to complete",
                    )
                )
            else:
                skipped.append(deployment)

        return DeploymentExecutionResult(successes=successes, failures=failures, skipped=skipped)

    @override
    async def post_process(self, result: DeploymentExecutionResult) -> None:
        await self._deployment_controller.mark_lifecycle_needed(
            DeploymentLifecycleType.DEPLOYING,
            sub_step=DeploymentLifecycleSubStep.DEPLOYING_PROVISIONING,
        )
        await self._deployment_controller.mark_lifecycle_needed(
            DeploymentLifecycleType.DEPLOYING,
            sub_step=DeploymentLifecycleSubStep.DEPLOYING_ROLLING_BACK,
        )
        await self._route_controller.mark_lifecycle_needed(RouteLifecycleType.PROVISIONING)


class DeployingRollingBackHandler(DeploymentHandler):
    """Handler for DEPLOYING / ROLLING_BACK sub-step.

    Clears ``deploying_revision`` and transitions directly to READY.
    """

    def __init__(
        self,
        deployment_controller: DeploymentController,
        route_controller: RouteController,
        deployment_repo: DeploymentRepository,
    ) -> None:
        self._deployment_controller = deployment_controller
        self._route_controller = route_controller
        self._deployment_repo = deployment_repo

    @classmethod
    @override
    def name(cls) -> str:
        return "deploying-rolling-back"

    @property
    @override
    def lock_id(self) -> LockID | None:
        return LockID.LOCKID_DEPLOYMENT_DEPLOYING

    @classmethod
    @override
    def target_statuses(cls) -> list[DeploymentLifecycleStatus]:
        return [
            DeploymentLifecycleStatus(
                lifecycle=EndpointLifecycle.DEPLOYING,
                sub_step=DeploymentLifecycleSubStep.DEPLOYING_ROLLING_BACK,
            )
        ]

    @classmethod
    @override
    def status_transitions(cls) -> DeploymentStatusTransitions:
        return DeploymentStatusTransitions(
            success=DeploymentLifecycleStatus(
                lifecycle=EndpointLifecycle.READY,
                sub_step=None,
            ),
        )

    @override
    async def execute(
        self, deployments: Sequence[DeploymentWithHistory]
    ) -> DeploymentExecutionResult:
        all_deployment_ids = {deployment.deployment_info.id for deployment in deployments}
        await self._deployment_repo.clear_deploying_revision(all_deployment_ids)
        log.info(
            "Cleared deploying_revision for {} rolling-back deployments",
            len(all_deployment_ids),
        )

        return DeploymentExecutionResult(successes=list(deployments))

    @override
    async def post_process(self, result: DeploymentExecutionResult) -> None:
        await self._deployment_controller.mark_lifecycle_needed(
            DeploymentLifecycleType.DEPLOYING,
            sub_step=DeploymentLifecycleSubStep.DEPLOYING_ROLLING_BACK,
        )
        await self._route_controller.mark_lifecycle_needed(RouteLifecycleType.PROVISIONING)