[アップデート] Lambda durable functions の AWS 公式 Kiro Power が追加されていたので使ってみた

AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: Durable Lambda Function - Parallel Calculation Workflow

Globals:
  Function:
    Timeout: 900
    MemorySize: 512
    Runtime: python3.13

Resources:
  CalculationFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: parallel-calculation-function
      Handler: handler.handler
      CodeUri: ./src
      DurableConfig:
        ExecutionTimeout: 3600
        RetentionPeriodInDays: 7
      Policies:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicDurableExecutionRolePolicy
      AutoPublishAlias: prod
      Environment:
        Variables:
          LOG_LEVEL: INFO

Outputs:
  FunctionArn:
    Description: Function ARN
    Value: !GetAtt CalculationFunction.Arn
  AliasArn:
    Description: Function Alias ARN
    Value: !Ref CalculationFunction.Alias

from aws_durable_execution_sdk_python import durable_execution, DurableContext, durable_step, StepContext
from aws_durable_execution_sdk_python.config import MapConfig
from typing import Dict, List


@durable_step
def calculate_square(step_ctx: StepContext, number: int) -> int:
    """数字の二乗を計算"""
    return number * number


@durable_step
def calculate_cube(step_ctx: StepContext, number: int) -> int:
    """数字の三乗を計算"""
    return number * number * number


@durable_execution
def handler(event: Dict, context: DurableContext) -> Dict:
    """
    並列計算ワークフロー
    
    入力例:
    {
        "numbers": [1, 2, 3, 4, 5],
        "operation": "square"  # または "cube"
    }
    """
    numbers = event.get('numbers', [1, 2, 3, 4, 5])
    operation = event.get('operation', 'square')
    
    context.logger.info(f"Starting parallel calculation for {len(numbers)} numbers")
    
    # 計算関数を選択
    if operation == 'cube':
        def process_number(ctx: DurableContext, num: int, index: int, items: List[int]):
            return ctx.step(calculate_cube(num), name=f'calculate-cube-{index}')
    else:
        def process_number(ctx: DurableContext, num: int, index: int, items: List[int]):
            return ctx.step(calculate_square(num), name=f'calculate-square-{index}')
    
    # 並列実行（最大3つ同時）
    results = context.map(
        inputs=numbers,
        func=process_number,
        name='parallel-calculations',
        config=MapConfig(max_concurrency=3)
    )
    
    # エラーチェック
    results.throw_if_error()
    
    # 結果を取得
    calculated_values = results.get_results()
    
    context.logger.info(f"Completed {results.success_count} calculations")
    
    return {
        'success': True,
        'operation': operation,
        'input': numbers,
        'results': calculated_values,
        'total_count': results.total_count,
        'success_count': results.success_count
    }