見出し画像

AWS SageMakerのJupyterLab環境を全自動で作成(したかった)

こんにちは、DevOpsエンジニアのSham六です。
普段はインフラ業務をしているのですが、機械学習を触ってみようと思い、AWS SageMaker StudioでJupyterLabを準備してみました。
今回CloudFromationで構築してみたのですが、ライフサイクル設定が思うようにいかなかったので共有したいと思います。

コード全体

CloudFormaitonでJupyterLab起動までできました。
コード全体を見たい方は以下のGithubをご参照ください。

ライフサイクル設定

SageMakerスペースではLifecycle設定を入れることで、インスタンスの停止を自動化することができます。

公式ドキュメントを探してみたのですが、LifeCycle設定はCFnテンプレートでは定義できないようです。
仕方ないのでaws cliでLifecycle設定を作成してみることにしました。
今回は環境構築が目的なので、AWS公式のサンプルのスクリプトをそのまま使用することにしました。

環境構築

環境構築は以下のような流れになります。

ステップ1:スクリプトでSageMakerでライフサイクル設定を作成する
ステップ2:CloudFormationで作成したライフサイクル設定をSageMaker Domainに紐づける

CI/CDパイプラインへの組み込みを想定しステップ1,2ともにスクリプトで実行します。
それでは各ステップを説明していきます。

ステップ1:ライフサイクル設定を作成する

公式 のサンプルを参考にライフサイクル設定を作成します。

#!/bin/bash
ASI_VERSION=0.3.1
curl -LO https://github.com/aws-samples/sagemaker-studio-apps-lifecycle-config-examples/releases/download/v$ASI_VERSION/jupyterlab-lccs-$ASI_VERSION.tar.gz
tar -xvzf jupyterlab-lccs-$ASI_VERSION.tar.gz

cd auto-stop-idle/

LCC_NAME=jupyter-lab-auto-stop-idle
LCC_CONTENT=`openssl base64 -A -in on-start.sh`

aws sagemaker create-studio-lifecycle-config \
    --studio-lifecycle-config-name $LCC_NAME \
    --studio-lifecycle-config-content $LCC_CONTENT \
    --studio-lifecycle-config-app-type JupyterLab \
    --query 'StudioLifecycleConfigArn'

ライフサイクル設定名(LCC_NAME)は任意ですが、CFnテンプレートからの指定に使うのでメモしておきます。

独自のライフサイクル設定ファイルを適用させたい場合はスクリプトの以下の部分でファイル名を変更します。

LCC_CONTENT=`openssl base64 -A -in on-start.sh`

スクリプトを実行すると以下のようにライフサイクル設定が作成できました。

ステップ2:ライフサイクル設定をSageMaker Domainに紐づける

SageMakerドメインのCFnテンプレート全体は以下になります。

  SageMakerStudioDomain:
    Type: AWS::SageMaker::Domain
    Properties:
      DomainName: !Sub "${EnvPrefix}-sagemaker-domain"
      AppNetworkAccessType: PublicInternetOnly
      AuthMode: IAM
      VpcId: !Ref Vpc
      SubnetIds:
        - !Ref SubnetAzA
        - !Ref SubnetAzC
      DefaultUserSettings:
        ExecutionRole: !GetAtt SageMakerExecutionRole.Arn
        StudioWebPortal: ENABLED
        DefaultLandingUri: "studio::"
      DefaultSpaceSettings:
        ExecutionRole: !GetAtt SageMakerExecutionRole.Arn
        JupyterLabAppSettings:
          LifecycleConfigArns:
            - !Sub "arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:studio-lifecycle-config/jupyter-lab-auto-stop-idle"

以下の部分でライフサイクル設定のARNを指定します。

DefaultSpaceSettings:
  JupyterLabAppSettings:
    LifecycleConfigArns:
      - !Sub "arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:studio-lifecycle-config/jupyter-lab-auto-stop-idle"

Jupyterlabスペース側でライフサイクル設定を紐付けできなかったので(2024年7月時点) 、仕方なくドメイン側で定義しました。

ちなみにJupyterLabスペースのCFnテンプレートは以下のようになります。

JupyterLabSpace:
    Type: AWS::SageMaker::Space
    Properties:
      DomainId: !Ref SageMakerStudioDomain
      SpaceName: !Sub "${EnvPrefix}-jupyterlab-space"
      SpaceDisplayName: "JupyterLab"
      OwnershipSettings:
        OwnerUserProfileName: !Sub "${EnvPrefix}-sagemaker-user-profile"
      SpaceSharingSettings:
        SharingType: Private
      SpaceSettings:
        AppType: JupyterLab
        JupyterLabAppSettings:
          DefaultResourceSpec:
            InstanceType: "ml.t3.medium"
            SageMakerImageVersionArn: !Sub
              - "arn:aws:sagemaker:${AWS::Region}:{SageMakerDistImg}:image-version/sagemaker-distribution-cpu/1"
              - { SageMakerDistImg: !FindInMap [SageMakerDistImgMap, !Ref AWS::Region, ImageNumber] }
        SpaceStorageSettings:
          EbsStorageSettings:
            EbsVolumeSizeInGb: 5


CFnテンプレートからSageMakerドメインとJupyterLabスペースを作成してみます。

まずSageMakerドメイン。問題なく作成できています。

続いてJupyterLabスペースを見てみると、

おや?ライフサイクル設定が紐づいていませんね。

一応、手動で設定できました。
しかし、ライフサイクル設定まで全自動にしたかったので、残念です。

自動停止確認

最後にインスタンスが自動で停止されるのかを確認してみます。

JupyterLabスペースを起動して、アイドルタイム(サンプルでは1時間)が経過するのを待ちます。

1時間後、問題なく停止しました。

CloudWatch Logsにもログが残っていました。

まとめ

SageMakerの環境設定をほぼ全自動にすることができました。
一部手動で設定する部分があるのが残念ですが、今後のアップデートに期待したいと思います。
今後は、今回作成したJupyterLabを使ってAWS Bedrockを触ってみたいです。


この記事が気に入ったらサポートをしてみませんか?