01: High-Performance FM Architecture and Throughput Management

mindmap
  root((High-Performance
    FM Systems))
    Batching
      Micro-Batching by Intent
      Request Coalescing
      Batch Window Tuning
      Priority Queue Separation
    Capacity Planning
      Token Throughput Targets
      Peak Hour Analysis
      Event Provisioning
      Growth Forecasting
    Utilization Monitoring
      Invocation Metrics
      Tokens-per-Second Tracking
      Throttle Rate Monitoring
      Waste Detection
    Auto-Scaling
      ECS Fargate Scaling
      Queue-Depth Triggers
      Step Scaling Policies
      Target Tracking
    Provisioned Throughput
      Model Units Sizing
      On-Demand vs Provisioned
      Breakeven Analysis
      Dynamic Adjustment

sequenceDiagram
    participant U1 as User A (product_search)
    participant U2 as User B (product_search)
    participant U3 as User C (order_status)
    participant Q as Intent Queue
    participant MB as Micro-Batcher
    participant B as Bedrock Claude

    U1->>Q: "Find shonen manga under 800 yen"
    U2->>Q: "Any new isekai releases?"
    U3->>Q: "Where is my order 12345?"

    Q->>MB: Batch [U1, U2] (same intent: product_search)
    Q->>MB: Single [U3] (order_status, different template)

    MB->>B: Batched invocation for product_search (2 queries)
    MB->>B: Single invocation for order_status

    B-->>MB: Batched response [resp_A, resp_B]
    B-->>MB: Single response [resp_C]

    MB-->>U1: resp_A
    MB-->>U2: resp_B
    MB-->>U3: resp_C

Tokens/sec
35,000 |                                          * (manga release spike)
30,000 |
25,000 |                                    ***
20,000 |                                  **   **
15,000 |                                **       **
10,000 |                    ********** *           ***
 7,000 |--- avg ------*****---------------------------------****-------
 5,000 |           ***                                          ***
 3,000 |        ***                                                ***
 1,000 | *******                                                      **
        |_____|_____|_____|_____|_____|_____|_____|_____|_____|_____|____
        00    03    06    09    12    15    18    21    00    03    06  JST
              trough        lunch     peak evening       overnight

TPS = (InputTokenCount + OutputTokenCount) / period_seconds

Waste Ratio = 1 - (actual_TPS / provisioned_TPS)

graph TD
    subgraph "Auto-Scaling Loop"
        A[CloudWatch Custom Metric<br>bedrock_queue_depth] --> B{Queue Depth<br>Evaluation}
        B -->|depth > 50 for 2 min| C[Step Scaling: +2 tasks]
        B -->|depth > 100 for 1 min| D[Step Scaling: +5 tasks]
        B -->|depth > 200 for 30s| E[Step Scaling: +10 tasks<br>Emergency]
        B -->|depth < 20 for 10 min| F[Scale-in: -1 task]
        B -->|depth < 5 for 15 min| G[Scale-in: -2 tasks]
    end

    subgraph "Metric Publisher"
        H[ECS Orchestrator] --> I[Publish queue depth<br>every 10 seconds]
        I --> A
    end

    subgraph "Bedrock"
        H --> J[Bedrock Invocations]
        J --> K[InvocationThrottles metric]
        K --> A
    end

Target TPS per task = 500
Desired task count  = ceil(current_total_TPS / 500)

graph TB
    subgraph "Client Layer"
        WS[WebSocket Clients<br>1M messages/day]
    end

    subgraph "Ingress"
        AG[API Gateway WebSocket]
    end

    subgraph "Orchestration Layer (ECS Fargate)"
        IC[Intent Classifier]
        MB[Micro-Batcher]
        RC[Request Coalescer]
        PQ[Priority Queue<br>High / Normal / Low]
        RL[Rate Limiter]
    end

    subgraph "Cache Layer"
        REDIS[ElastiCache Redis<br>Coalesced Response Cache<br>Embedding Cache]
    end

    subgraph "FM Layer"
        BPT[Bedrock Provisioned Throughput<br>Sonnet: 15 units<br>Haiku: 5 units]
        BOD[Bedrock On-Demand<br>Overflow / Burst]
    end

    subgraph "Monitoring"
        CW[CloudWatch Metrics<br>TPS, Queue Depth, Throttles]
        AS[Auto-Scaling Policies<br>Step + Target Tracking]
        ALARM[Alarms + SNS]
    end

    WS --> AG --> IC
    IC --> MB
    MB --> RC
    RC -->|cache hit| REDIS
    RC -->|cache miss| PQ
    PQ --> RL
    RL -->|within provisioned capacity| BPT
    RL -->|overflow| BOD
    BPT --> REDIS
    BOD --> REDIS

    CW --> AS
    AS -->|scale ECS| IC
    AS -->|adjust provisioned units| BPT
    CW --> ALARM

import asyncio
import time
import hashlib
import json
from dataclasses import dataclass, field
from typing import Optional
from collections import defaultdict

import boto3


@dataclass
class InferenceRequest:
    """A single MangaAssist inference request."""
    request_id: str
    user_id: str
    intent: str
    prompt: str
    model_id: str
    priority: str = "normal"  # high, normal, low
    created_at: float = field(default_factory=time.time)
    normalized_query_hash: Optional[str] = None


@dataclass
class InferenceResponse:
    """Response from Bedrock invocation."""
    request_id: str
    output_text: str
    input_tokens: int
    output_tokens: int
    latency_ms: float
    from_cache: bool = False
    from_coalesce: bool = False


class ThroughputManager:
    """
    Manages Bedrock throughput for MangaAssist with:
    - Micro-batching by intent
    - Request coalescing for identical queries
    - Priority-based rate limiting
    - Provisioned vs on-demand routing
    """

    def __init__(
        self,
        provisioned_tps: int = 20_000,
        batch_window_ms: int = 150,
        coalesce_window_ms: int = 1_000,
        max_batch_size: int = 8,
        max_concurrent_invocations: int = 100,
        redis_client=None,
    ):
        self.provisioned_tps = provisioned_tps
        self.batch_window_ms = batch_window_ms
        self.coalesce_window_ms = coalesce_window_ms
        self.max_batch_size = max_batch_size
        self.redis_client = redis_client

        # Semaphore to limit concurrent Bedrock invocations
        self._invocation_semaphore = asyncio.Semaphore(max_concurrent_invocations)

        # Intent-based batch queues: intent -> list of (request, future)
        self._batch_queues: dict[str, list] = defaultdict(list)
        self._batch_locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)

        # Coalescing: query_hash -> (future, timestamp)
        self._coalesce_map: dict[str, asyncio.Future] = {}
        self._coalesce_lock = asyncio.Lock()

        # Metrics tracking
        self._tokens_this_second = 0
        self._tokens_lock = asyncio.Lock()
        self._last_token_reset = time.time()

        # Rate limiting buckets
        self._rate_limits = {
            "high": asyncio.Semaphore(50),
            "normal": asyncio.Semaphore(40),
            "low": asyncio.Semaphore(10),
        }

        # Bedrock client
        self._bedrock_runtime = boto3.client("bedrock-runtime", region_name="ap-northeast-1")

    async def process_request(self, request: InferenceRequest) -> InferenceResponse:
        """
        Main entry point. Routes through coalescing, batching, and rate limiting.
        """
        # Step 1: Check coalescing (skip for order-status, unique per user)
        if request.intent not in ("order_status", "account_info"):
            coalesced = await self._try_coalesce(request)
            if coalesced is not None:
                return coalesced

        # Step 2: Check Redis cache
        cached = await self._check_cache(request)
        if cached is not None:
            return cached

        # Step 3: Enqueue for micro-batching
        response = await self._enqueue_for_batch(request)
        return response

    def _compute_query_hash(self, request: InferenceRequest) -> str:
        """Compute a normalized hash for coalescing identical queries."""
        normalized = json.dumps({
            "intent": request.intent,
            "prompt_hash": hashlib.sha256(
                request.prompt.strip().lower().encode()
            ).hexdigest()[:16],
            "model_id": request.model_id,
        }, sort_keys=True)
        return hashlib.sha256(normalized.encode()).hexdigest()[:32]

    async def _try_coalesce(self, request: InferenceRequest) -> Optional[InferenceResponse]:
        """
        If an identical query is already in-flight, wait for its result
        instead of making a duplicate Bedrock call.
        """
        query_hash = self._compute_query_hash(request)
        request.normalized_query_hash = query_hash

        async with self._coalesce_lock:
            if query_hash in self._coalesce_map:
                existing_future, created_at = self._coalesce_map[query_hash]
                age_ms = (time.time() - created_at) * 1000
                if age_ms < self.coalesce_window_ms:
                    # Wait for the in-flight request to complete
                    result = await existing_future
                    return InferenceResponse(
                        request_id=request.request_id,
                        output_text=result.output_text,
                        input_tokens=result.input_tokens,
                        output_tokens=result.output_tokens,
                        latency_ms=result.latency_ms,
                        from_coalesce=True,
                    )

            # No in-flight duplicate; register this request for coalescing
            future = asyncio.get_event_loop().create_future()
            self._coalesce_map[query_hash] = (future, time.time())

        return None  # Proceed to batching/invocation

    async def _check_cache(self, request: InferenceRequest) -> Optional[InferenceResponse]:
        """Check Redis for a cached response."""
        if self.redis_client is None:
            return None

        cache_key = f"manga:resp:{request.normalized_query_hash or self._compute_query_hash(request)}"
        cached_data = self.redis_client.get(cache_key)
        if cached_data:
            data = json.loads(cached_data)
            return InferenceResponse(
                request_id=request.request_id,
                output_text=data["output_text"],
                input_tokens=data["input_tokens"],
                output_tokens=data["output_tokens"],
                latency_ms=0.0,
                from_cache=True,
            )
        return None

    async def _enqueue_for_batch(self, request: InferenceRequest) -> InferenceResponse:
        """
        Add request to the intent-specific batch queue.
        If the batch fills up or the window expires, flush the batch.
        """
        future = asyncio.get_event_loop().create_future()

        async with self._batch_locks[request.intent]:
            self._batch_queues[request.intent].append((request, future))

            if len(self._batch_queues[request.intent]) >= self.max_batch_size:
                batch = self._batch_queues[request.intent][:self.max_batch_size]
                self._batch_queues[request.intent] = self._batch_queues[request.intent][self.max_batch_size:]
                asyncio.create_task(self._flush_batch(request.intent, batch))
            elif len(self._batch_queues[request.intent]) == 1:
                # First item in batch -- start the window timer
                asyncio.create_task(self._batch_timer(request.intent))

        return await future

    async def _batch_timer(self, intent: str):
        """Wait for the batch window, then flush whatever is queued."""
        await asyncio.sleep(self.batch_window_ms / 1000.0)
        async with self._batch_locks[intent]:
            if self._batch_queues[intent]:
                batch = self._batch_queues[intent][:]
                self._batch_queues[intent] = []
                asyncio.create_task(self._flush_batch(intent, batch))

    async def _flush_batch(self, intent: str, batch: list):
        """
        Invoke Bedrock for a batch of requests.
        Uses rate limiting and routes to provisioned or on-demand.
        """
        for request, future in batch:
            priority_sem = self._rate_limits.get(request.priority, self._rate_limits["normal"])
            await priority_sem.acquire()
            try:
                async with self._invocation_semaphore:
                    response = await self._invoke_bedrock(request)
                    future.set_result(response)

                    # Update coalescing map
                    if request.normalized_query_hash and request.normalized_query_hash in self._coalesce_map:
                        coalesce_future, _ = self._coalesce_map.pop(request.normalized_query_hash)
                        if not coalesce_future.done():
                            coalesce_future.set_result(response)

                    # Cache the result
                    await self._cache_response(request, response)

                    # Track tokens
                    await self._track_tokens(response.input_tokens + response.output_tokens)
            except Exception as e:
                if not future.done():
                    future.set_exception(e)
            finally:
                priority_sem.release()

    async def _invoke_bedrock(self, request: InferenceRequest) -> InferenceResponse:
        """Call Bedrock and return the response."""
        start = time.time()

        body = json.dumps({
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": 1024,
            "messages": [{"role": "user", "content": request.prompt}],
        })

        # Route: provisioned throughput if within capacity, else on-demand
        current_tps = await self._get_current_tps()
        model_id = request.model_id
        if current_tps > self.provisioned_tps * 0.9:
            # Near provisioned capacity -- use on-demand endpoint as overflow
            model_id = f"{request.model_id}"  # on-demand uses same model ID

        response = self._bedrock_runtime.invoke_model(
            modelId=model_id,
            contentType="application/json",
            accept="application/json",
            body=body,
        )

        response_body = json.loads(response["body"].read())
        latency_ms = (time.time() - start) * 1000

        return InferenceResponse(
            request_id=request.request_id,
            output_text=response_body["content"][0]["text"],
            input_tokens=response_body["usage"]["input_tokens"],
            output_tokens=response_body["usage"]["output_tokens"],
            latency_ms=latency_ms,
        )

    async def _track_tokens(self, token_count: int):
        """Track tokens-per-second for monitoring and scaling decisions."""
        async with self._tokens_lock:
            now = time.time()
            if now - self._last_token_reset >= 1.0:
                # Publish the previous second's count to CloudWatch
                await self._publish_tps_metric(self._tokens_this_second)
                self._tokens_this_second = 0
                self._last_token_reset = now
            self._tokens_this_second += token_count

    async def _get_current_tps(self) -> int:
        """Return the current tokens-per-second rate."""
        async with self._tokens_lock:
            return self._tokens_this_second

    async def _publish_tps_metric(self, tps: int):
        """Publish tokens-per-second to CloudWatch."""
        cloudwatch = boto3.client("cloudwatch", region_name="ap-northeast-1")
        cloudwatch.put_metric_data(
            Namespace="MangaAssist",
            MetricData=[{
                "MetricName": "TokensPerSecond",
                "Value": tps,
                "Unit": "Count/Second",
                "Dimensions": [
                    {"Name": "Service", "Value": "ThroughputManager"},
                    {"Name": "Environment", "Value": "production"},
                ],
            }],
        )

    async def _cache_response(self, request: InferenceRequest, response: InferenceResponse):
        """Cache the response in Redis with a short TTL."""
        if self.redis_client is None:
            return
        if request.intent in ("order_status", "account_info"):
            return  # Do not cache user-specific responses

        cache_key = f"manga:resp:{request.normalized_query_hash}"
        self.redis_client.setex(
            cache_key,
            30,  # 30-second TTL
            json.dumps({
                "output_text": response.output_text,
                "input_tokens": response.input_tokens,
                "output_tokens": response.output_tokens,
            }),
        )

    def get_metrics_snapshot(self) -> dict:
        """Return a snapshot of current throughput metrics."""
        return {
            "current_tps": self._tokens_this_second,
            "provisioned_tps": self.provisioned_tps,
            "utilization_pct": round(
                (self._tokens_this_second / self.provisioned_tps) * 100, 1
            ) if self.provisioned_tps > 0 else 0,
            "batch_queue_depths": {
                intent: len(queue)
                for intent, queue in self._batch_queues.items()
            },
            "coalesce_map_size": len(self._coalesce_map),
        }

"""
CDK-style auto-scaling configuration for MangaAssist ECS Fargate service
with Bedrock-aware custom metrics.
"""

from aws_cdk import (
    Stack,
    Duration,
    aws_ecs as ecs,
    aws_applicationautoscaling as appscaling,
    aws_cloudwatch as cloudwatch,
)
from constructs import Construct


class MangaAssistAutoScalingStack(Stack):
    """
    Configures auto-scaling for the MangaAssist orchestrator service
    running on ECS Fargate, using Bedrock queue depth and tokens-per-second
    as scaling signals instead of CPU/memory alone.
    """

    def __init__(self, scope: Construct, id: str, ecs_service: ecs.FargateService, **kwargs):
        super().__init__(scope, id, **kwargs)

        # --- Scalable Target ---
        scaling = ecs_service.auto_scale_task_count(
            min_capacity=5,    # Minimum tasks even at overnight trough
            max_capacity=100,  # Hard ceiling for cost safety
        )

        # --- Policy 1: Target Tracking on Tokens-per-Second ---
        # Each ECS task should handle roughly 500 tokens/s.
        # If total TPS rises, add tasks proportionally.
        tps_metric = cloudwatch.Metric(
            namespace="MangaAssist",
            metric_name="TokensPerSecond",
            dimensions_map={
                "Service": "ThroughputManager",
                "Environment": "production",
            },
            statistic="Average",
            period=Duration.seconds(60),
        )

        scaling.scale_on_metric(
            "ScaleOnTPS",
            metric=tps_metric,
            scaling_steps=[
                appscaling.ScalingInterval(change=-2, lower=0, upper=2_000),
                appscaling.ScalingInterval(change=0, lower=2_000, upper=5_000),
                appscaling.ScalingInterval(change=2, lower=5_000, upper=10_000),
                appscaling.ScalingInterval(change=5, lower=10_000, upper=20_000),
                appscaling.ScalingInterval(change=10, lower=20_000),
            ],
            adjustment_type=appscaling.AdjustmentType.CHANGE_IN_CAPACITY,
            cooldown=Duration.seconds(120),
        )

        # --- Policy 2: Step Scaling on Bedrock Queue Depth ---
        queue_depth_metric = cloudwatch.Metric(
            namespace="MangaAssist",
            metric_name="BedrockQueueDepth",
            dimensions_map={
                "Service": "Orchestrator",
                "Environment": "production",
            },
            statistic="Maximum",
            period=Duration.seconds(30),
        )

        # Scale-out steps (aggressive)
        scaling.scale_on_metric(
            "ScaleOutOnQueueDepth",
            metric=queue_depth_metric,
            scaling_steps=[
                appscaling.ScalingInterval(change=2, lower=50, upper=100),
                appscaling.ScalingInterval(change=5, lower=100, upper=200),
                appscaling.ScalingInterval(change=10, lower=200),
            ],
            adjustment_type=appscaling.AdjustmentType.CHANGE_IN_CAPACITY,
            cooldown=Duration.seconds(60),
        )

        # Scale-in steps (conservative)
        scaling.scale_on_metric(
            "ScaleInOnQueueDepth",
            metric=queue_depth_metric,
            scaling_steps=[
                appscaling.ScalingInterval(change=-1, upper=20),
                appscaling.ScalingInterval(change=-2, upper=5),
            ],
            adjustment_type=appscaling.AdjustmentType.CHANGE_IN_CAPACITY,
            cooldown=Duration.seconds(300),
        )

        # --- Policy 3: Scheduled Scaling for Known Peaks ---
        # Evening peak: 17:00-23:59 JST (08:00-14:59 UTC)
        scaling.scale_on_schedule(
            "EveningPeakScaleUp",
            schedule=appscaling.Schedule.cron(hour="8", minute="0"),
            min_capacity=30,
        )
        scaling.scale_on_schedule(
            "EveningPeakScaleDown",
            schedule=appscaling.Schedule.cron(hour="15", minute="0"),
            min_capacity=10,
        )

        # Overnight trough: 02:00-06:00 JST (17:00-21:00 UTC)
        scaling.scale_on_schedule(
            "OvernightTrough",
            schedule=appscaling.Schedule.cron(hour="17", minute="0"),
            min_capacity=5,
        )

        # Monday manga release: Scale up Sunday 22:00 JST (13:00 UTC)
        scaling.scale_on_schedule(
            "MondayMangaReleasePrepare",
            schedule=appscaling.Schedule.cron(
                week_day="SUN", hour="13", minute="0"
            ),
            min_capacity=50,
        )
        scaling.scale_on_schedule(
            "MondayMangaReleaseEnd",
            schedule=appscaling.Schedule.cron(
                week_day="MON", hour="6", minute="0"  # 15:00 JST
            ),
            min_capacity=10,
        )

        # --- Policy 4: Throttle-Based Emergency Scaling ---
        throttle_alarm = cloudwatch.Alarm(
            self,
            "BedrockThrottleAlarm",
            metric=cloudwatch.Metric(
                namespace="AWS/Bedrock",
                metric_name="InvocationThrottles",
                statistic="Sum",
                period=Duration.seconds(60),
            ),
            threshold=5,
            evaluation_periods=1,
            comparison_operator=cloudwatch.ComparisonOperator.GREATER_THAN_OR_EQUAL_TO_THRESHOLD,
            alarm_description="Bedrock throttling detected -- emergency scale-out",
        )

        scaling.scale_on_metric(
            "EmergencyScaleOnThrottle",
            metric=cloudwatch.Metric(
                namespace="AWS/Bedrock",
                metric_name="InvocationThrottles",
                statistic="Sum",
                period=Duration.seconds(60),
            ),
            scaling_steps=[
                appscaling.ScalingInterval(change=5, lower=5, upper=20),
                appscaling.ScalingInterval(change=15, lower=20),
            ],
            adjustment_type=appscaling.AdjustmentType.CHANGE_IN_CAPACITY,
            cooldown=Duration.seconds(60),
        )