最近の砂場活動その17: CloudWatch Eventsから直接AWS Batchのジョブをキックする

前置き

  • 機械学習のモデルの定期的な再学習、CloudWatch Eventsを使うと便利です
    • Cron的なスケジュール管理ができる
    • 機械学習のモデルの学習はAWS Batchで行なっていた(今だとSageMakerな人も多いのかもしれない)
  • しかし、以前はCloudWatch Eventsから直接AWS Batchのジョブを起動させることができなかった
    • LambdaからAPIを叩いてsubmitして〜という世界観
    • できるのは分かるけど、率直に言ってダルかった...
  • Lambdaを起動させるためのIAMの管理などもあって、AWS側でどうにかして欲しかった
Parameters:
  LambdaCodeStringSubmitBatchJob:
    Type: String
    Default: |
      import os
      from datetime import datetime as dt
      import boto3
      batch = boto3.client('batch')
      def lambda_handler(event, context):
          try:
              response = batch.submit_job(
                  jobName=os.environ['JOB_NAME'],
                  jobQueue=os.environ['JOB_QUEUE_NAME'],
                  jobDefinition=os.environ['JOB_DEFINITION_ARN']
              )
              jobId = response['jobId']
              return {
                  'jobId': jobId
              }
          except Exception as e:
              raise e
Resources:
  UpdateModelBatchJobTrigger:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: UpdateModelBatchJobTriggerResource
      Description: "学習データから分類器を構築"
      Role: 
        Fn::ImportValue: !Sub "${IAMStackName}:BatchJobTriggerFunctionRole"
      Handler: index.lambda_handler
      Runtime: python3.6
      MemorySize: 128
      Timeout: 30
      Environment:
        Variables:
          JOB_NAME: update_model
          JOB_QUEUE_NAME: go-active-learning
          JOB_DEFINITION_ARN: !Ref UpdateModelJobDefinition
      Code:
        ZipFile: !Sub "${LambdaCodeStringSubmitBatchJob}"
   LambdaUpdateModelRule:
    Type: AWS::Events::Rule
    Properties:
      Name: LambdaUpdateModelRule
      ScheduleExpression: rate(2 days)
      Targets:
        - Id: UpdateModelBatchJobTrigger
          Arn: !GetAtt UpdateModelBatchJobTrigger.Arn
      State: "ENABLED"
  LambdaUpdateModelPermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !Ref UpdateModelBatchJobTrigger
      SourceArn: !GetAtt LambdaUpdateModelRule.Arn
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com

現在: CloudWatch Eventsからジョブを直接キックできるようになった!

そうそう、これが欲しかった! 大分昔のアップデートなので、ちゃんとアップデートを追っていなかったことが明らかになってしまった!

CFnも欲しかったものを直接的に書けるようになったので、最高です。

Resources:
  UpdateModelRule:
    Type: AWS::Events::Rule
    Properties:
      Name: UpdateModelRule
      ScheduleExpression: rate(2 days)
      Targets:
        - Id: UpdateModelBatch
          Arn: !Ref JobQueue
          BatchParameters:
            JobDefinition: !Ref UpdateModelJobDefinition
            JobName: update_model
          RoleArn:
            Fn::ImportValue: !Sub "${IAMStackName}:SubmitBatchJobRuleRole"
      State: "ENABLED"

CloudWatch EventsからジョブをキックするためのIAMはこんな感じ。

Resources:
  SubmitBatchJobRuleRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: events.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: batch-job-trigger-policy
          PolicyDocument:
            Statement:
              - Effect: Allow
                Action:
                  - batch:SubmitJob
                Resource: 
                  - "*"
      Path: "/"

クラウド側のアップデートに合わせて、IaCもリファクタリングしていこうねという話でした。