こんな人向けの記事

AWSを使って機械学習をやりたいけど、一度説明を受けたSageMakerはジョブを作って学習や推論をするにはいいけど、もっと基礎的な機械学習の動きを知ったり、Trial & Errorをするには使いにくいなぁ。
SageMaker Notebook InstanceでJupyter環境は作れるけど、いやVSCODEを使いたいんだよね。
環境をいちいちマネジメントコンソールからぽちぽち作ると、ちょっと環境消したいとか、変更したいというのが大変なので、IaCでやりたいなぁ

っていう人向け。
技術的にはEC2を機械学習用AMI（DeepLearningAMI）で立ち上げて、そこにSSHトンネル接続するということです。

IaCの選択

今回のIaCにはCDKを利用します。余計な記述をしなくてもAWSのベストプラクティスが盛り込まれているし、シンプルに書けるため。CDKについてはこういうサイトが詳しいので、参考にすべし。
またCDKの言語としてはPythonやTypeScriptなど色々選べるが、公式ドキュメントやサンプルの多い、TypeScriptを利用します。

アーキテクチャ

構成は非常によくある構成。セッションマネージャー経由でPrivate Subnet配下にあるEC2に接続します。Public Subnetにも踏み台サーバーを置いて接続する形はセキュリティ的にもベストプラクティスではないので、今ならこの方法を使いましょう。

IaCの解説

全体のコード

全体が必要な方はこちら。

cdk.json

{
    "app": "npx ts-node ./cdk.ts",
    "context": {
        "region": "ap-northeast-1",
        "ami_id": "ami-082ecedb3ced1f080",
        "volume_size": 500
    }
}

cdk.ts

#!/usr/bin/env node
import { App, Stack } from 'aws-cdk-lib';
import { MLBaseStack } from './lib/mlbase_stack';

const app = new App();
const region = app.node.tryGetContext("region") as string;
const ami_id = app.node.tryGetContext("ami_id") as string;
const volume_size = app.node.tryGetContext("volume_size") as number;

new MLBaseStack(app, 'MLBaseStack', {
    env: {
        region: region,
    },
    ami_id: ami_id,
    volume_size: volume_size
});

mlbase_stack.ts

import * as cdk from "aws-cdk-lib";
import * as ec2 from "aws-cdk-lib/aws-ec2";
import * as iam from "aws-cdk-lib/aws-iam";

import { Construct } from 'constructs';
import * as keypair from "cdk-ec2-key-pair";

interface MLbaseStackProps extends cdk.StackProps {
    ami_id: string;
    volume_size: number;
}

export class MLBaseStack extends cdk.Stack {
    constructor(scope: Construct, id: string, props: MLbaseStackProps) {
        super(scope, id, props);

        // VPC for EC2
        const vpc = new ec2.Vpc(this, 'myVPC', {
            cidr: "10.0.0.0/24",
            natGateways: 1,
            subnetConfiguration: [
                {
                    cidrMask: 28,
                    name: "public",
                    subnetType: ec2.SubnetType.PUBLIC
                },
                {
                    cidrMask: 28,
                    name: "private",
                    subnetType: ec2.SubnetType.PRIVATE_WITH_NAT
                }
            ],
            maxAzs: 1
        });

        // Key Pair for EC2
        const key = new keypair.KeyPair(this, 'keyPair', {
            name: 'ml-ec2-keypair',
            description: 'Key Pair created with CDK'
        });
        key.grantReadOnPublicKey;

        // Role for EC2
        const instanceRole = new iam.Role(this, 'instanceRole', {
            assumedBy: new iam.ServicePrincipal('ec2.amazonaws.com'),
            managedPolicies: [
                iam.ManagedPolicy.fromAwsManagedPolicyName('AmazonS3FullAccess'),
                iam.ManagedPolicy.fromAwsManagedPolicyName("AmazonSSMManagedInstanceCore")
            ],
        });

        // EC2 Instance
        const ec2Instance = new ec2.Instance(this, 'mlInstance', {
            vpc,
            role: instanceRole,
            keyName: key.keyPairName,
            instanceType: ec2.InstanceType.of(
                ec2.InstanceClass.P3,
                ec2.InstanceSize.XLARGE2,
            ),
            machineImage: new ec2.GenericLinuxImage({ [this.region]: props.ami_id }),
            blockDevices: [
                {
                    deviceName: '/dev/sda1',
                    volume: {
                        ebsDevice: {
                            volumeSize: props.volume_size
                        }
                    }
                }
            ]
        });

        new cdk.CfnOutput(this, "Download Key Command", {
            value: `aws secretsmanager get-secret-value --secret-id ec2-ssh-key/${key.keyPairName}/private --query SecretString --output text > cdk-key.pem && chmod 400 cdk-key.pem`
        });

        new cdk.CfnOutput(this, "EC2 Instance ID", {
            value: `${ec2Instance.instanceId}`
        });
    }
}

事前準備

基本的にはここに書いてあるので、aws-cdkを使えるようにしておきます。

ディレクトリ構成

infra/
 ├ cdk.json
 ├ cdk.ts
 └ lib/
    └ mlbase_stack.ts

cdk.json

パラメータをいくつか設定しているところ。コードに直接埋め込んでもいいし、いろいろな方法で渡すことが可能ですが。今回はcdk.jsonでcontextとして定義しています。
AMIのIDは例えばここに書いてあるように取得します。あるいはマネジメントコンソールのEC2でインスタンスを作ろうとすると、確認することも可能です。今回サンプルで使用しているAMIはDeepLearningAMI Ubuntu18.04 Version61.0 64-bit(x86)です。ボリュームサイズはEC2にくっつけるEBSのサイズで、とりあえず500GBとしていますが、好きなサイズにして良いです。

cdk.json

{
    "app": "npx ts-node ./cdk.ts",
    "context": {
        "region": "ap-northeast-1",
        "ami_id": "ami-0d69b2fd9641af433",
        "volume_size": 500
    }
}

cdk.ts

メイン部分。region, ami_id, volume_sizeはcdk.jsonに格納しているのでそれをロードします。
MLBaseStackは今回構築するVPCやEC2などのリソースの設定が書かれているところです。

cdk.ts

#!/usr/bin/env node
import { App, Stack } from 'aws-cdk-lib';
import { MLBaseStack } from './lib/mlbase_stack';

const app = new App();
const region = app.node.tryGetContext("region") as string;
const ami_id = app.node.tryGetContext("ami_id") as string;
const volume_size = app.node.tryGetContext("volume_size") as number;

new MLBaseStack(app, 'MLBaseStack', {
    env: {
        region: region,
    },
    ami_id: ami_id,
    volume_size: volume_size
});

mlbase_stack.ts

ここで今回構築しているリソースを定義している。下記詳細を説明する。

VPC

可読性が高いので、見て分かる通りだと思いますが、VPCを定義して、SubnetのCIDRを設定しています。
たったこれだけで、VPC、Private, Public, Gatewayなどのリソース定義ができます。

mlbase_stack.ts

const vpc = new ec2.Vpc(this, 'myVPC', {
    cidr: "10.0.0.0/24",
    natGateways: 1,
    subnetConfiguration: [
        {
            cidrMask: 28,
            name: "public",
            subnetType: ec2.SubnetType.PUBLIC
        },
        {
            cidrMask: 28,
            name: "private",
            subnetType: ec2.SubnetType.PRIVATE_WITH_NAT
        }
    ],
    maxAzs: 1
});

Key Pair

SSH接続するための鍵を作成。鍵の名前はなんでも良いです。

mlbase_stack.ts

const key = new keypair.KeyPair(this, 'keyPair', {
    name: 'ml-ec2-keypair',
    description: 'Key Pair created with CDK'
});
key.grantReadOnPublicKey;

EC2

EC2のインスタンスロールと、EC2そのものについての定義を記載しています。
インスタンスタイプとインスタンスサイズがハードコーディングされているので注意。ここではGPUを使う想定でP3.2xlargeとしています。EC2自体はcdk.jsonで設定しているDeepLearningAMIを使って立ち上がるので、立ち上がったEC2はすぐにでもGPUを使った学習や検討を行うことができます。
また、権限としてS3全体へのアクセス権と、SSMのSessionManager経由でSSH接続をするためのAmazonSSMManagedInstanceCoreの権限を付与しています。DBなどの接続が必要であれば、必要に応じてアクセス権をここに追加することが可能です。

mlbase_stack.ts

// Role for EC2
const instanceRole = new iam.Role(this, 'instanceRole', {
    assumedBy: new iam.ServicePrincipal('ec2.amazonaws.com'),
    managedPolicies: [
        iam.ManagedPolicy.fromAwsManagedPolicyName('AmazonS3FullAccess'),
        iam.ManagedPolicy.fromAwsManagedPolicyName("AmazonSSMManagedInstanceCore")
    ],
});

// EC2 Instance
const ec2Instance = new ec2.Instance(this, 'mlInstance', {
    vpc,
    role: instanceRole,
    keyName: key.keyPairName,
    instanceType: ec2.InstanceType.of(
        ec2.InstanceClass.P3,      // 任意のタイプに書き換えること
        ec2.InstanceSize.XLARGE2,  // 任意のサイズに書き換えること
    ),
    machineImage: new ec2.GenericLinuxImage({ [this.region]: props.ami_id }),
    blockDevices: [
        {
            deviceName: '/dev/sda1',
            volume: {
                ebsDevice: {
                    volumeSize: props.volume_size
                }
            }
        }
    ]
});

CfnOutput

リソースを作ったあとEC2に接続するためには、SSH接続のための鍵とインスタンスIDが必要となります。
その鍵の取得コマンドと、EC2のインスタンスIDを表示するために、実行後に出力されるようにしています。

mlbase_stack.ts

new cdk.CfnOutput(this, "Download Key Command", {
    value: `aws secretsmanager get-secret-value --secret-id ec2-ssh-key/${key.keyPairName}/private --query SecretString --output text > cdk-key.pem && chmod 400 cdk-key.pem`
});

new cdk.CfnOutput(this, "EC2 Instance ID", {
    value: `${ec2Instance.instanceId}`
});

デプロイ

以上のようなコードを準備したら、Deployを行います。
詳細はこういうページを参照していただくとして、基本的に下記コマンドを叩けば良いです。

npm install
cdk bootstrap
cdk deploy

EC2に接続

鍵のダウンロード

deployが完了すると、鍵の入手コマンドが現れるので、1行目のコマンドを実行して鍵をダウンロードします。

aws secretsmanager get-secret-value --secret-id ec2-ssh-key/ml-ec2-keypair/private --query SecretString --output text > cdk-key.pem && chmod 400 cdk-key.pem

SSH接続

VSCODEを使っている場合は、session manager経由で接続するために、下記のように~/.ssh/configに記載を行います。
<YOUR_PROFILE_NAME>には、AWS_PROFILEを、<HOST_NAME>には、cdk deployが完了すると表示されるEC2のインスタンスIDを記載する。

Host <任意の名前>
  ProxyCommand sh -c "aws --region ap-northeast-1 --profile <YOUR_PROFILE_NAME> ssm start-session --target %h --document-name AWS-StartSSHSession --parameters 'portNumber=%p'"
  User ubuntu
  HostName <HOST_NAME>
  IdentityFile <PATH_TO_THE_KEY>/cdk-key.pem

これでめでたく接続ができます。

削除

環境を消したい場合は下記でOK。作ったり消したりが簡単です。

$ cdk destroy
Are you sure you want to delete: MLBaseStack (y/n)?  y
MLBaseStack: destroying...
 ✅  MLBaseStack: destroyed

まとめ

VPC、EC2の機械学習環境をCDKでパッとデプロイして、すぐにSSHで接続をして開発を行うような環境を整えました。データセット用にS3のバケットを用意し、どこかのGitにソースコードを置いておけば、便利な開発環境の出来上がりです。5分程度ですぐに使えるようになると思います。
AWS上で機械学習の超序盤の研究開発をしたいんだけど、Jupyter Notebookだけではちょっと、、という方にはこんな使い方もいかがでしょうか。

ちなみに今回は機械学習用に作りましたが、もちろん別になんでも良いわけです。CDKを気軽に使って、自分専用の検討環境を作ったり消したり、使いやすいようにしておきましょう。