AWS SageMakerローカルモードで学ぶコンテナベース機械学習

R&Dチームの徳田(@dakuton)です。
時系列分析に携わるエンジニア目線でコロナウイルスに関する動向を見ていると、短い期間で当てはまる予測モデルをつくるというよりは、カットオフ発生(データの分断点)を調べることが重要になるだろうと思いながら自宅で静観しています。(どういうことか?気になる方は効果検証入門をお読みください)

本題です。AWS SageMakerについて調べていたら、AWSのリソースを全く消費せずに動作する謎のモード(local mode)があることに気がつきました。今回は、そちらの謎機能をもとに、AWS SageMakerの技術的な部分の解説をしようと思います。

SageMakerとは?

広義としては、AWSのホスティング環境で機械学習モデルの訓練とデプロイを行うためのサービス郡のことを指します。
SageMakerで提供されているマネージドサービスは、機械学習プロセス(アノテーション -> モデル開発 -> 訓練 -> モデル変換 -> デプロイ)の1プロセスを単純に構成するというよりは、より機能レベルの高いサービスとして提供されているサービスが数多くあります。

上記のような高機能サービスを提供できている背景はなんだろう?と思い、技術的にSageMakerの中心要素であるSageMakerコンテナ(モデル実行環境)に着目してみました。すると、各モデルの実行環境がそれぞれ「共通化されたDockerのコンテナデザインパターンで実現されている」点が重要な役目を果たしている、ということがわかりました。

一方で、Amazon SageMaker Studioを利用すれば、Google ColabのようにNotebookベースでインスタンスの起動やデータの保管などまとめて管理されているため、Dockerコンテナの存在を意識することなく利用できるようになっています。

実行インスタンス

Amazon SageMakerドキュメントだと、Notebookの使えるML専用インスタンスが案内されています。Notebookを使わないのであれば、それ以外のAmazon EC2インスタンスでも訓練・推論させることは可能です。

データ保管方法

SageMakerコンテナの訓練・推論にかかわるデータ(以下の5種類)は、S3 BucketやContainer Registry(コンテナイメージ保管用、Amazon ECRなど)にて保管されます。

  1. 訓練・推論環境(コンテナイメージ)
  2. データセット
  3. ハイパーパラメータ
  4. 学習済みモデル
  5. 評価指標

SageMakerコンテナの起動モード

SageMakerコンテナの振る舞いとしては、大きく分けて2つの起動モード(訓練モードと推論モード)があり、1つのSageMakerコンテナ内に同居しています。 どちらのモードで起動するかは、コンテナ外部から訓練・推論スクリプトを実行するために用意するエンドポイント2種類のうち、どちらのエンドポイントが呼び出されたか?によって決定されます。

https://cdn-ak.f.st-hatena.com/images/fotolife/o/optim-tech/20200409/20200409081335.png

デザインパターン

Docker公式にて紹介されている名称ではありませんが、私家版 Dockerfile Patternで当てはめると

  1. 訓練モード: 「Builder パターン」に該当、Dockerホストから受け渡されたデータセットをもとに学習済みモデルを生成
  2. 推論モード: 「Service パターン」に該当、常駐daemonとして起動し推論可能なWebAPIを払い出す

という2つのDockerデザインパターンを組み合わせて構成されていることがわかります。 先ほどの図をもとに--volume経由でのSageMakerコンテナ-Dockerホスト間でのデータの受け渡しに関して少し具体化すると、以下のように表されます。

https://cdn-ak.f.st-hatena.com/images/fotolife/o/optim-tech/20200409/20200409081329.png

ローカルでSageMakerを起動する

local modeの存在意義

利用例がAWS公式で紹介されています。(オンプレミス環境から Amazon SageMaker を利用する)

訓練・推論にかかるリソース(インスタンス、ストレージ)をクラウド管理外のサーバー含めて選択できるという点がメリットです。

local modeで起動する

SageMaker​ Python SDKを利用して、AWSのインスタンスではなく手元のPCで動かすことができます。 スクリプトは以下のようになります。

  • 私の動かした環境だと、exampleとして示されている、リージョンのlocal指定に加えて、ダミーのARNを挿入する必要がありました。
  • GPUコンテナ利用可能であれば、リージョンとしてlocal_gpuを指定することでGPUによる訓練を行うことができる場合もあります。(--runtime=nvidiaオプション起動となりますが、NVIDIA Container Toolkitバージョンによっては起動しないので注意)

Dockerfileサンプル(カスタムイメージ)

FROM python:3.7

ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE
ENV PATH="/opt/program:${PATH}"

# requirements serving
RUN pip install flask

# copy your custom endpoints
COPY train.py /opt/program/train
COPY serve.py /opt/program/serve

RUN chmod +x /opt/program/train
RUN chmod +x /opt/program/serve

訓練サンプル(Python)

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from pathlib import Path

from boto3 import Session
from sagemaker.estimator import Estimator
from sagemaker.local import LocalSession, LocalSagemakerRuntimeClient, LocalSagemakerClient
from sagemaker.model import Model

REGION_LOCAL = "local"
ROLE_LOCAL = "SageMakerRole/dummy"

class DummySession(LocalSession):
    def __init__(self, port):
        self.port = port
        super().__init__()

    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
        self.boto_session = Session(region_name=REGION_LOCAL)
        self.config = {
            "local": {
                "local_code": True,
                "region_name": REGION_LOCAL,
                "serving_port": self.port,
            }
        }
        self._region_name = REGION_LOCAL
        self.sagemaker_client = LocalSagemakerClient(self)
        self.sagemaker_runtime_client = LocalSagemakerRuntimeClient(self.config)
        self.local_mode = True

class LocalEstimator(Estimator):
    def __init__(self, image_name, output_dir, port):
        super().__init__(
            image_name,
            ROLE_LOCAL,
            1,
            REGION_LOCAL,
            REGION_LOCAL,
            output_path=Path(output_dir).resolve().as_uri(),
            sagemaker_session=DummySession(port)
        )


def main():
    estimator = LocalEstimator('my_sagemaker_image:latest', './train', 8000)
    estimator.fit(Path('./data').resolve().as_uri(), wait=False)


if __name__ == '__main__':
    main()

上記訓練スクリプトを実行すると、SDK内部的には下記のような自動生成されたdocker-compose.yamlに基づき、Dockerコンテナベースにて訓練実行されます。

SDKにより自動生成されるdocker-compose.yaml

networks:
  sagemaker-local:
    name: sagemaker-local
services:
  algo-1-xxxxx:
    command: train
    environment:
    - AWS_REGION=local
    - TRAINING_JOB_NAME=my_sagemaker_image-YYYY-mm-DD-HH-MM-SS-fff
    image: my_sagemaker_image:latest
    networks:
      sagemaker-local:
        aliases:
        - algo-1-xxxxx
    stdin_open: true
    tty: true
    volumes:
    - /private/var/folders/nn/xxxx/T/xxxx/algo-1-xxxxx/output/data:/opt/ml/output/data
    - /private/var/folders/nn/xxxx/T/xxxx/algo-1-xxxxx/output:/opt/ml/output
    - /private/var/folders/nn/xxxx/T/xxxx/algo-1-xxxxx/input:/opt/ml/input
    - /private/var/folders/nn/xxxx/T/xxxx/model:/opt/ml/model
    - /path/to/project_dir:/opt/ml/input/data/training
version: '2.3'

SDKにより間接的に実行される訓練スクリプト

docker-compose \
    -f /private/var/folders/nn/xxxx/T/xxxx/docker-compose.yaml \
    up --build --abort-on-container-exit

AWS S3以外でデータを保管する方法

SageMakerのデータ保管先として、AWS S3以外を選択することも技術的には可能です。
今回はAWS S3互換のオブジェクトストレージ(MinIO)を使って、S3 Bukcet とContainer Registryを同居させる例をご紹介します。

docker-compose.yaml

version: "3"
services:
  registry:
    image: registry:2.7.1
    depends_on:
      - minio
      - defaultbucket
    ports:
      - 9001:5000
    volumes:
      - ./docker-registry-config.yml:/etc/docker/registry/config.yml

  minio:
    image: minio/minio
    ports:
      - 9000:9000
    volumes:
     - ./minio:/export
    environment:
      MINIO_ACCESS_KEY: access_key
      MINIO_SECRET_KEY: secret
    command: server /export
    restart: always

  defaultbucket:
    image: minio/mc
    depends_on:
      - minio
    entrypoint: >
      /bin/sh -c "
      until (/usr/bin/mc config host add myminio http://minio:9000 access_key secret) do echo 'try to create buckets...' && sleep 1; done;
      /usr/bin/mc mb myminio/your_bucket;
      /usr/bin/mc policy download myminio/your_bucket;
      /usr/bin/mc mb myminio/docker-registry;
      /usr/bin/mc policy download myminio/docker-registry;
      exit 0;
      "

docker-registry-config.yml

version: 0.1
log:
  formatter: text
  fields:
    service: registry
    environment: staging
http:
  addr: :5000
storage:
  s3:
    accesskey: access_key
    secretkey: secret
    region: us-east-1
    regionendpoint: http://minio:9000
    bucket: docker-registry
    encrypt: false
    secure: false
    v4auth: true
    chunksize: 5242880
    rootdirectory: /images
  cache:
    blobdescriptor: 'inmemory'
  delete:
    enabled: true
  maintenance:
    uploadpurging:
      enabled: true
      age: 168h
      interval: 24h
      dryrun: false
    readonly:
      enabled: false
  redirect:
    disable: true

データ保存サンプル(Python)

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import boto3

s3 = boto3.resource('s3',
                    endpoint_url='http://minio:9000',
                    aws_access_key_id='access_key',
                    aws_secret_access_key='secret',
                    config=Config(signature_version='s3v4'),
                    region_name='us-east-1')
s3.Bucket('your_bucket').upload_file('model.tar.gz','model_latest.tar.gz')

コンテナpushサンプルコマンド

docker tag \
    my_sagemaker_image:latest \
    localhost:9001/<image_name>
docker push \
    localhost:9001/<image_name>

訓練・推論にかかわるスクリプトをGitHubやGitLabなどで管理している場合は、リポジトリに付随するContainer Registryを利用してコンテナイメージを保管してもよいでしょう。

# GitHub Package Registryの場合
docker login \
    docker.pkg.github.com -u <user> -p <password>
docker tag \
    my_sagemaker_image:latest \
    docker.pkg.github.com/<user または group>/<repository_name>/<image_name>
docker push \
    docker.pkg.github.com/<user または group>/<repository_name>/<image_name>

# GitLab Container Registryの場合
docker login \
    <GitLab CE host> -u <user> -p <password>
docker tag \
    my_sagemaker_image:latest \
    <GitLab CE host>/<user または group/subgroup>/<repository_name>/<image_name>
docker push \
    <GitLab CE host>/<user または group/subgroup>/<repository_name>/<image_name>

次に検討すべきことは?

データセットやハイパーパラメータを変えながら訓練を何度も試していくなかで、どの訓練結果が最良か?の比較を行いたくなることが想定されます。
今回紹介した範囲だと、S3 Bucketにファイルとしては結果集約されるものの、S3だけでは結果の比較管理まではサポートしていません。
必要な場合はSageMaker Experimentsなど結果トラッキング可能な仕組みと組み合わせて利用するのがよいでしょう。

オプティムではコンテナベースのAIサービス開発など様々な技術を取り扱っています。興味がある方は下記をご確認ください。

www.optim.co.jp