StepFunctionsでECSタスクのMemoryを倍々にしながら再実行する

概要

表題の通り、StepFunctionsでECSタスクがOutOfMemory起因で失敗した場合に限り、そのメモリを倍々にしながら再実行する方法についてです。要するに ContainerOverrides を使っているだけなのですが、備忘をかねてまとめてみました。どなたかの参考になれば幸いです。

実装

早速ですが、以下がStateMachineのサンプルです。

実装（json）は以下の通り。

{
  "Comment": "sample",
  "StartAt": "initialize_memory",
  "States": {
    "check_memory_error": {
      "Choices": [
        {
          "Next": "increase_memory",
          "StringMatches": "*OutOfMemoryError*",
          "Variable": "$.error_info.Cause"
        }
      ],
      "Default": "retry_without_memory_increase",
      "Type": "Choice"
    },
    "check_retry_count": {
      "Choices": [
        {
          "Next": "run_task",
          "NumericLessThan": 3,
          "Variable": "$.task_info.retry_count"
        }
      ],
      "Default": "fail_state",
      "Type": "Choice"
    },
    "fail_state": {
      "Type": "Fail"
    },
    "increase_memory": {
      "Next": "check_retry_count",
      "Parameters": {
        "cpu.$": "States.MathAdd($.task_info.cpu, $.task_info.cpu)",
        "memory.$": "States.MathAdd($.task_info.memory, $.task_info.memory)",
        "retry_count.$": "States.MathAdd($.task_info.retry_count, 1)"
      },
      "ResultPath": "$.task_info",
      "Type": "Pass"
    },
    "initialize_memory": {
      "Next": "run_task",
      "Result": {
        "cpu": 1024,
        "memory": 4096,
        "retry_count": 0
      },
      "ResultPath": "$.task_info",
      "Type": "Pass"
    },
    "retry_without_memory_increase": {
      "Next": "check_retry_count",
      "Parameters": {
        "retry_count.$": "States.MathAdd($.task_info.retry_count, 1)"
      },
      "ResultPath": "$.task_info",
      "Type": "Pass"
    },
    "run_task": {
      "Catch": [
        {
          "ErrorEquals": [
            "States.ALL"
          ],
          "Next": "check_memory_error",
          "ResultPath": "$.error_info"
        }
      ],
      "Next": "success_state",
      "Parameters": {
        "Cluster": "arn:aws:ecs:ap-northeast-1:123456789000:cluster/sample",
        "LaunchType": "FARGATE",
        "NetworkConfiguration": {
          "AwsvpcConfiguration": {
            "AssignPublicIp": "DISABLED",
            "SecurityGroups": [
              "sg-sample"
            ],
            "Subnets": [
              "subnet-sample1a",
              "subnet-sample1c",
              "subnet-sample1d"
            ]
          }
        },
        "Overrides": {
          "ContainerOverrides": [
            {
              "Name": "app"
            }
          ],
          "Cpu.$": "States.JsonToString($.task_info.cpu)",
          "Memory.$": "States.JsonToString($.task_info.memory)"
        },
        "PlatformVersion": "1.4.0",
        "TaskDefinition": "arn:aws:ecs:ap-northeast-1:123456789000:task-definition/sample"
      },
      "Resource": "arn:aws:states:::ecs:runTask.sync",
      "Type": "Task"
    },
    "success_state": {
      "Type": "Succeed"
    }
  },
  "TimeoutSeconds": 86400
}

ECSはMemoryとvCPUの組み合わせについて規定があるため、上記サンプルではMemoryを倍にする際、一緒にvCPUも倍にするようにしています。単純な倍々にするのではなく、よりきめ細やかな制御が必要な場合は、例えばvCPUやMemoryの値について事前に配列を定義しておいて、Retry回数に応じて、配列からN番目の要素を引っ張ってくるような手法で実現可能ではないかと思っています。
また、一旦再試行回数は３回までとしています。それを超えたら（OOMだろうと、それ以外のエラーだろうと）再実行はしない設定です。この辺りの調整については、皆様のユースケースに合わせて変更していただければと思います。