経緯

自分がAWS初心者ということもあり、ECSでGPUを使うタスクを実行するまでに、いくつか詰まるところがあったので共有します。

やりたいこと

自分がやりたかったことは、下の画像のような流れを作ることです。画像データがS3にアップロードされたタイミングで、LambdaがECSの機械学習関連のタスクを実行して、処理の結果をs3にアップロードするというものです。

今回は、画像の赤枠で囲まれた、ECSでGPUを使うタスクを実行する方法を説明します。

手順

以下の手順では、起動タイプがEC2のECSのタスクのやり方です。こちらのやりとりを見てみたところ、現時点ではFargateはGPUサポートされていないみたいです。

1. GPU対応AMIの選択

今回は、amzn2-ami-ecs-gpu-hvm-2.0.20220209-x86_64-ebsというGPU対応のAMIを使用しました。
なるべく新しいバージョンのAMIを使用してください。
古いAMIだと、コンテナエージェントのバージョンが低く、タスクの実行時にATTRIBUTEエラーが発生する場合があります。

2. インスタンス作成

インスタンスを作成する際に、注意する点は以下の3つです。

公式通り、GPU対応のp2、p3、g3、g4インスタンスのいずれかをインスタンスタイプとして選択してください。
IAMロールにecsInstanceRoleを選択してください。ロールの作成方法はこちら。
ユーザーデータに以下の行を追加してください。
cluster-nameには、紐づけたいクラスター名を入れてください。

#!/bin/bash
echo ECS_CLUSTER=cluster-name >> /etc/ecs/ecs.config

3. クラスターを作成

クラスターは前の手順でのクラスター名で作成してください。

4. タスク定義

起動タイプにはEC2を選択し、タスクとコンテナの定義の設定を行います。
コンテナの定義のコンテナの追加ボタンを押し、タスクで動かすコンテナイメージを設定します。

[環境]->[GPU]の項目で、タスクに使うGPU数を入力してください。インスタンスタイプのGPU数を超えない数で設定します。

GPU数の選択をしないと、タスクでGPUを使用できないのですが、新しいAWSのUIでは、GPU数の選択の項目が見つけられませんでした。なので、古いUIで設定することをおすすめします。

まとめ

設定は以上です。あとは、インスタンス起動してタスク実行してみてください。
AWSの日本語の説明、時々わけわからん(笑)。

参考文献