AWS Cloud Development Kit(CDK)(Python)を使って「VPC+Sagemaker+Lambda」環境構築してみた(Sagemaker編)

2020.08.18

どーもsutoです。

前回の記事でAWS Cloud Development Kit(CDK)を使ってPythonコードでVPC構築自動化をやってみました。

今回はこちらにsagemakerのノートブックインスタンスを作成するコードを追記していきたいと思います。

はじめに

今回、自分のアカウントで機械学習検証用環境のインフラを構築し、いつでも環境のデプロイ/破棄ができるように自動化&コード管理したい、というモチベからAWS CDKを本格的に触ってみることにしました。

本検証の目標

  1. AWS CDKのインストールと初期設定、VPC構築
  2. Sagemakerのノートブックインスタンスを自動構築、VPCとの関連付け(←本記事)
  3. ノートブックインスタンスのスケジュールによる自動停止を行うLambdaの追加設定

でやっていきます。

1、3は別記事に掲載しています。

CDK環境にSagemakerモジュール追加

前回作成したsage-nwプロジェクトの仮想環境ディレクトリ内に入っている状態からスタートします。

~ sage-nw % source .env/bin/activate
(.env) ~ sage-nw %

setup.pyの「install_requires」部分を編集(aws-cdk.aws-sagemaker)して必要なモジュールをインストールします。

    install_requires=[
        "aws-cdk.core",
        "aws-cdk.aws-ec2",
        "aws-cdk.aws-sagemaker",
        ],
pip install -r requirements.txt

コードを記述

今回のSagemaker構築内容は以下のとおりです。

  • インスタンスタイプを指定したノートブックインスタンス1台新規作成
  • ライフサイクル設定を追加し、インスタンス起動時に実行するスクリプトを起動するようにする
  • ネットワークオプション設定を追加し、プロジェクトで前回デプロイしたVPCと関連付け

まず「sage_nw」フォルダ内に移動し、ライフサイクル設定用のスクリプトファイルを作成し、保存します。

onstart.sh

#!/bin/bash
set -e

nohup pip install --upgrade pip

sagemaker_stack.pyを新規作成し、以下のコードを記述します。

from aws_cdk import core
from aws_cdk import aws_sagemaker as sagemaker
import base64


class SagemakerStack(core.Construct):

    def __init__(self, scope: core.Construct, id: str, subnetid: str, sgid: str) -> None:
        super().__init__(scope, id,)
    
            # The code that defines your stack goes here
            # 各種パラメータの格納する変数を宣言
            instance_type = 'ml.t2.medium'
            role_arn = 'arn:aws:iam::111122223333:role/service-role/AmazonSageMaker-ExecutionRole-YYYYYYYYY'
            internet_access = 'Enabled'
            LifeCycleConfigName = 'notebook-lf'
            LifecycleScriptStr = open("sage_nw/onstart.sh", "r").read()
    
            content = [
                {"content": core.Fn.base64(LifecycleScriptStr)}
            ]
            sagemaker_lifecycle = sagemaker.CfnNotebookInstanceLifecycleConfig(
                self,
                LifeCycleConfigName,
                notebook_instance_lifecycle_config_name=LifeCycleConfigName,
                on_start=content,
            )
            sagemaker_notebook = sagemaker.CfnNotebookInstance(
                self,
                id='notebook',
                instance_type=instance_type,
                role_arn=role_arn,
                security_group_ids=sgid,
                subnet_id=subnetid,
                direct_internet_access=internet_access,
                lifecycle_config_name=LifeCycleConfigName,
            )

処理ごとにコードを管理しやすいようにファイルを分けました。この処理をメインであるsage_nw_stack.pyで呼び出すようにします。

sage_nw_stack.pyにsagemaker_stackの処理を呼び出す内容を追記して保存します。

from aws_cdk import core
from aws_cdk import aws_ec2 as ec2
from sagemaker_stack import SagemakerStack

class SageNwStack(core.Stack):
    def __init__(self, scope: core.Construct, id: str, **kwargs) -> None:
        super().__init__(scope, id, **kwargs)

    # The code that defines your stack goes here
    # 変数の宣言
    vpc_cidr = '10.1.0.0/16'
    subnet_mask = 24
    
    # 新規VPC作成
    vpc = ec2.Vpc(
        self,
        id="Sage-vpc",
        cidr=vpc_cidr,
        nat_gateways=0,
        subnet_configuration=[
            ec2.SubnetConfiguration(
                cidr_mask=subnet_mask, name='public', subnet_type=ec2.SubnetType.PUBLIC,
            ),
            ec2.SubnetConfiguration(
                cidr_mask=subnet_mask, name='private', subnet_type=ec2.SubnetType.ISOLATED,
            ),
        ],
    )
    
    security_group = ec2.SecurityGroup(
        self,
        id='Sage-sg',
        vpc=vpc,
        security_group_name='test-sg'
    )
    security_group.add_ingress_rule(
        peer=ec2.Peer.ipv4(vpc_cidr),
        connection=ec2.Port.all_traffic(),
    )
    # vpcメソッドの結果を変数に格納
    subnet = vpc.isolated_subnets[0].subnet_id
    sg = [security_group.security_group_id]
    
    # ノートブックインスタンス作成
    sagemaker = SagemakerStack(
        self,
        "Notebook",
        subnetid = subnet,
        sgid = sg,
    )

スタックの確認とデプロイ

cdk diffコマンドでデプロイ済の Stack との差異を表示することができます。

cdk diff sage-nw --profile suto

確認が終わったら以下コマンドでスタックを更新してみます。

(.env) ~ sage-nw % cdk deploy sage-nw --profile suto

スタックが更新されていることを確認できました。

まとめ

今回のSagemakerモジュールについてはCfnがついているためCloudFormationリソースと1対1で対応していることがわかります。PythonとCloudFormationで比較はしやすいメリットがありますが、コード短縮はVPCのときほどの恩恵はないかな〜。それでもコードを書く量は削減できますし、スタック更新が同じコマンドで簡単にできるのが大きいですね。

次回はこれにLambda関数を追加してみようと思います。

<次回記事>