[小ネタ] AWS Batch のジョブ数を CloudWatch メトリクスに Put する仕組みを CloudFormation で構築してみた

AWSTemplateFormatVersion: 2010-09-09
Description: AWS Batch Cloudwatch Metrics Job Count

Parameters:
  QueueName:
    Type: String
    Description: Name of the AWS Batch job queue to monitor

Resources:
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: BatchAndCloudWatchAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - batch:ListJobs
                  - cloudwatch:PutMetricData
                Resource: '*'

  BatchMetricsFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: index.lambda_handler
      Role: !GetAtt LambdaExecutionRole.Arn
      Code:
        ZipFile: |
          import boto3
          import os

          def lambda_handler(event, context):
              batch = boto3.client('batch')
              cloudwatch = boto3.client('cloudwatch')

              queue_name = os.environ['QUEUE_NAME']

              # ジョブキューの状態別のジョブ数を取得
              submitted_count = len(batch.list_jobs(
                  jobQueue=queue_name,
                  jobStatus='SUBMITTED'
              )['jobSummaryList'])

              pending_count = len(batch.list_jobs(
                  jobQueue=queue_name,
                  jobStatus='PENDING'
              )['jobSummaryList'])

              runnable_count = len(batch.list_jobs(
                  jobQueue=queue_name,
                  jobStatus='RUNNABLE'
              )['jobSummaryList'])

              starting_count = len(batch.list_jobs(
                  jobQueue=queue_name,
                  jobStatus='STARTING'
              )['jobSummaryList'])

              running_count = len(batch.list_jobs(
                  jobQueue=queue_name,
                  jobStatus='RUNNING'
              )['jobSummaryList'])

              # CloudWatchにメトリクスを送信
              cloudwatch.put_metric_data(
                  Namespace='BatchJobMetrics',
                  MetricData=[
                      {
                          'MetricName': 'SubmittedCount',
                          'Value': submitted_count,
                          'Unit': 'Count',
                          'Dimensions': [
                              {
                                  'Name': 'QueueName',
                                  'Value': queue_name
                              }
                          ]
                      },
                      {
                          'MetricName': 'PendingCount',
                          'Value': pending_count,
                          'Unit': 'Count',
                          'Dimensions': [
                              {
                                  'Name': 'QueueName',
                                  'Value': queue_name
                              }
                          ]
                      },
                      {
                          'MetricName': 'RunnableCount',
                          'Value': runnable_count,
                          'Unit': 'Count',
                          'Dimensions': [
                              {
                                  'Name': 'QueueName',
                                  'Value': queue_name
                              }
                          ]
                      },
                      {
                          'MetricName': 'StartingCount',
                          'Value': starting_count,
                          'Unit': 'Count',
                          'Dimensions': [
                              {
                                  'Name': 'QueueName',
                                  'Value': queue_name
                              }
                          ]
                      },
                      {
                          'MetricName': 'RunningCount',
                          'Value': running_count,
                          'Unit': 'Count',
                          'Dimensions': [
                              {
                                  'Name': 'QueueName',
                                  'Value': queue_name
                              }
                          ]
                      }
                  ]
              )

              return {
                  'statusCode': 200,
                  'body': {
                      'submitted': submitted_count,
                      'pending': pending_count,
                      'runnable': runnable_count,
                      'starting': starting_count,
                      'running': running_count
                  }
              }
      Runtime: python3.13
      Timeout: 30
      Environment:
        Variables:
          QUEUE_NAME: !Ref QueueName

  MetricsScheduleRule:
    Type: AWS::Events::Rule
    Properties:
      Description: "Schedule for collecting Batch queue metrics"
      ScheduleExpression: "rate(1 minute)"
      State: ENABLED
      Targets:
        - Arn: !GetAtt BatchMetricsFunction.Arn
          Id: "BatchMetricsTarget"

  LambdaInvokePermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !Ref BatchMetricsFunction
      Action: "lambda:InvokeFunction"
      Principal: "events.amazonaws.com"
      SourceArn: !GetAtt MetricsScheduleRule.Arn