02: Auto-Scaling and Utilization for GenAI Workloads

graph LR
    subgraph "What Request-Count Scaling Sees"
        A[5x request spike] --> B[Scale to 5x capacity]
    end

    subgraph "What Actually Happens"
        C[5x request spike] --> D[Heavier intent mix<br>+2x Sonnet ratio]
        D --> E[Longer token counts<br>+3x avg tokens]
        E --> F[Effective demand: 15x]
    end

    B --> G[Result: Severe<br>under-provisioning]
    F --> G

gantt
    title MangaAssist Annual Capacity Calendar
    dateFormat  YYYY-MM-DD
    axisFormat  %b

    section Sustained High
    New Year Sale (5x)          :crit, 2026-01-01, 7d
    Golden Week (3x)            :crit, 2026-04-29, 7d
    Amazon Prime Day (4x)       :crit, 2026-07-15, 3d
    Summer Comiket (3x)         :crit, 2026-08-10, 4d
    Black Friday (5x)           :crit, 2026-11-27, 4d
    Cyber Monday (4x)           :crit, 2026-11-30, 2d
    Year-End Sale (4x)          :crit, 2026-12-20, 12d

    section Burst Events
    Jump Festa (3x)             :active, 2026-12-19, 2d
    Manga Award Season (2x)     :active, 2026-03-01, 14d
    Anime Season Premieres (2x) :active, 2026-01-10, 7d
    Anime Season Premieres (2x) :active, 2026-04-10, 7d
    Anime Season Premieres (2x) :active, 2026-07-10, 7d
    Anime Season Premieres (2x) :active, 2026-10-10, 7d

graph TD
    subgraph "ECS Task Metrics Published Every 10s"
        M1[bedrock_queue_depth<br>In-flight Bedrock calls waiting]
        M2[active_bedrock_invocations<br>Concurrent calls to Bedrock]
        M3[request_latency_p99<br>99th percentile response time]
    end

    subgraph "CloudWatch Aggregation"
        M1 --> A1[Aggregate: Maximum across tasks]
        M2 --> A2[Aggregate: Sum across tasks]
        M3 --> A3[Aggregate: Maximum across tasks]
    end

    subgraph "Scaling Policies"
        A1 --> P1[Step Scaling Policy<br>Queue Depth]
        A2 --> P2[Target Tracking Policy<br>Invocations per Task = 40]
        A3 --> P3[Step Scaling Policy<br>Latency Breach]
    end

    subgraph "ECS Auto Scaling"
        P1 --> S[Desired Task Count]
        P2 --> S
        P3 --> S
        S --> T[ECS Fargate Tasks<br>min=5, max=100]
    end

flowchart TD
    A[Bedrock InvocationThrottles > 0] --> B{Provisioned<br>throughput active?}

    B -->|Yes| C[Increase provisioned<br>model units via API]
    B -->|No| D[Enable provisioned<br>throughput immediately]

    A --> E[Activate request<br>coalescing aggressively]
    A --> F[Shift eligible queries<br>from Sonnet to Haiku]
    A --> G[Extend Redis cache TTL<br>from 30s to 120s]
    A --> H[Enable degraded mode:<br>return cached/template<br>responses for FAQ intent]

    C --> I[Monitor: Throttles<br>should drop within 5 min]
    D --> I
    E --> I
    F --> I
    G --> I
    H --> I

    I -->|Throttles persist| J[Escalate: Page on-call,<br>consider cross-region failover]
    I -->|Throttles resolved| K[Gradually revert<br>to normal settings]

Predicted TPS in 5 minutes = current_TPS + (TPS_slope * 300)

sequenceDiagram
    participant CAL as Event Calendar
    participant SCHED as Throughput Scheduler
    participant BPT as Bedrock Provisioned<br>Throughput API
    participant CW as CloudWatch
    participant OPS as On-Call Engineer

    CAL->>SCHED: Monday manga release in 3 hours
    SCHED->>BPT: Increase to event capacity (25 Sonnet + 8 Haiku)
    BPT-->>SCHED: Provisioning in progress (ETA 15 min)
    SCHED->>CW: Publish "ProvisioningState: SCALING" metric

    Note over BPT: 15 minutes later...
    BPT-->>SCHED: Provisioning complete
    SCHED->>CW: Publish "ProvisioningState: ACTIVE" metric

    Note over CAL: Manga release begins
    CW-->>SCHED: TPS climbing: 15K -> 25K -> 32K

    alt TPS exceeds provisioned capacity
        SCHED->>BPT: Emergency increase request
        SCHED->>OPS: Alert: approaching provisioned ceiling
    else TPS within capacity
        SCHED->>CW: All nominal
    end

    Note over CAL: 4 hours later, event subsides
    SCHED->>BPT: Scale back to evening baseline (15 Sonnet + 5 Haiku)

"""
GenAI-aware auto-scaler for MangaAssist.
Combines ECS task scaling with Bedrock provisioned throughput adjustment.
"""

import time
import json
import logging
from dataclasses import dataclass
from datetime import datetime, timezone, timedelta
from enum import Enum
from typing import Optional

import boto3

logger = logging.getLogger("mangaassist.autoscaler")

JST = timezone(timedelta(hours=9))


class ScalingState(Enum):
    STABLE = "stable"
    SCALING_OUT = "scaling_out"
    SCALING_IN = "scaling_in"
    EMERGENCY = "emergency"
    EVENT_OVERRIDE = "event_override"


@dataclass
class TrafficSnapshot:
    """Point-in-time traffic measurements."""
    timestamp: float
    tokens_per_second: float
    bedrock_queue_depth: int
    active_invocations: int
    throttle_count: int
    p99_latency_ms: float
    sonnet_ratio: float  # Fraction of invocations going to Sonnet vs Haiku
    ecs_task_count: int
    provisioned_sonnet_units: int
    provisioned_haiku_units: int


@dataclass
class ScalingDecision:
    """Output of the scaling evaluation."""
    state: ScalingState
    ecs_desired_count: Optional[int]
    sonnet_units: Optional[int]
    haiku_units: Optional[int]
    reason: str
    actions: list[str]


class GenAIAutoScaler:
    """
    Auto-scaler for MangaAssist that understands FM workload characteristics.

    Key differences from a generic auto-scaler:
    1. Scales on tokens-per-second and queue depth, not CPU.
    2. Adjusts Bedrock provisioned throughput alongside ECS tasks.
    3. Uses traffic slope prediction to scale proactively.
    4. Integrates with an event calendar for pre-provisioning.
    5. Implements throttle-response strategies beyond just adding capacity.
    """

    # Capacity constants
    SONNET_TPS_PER_UNIT = 1_000
    HAIKU_TPS_PER_UNIT = 2_000
    TPS_PER_ECS_TASK = 500
    ECS_MIN_TASKS = 5
    ECS_MAX_TASKS = 100

    # Thresholds
    TPS_UTILIZATION_SCALE_OUT = 0.80  # Scale out at 80% utilization
    TPS_UTILIZATION_SCALE_IN = 0.40   # Scale in at 40% utilization
    TPS_UTILIZATION_EMERGENCY = 0.95  # Emergency at 95%
    QUEUE_DEPTH_SCALE_OUT = 50
    QUEUE_DEPTH_EMERGENCY = 200
    LATENCY_SCALE_OUT_MS = 4_000
    THROTTLE_EMERGENCY = 1  # Any throttle is an emergency

    # Prediction
    PREDICTION_WINDOW_SECONDS = 300  # Look ahead 5 minutes

    def __init__(
        self,
        ecs_cluster: str,
        ecs_service: str,
        region: str = "ap-northeast-1",
        event_calendar: Optional[dict] = None,
    ):
        self.ecs_cluster = ecs_cluster
        self.ecs_service = ecs_service
        self.region = region
        self.event_calendar = event_calendar or {}

        self._ecs_client = boto3.client("ecs", region_name=region)
        self._bedrock_client = boto3.client("bedrock", region_name=region)
        self._cloudwatch_client = boto3.client("cloudwatch", region_name=region)
        self._autoscaling_client = boto3.client("application-autoscaling", region_name=region)

        # Rolling window of snapshots for slope calculation
        self._snapshot_history: list[TrafficSnapshot] = []
        self._max_history = 60  # Keep last 60 snapshots (10 minutes at 10s intervals)

        self._last_scale_action_time = 0.0
        self._cooldown_seconds = 60.0

    def evaluate(self, snapshot: TrafficSnapshot) -> ScalingDecision:
        """
        Evaluate current traffic and return a scaling decision.
        This is the main entry point, called every 10 seconds.
        """
        self._snapshot_history.append(snapshot)
        if len(self._snapshot_history) > self._max_history:
            self._snapshot_history = self._snapshot_history[-self._max_history:]

        # Check for event override first
        event_decision = self._check_event_override(snapshot)
        if event_decision is not None:
            return event_decision

        # Check for emergency conditions
        if snapshot.throttle_count >= self.THROTTLE_EMERGENCY:
            return self._handle_throttle_emergency(snapshot)

        if snapshot.bedrock_queue_depth >= self.QUEUE_DEPTH_EMERGENCY:
            return self._handle_queue_emergency(snapshot)

        # Calculate provisioned capacity
        provisioned_tps = (
            snapshot.provisioned_sonnet_units * self.SONNET_TPS_PER_UNIT
            + snapshot.provisioned_haiku_units * self.HAIKU_TPS_PER_UNIT
        )
        utilization = snapshot.tokens_per_second / provisioned_tps if provisioned_tps > 0 else 1.0

        # Predictive check: will we breach in 5 minutes?
        predicted_tps = self._predict_tps()
        predicted_utilization = predicted_tps / provisioned_tps if provisioned_tps > 0 else 1.0

        # Scale-out evaluation
        if (
            utilization > self.TPS_UTILIZATION_SCALE_OUT
            or predicted_utilization > self.TPS_UTILIZATION_EMERGENCY
            or snapshot.bedrock_queue_depth > self.QUEUE_DEPTH_SCALE_OUT
            or snapshot.p99_latency_ms > self.LATENCY_SCALE_OUT_MS
        ):
            return self._compute_scale_out(snapshot, utilization, predicted_tps)

        # Scale-in evaluation
        if (
            utilization < self.TPS_UTILIZATION_SCALE_IN
            and snapshot.bedrock_queue_depth < 10
            and snapshot.p99_latency_ms < 2_000
            and self._in_cooldown() is False
        ):
            return self._compute_scale_in(snapshot, utilization)

        # Stable -- no action needed
        return ScalingDecision(
            state=ScalingState.STABLE,
            ecs_desired_count=None,
            sonnet_units=None,
            haiku_units=None,
            reason=f"Stable: utilization={utilization:.1%}, queue={snapshot.bedrock_queue_depth}, p99={snapshot.p99_latency_ms:.0f}ms",
            actions=[],
        )

    def _predict_tps(self) -> float:
        """
        Predict tokens-per-second in 5 minutes using linear regression
        on the recent snapshot history.
        """
        if len(self._snapshot_history) < 6:
            # Not enough data; return current value
            return self._snapshot_history[-1].tokens_per_second if self._snapshot_history else 0

        recent = self._snapshot_history[-30:]  # Last 5 minutes of 10s snapshots
        if len(recent) < 2:
            return recent[-1].tokens_per_second

        # Simple linear regression: TPS = a + b * t
        t_values = [s.timestamp - recent[0].timestamp for s in recent]
        tps_values = [s.tokens_per_second for s in recent]

        n = len(t_values)
        sum_t = sum(t_values)
        sum_tps = sum(tps_values)
        sum_t_tps = sum(t * tps for t, tps in zip(t_values, tps_values))
        sum_t_sq = sum(t * t for t in t_values)

        denominator = n * sum_t_sq - sum_t * sum_t
        if denominator == 0:
            return tps_values[-1]

        slope = (n * sum_t_tps - sum_t * sum_tps) / denominator
        intercept = (sum_tps - slope * sum_t) / n

        # Predict 5 minutes into the future
        future_t = t_values[-1] + self.PREDICTION_WINDOW_SECONDS
        predicted = intercept + slope * future_t

        return max(predicted, 0)  # TPS can't be negative

    def _handle_throttle_emergency(self, snapshot: TrafficSnapshot) -> ScalingDecision:
        """Respond to Bedrock throttling with multi-layered mitigation."""
        additional_sonnet = max(2, snapshot.throttle_count)  # At least 2 extra units
        additional_haiku = max(1, snapshot.throttle_count // 2)

        return ScalingDecision(
            state=ScalingState.EMERGENCY,
            ecs_desired_count=min(snapshot.ecs_task_count + 10, self.ECS_MAX_TASKS),
            sonnet_units=snapshot.provisioned_sonnet_units + additional_sonnet,
            haiku_units=snapshot.provisioned_haiku_units + additional_haiku,
            reason=f"EMERGENCY: {snapshot.throttle_count} throttles detected",
            actions=[
                f"Increase provisioned Sonnet by {additional_sonnet} units",
                f"Increase provisioned Haiku by {additional_haiku} units",
                "Activate aggressive request coalescing",
                "Shift eligible queries from Sonnet to Haiku",
                "Extend Redis cache TTL to 120s",
                "Page on-call engineer",
            ],
        )

    def _handle_queue_emergency(self, snapshot: TrafficSnapshot) -> ScalingDecision:
        """Respond to extreme queue depth."""
        return ScalingDecision(
            state=ScalingState.EMERGENCY,
            ecs_desired_count=min(snapshot.ecs_task_count + 10, self.ECS_MAX_TASKS),
            sonnet_units=snapshot.provisioned_sonnet_units + 3,
            haiku_units=snapshot.provisioned_haiku_units + 2,
            reason=f"EMERGENCY: Queue depth {snapshot.bedrock_queue_depth} exceeds {self.QUEUE_DEPTH_EMERGENCY}",
            actions=[
                "Add 10 ECS tasks immediately",
                "Increase provisioned throughput",
                "Enable degraded mode for FAQ queries (template responses)",
                "Alert on-call engineer",
            ],
        )

    def _compute_scale_out(
        self, snapshot: TrafficSnapshot, utilization: float, predicted_tps: float
    ) -> ScalingDecision:
        """Calculate how much to scale out."""
        # ECS task scaling: target TPS_PER_ECS_TASK per task
        desired_tasks = max(
            snapshot.ecs_task_count + 2,
            int(predicted_tps / self.TPS_PER_ECS_TASK) + 1,
        )
        desired_tasks = min(desired_tasks, self.ECS_MAX_TASKS)

        # Provisioned throughput: add units if utilization is high
        sonnet_units = snapshot.provisioned_sonnet_units
        haiku_units = snapshot.provisioned_haiku_units

        if utilization > self.TPS_UTILIZATION_SCALE_OUT:
            # Add units proportional to the gap
            needed_tps = predicted_tps * 1.3  # 30% headroom
            current_capacity = (
                sonnet_units * self.SONNET_TPS_PER_UNIT
                + haiku_units * self.HAIKU_TPS_PER_UNIT
            )
            gap = needed_tps - current_capacity
            if gap > 0:
                # Split additional capacity between Sonnet and Haiku
                # based on current Sonnet ratio
                sonnet_gap = int(gap * snapshot.sonnet_ratio / self.SONNET_TPS_PER_UNIT) + 1
                haiku_gap = int(gap * (1 - snapshot.sonnet_ratio) / self.HAIKU_TPS_PER_UNIT) + 1
                sonnet_units += sonnet_gap
                haiku_units += haiku_gap

        actions = []
        if desired_tasks > snapshot.ecs_task_count:
            actions.append(
                f"Scale ECS from {snapshot.ecs_task_count} to {desired_tasks} tasks"
            )
        if sonnet_units > snapshot.provisioned_sonnet_units:
            actions.append(
                f"Increase Sonnet provisioned from {snapshot.provisioned_sonnet_units} to {sonnet_units} units"
            )
        if haiku_units > snapshot.provisioned_haiku_units:
            actions.append(
                f"Increase Haiku provisioned from {snapshot.provisioned_haiku_units} to {haiku_units} units"
            )

        return ScalingDecision(
            state=ScalingState.SCALING_OUT,
            ecs_desired_count=desired_tasks,
            sonnet_units=sonnet_units,
            haiku_units=haiku_units,
            reason=f"Scale-out: utilization={utilization:.1%}, predicted_tps={predicted_tps:.0f}, queue={snapshot.bedrock_queue_depth}",
            actions=actions,
        )

    def _compute_scale_in(self, snapshot: TrafficSnapshot, utilization: float) -> ScalingDecision:
        """Calculate how much to scale in (conservative)."""
        # Never scale in by more than 10% at a time
        reduction = max(1, int(snapshot.ecs_task_count * 0.10))
        desired_tasks = max(snapshot.ecs_task_count - reduction, self.ECS_MIN_TASKS)

        # Reduce provisioned throughput only if utilization is very low
        sonnet_units = snapshot.provisioned_sonnet_units
        haiku_units = snapshot.provisioned_haiku_units

        if utilization < 0.25:
            sonnet_units = max(3, sonnet_units - 1)
            haiku_units = max(2, haiku_units - 1)

        actions = []
        if desired_tasks < snapshot.ecs_task_count:
            actions.append(
                f"Scale ECS from {snapshot.ecs_task_count} to {desired_tasks} tasks"
            )
        if sonnet_units < snapshot.provisioned_sonnet_units:
            actions.append(
                f"Reduce Sonnet provisioned from {snapshot.provisioned_sonnet_units} to {sonnet_units} units"
            )

        return ScalingDecision(
            state=ScalingState.SCALING_IN,
            ecs_desired_count=desired_tasks,
            sonnet_units=sonnet_units,
            haiku_units=haiku_units,
            reason=f"Scale-in: utilization={utilization:.1%}, queue={snapshot.bedrock_queue_depth}",
            actions=actions,
        )

    def _check_event_override(self, snapshot: TrafficSnapshot) -> Optional[ScalingDecision]:
        """Check if an event calendar entry requires override scaling."""
        now_jst = datetime.now(JST)
        for event_name, event_config in self.event_calendar.items():
            start = event_config["start"]
            end = event_config["end"]
            if start <= now_jst <= end:
                return ScalingDecision(
                    state=ScalingState.EVENT_OVERRIDE,
                    ecs_desired_count=min(event_config["ecs_tasks"], self.ECS_MAX_TASKS),
                    sonnet_units=event_config["sonnet_units"],
                    haiku_units=event_config["haiku_units"],
                    reason=f"Event override: {event_name} ({start} to {end})",
                    actions=[
                        f"Event '{event_name}' active -- using override capacity",
                        f"ECS: {event_config['ecs_tasks']} tasks",
                        f"Sonnet: {event_config['sonnet_units']} units",
                        f"Haiku: {event_config['haiku_units']} units",
                    ],
                )
        return None

    def _in_cooldown(self) -> bool:
        """Check if we are in a cooldown period after a recent scale action."""
        return (time.time() - self._last_scale_action_time) < self._cooldown_seconds

    async def apply_decision(self, decision: ScalingDecision):
        """Apply a scaling decision to AWS resources."""
        if decision.state == ScalingState.STABLE:
            return

        logger.info(
            "Applying scaling decision: state=%s reason=%s actions=%s",
            decision.state.value, decision.reason, decision.actions,
        )

        # Apply ECS scaling
        if decision.ecs_desired_count is not None:
            self._ecs_client.update_service(
                cluster=self.ecs_cluster,
                service=self.ecs_service,
                desiredCount=decision.ecs_desired_count,
            )
            logger.info("ECS desired count set to %d", decision.ecs_desired_count)

        # Apply provisioned throughput changes
        if decision.sonnet_units is not None:
            # Bedrock provisioned throughput API call would go here.
            # The exact API depends on the provisioned throughput model
            # (e.g., CreateProvisionedModelThroughput or UpdateProvisionedModelThroughput).
            logger.info(
                "Provisioned throughput update: Sonnet=%d units, Haiku=%d units",
                decision.sonnet_units,
                decision.haiku_units or 0,
            )

        # Publish scaling event to CloudWatch
        self._cloudwatch_client.put_metric_data(
            Namespace="MangaAssist",
            MetricData=[
                {
                    "MetricName": "ScalingAction",
                    "Value": 1,
                    "Unit": "Count",
                    "Dimensions": [
                        {"Name": "State", "Value": decision.state.value},
                        {"Name": "Service", "Value": "GenAIAutoScaler"},
                    ],
                },
            ],
        )

        self._last_scale_action_time = time.time()

sequenceDiagram
    participant CAL as Event Calendar<br>(Lambda CronJob)
    participant SCALER as GenAIAutoScaler
    participant ECS as ECS Fargate<br>Service
    participant BPT as Bedrock Provisioned<br>Throughput
    participant REDIS as ElastiCache Redis
    participant CW as CloudWatch
    participant OPS as On-Call Eng

    Note over CAL: Sunday 21:00 JST<br>3 hours before release

    CAL->>SCALER: Event: "weekly_shonen_jump"<br>starts in 3h
    SCALER->>BPT: UpdateProvisionedThroughput<br>Sonnet: 15 → 25 units<br>Haiku: 5 → 8 units
    BPT-->>SCALER: Status: PROVISIONING
    SCALER->>ECS: UpdateService<br>desiredCount: 30 → 50
    ECS-->>SCALER: Tasks launching...
    SCALER->>REDIS: Pre-warm popular queries<br>(top 100 manga titles)
    SCALER->>CW: Publish ScalingState=EVENT_OVERRIDE

    Note over ECS: 5 minutes later<br>50 tasks running

    Note over BPT: 15 minutes later<br>Provisioned throughput active

    SCALER->>CW: Publish ProvisionedState=ACTIVE
    SCALER->>OPS: INFO: Event capacity ready

    Note over CAL: Monday 00:00 JST<br>Release goes live

    CW-->>SCALER: TPS: 7K → 15K → 22K → 30K
    Note over SCALER: TPS within provisioned<br>capacity (33K). Stable.

    CW-->>SCALER: TPS: 30K → 33K → 36K
    Note over SCALER: TPS exceeding provisioned!

    SCALER->>BPT: Emergency: Sonnet 25 → 28 units
    SCALER->>ECS: UpdateService: 50 → 65 tasks
    SCALER->>REDIS: Extend cache TTL: 30s → 90s
    SCALER->>OPS: WARN: Approaching capacity ceiling

    Note over CW: 30 minutes later<br>Surge subsides

    CW-->>SCALER: TPS: 36K → 28K → 20K
    SCALER->>SCALER: Predicted TPS in 5m: 16K<br>Safe to scale in

    Note over CAL: Monday 04:00 JST<br>Event window ends

    CAL->>SCALER: Event "weekly_shonen_jump" ended
    SCALER->>BPT: Reduce to evening baseline<br>Sonnet: 28 → 15 units
    SCALER->>ECS: UpdateService: 65 → 30 tasks
    SCALER->>REDIS: Restore cache TTL: 90s → 30s
    SCALER->>CW: Publish ScalingState=STABLE