😺

SageMaker StudioのCode Editorにカスタムイメージを適用させる方法

に公開
AWS
Docker
SageMaker
image
custom
tech

概要

最近、SageMaker Studioのインスタンスのビルドが遅くなっている気がします。
それに伴って、今までは使い捨てが容易なライフサイクル設定で気楽に運用していたんですが、既存のライフサイクル設定だと上限時間である五分を超えることが多くなりました。これは困る。

ということで、今回はライフサイクル設定から脱却するために、SageMaker StudioのCode Editor環境でカスタムイメージを適用するための作成方法をまとめました。

Dockerfileを作成

AWSで提供されてるGPUイメージをベースにするのですが、注意しないといけないのは、

JupyterLab とコードエディタのアプリケーションは、UID:1001 および GID:101 の sagemaker-user として実行されます。デフォルトで、sagemaker-user には sudo/root のアクセス許可が付与されます。KernelGateway アプリケーションはデフォルトで root として実行されます。

と公式で述べられてる点です。

https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/studio-updated-migrate-lcc.html

こちらについては、dockerfileにある「コンテナ内ユーザー設定」をすることで解消できます。これがないと、イメージをアタッチできても起動する際にエラーになってしまいます。他にも、root権限でOSレベルの依存を追加したり、Studio が期待するエントリポイントを適切に設定する必要があります。

下記がdockerfileの例になります。

FROM public.ecr.aws/sagemaker/sagemaker-distribution:3.1-gpu

# ---- コンテナ内ユーザー設定 (Studio の既定 UID/GID と合わせる) ----
ARG NB_USER="sagemaker-user"
ARG NB_UID=1001
ARG NB_GID=101
ENV MAMBA_USER=${NB_USER}

# ---- root 権限で OS レベルの依存を追加 ----
USER root

RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        libopencv-dev \
        libgl1-mesa-dev \
        zip \
        ca-certificates \
        curl && \
    rm -rf /var/lib/apt/lists/*

# Code-Editor側でDockerやDocker-Composeを取り扱うためにインストール
RUN mkdir -p /etc/apt/keyrings && \
    curl -fsSL https://download.docker.com/linux/ubuntu/gpg | tee /etc/apt/keyrings/docker.asc > /dev/null && \
    echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu jammy stable" \
    > /etc/apt/sources.list.d/docker.list && \
    apt-get update && apt-get install -y --allow-downgrades \
    docker-ce=5:20.10.24~3-0~ubuntu-jammy \
    docker-ce-cli=5:20.10.24~3-0~ubuntu-jammy \
    docker-compose-plugin && \
    rm -rf /var/lib/apt/lists/*
RUN apt-get update && apt-get install -y docker-compose-plugin

# pip でライブラリをインストール
RUN pip install --no-cache-dir \
    numpy \
    pandas

# ---- 権限調整 & 作業ディレクトリ ----
RUN chown -R ${NB_UID}:${NB_GID} /opt/conda /home/${NB_USER}
USER ${MAMBA_USER}
WORKDIR /home/${NB_USER}

# ---- Studio が期待するエントリポイント (VS Code-OSS 起動) ----
ENTRYPOINT ["entrypoint-code-editor"]

ECRへプッシュするためのスクリプトを作成

# 変数定義
AWS_ACCOUNT_ID=<AWS_ACCOUNT_ID>
REGION=ap-northeast-1
REPO_NAME=<ECR Repository Name>

# ECR リポジトリ作成（初回のみ）
aws ecr create-repository --repository-name ${REPO_NAME} --region ${REGION}

# ECR にログイン
aws ecr get-login-password --region ${REGION} \
  | docker login --username AWS --password-stdin \
    ${AWS_ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com

# ビルド & タグ付け(code-editorで進める場合、--network sagemakerが必須)
docker build -t ${REPO_NAME}:latest . --network sagemaker
docker tag ${REPO_NAME}:latest \
  ${AWS_ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/${REPO_NAME}:v1

# Push
docker push ${AWS_ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/${REPO_NAME}:v1

Code-Editorから実行したい場合は、create(or update) domainで下記のコマンドを実行してください。これで、SageMaker Studio内の環境でdockerコマンドを呼び出すことができるようになります。

aws --region region \
    sagemaker update-domain --domain-id domain-id \
    --domain-settings-for-update '{"DockerSettings": {"EnableDockerAccess": "ENABLED",}}'

その後、Code-Editorでdocker環境をインストールする必要があるので、下記を実行します。

# Install dependencies
sudo apt-get update && sudo apt-get install -y --no-install-recommends \
  libopencv-dev libgl1-mesa-dev zip ffmpeg \
  ca-certificates curl nvidia-cuda-toolkit

# Docker install
sudo mkdir -p /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo tee /etc/apt/keyrings/docker.asc > /dev/null
echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] \
  https://download.docker.com/linux/ubuntu $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \
  sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

sudo apt-get update && sudo apt-get install -y --allow-downgrades \
  docker-ce=5:20.10.24~3-0~ubuntu-jammy \
  docker-ce-cli=5:20.10.24~3-0~ubuntu-jammy

# CUDA
export CUDA_HOME=/opt/conda

sudo apt-get update && sudo apt-get install -y docker-compose-plugin

上記のシェルを実行すると、ECRでリポジトリが作成されて、イメージがECRへプッシュされます。今回使用してるイメージサイズは30GB以上なので、30分前後処理がかかります。

後は、sagemaker studioで左のサイドバーにあるイメージを選択して、ECRにあるイメージを登録します。そして、「ドメイン」→「環境」へ移動して、登録したカスタムイメージをアタッチすれば使用できるようになります。

まとめ

今回は、SageMaker StudioのCode Editorにカスタムイメージを適用させる方法を紹介しました。

意外ときちんとカスタムイメージをCode Editor等で扱うためには注意しないといけない項目が多いため、あまり記事がなく苦戦しました。公式マニュアルと格闘する羽目になりました。

ウェブ記事にない場合って、大抵githubリポジトリの中身を見たり、公式マニュアルと格闘しがちですよね。

参考資料

https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/studio-updated-migrate-lcc.html

https://qiita.com/kazuneet/items/ecfd40c67d4ad83f4855

Discussion