AWS Step Functions でジョブ管理システムを組む

description: 指定したタグを持つインスタンス群のうち、1台でコマンドを実行する。
schemaVersion: '0.3'
assumeRole: 'arn:aws:iam::********:role/********'
outputs:
  - runCommand.OutputPayload
parameters:
  Tags:
    description: |-
      タグ
      (形式: {"<TAG_NAME>": "<VALUE>", ...})      
    type: StringMap
  Commands:
    description: コマンド
    type: StringList
    displayType: textarea
  WorkingDirectory:
    description: ワーキング・ディレクトリ
    type: String
    default: ''
  ExecutionTimeout:
    description: タイムアウト (秒)
    type: String
    default: '3600'
  TaskToken:
    description: |-
      Step Functions タスク・トークン    
      https://docs.aws.amazon.com/step-functions/latest/dg/connect-to-resource.html#connect-wait-token      
    type: String
mainSteps:
  - name: getInstanceId
    description: 実行対象のインスタンスを 1台選択する。
    action: 'aws:executeScript'
    inputs:
      Runtime: python3.7
      Handler: script_handler
      Script: |-
        import random

        import boto3
        
        def script_handler(events, context):
          ids = instance_ids(
            describe_instances(events['tags'])
          )
        
          print(ids)
        
          return select_instance(ids)
        
        def instance_ids(payload):
          return [
            instances['InstanceId'] for launch_request in payload['Reservations'] for instances in launch_request['Instances']
          ]
        
        def describe_instances(tags):
          return boto3.client('ec2').describe_instances(
            Filters=(
              cond_tags(tags) + cond_instance_status()
            )
          )
        
        def cond_tags(tags):
          return [
            {
              'Name': f'tag:{ k }',
              'Values': [
                v
              ]
            } for k, v in tags.items()
          ]
        
        def cond_instance_status():
          return [
            {
              'Name': 'instance-state-name',
              'Values': [
                'running'
              ]
            }
          ]
        
        def select_instance(ids):
          return random.choice(ids)        
      InputPayload:
        tags: '{{ Tags }}'
    outputs:
      - Name: instanceId
        Selector: $.Payload
        Type: String
    onFailure: 'step:callbackWorkerBusy'
    onCancel: 'step:callbackFailure'
  - name: runCommand
    description: 選択した 1台でコマンドを実行する。
    action: 'aws:runCommand'
    inputs:
      DocumentName: AWS-RunShellScript
      InstanceIds:
        - '{{ getInstanceId.instanceId }}'
      Parameters:
        commands: '{{ Commands }}'
        workingDirectory: '{{ WorkingDirectory }}'
        executionTimeout: '{{ ExecutionTimeout }}'
      CloudWatchOutputConfig:
        CloudWatchOutputEnabled: true
      OutputS3BucketName: quickguard-example
      OutputS3KeyPrefix: run_command/logs
      TimeoutSeconds: 600
    onFailure: Continue
    isCritical: true
  - name: chooseTaskResult
    description: Step Functions にコールバックする結果を選択する。
    action: 'aws:branch'
    inputs:
      Choices:
        - NextStep: callbackSuccess
          Variable: '{{ runCommand.Status }}'
          StringEquals: Success
      Default: callbackFailure
  - name: callbackSuccess
    description: Step Functions にタスクが成功した旨をコールバックする。
    action: 'aws:executeAwsApi'
    inputs:
      Service: stepfunctions
      Api: SendTaskSuccess
      output: |-
        {
          "AutomationExecutionURL": "https://ap-northeast-1.console.aws.amazon.com/systems-manager/automation/execution/{{ automation:EXECUTION_ID }}?region=ap-northeast-1",
          "CommandResult": {{ runCommand.OutputPayload }}
        }        
      taskToken: '{{ TaskToken }}'
    isEnd: true
  - name: callbackFailure
    description: Step Functions にタスクが失敗した旨をコールバックする。
    action: 'aws:executeAwsApi'
    inputs:
      Service: stepfunctions
      Api: SendTaskFailure
      cause: |-
        {
          "AutomationExecutionURL": "https://ap-northeast-1.console.aws.amazon.com/systems-manager/automation/execution/{{ automation:EXECUTION_ID }}?region=ap-northeast-1",
          "CommandResult": {{ runCommand.OutputPayload }}
        }        
      error: ExecutionFailed
      taskToken: '{{ TaskToken }}'
    isEnd: true
  - name: callbackWorkerBusy
    description: Step Functions にワーカが Busy である旨をコールバックする。
    action: 'aws:executeAwsApi'
    inputs:
      Service: stepfunctions
      Api: SendTaskFailure
      cause: |-
        {
          "AutomationExecutionURL": "https://ap-northeast-1.console.aws.amazon.com/systems-manager/automation/execution/{{ automation:EXECUTION_ID }}?region=ap-northeast-1"
        }        
      error: WorkerBusy
      taskToken: '{{ TaskToken }}'
    isEnd: true

{
  "StartAt": "Run job",
  "States": {
    "Run job": {
      "Type": "Task",
      "Resource": "arn:aws:states:::aws-sdk:ssm:startAutomationExecution.waitForTaskToken",
      "Parameters": {
        "DocumentName": "RunShellScriptOnOneInstance",
        "Parameters": {
          "TaskToken.$": "States.Array($$.Task.Token)",
          "Tags.$": "States.Array($.Tags)",
          "Commands.$": "$.Commands",
          "WorkingDirectory.$": "States.Array($.WorkingDirectory)",
          "ExecutionTimeout.$": "States.Array($.ExecutionTimeout)"
        }
      },
      "Next": "Callback success",
      "Catch": [
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "Next": "Callback failure"
        }
      ],
      "Retry": [
        {
          "ErrorEquals": [
            "WorkerBusy"
          ],
          "BackoffRate": 2,
          "IntervalSeconds": 100,
          "MaxAttempts": 2
        }
      ]
    },
    "Callback success": {
      "Type": "Task",
      "Resource": "arn:aws:states:::aws-sdk:sfn:sendTaskSuccess",
      "Parameters": {
        "Output.$": "$",
        "TaskToken.$": "$$.Execution.Input.TaskToken"
      },
      "Retry": [
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "BackoffRate": 1,
          "IntervalSeconds": 5,
          "MaxAttempts": 2
        }
      ],
      "Next": "Success"
    },
    "Success": {
      "Type": "Succeed"
    },
    "Callback failure": {
      "Type": "Task",
      "Resource": "arn:aws:states:::aws-sdk:sfn:sendTaskFailure",
      "Parameters": {
        "Error.$": "$.Error",
        "Cause.$": "$.Cause",
        "TaskToken.$": "$$.Execution.Input.TaskToken"
      },
      "Retry": [
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "BackoffRate": 1,
          "IntervalSeconds": 5,
          "MaxAttempts": 2
        }
      ],
      "Next": "Fail"
    },
    "Fail": {
      "Type": "Fail"
    }
  }
}

{
  "StartAt": "Send",
  "States": {
    "Send": {
      "Type": "Task",
      "Resource": "arn:aws:states:::events:putEvents",
      "Parameters": {
        "Entries": [
          {
            "Source": "net.quickguard.states",
            "DetailType": "Job Network Execution Fail",
            "Detail": {
              "executionURL.$": "States.Format('https://ap-northeast-1.console.aws.amazon.com/states/home?region=ap-northeast-1#/executions/details/{}', $$.Execution.Input.AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID)",
              "error.$": "$.Error",
              "cause.$": "States.StringToJson($.Cause)",
              "channel.$": "$.Channel"
            }
          }
        ]
      },
      "End": true,
      "Retry": [
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "BackoffRate": 1,
          "IntervalSeconds": 5,
          "MaxAttempts": 2
        }
      ]
    }
  }
}

{
  "source": [
    "net.quickguard.states"
  ],
  "detail-type": [
    "Job Network Execution Fail"
  ],
  "detail": {
    "channel": [
      "{{ チャネル }}"
    ]
  }
}

{
  "StartAt": "Job 1",
  "States": {
    "Job 1": {
      "Type": "Task",
      "Resource": "arn:aws:states:::states:startExecution.waitForTaskToken",
      "Parameters": {
        "StateMachineArn": "arn:aws:states:ap-northeast-1:********:stateMachine:RunJob",
        "Input": {
          "AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID.$": "$$.Execution.Id",
          "TaskToken.$": "$$.Task.Token",
          "Tags.$": "$$.Execution.Input",
          "Commands": [
            "echo 'SUCCESS'",
            "true"
          ],
          "WorkingDirectory": "/tmp",
          "ExecutionTimeout": "300"
        }
      },
      "Next": "Job 2",
      "Catch": [
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "Next": "Send fail event"
        }
      ]
    },
    "Job 2": {
      "Type": "Task",
      "Resource": "arn:aws:states:::states:startExecution.waitForTaskToken",
      "Parameters": {
        "StateMachineArn": "arn:aws:states:ap-northeast-1:********:stateMachine:RunJob",
        "Input": {
          "AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID.$": "$$.Execution.Id",
          "TaskToken.$": "$$.Task.Token",
          "Tags.$": "$$.Execution.Input",
          "Commands": [
            "echo 'FAILURE'",
            "false"
          ],
          "WorkingDirectory": "/tmp",
          "ExecutionTimeout": "300"
        }
      },
      "Next": "Job 3",
      "Catch": [
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "Next": "Send fail event"
        }
      ]
    },
    "Job 3": {
      "Type": "Task",
      "Resource": "arn:aws:states:::states:startExecution.waitForTaskToken",
      "Parameters": {
        "StateMachineArn": "arn:aws:states:ap-northeast-1:********:stateMachine:RunJob",
        "Input": {
          "AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID.$": "$$.Execution.Id",
          "TaskToken.$": "$$.Task.Token",
          "Tags.$": "$$.Execution.Input",
          "Commands": [
            "echo 'NEVER EXECUTED'",
            "true"
          ],
          "WorkingDirectory": "/tmp",
          "ExecutionTimeout": "300"
        }
      },
      "Catch": [
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "Next": "Send fail event"
        }
      ],
      "End": true
    },
    "Send fail event": {
      "Type": "Task",
      "Resource": "arn:aws:states:::states:startExecution.sync:2",
      "Parameters": {
        "StateMachineArn": "arn:aws:states:ap-northeast-1:********:stateMachine:SendFailEvent",
        "Input": {
          "AWS_STEP_FUNCTIONS_STARTED_BY_EXECUTION_ID.$": "$$.Execution.Id",
          "Error.$": "$.Error",
          "Cause.$": "$.Cause",
          "Channel": "high"
        }
      },
      "Next": "Fail"
    },
    "Fail": {
      "Type": "Fail"
    }
  }
}

{
  "role": "job_worker",
  "env": "production"
}

{
  "schemaVersion": "1.2",
  "description": "Run a shell script or specify the commands to run.",
  "parameters": {
    ...(略)...
    "executionTimeout": {
      "type": "String",
      "default": "3600",
      "description": "(Optional) The time in seconds for a command to complete before it is considered to have failed. Default is 3600 (1 hour). Maximum is 172800 (48 hours).",
      "allowedPattern": "([1-9][0-9]{0,4})|(1[0-6][0-9]{4})|(17[0-1][0-9]{3})|(172[0-7][0-9]{2})|(172800)"
    }
  },
  "runtimeConfig": {
    "aws:runShellScript": {
      "properties": [
        {
          "id": "0.aws:runShellScript",
          "runCommand": "{{ commands }}",
          "workingDirectory": "{{ workingDirectory }}",
          "timeoutSeconds": "{{ executionTimeout }}"
        }
      ]
    }
  }
}