EKSユーザーもSageMakerを使うと楽に機械学習に取り組める話

Amazon EKS #2 Advent Calendar 2019の24日目に出す予定だったの記事です。

はじめに

re:Invent 2019が終わってはや3週間。

EKSのFargate対応やFargate Spotの登場など華々しいアップデートの影に隠れて、
Amazon SageMaker Operators for Kubernetesというアップデートがあったことは皆さんご存知でしょうか？

普段はSageMaker[1]島に生息している身ですが、この度はEKS Advent Calenderの1枠を頂いて、まだまだ記事も少ない[2]こちらのアップデートについて少し深堀りしてみることにします。

Amazon SageMaker Operators for Kubernetesとは

ドキュメントを引用すると、

Amazon SageMaker Operators for Kubernetes make it easier for developers and data scientists using Kubernetes to train, tune, and deploy machine learning (ML) models in Amazon SageMaker. You can install these SageMaker Operators on your Kubernetes cluster in Amazon Elastic Kubernetes Service (EKS) to create SageMaker jobs natively using the Kubernetes API and command-line Kubernetes tools such as ‘kubectl’.

つまりは、

• kubectlを利用してEKS上からSageMakerを操作することができるようになった
• これにより、EKS上で構築されたシステムの一部としてMLワークフローを構築・管理することができ、且つ裏側ではSageMakerを利用することでその高い可用性・セキュリティ・コスト効率などの恩恵を得ることができる

というアップデートです。

Kubeflowと何が違うの？

k8sでMLワークフローを構築する試みとしてはKubeflowなどが挙げられますが、Kubeflowの様に全てをk8s上で再構築するのではなく、AWS上で最適化されたML基盤ととして既に実績のあるAmazon SageMakerを最大限活用していこうという思想の機能なのだと自分は理解しています。ですので、機械学習ジョブや推論APIのコンテナはEKS上には立ち上がらず、SageMaker基盤上に立ち上がります。

どう使うの？

このSageMaker OperatorをEKSクラスタにインストールすることで、例えば学習ジョブの実行やハイパーパラメータ最適化ジョブの実行、推論APIの起動などがkubectlのインターフェースを利用して実行できるようになります。

$ kubectl apply -f traing.yaml
$ kubectl apply -f hpo-training.yaml
$ kubectl apply -f hosting.yaml

例えばtraining.yamlには利用するコンテナイメージや、インスタンスタイプ、学習データのS3パス、ハイパーパラメータなどを記述します。

traing.yaml

apiVersion: sagemaker.aws.amazon.com/v1
kind: TrainingJob
metadata:
  name: traing-sample
spec:
    hyperParameters:
        - name: epoch
          value: "10"
    algorithmSpecification:
        trainingImage: 123456789012.dkr.ecr.us-west-2.amazonaws.com/sample:1
        trainingInputMode: File
    roleArn: arn:aws:iam::123456789012:role/service-role/AmazonSageMaker-ExecutionRole
    region: ap-northeast-1
    outputDataConfig:
        s3OutputPath: s3://s3-bucket/output
    resourceConfig:
        instanceCount: 1
        instanceType: ml.m4.xlarge
        volumeSizeInGB: 5
    stoppingCondition:
        maxRuntimeInSeconds: 86400
    inputDataConfig:
        - channelName: train
          dataSource:
            s3DataSource:
                s3DataType: S3Prefix
                s3Uri: https://s3-bucket/train
                s3DataDistributionType: FullyReplicated
          contentType: text/csv

機械学習プロジェクトのフェーズごとのツールの使い分け

個人的には下記のような使い分けがハマりそうなイメージを持っています。

実験や検証のフェーズ： MLエンジニアは、SageMaker Studioと呼ばれる機械学習用の統合開発環境(JupyterLabの拡張版のようなもの)上で特徴量・アルゴリズムの開発、学習、推論APIの動作確認などを行う。実際の学習ジョブや推論APIの起動にはSageMaker Python SDKを利用。

EKS主体のProduction環境に組み込むフェーズ： インフラエンジニアは学習ジョブの再実行や推論APIの作成・更新などにSageMaker Operatorを利用。

詳細な機能紹介はこちらをご参照下さい。

• AWS Blog: Amazon SageMaker Operators for Kubernetes のご紹介
• GitHub
• ドキュメント[3]

利用方法

利用方法は上記に挙げたAWS Blogやドキュメントにも記載があるのですが、権限周りで手間取ったので、AWS Blogに加筆修正する形で以下に記載します。

大まかに以下の流れで進めていきます。

• 前準備
• EKS Cluster作成
• OIDCプロバイダーの設定とOperator用IAM Roleの作成
• Kubernetes クラスターへの Operator のセットアップ
• 学習ジョブの準備
• 学習ジョブの実行
• 推論APIの作成

前準備

kubectl、AWS CLI及びeksctlを最新のものにアップデート。自分の環境ではkubectlのバージョンが1.16.2、AWS CLIが1.16.298、eksctlが0.11.1でした。

EKS Cluster作成

検証用EKSクラスタをeksctlで作成します。

$ eksctl create cluster -f eks-sagemaker.yaml

eks-sagemaker.yaml

apiVersion: eksctl.io/v1alpha5
kind: ClusterConfig

metadata:
  name: eks-sagemaker
  region: ap-northeast-1

managedNodeGroups:
  - name: managed-ng
    instanceType: t2.large
    desiredCapacity: 1

本ブログでは基本チュートリアルの流れに寄せて書いていますが、代わりに以下のeks-sagemaker.yamlを利用してクラスタを立ち上げると次の「OIDCプロバイダーの設定とOperator用IAM Roleの作成」の部分をスキップすることもできます。Amazon EKS Advent Calendar 2019の @taishinさんの記事はこちらの方法で環境構築されていました。

eks-sagemaker.yaml

apiVersion: eksctl.io/v1alpha5
kind: ClusterConfig

metadata:
  name: eks-sagemaker
  region: ap-northeast-1

iam:
  withOIDC: true
  serviceAccounts:
  - metadata:
      name: sagemaker-k8s-operator-default
      namespace: sagemaker-k8s-operator-system
    attachPolicyARNs:
    - "arn:aws:iam::aws:policy/AmazonSageMakerFullAccess"

managedNodeGroups:
  - name: managed-ng
    instanceType: t2.large
    desiredCapacity: 1

OIDCプロバイダーの設定とOperator用IAM Roleの作成

IAMロール周りのセットアップを行う前に、Amazon SageMaker Operators for Kubernetesを利用する場合、どの層でどのように認証認可が行われているのかを整理してみます。ここを意識せずにチュートリアルを進めて自分は沼にハマりました。

大きく分けると以下の3つの認証認可が関わってきます。

1. EKSクラスタの認証認可。手元のPCからkubectlを叩く場合、IAMを利用して認証し、k8sのRBACを利用し認可します。
2. EKS上のSageMaker Operator -> SageMakerの認証認可。 SageMaker OperatorがSageMakerを操作(学習ジョブ実行、推論API作成など)するための権限を与える必要があります。ここで利用するIAMロールをOperator用IAM Roleと呼ぶことにします。
3. SageMaker -> 他のAWSのサービスの認証認可。SageMakerの各種リソース(学習ジョブや推論APIなど)が他のAWSサービスにアクセスするための権限を与える必要があります。例えば、SageMakerの各学習ジョブがS3やCloudWatchにアクセスする際の権限などです。ここで利用するIAMロールをSageMaker用IAM Roleと呼ぶことにします。

このステップで中心に説明するのは2:「SageMaker Operator -> SageMakerの認証認可」の部分になります。

1:「EKSクラスタへの認証認可」の部分の説明に関してはドキュメントやトラブシュートのページを参考にして下さい。デフォルトではクラスタ作成で利用したIAMユーザー・ロールにsystem:masters権限が付与されるので、今回は特に設定せずとも動作しました。

3:「SageMaker -> 他のAWSのサービスの認証認可」の部分の「SageMaker用IAM Role」についてはこちらのドキュメントを参考に作成しておいて下さい。SageMakerへの信頼ポリシーを設定し、AmazonSageMakerFullAccessポリシーをアタッチすれば良いです。自分は以前SageMakerを利用した時に作成されたAmazonSageMaker-ExecutionRoleを再利用して検証しました。

それでは、2:「SageMaker Operator -> SageMakerの認証認可」の話に入りますが、ここでは「IAM Roles for Service Accounts (IRSA)」という仕組みが使われます。この仕組みを用いることで、k8sのService AccountとIAM Roleを関連付け、そのService Accountを利用するPodに対してIAM Roleに紐づく一時的な認証情報をセキュアに渡すことができます。

IAM Roles for Service Accountsを利用するために以下の手順で設定を行います。

• EKSクラスターへOIDC プロバイダーを関連付ける
• 信頼ポリシーにOIDCプロバイダーやService Accountを指定してIAMロールを作成
• Service AccountのアノテーションにIAM RoleのARNを指定(次のステップで実施)

まず、以下のコマンドでEKSクラスターに OIDC プロバイダーを関連付けます。

$ eksctl utils associate-iam-oidc-provider --cluster ${CLUSTER_NAME} \
    --region ${AWS_REGION} --approve

[ℹ]  eksctl version 0.11.1
[ℹ]  using region ap-northeast-1
[ℹ]  will create IAM Open ID Connect provider for cluster "eks-sagemaker" in "ap-northeast-1"
[✔]  created IAM Open ID Connect provider for cluster "eks-sagemaker" in "ap-northeast-1"

次に、IAM Role作成時に必要なOIDC IDを取得します。

$ aws eks describe-cluster --name ${CLUSTER_NAME} --region ${AWS_REGION} \
    --query cluster.identity.oidc.issuer --output text

https://oidc.eks.${AWS_REGION}.amazonaws.com/id/{Your OIDC ID}

次に、以下のような"trust.json"ファイルを新規作成します。プレースホルダーは、ご自身の OIDC ID、AWS アカウント番号、EKS クラスターのリージョンにそれぞれ置き換えてください。

trust.json

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Federated": "arn:aws:iam::<AWS account number>:oidc-provider/oidc.eks.<EKS Cluster region>.amazonaws.com/id/<OIDC ID>"
            },
            "Action": "sts:AssumeRoleWithWebIdentity",
            "Condition": {
                "StringEquals": {
                    "oidc.eks.<EKS Cluster region>.amazonaws.com/id/<OIDC ID>:aud": "sts.amazonaws.com",
                    "oidc.eks.<EKS Cluster region>.amazonaws.com/id/<OIDC ID>:sub": "system:serviceaccount:sagemaker-k8s-operator-system:sagemaker-k8s-operator-default"
                }
            }
        }
    ]
}

それではOperator用のIAM Roleを作成します。次のステップで利用するので、作成した出力されたIAM Roleの情報の中からIAM RoleのARN(arn:aws:iam::000000000000:role/SageMakerOperatorRole)をメモしておいて下さい。

$ aws iam create-role --role-name SageMakerOperatorRole --assume-role-policy-document file://trust.json --output=text

ROLE    arn:aws:iam::000000000000:role/SageMakerOperatorRole   2019-12-22T15:44:02Z    /       AAAAAAAAAAAAAAAA   SageMakerOperatorRole
ASSUMEROLEPOLICYDOCUMENT        2012-10-17
STATEMENT       sts:AssumeRoleWithWebIdentity   Allow
STRINGEQUALS    sts.amazonaws.com       system:serviceaccount:sagemaker-k8s-operator-system:sagemaker-k8s-operator-default
PRINCIPAL       arn:aws:iam::000000000000:oidc-provider/oidc.eks.ap-northeast-1.amazonaws.com/id/AAAAAAAAAAAAAAAAAAAAAAAAAA

最後に、作成したIAM Roleに Amazon SageMaker へのアクセスするためのAmazonSageMakerFullAccess ポリシーをアタッチします。

$ aws iam attach-role-policy --role-name SageMakerOperatorRole  --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerFullAccess

Kubernetes クラスターへの Operator のセットアップ

ここからはEKSクラスタにOperatorをインストールしていきます。

まず、Operator をインストールするための YAMLファイルをダウンロードします。

$ wget https://raw.githubusercontent.com/aws/amazon-sagemaker-operator-for-k8s/master/release/rolebased/installer.yaml

installer.yaml ファイルの eks.amazonaws.com/role-arn を、前のステップで作成したOperator用IAM Roleの ARN で更新します。 以下に該当箇所を一部抜粋して示しました。

installer.yaml

---
apiVersion: v1
kind: ServiceAccount
metadata:
  annotations:
    eks.amazonaws.com/role-arn: arn:aws:iam::000000000000:role/SageMakerOperatorRole
  name: sagemaker-k8s-operator-default
  namespace: sagemaker-k8s-operator-system
---

Kubernetes クラスターに Amazon SageMaker Custom Resource Definition (CRD) をインストールし、使用する Operator をセットアップします。

$ kubectl apply -f installer.yaml

以下のコマンドでCRDが正常にインストールされているかを確認します。

$ kubectl get crd | grep sagemaker

batchtransformjobs.sagemaker.aws.amazon.com         2019-11-20T17:12:34Z
endpointconfigs.sagemaker.aws.amazon.com            2019-11-20T17:12:34Z
hostingdeployments.sagemaker.aws.amazon.com         2019-11-20T17:12:34Z
hyperparametertuningjobs.sagemaker.aws.amazon.com   2019-11-20T17:12:34Z
models.sagemaker.aws.amazon.com                     2019-11-20T17:12:34Z
trainingjobs.sagemaker.aws.amazon.com               2019-11-20T17:12:34Z

ここまでで、SageMaker Operatorが正常にインストールされました🎉

必須ではありませんが、Amazon SageMaker からCloudWatch Logsに出力されたログを kubetl を使用してコマンドラインで確認するためのクライアント(smlogs)をインストールするために以下のコマンドを実行します。Macの場合はexport os="darwin"に変更して実行して下さい。

export os="linux"

wget https://amazon-sagemaker-operator-for-k8s-us-east-1.s3.amazonaws.com/kubectl-smlogs-plugin/latest/${os}.amd64.tar.gz
tar xvzf ${os}.amd64.tar.gz

# Move binaries to a directory in your homedir.
mkdir ~/sagemaker-k8s-bin
cp ./kubectl-smlogs.${os}.amd64/kubectl-smlogs ~/sagemaker-k8s-bin/.

# This line will add the binaries to your PATH in your .bashrc.

echo 'export PATH=$PATH:~/sagemaker-k8s-bin' >> ~/.bashrc

# Source your .bashrc to update environment variables:
source ~/.bashrc

学習ジョブの実行

さてここからは、機械学習タスクの1例としてPytorchを利用したCIFAR-10の画像分類を行います。

「実験・検証のフェーズはSageMaker Studioなどを利用して一通り済ませており、必要なデータセット、Pytorchのスクリプトが用意できたので、ここからはインフラエンジニアが学習ジョブを定期的に実行し、推論APIを起動するぞ」というシナリオを想定しています。

SageMakerの学習ジョブを利用する場合は、その都度、学習・推論用のコンテナイメージをECRに登録し、学習データをS3に置いておく必要があります。ここではpytorchの学習・推論用スクリプトを含んだコンテナイメージを作成し、ECRに登録します。次にCIFAR-10のデータセットをS3バケットにアップロードします。
ここまでの手順はこちらのレポジトリをクローンして、setup.shを実行することで実現できます。

次に、先程のレポジトリ中のtraining.yamlを編集します。

• trainingImage
  • {account_id}.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-cifar10:latest
• roleArn
  • SageMaker用IAM RoleのARN
• s3OutputPath
  • 任意のモデルを出力する先のS3パス
  • 今回は s3://sagemaker-ap-northeast-1-{account_id}/
• s3Uri
  • 学習データがアップロードされているS3パス
  • 今回は s3://sagemaker-ap-northeast-1-{account_id}/pytorch-cifar10

各パラメータの詳細はこちらに記載されています。

最後に以下のコマンドを実行することで、Operatorを介して学習ジョブがSageMakerの環境上で実行されます。SageMakerはジョブの実行時のみリソースを用意するため、インスタンスの費用はジョブの実行時間に対してのみ発生します。

$ kubectl apply -f training.yaml

全学習ジョブの状況の確認には下記のコマンドを利用します。学習ジョブが正常終了すると、STATUSがCompletedに変化します。この時点で、s3OutputPathに学習済みのモデルが格納されています。

$ kubectl get trainingjobs
NAME                STATUS      SECONDARY-STATUS   CREATION-TIME          SAGEMAKER-JOB-NAME
pytorch-cifar10-training   Completed   Completed          2019-12-27T06:20:18Z   pytorch-cifar10-training-aaaaaaaaaaa

$ kubectl describe trainingjobs

smlogsを利用すると、CloudWatch Logsに格納された学習ジョブコンテナの標準出力を確認できます。

$ kubectl smlogs trainingjob pytorch-cifar10-training

トラブルシュート時は、SageMaker OperatorのControllerのログを確認するのも有用です。

$ kubectl get pod --namespace sagemaker-k8s-operator-system
NAME                                                         READY   STATUS    RESTARTS   AGE
sagemaker-k8s-operator-controller-manager-5cd7df4d74-2xtss   2/2     Running   0          12d

$ kubectl logs sagemaker-k8s-operator-controller-manager-5cd7df4d74-2xtss  -c manager --namespace sagemaker-k8s-operator-system

推論APIの作成

ここでは学習済みモデルをホスティングしたAPIサーバーを作成します。

まずクローンしたレポジトリ下のhosting.yamlを編集します。

• executionRoleArn
  • SageMaker用IAM RoleのARN
• modelDataUrl
  • 以下のコマンドで取得したURL
  • kubectl get trainingjob pytorch-cifar10-training --output=json | jq '.status.modelPath'
• image
  • {account_id}.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-cifar10:latest

以下のコマンドを実行することで、Operatorを介して推論APIがSageMakerの環境上で作成されます。こちらは明示的に削除するまで料金が発生する点に注意して下さい。

$ kubectl apply -f hosting.yaml

起動状況の確認には以下のコマンドを利用します。STATUSがCreatingからInServiceに変化するのに数分から10分程掛かりました。

$ kubectl get hostingdeployments
NAME                 STATUS      SAGEMAKER-ENDPOINT-NAME
hosting-deployment   InService   hosting-deployment-b5353d03287311eab98106bce8b274a8

$ kubectl describe hostingdeployments

STATUSについての詳細はこちらに記載されています。

最後に、シンプルですがAPIが正常に動作しているか確認しましょう。
適当なテスト画像を用意し、SageMakerのInvokeEndpointを実行する権限のある環境で、inference-sample.pyを実行すると推論結果が表示されます。

$ python inference-sample.py 
dog

無事推論結果が取得できました🎉

推論APIの設定変更は、yamlファイルを編集し、再度kubectl applyを実行することで実現できます。
https://sagemaker.readthedocs.io/en/stable/amazon_sagemaker_operators_for_kubernetes_jobs.html#update-hostingdeployment

また、今回は検証できませんでしたが、hosting用のCRDに複数のProductionVariantを設定することで、複数のモデルに対してある割合でトラフィックを流すなど、A/Bテストに利用することもできそうですね。

おわりに

この記事では、Amazon SageMaker Operators for Kubernetesについて、その概要や利用法を紹介しました。

一方で、k8sのリソースとして管理する必要がない場合、手動でSageMakerを操作している例や、AWS Step FunctionsやAirflowとSageMakerを組み合わせてMLワークフローを構築している例[4]もよく見かけます。細かいところですとKubeflow PipelineとSageMakerの統合[5]も徐々に進んでいますし、ありきたりな話ですが、チームのスキルセットや実現したいことに合わせて適切な方法を選ぶのが良さそうです💪

--
[1] Amazon SageMakerとは機械学習エンジニア向けのサービスで、開発環境構築や学習ジョブ実行、推論API作成周りのインフラの手間をざっと肩代わりしてくれるというものです。詳細はこちら
[2] 2019/12/27現在では日本語の記事は @taishin さんの記事 だけしか見つかりませんでしたね。
[3] 余談ですがAmazon SageMaker Operators for KubernetesのドキュメントはSageMakerドキュメントではなくSageMaker Python SDKのドキュメントに入っているという不思議。ただこのSageMaker Python SDKのドキュメントはSageMakerユーザーであれば一度目を通すことをオススメします。とにかく情報量が多いので。
[4] https://www.infoq.com/jp/news/2019/01/amazon-step-function-integration/
[5] https://github.com/kubeflow/pipelines/tree/master/components/aws/sagemaker