AWS CDKのStepFunctionsでEcsRunTaskにcpuパラメータを渡す

Musubi AI 在庫管理で DevOps エンジニアをしている kacky です。

この記事では AWS CDK 上でのジョブ開発でのとある困り事を解決した話をしたいと思います。

AWS CDK の StepFunctions には、EcsRunTask という API が存在します。この API を使用すると、StepFunctions 上で ECS の Task を指定して実行することが可能です。この機能は Fargate や EC2 に対応しており、バッチ処理などには必要不可欠なものです。しかし、この API には CPU や Memory などのリソース割り当てを実行時に上書きできないという制限があります。

AWS の REST API や boto3 などの SDK では、この機能を実現できますが、CDK の StepFunctions には含まれていません。これは非常に面倒な問題です。

これらの問題を解決すべく、StepFunctions の EcsRunTask でリソース割り当てを上書きできないかを模索しました。

調査の結果、CDK の stepfunctions のソースコード内で以下のような API の実装を発見しました。

https://github.com/aws/aws-cdk/blob/4f1c94b27ef7f4ceccea0ff39625c0e8add31c9f/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/ecs/run-task.ts#L315

  protected _renderTask(): any {
    return {
      Resource: integrationResourceArn('ecs', 'runTask', this.integrationPattern),
      Parameters: sfn.FieldUtils.renderObject({
        Cluster: this.props.cluster.clusterArn,
        TaskDefinition: this.props.revisionNumber === undefined ? this.props.taskDefinition.family : `${this.props.taskDefinition.family}:${this.props.revisionNumber.toString()}`,
        NetworkConfiguration: this.networkConfiguration,
        Overrides: renderOverrides(this.props.containerOverrides),
        PropagateTags: this.props.propagatedTagSource,
        ...this.props.launchTarget.bind(this, { taskDefinition: this.props.taskDefinition, cluster: this.props.cluster }).parameters,
        EnableExecuteCommand: this.props.enableExecuteCommand,
      }),
    };
  }

この関数は JavaScript の Object を sfn.FieldUtils.renderObject で JSON に変換して ECS Run Task の パラメータ として実行します。この実装ゆえに ECS Run Task のパラメータに固定のキーしか指定できないことがわかります。一方で、注目すべきは、launchTarget.bind の戻り値で上書きする部分です。つまり、launchTarget.bind の戻り値に cpu や memory などのパラメータを返せば実現できそうに見えます。

つぎに、cpu や memory などのパラメータを受け付けるカスタムの LaunchTarget を作成してみました。以下にそのサンプルコードを示します。

/**
 * Properties to define an ECS service
 */
export interface EcsFargateCustomLaunchTargetOptions {
  readonly platformVersion?: ecs.FargatePlatformVersion;

  readonly cpu?: string;

  readonly memory?: string;
}

/**
 * FargateTaskでCPU、Memoryを上書き実行する LaunchTargetConfig
 */
export class EcsFargateCustomLaunchTarget implements IEcsLaunchTarget {
  constructor(private readonly options?: EcsFargateCustomLaunchTargetOptions) {}

  bind(
    task: EcsRunTask,
    launchTargetOptions: LaunchTargetBindOptions
  ): EcsLaunchTargetConfig {
    if (!launchTargetOptions.taskDefinition.isFargateCompatible) {
      throw new Error("Supplied TaskDefinition is not compatible with Fargate");
    }

    return {
      parameters: {
        LaunchType: "FARGATE",
        PlatformVersion: this.options?.platformVersion,
        Overrides: {
          Cpu: this.options?.cpu,
          Memory: this.options?.memory,
        },
      },
    };
  }
}

StepFunctions で ECS Task を実行する際には、上記のカスタム LaunchTarget を使用します。

以下のように使用することで、CDK の StepFunctions から CPU や Memory などのパラメータを上書きして、ECSTask を実行することができます。

import { Stack } from "aws-cdk-lib";
import * as ecs from "aws-cdk-lib/aws-ecs";
import * as sfnt from "aws-cdk-lib/aws-stepfunctions-tasks";
import { Construct } from "constructs";

interface CustomRunTaskProps {
  readonly cluster: ecs.ICluster;
  readonly task: ecs.TaskDefinition;
}

class CustomRunTask extends Construct {
  readonly task: sfnt.EcsRunTask;
  constructor(scope: Construct, id: string, props: CustomRunTaskProps) {
    super(scope, id);

    this.task = new sfnt.EcsRunTask(scope, "RunTask", {
      cluster: props.cluster,
      taskDefinition: props.task,
      launchTarget: new EcsFargateCustomLaunchTarget({
        platformVersion: ecs.FargatePlatformVersion.VERSION1_4,
        cpu: "1024",
        memory: "2048",
      }),
      integrationPattern: sfn.IntegrationPattern.RUN_JOB,
    });
  }
}

この記事では、カスタムの LaunchTarget を作成することでこの問題を回避する方法を示しました。しかし、この解決策は一時的なものであり、本質的な解決策ではありません。この機能は CDK の SDK の標準機能として実装されるべきです。そのため、私はこの問題を修正するための Pull Request を作成しました。 https://github.com/aws/aws-cdk/pull/30140 これにより、AWS CDK の StepFunctions の EcsRunTask API がより柔軟で使いやすいものになることを期待しています。