バッチ処理のコスト削減に EC2 インスタンスプール方式を実装してみた

[DynamoDB] available → in-use → (処理完了後) → available
[EC2 実機] stopped → running → (処理完了後) → stopped

import json
import boto3
from datetime import datetime
from botocore.exceptions import ClientError

ec2 = boto3.client('ec2')
dynamodb = boto3.resource('dynamodb')
table = dynamodb.Table('InstancePool')

def lambda_handler(event, context):
    # SQS イベントソースから Records を取得
    for record in event.get('Records', []):
        message_body = json.loads(record['body'])
        tenant_id = message_body.get('tenant_id')
        job_id = message_body.get('job_id')

        instance_id = find_and_allocate_instance(tenant_id, job_id)

        if instance_id is None:
            # インスタンスが空くまで待機（SQS が自動リトライ）
            raise RuntimeError("No available instances in pool")

        start_instance(instance_id)

        execute_batch_job(instance_id, message_body)

def find_and_allocate_instance(tenant_id, job_id):
    response = table.scan(
        FilterExpression='#status = :available',
        ExpressionAttributeNames={'#status': 'status'},
        ExpressionAttributeValues={':available': 'available'}
    )

    instances = response.get('Items', [])
    if not instances:
        return None

    # 複数見つかった場合は順番に試すことで、競合時の成功率を上げる
    for instance in instances:
        try:
            # version を使うことで、読み取りから更新までの間に他のプロセスが
            # 変更していないことを保証する
            table.update_item(
                Key={'instance_id': instance['instance_id']},
                UpdateExpression='SET #status = :in_use, tenant_id = :tenant, job_id = :job, version = version + :inc, updated_at = :time',
                ConditionExpression='#status = :available AND version = :version',
                ExpressionAttributeNames={
                    '#status': 'status'  # status は予約語のため、ExpressionAttributeNames を使用
                },
                ExpressionAttributeValues={
                    ':in_use': 'in-use',
                    ':available': 'available',
                    ':tenant': tenant_id,
                    ':job': job_id,  # トレーサビリティのため、実行中のジョブ ID を記録
                    ':version': instance['version'],  # 読み取り時の version と一致することを確認
                    ':inc': 1,  # 更新のたびに version をインクリメントして競合を検知可能に
                    ':time': datetime.now().isoformat()
                }
            )
            return instance['instance_id']
        except ClientError as e:
            if e.response['Error']['Code'] == 'ConditionalCheckFailedException':
                # これにより、全インスタンスが使用中になるまで試行を続ける
                continue
            raise

    return None

def start_instance(instance_id):
    ec2.start_instances(InstanceIds=[instance_id])

    # 起動完了を待たないと、バッチ処理が実行できない
    waiter = ec2.get_waiter('instance_running')
    waiter.wait(
        InstanceIds=[instance_id],
        WaiterConfig={
            'Delay': 15,  # 15 秒ごとにチェック
            'MaxAttempts': 40  # 最大 10 分待機
        }
    )

    print(f"Instance {instance_id} is now running")

{
  "source": ["aws.ec2"],
  "detail-type": ["EC2 Instance State-change Notification"],
  "detail": {
    "state": ["stopped", "running"]
  }
}

import boto3
from datetime import datetime

dynamodb = boto3.resource('dynamodb')
table = dynamodb.Table('InstancePool')

def lambda_handler(event, context):
    """EC2 状態変更イベントを処理して DynamoDB を更新"""
    detail = event.get('detail', {})
    instance_id = detail.get('instance-id')
    state = detail.get('state')

    if not instance_id or not state:
        print(f"Invalid event: missing instance-id or state")
        return {'statusCode': 400}

    if state == 'stopped':
        # tenant_id と job_id をクリアすることで、次のジョブで使用可能にする
        table.update_item(
            Key={'instance_id': instance_id},
            UpdateExpression='SET #status = :available, tenant_id = :null, job_id = :null, updated_at = :time',
            ExpressionAttributeNames={'#status': 'status'},
            ExpressionAttributeValues={
                ':available': 'available',
                ':null': None,
                ':time': datetime.now().isoformat()
            }
        )
        print(f"Instance {instance_id} marked as available")

    elif state == 'running':
        # 実際の状態更新は割り当て時に完了しているため、ここでは何もしない
        print(f"Instance {instance_id} is now running")