Elixir Phoenix と Bumblebee による画像分類を Amazon SageMaker 上でマイクロサービス化する

はじめに

前回、前々回で Phoenix と Bumblebee による画像分類 REST API をコンテナで実装しました

今回はこのコンテナを Amazon SageMaker にデプロイし、 AWS 上のシステムに組み込めるマイクロサービスにします

ただし、今回は CPU インスタンスで、 GPU は使用しません

Amazon SageMaker とは

SageMaker はソーセージを作る機械のように、簡単に AI システムを実装できてしまう AWS のサービスです

ノートブックインスタンスやトレーニングなど、さまざまな機能がありますが、今回はリアルタイム推論のエンドポイントを作ります

実装はこちら

コンテナ定義

SageMaker 操作

実装の流れ

• ECR にコンテナのイメージを保存する

  • ECR にリポジトリーを作成
  • Dockerfile で　 REST API として機能するコンテナを定義
  • コンテナのイメージを ECR のリポジトリーにプッシュ

• S3 にモデルファイルを保存する

  • モデルファイルを .tar.gz に圧縮する
  • S3 上にアップロードする

• IAM にコンテナを実行する権限を持ったロールを作成する

• SageMaker でモデルを作成する

  • ECR のイメージを指定
  • S3 のモデルファイルを指定
  • IAM の実行ロールを指定

• SageMaker でエンドポイント設定を作成する

  • SageMaker のモデルを指定
  • インスタンスタイプを指定

• SageMaker でエンドポイントを作成する

  • SageMaker のエンドポイント設定を指定

本来はこの流れを Terraform で定義しますが、今回はあえて Livebook で実装してみます

事前作業

AWS のアカンウトと、管理者権限を持った IAM ユーザーと、その認証情報（ACCESS_KEY_ID と SECRET_ACCESS_KEY）が必要です

セットアップ

Livebook を開き、以下のコードを実行します

Mix.install([
  {:aws, "~> 0.13"},
  {:ex_aws, "~> 2.3"},
  {:ex_aws_sts, "~> 2.3"},
  {:ex_aws_s3, "~> 2.3"},
  {:req, "~> 0.3"},
  {:nx, "~> 0.4"},
  {:stb_image, "~> 0.6"},
  {:hackney, "~> 1.18"},
  {:sweet_xml, "~> 0.7"},
  {:kino, "~> 0.8"}
])

ex_aws だけでは SageMaker 周りの実装が不十分だったため、 aws もインストールしています

認証

入力エリアを用意し、そこに IAM ユーザーの認証情報を入力します

ACCESS_KEY_ID　と SECRET_ACCESS_KEY は秘密情報なので、値が見えないように Kino.Input.password を使います

access_key_id_input = Kino.Input.password("ACCESS_KEY_ID")

secret_access_key_input = Kino.Input.password("SECRET_ACCESS_KEY")

リージョンもここで入力しておきましょう

region_input = Kino.Input.text("REGION")

各認証情報を使って、 aws-elixir で AWS にリクエストを投げるためのクライアントを生成します

client =
  AWS.Client.create(
    Kino.Input.read(access_key_id_input),
    Kino.Input.read(secret_access_key_input),
    Kino.Input.read(region_input)
  )

また、 ex_aws 用の認証情報も作っておきます

認証情報が出力されないよう、最終行は "dummy" にしています

auth_config = [
  access_key_id: Kino.Input.read(access_key_id_input),
  secret_access_key: Kino.Input.read(secret_access_key_input),
  region: Kino.Input.read(region_input)
]

"dummy"

アカウントID取得

AWS　アカウントのアカウント ID （12桁の数字）を使うため、取得しておきます

account_id =
  ExAws.STS.get_caller_identity()
  |> ExAws.request!(auth_config)
  |> then(& &1.body.account)

ECR リポジトリー作成

アカウント ID 、リージョン、イメージ名を使ってい、リポジトリーのフルネーム（外部から指定するときの名前）を設定します

region = Kino.Input.read(region_input)
image = "sagemaker-phoenix"
fullname = "#{account_id}.dkr.ecr.#{region}.amazonaws.com/#{image}:latest"

ECR 上にリポジトリーを作成します

client
|> AWS.ECR.create_repository(%{
  "repositoryName" => image
})

作成したリポジトリーを確認します

client
|> AWS.ECR.describe_repositories(%{})

リポジトリーを SageMaker から読み込めるようにポリシー設定します

client
|> AWS.ECR.set_repository_policy(%{
  "repositoryName" => image,
  "policyText" =>
    Jason.encode!(%{
      "Statement" => [
        %{
          "Sid" => "ECR202201051440",
          "Effect" => "Allow",
          "Principal" => %{
            "AWS" => ["*"]
          },
          "Action" => "ecr:*"
        }
      ]
    })
})

ECR リポジトリーへのイメージプッシュ

ECR 上のリポジトリーはプライベート（認証しないとプッシュできない場所）なので、認証情報を取得します

token =
  client
  |> AWS.ECR.get_authorization_token(%{})
  |> elem(1)
  |> then(& &1["authorizationData"])
  |> Enum.at(0)
  |> then(& &1["authorizationToken"])

このトークンは ユーザー名：パスワード のBAE64文字列になっているので、これをデコードして docker login します

[username, password] =
  token
  |> Base.decode64!()
  |> String.split(":")

Docker 操作の Elixir モジュールはどれもイマイチだったので、直接 docker コマンドを実行します

System.cmd(
  "docker",
  [
    "login",
    "--username",
    username,
    "--password",
    password,
    fullname
  ]
)

前回定義したコンテナに sagemaker-phoenix のタグを付けてビルドします

/sagemaker/serve の部分が Dockerfile の配置ディレクトリーです

System.cmd(
  "docker",
  [
    "build",
    "-t",
    image,
    "/sagemaker/serve"
  ]
)

更にこのイメージにフルネームのタグを付けます

System.cmd(
  "docker",
  [
    "tag",
    image,
    fullname
  ]
)

このイメージを ECR にプッシュします

System.cmd(
  "docker",
  [
    "push",
    fullname
  ]
)

ECR 上のイメージ一覧でプッシュできたことを確認します

client
|> AWS.ECR.list_images(%{
  "repositoryName" => image
})
|> elem(1)
|> then(& &1["imageIds"])
|> Kino.DataTable.new()

imageTag が latest のイメージが1件できています

S3 へのモデルファイルアップロード

アップロード先の指定

バケットの一覧を確認します

ExAws.S3.list_buckets()
|> ExAws.request!(auth_config)
|> then(& &1.body.buckets)
|> Kino.DataTable.new()

アップロード先のバケットを入力します

bucket_name_input = Kino.Input.text("BUCKET_ANME")

バケット名とバケット上のパスを指定します

bucket_name = Kino.Input.read(bucket_name_input)
model_prefix = "models/"

モデルファイルの圧縮

モデルファイルをダウンロードしておくため、先にコンテナで1回動かしておいてください

SageMaker で読み込む場合、モデルファイルは .tar.gz の形式で圧縮する必要があります

圧縮対象のファイル一覧を準備します

/sagemaker/serve/models が Bumblebee のキャッシュディレクトリー（モデルファイルの保存先）です

このとき、 {<相対パス>、<絶対パス>} のタプル配列にしておきます

models_path = "/sagemaker/serve/models"

filenames =
  models_path
  # 指定ディレクトリー配下のファイル一覧
  |> File.ls!()
  |> Enum.map(fn filename ->
    {
      # 相対パス
      to_charlist(filename),
      # 絶対パス
      to_charlist(Path.join(models_path, filename))
    }
  end)

:erl_tar.create で指定したファイル一覧を models.tar.gz に圧縮します

このとき、 [:compressed] を指定しないと .tar の形式になってしまいます

tar_filename = "models.tar.gz"

:erl_tar.create(tar_filename, filenames, [:compressed])

アップロードの実行

ExAws.S3.upload を使って圧縮ファイルを S3 上にアップロードします

tar_filename
|> ExAws.S3.Upload.stream_file()
|> ExAws.S3.upload(bucket_name, model_prefix <> tar_filename)
|> ExAws.request!(auth_config)

アップロードされたことを確認します

bucket_name
|> ExAws.S3.list_objects(prefix: model_prefix)
|> ExAws.request!(auth_config)
|> then(& &1.body.contents)
|> Kino.DataTable.new()

IAM ロールの定義

SageMaker 実行用ロールを定義します

client
|> AWS.IAM.create_role(%{
  "RoleName" => "sagemaker-phoenix-role",
  "AssumeRolePolicyDocument" =>
    Jason.encode!(%{
      "Statement" => [
        %{
          "Sid" => "STS202201051440",
          "Effect" => "Allow",
          "Principal" => %{
            "Service" => ["sagemaker.amazonaws.com"]
          },
          "Action" => "sts:AssumeRole"
        }
      ]
    })
})

ログ出力や ECR リポジトリー、S3バケットへのアクセス権限をポリシーとして定義します

client
|> AWS.IAM.create_policy(%{
  "PolicyName" => "sagemaker-phoenix-role-policy",
  "PolicyDocument" =>
    Jason.encode!(%{
      "Version" => "2012-10-17",
      "Statement" => [
        %{
          "Effect" => "Allow",
          "Action" => [
            "cloudwatch:PutMetricData",
            "logs:CreateLogStream",
            "logs:PutLogEvents",
            "logs:CreateLogGroup",
            "logs:DescribeLogStreams",
            "ecr:GetAuthorizationToken"
          ],
          "Resource" => ["*"]
        },
        %{
          "Effect" => "Allow",
          "Action" => [
            "s3:GetObject"
          ],
          "Resource" => [
            "arn:aws:s3:::#{bucket_name}/*"
          ]
        },
        %{
          "Effect" => "Allow",
          "Action" => [
            "ecr:BatchCheckLayerAvailability",
            "ecr:GetDownloadUrlForLayer",
            "ecr:BatchGetImage"
          ],
          "Resource" => [
            "arn:aws:ecr:::repository/#{image}"
          ]
        }
      ]
    })
})

ポリシーが正しく作成されたことを確認します

client
|> AWS.IAM.get_policy_version(%{
  "PolicyArn" => "arn:aws:iam::#{account_id}:policy/sagemaker-phoenix-role-policy",
  "VersionId" => "v1"
})
|> elem(1)
|> then(& &1["GetPolicyVersionResponse"]["GetPolicyVersionResult"]["PolicyVersion"]["Document"])
|> URI.decode()
|> Jason.decode!()

ポリシーをロールに適用します

client
|> AWS.IAM.attach_role_policy(%{
  "RoleName" => "sagemaker-phoenix-role",
  "PolicyArn" => "arn:aws:iam::#{account_id}:policy/sagemaker-phoenix-role-policy"
})

SageMaker モデル作成

SageMaker のモデルを作成します

作成した IAM ロール、 ECR リポジトリー、 S3 上のモデルファイルパスを指定します

client
|> AWS.SageMaker.create_model(%{
  "ModelName" => "sagemaker-phoenix-model",
  "ExecutionRoleArn" => "arn:aws:iam::#{account_id}:role/sagemaker-phoenix-role",
  "PrimaryContainer" => %{
    "Image" => fullname,
    "ModelDataUrl" => "s3://#{bucket_name}/#{model_prefix}#{tar_filename}"
  }
})

SageMaker エンドポイント設定作成

SageMaker のエンドポイント設定を作成します

先程作成したモデル名、インスタンスタイプなどを指定します

今回は ResNet なので小さいインスタンス（ml.t2.medium）で十分です

ある程度重いモデルを高速に動作させるためには GPU が必要になるため、 GPU 用のコンテナイメージを定義した上で ml.g4dn.xlarge などのインスタンスタイプを指定します

client
|> AWS.SageMaker.create_endpoint_config(%{
  "EndpointConfigName" => "sagemaker-phoenix-endpoint-config",
  "ProductionVariants" => [
    %{
      "VariantName" => "variant-1",
      "ModelName" => "sagemaker-phoenix-model",
      "InstanceType" => "ml.t2.medium",
      "InitialInstanceCount" => 1,
      "InitialVariantWeight" => 1
    }
  ]
})

SageMaker エンドポイント作成

エンドポイント設定を指定してエンドポイントを作成します

エンドポイント作成は非同期で行われるため、:ok が返ってきてもしばらくは作成中です

内部的には ECR からのコンテナイメージプル、コンテナの起動、 S3 からのモデルファイルダウンロード、展開などが実行されています

client
|> AWS.SageMaker.create_endpoint(%{
  "EndpointName" => "sagemaker-phoenix-endpoint",
  "EndpointConfigName" => "sagemaker-phoenix-endpoint-config"
})

エンドポイントの状況を確認します

この結果が Creating のうちはまだ作成中です InService になれば正常に起動できています

client
|> AWS.SageMaker.describe_endpoint(%{
  "EndpointName" => "sagemaker-phoenix-endpoint"
})
|> elem(1)
|> then(& &1["EndpointStatus"])

SageMaker エンドポイント呼出

SageMaker のエンドポイントは AWS のサービスのように、マイクロサービスとして呼び出すことができます

テスト画像を使ってエンドポイントを呼び出してみます

binary =
  "https://raw.githubusercontent.com/pjreddie/darknet/master/data/dog.jpg"
  |> Req.get!()
  |> then(& &1.body)

確認のため、バイナリーを画像として表示します

binary
|> StbImage.read_binary!()
|> StbImage.to_nx()
|> Kino.Image.new()

バイナリーを SageMaker のエンドポイントに渡して推論を実行します

client
|> AWS.SageMakerRuntime.invoke_endpoint(
  "sagemaker-phoenix-endpoint",
  %{
    "Body" => binary,
    "ContentType" => "image/jpeg"
  }
)
|> elem(1)
|> then(& &1["Body"])
|> Jason.decode!()
|> then(& &1["predictions"])
|> Kino.DataTable.new()

アラスカン・マラミュートやハスキー、チベタン・マスティフは犬種です

自転車も識別結果に出てきていますね

別の画像でも試してみましょう

image_input = Kino.Input.image("INPUT_IMAGE", format: :jpeg)

binary =
  image_input
  |> Kino.Input.read()
  |> then(& &1.data)

client
|> AWS.SageMakerRuntime.invoke_endpoint(
  "sagemaker-phoenix-endpoint",
  %{
    "Body" => binary,
    "ContentType" => "image/jpeg"
  }
)
|> elem(1)
|> then(& &1["Body"])
|> Jason.decode!()
|> then(& &1["predictions"])
|> Kino.DataTable.new()

ちゃんと分類できているようです

モデルを変えたりすれば、 Bumblebee で提供されている機能は全てマイクロサービス化できるわけです

この SageMaker エンドポイントを Lambda などから呼び出せばシステムに組み込めます

後片付け

SageMaker エンドポイントや ECR リポジトリー、 S3 上のファイルは存在している間課金されるので、使い終わったら削除しておきます

SageMaker エンドポイント削除

client
|> AWS.SageMaker.delete_endpoint(%{
  "EndpointName" => "sagemaker-phoenix-endpoint"
})

client
|> AWS.SageMaker.delete_endpoint_config(%{
  "EndpointConfigName" => "sagemaker-phoenix-endpoint-config"
})

client
|> AWS.SageMaker.delete_model(%{
  "ModelName" => "sagemaker-phoenix-model"
})

IAM ロール削除

client
|> AWS.IAM.detach_role_policy(%{
  "RoleName" => "sagemaker-phoenix-role",
  "PolicyArn" => "arn:aws:iam::#{account_id}:policy/sagemaker-phoenix-role-policy"
})

client
|> AWS.IAM.delete_policy(%{
  "PolicyArn" => "arn:aws:iam::#{account_id}:policy/sagemaker-phoenix-role-policy"
})

client
|> AWS.IAM.delete_role(%{
  "RoleName" => "sagemaker-phoenix-role"
})

ECR リポジトリー削除

client
|> AWS.ECR.batch_delete_image(%{
  "repositoryName" => image,
  "imageIds" => [
    %{"imageTag" => "latest"}
  ]
})

リポジトリーが空になったことの確認

client
|> AWS.ECR.list_images(%{
  "repositoryName" => image
})
|> elem(1)
|> then(& &1["imageIds"])
|> Kino.DataTable.new()

client
|> AWS.ECR.delete_repository(%{
  "repositoryName" => image
})

モデルファイル削除

ExAws.S3.delete_object(bucket_name, model_prefix <> tar_filename)
|> ExAws.request!(auth_config)

まとめ

Elixir で SageMaker というのは構想していたことではありますが、まさかこんなに早く簡単に実装できてしまうとは

Bumblebee の登場は想像を上回る出来事でした