Bumblebee による画像分類を Amazon SageMaker Serverless でサーバレスサービス化する

## はじめに

以前、 AWS の Amazon SageMaker 上に Elixir で画像分類サービスを構築しました

https://qiita.com/RyoWakabayashi/items/6cad17cd6047cc7113c0

SageMaker 上に構築する場合、通常の Web サーバーと同じように Phoenix フレームワークで構築可能です

また、画像分類を AWS Lambda でサーバレス化する記事も書いています

https://qiita.com/RyoWakabayashi/items/53d831e4e6ff3ad056c8

Lambda の場合、サーバレスの特殊な環境で実行するため、 FaasBase というモジュールを利用しました

https://github.com/imahiro-t/faas_base

サーバレス化することで、使用していない間のコスト抑制、アクセスピーク時の可用性向上に効果が期待できます

本記事では、基本的に通常の SageMaker と同じ実装で Phoenix を使いつつサーバレス化可能な SageMaker Serverless によるサーバレスサービスを Livebook から構築します

実装したノートブックはこちら

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/aws/sagemaker_serverless.livemd

## 事前準備

AWS のアカウント、 IAM ユーザーの認証情報を用意してください

Docker のビルドができるよう、 Docker Desktop や Rancher Desktop 等をインストールしてください

また、 Livebook をインストールし、起動してください

https://livebook.dev/

## コンテナのローカル起動

画像分類コンテナを AWS 上で動かす前にローカルでビルドしておきます

### コンテナ定義

全てここに記載するのは大変なので、以下の GitHub リポジトリーを参照してください

https://github.com/RyoWakabayashi/elixir-learning/tree/main/sagemaker/serve

基本的に[以前の記事](https://qiita.com/RyoWakabayashi/items/6cad17cd6047cc7113c0)と同じく、 Bumblebee で画像分類を実行する Phoenix サーバーになっています
（※以前の記事からモジュール等を更新しています）

以前の記事から変わったところを2点特筆しておきます

### Apple シリコン搭載 Mac でのビルド

以前の記事は古い Intel Mac でビルドしていましたが、今回は M2 Mac でビルドしました

Lambda の実行環境は Intel CPU なので、ビルド時にいくつか指定が必要です

```diff_docker
...
WORKDIR /app

+ ENV ERL_FLAGS="+JPperf true"
ENV HOME="/app"
ENV MIX_ENV="dev"

RUN mix deps.get
...
```

参考:

https://elixirforum.com/t/cant-compile-elixir-phoenix-api-in-docker-build-step/56863/7

Docker コマンドでビルドする場合は `--platform linux/amd64` の指定が必要です

また、 docker-compose でビルドする場合は `platform: linux/amd64` の指定が必要です

### 権限設定

詳細には分かっていませんが、通常の SageMaker だと問題ないのに、 SageMaker Serverless の場合のみ、実行時に権限不足のエラーが発生します
（おそらく、実行ユーザーが違うため）

そのため、ビルドディレクトリー配下の権限を追加する必要があります

```diff_docker
...
RUN mix compile.phoenix

RUN chmod +x /app/serve
+ RUN chmod -R 777 /app/_build
...
```

## Livebook からの AWS 環境構築

こちらも基本的に[以前の記事](https://qiita.com/RyoWakabayashi/items/6cad17cd6047cc7113c0)と同じです

コードは記載しますが、前回と同じ解説は割愛します

### セットアップ

ノートブックのセットアップセルで以下のコードを実行します

```elixir
Mix.install([
  {:aws, "~> 1.0"},
  {:ex_aws, "~> 2.5"},
  {:ex_aws_sts, "~> 2.3"},
  {:ex_aws_s3, "~> 2.5"},
  {:req, "~> 0.5"},
  {:nx, "~> 0.8"},
  {:stb_image, "~> 0.6"},
  {:hackney, "~> 1.20"},
  {:sweet_xml, "~> 0.7"},
  {:kino, "~> 0.14"},
  {:flow, "~> 1.2"}
])
```

今回はサーバーレス化されて自動スケーリングする状態を確認するため、並列処理用の Flow をインストールしています

https://github.com/dashbitco/flow

### 認証

テキスト入力を作り、 AWS の認証情報を入力します

```elixir
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")

[
  access_key_id_input,
  secret_access_key_input,
  region_input
]
|> Kino.Layout.grid(columns: 3)
```

![スクリーンショット 2024-09-22 17.51.42.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/0e4d4444-235d-4a57-8ec1-d37a43a2d1c4.png)

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

```elixir
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)
]

Kino.nothing()
```

### アカウントID取得

入力した AWS 認証情報のアカウントIDを取得しておきます

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

### ECRリポジトリー作成

リージョンを指定します

```elixir
region = Kino.Input.read(region_input)
```

任意の名前で ECR リポジトリーを用意します

```elixir
image = "sagemaker-phoenix"
```

```elixir
fullname = "#{account_id}.dkr.ecr.#{region}.amazonaws.com/#{image}:latest"
```

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

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

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

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

コンテナをビルドし、 ECR リポジトリーにプッシュします

```elixir
token =
  client
  |> AWS.ECR.get_authorization_token(%{})
  |> elem(1)
  |> Map.get("authorizationData")
  |> Enum.at(0)
  |> Map.get("authorizationToken")
```

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

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

```elixir
{result, 0} =
  System.cmd(
    "docker",
    [
      "build",
      "--platform",
      "linux/amd64",
      "-t",
      image,
      "/sagemaker/serve"
    ]
  )

Kino.Markdown.new(result)
```

:::note info
ここで `--platform linux/amd64` を指定することで、 Apple シリコンの Mac でも有効なビルドができます
:::

```elixir
{result, 0} =
  System.cmd(
    "docker",
    [
      "tag",
      image,
      fullname
    ]
  )
```

```elixir
{result, 0} =
  System.cmd(
    "docker",
    [
      "push",
      fullname
    ]
  )

Kino.Markdown.new(result)
```

```elixir
client
|> AWS.ECR.list_images(%{
  "repositoryName" => image
})
|> elem(1)
```

![スクリーンショット 2022-12-18 23.50.44.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/b4ce4b9c-adda-7d82-0fa2-4bc06a1ca675.png)

imageTag が latest のイメージが1件できていれば OK です

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

画像分類用のモデルファイルを圧縮して S3 にアップロードします

#### アップロード先の指定

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

```elixir
bucket_name_input = Kino.Input.text("BUCKET_NAME")
```

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

#### モデルファイルの圧縮

```elixir
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)
```

```elixir
tar_filename = "models.tar.gz"

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

#### アップロードの実行

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

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

### IAM ロールの定義

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

Serverless でも権限を変更する必要はありません

```elixir
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"
        }
      ]
    })
})
```

```elixir
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}"
          ]
        }
      ]
    })
})
```

```elixir
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!()
```

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

### SageMaker モデル作成

モデル作成は通常時と変わりません

```elixir
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 Serverless で変更するポイントです

```elixir
client
|> AWS.SageMaker.create_endpoint_config(%{
  "EndpointConfigName" => "sagemaker-phoenix-endpoint-config",
  "ProductionVariants" => [
    %{
      "VariantName" => "variant-1",
      "ModelName" => "sagemaker-phoenix-model",
      "ServerlessConfig" => %{ 
        "MaxConcurrency" => 4,
        "MemorySizeInMB" => 6144,
        "ProvisionedConcurrency" => 1
      }
    }
  ]
})
```

以下の `ServerlessConfig` を追加しするだけで SageMaker サーバレスに変更されます

`InstanceType` の指定は不要になります

```elixir
...
      "ServerlessConfig" => %{ 
        "MaxConcurrency" => 4,
        "MemorySizeInMB" => 6144,
        "ProvisionedConcurrency" => 1
      }
...
```

:::note warn
`InstanceType` で `ml.inf1` などのインスタンスタイプを指定していると、サーバレスに対応していない旨が表示されます

Inferentia や GPU インスタンスには対応していないため、その点は注意が必要です
:::

### SageMaker エンドポイント作成

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

以下のコードの結果が `InService` になるまで待ちます

```elixir
client
|> AWS.SageMaker.describe_endpoint(%{
  "EndpointName" => "sagemaker-phoenix-endpoint"
})
|> elem(1)
|> Map.get("EndpointStatus")
```

### SageMaker エンドポイント呼出

#### 単独画像

実際に画像を SageMaker に投げて画像分類を実行します

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

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

投げるのは以下の画像です

![スクリーンショット 2022-12-19 0.25.00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/be7fbddc-b0c4-f38c-c39a-8787ec06572e.png)

```elixir
client
|> AWS.SageMakerRuntime.invoke_endpoint(
  "sagemaker-phoenix-endpoint",
  %{
    "Body" => binary,
    "ContentType" => "image/jpeg",
    "Accept" => "application/json"
  },
  recv_timeout: 60_000
)
|> elem(1)
|> Map.get("Body")
|> Jason.decode!()
|> Map.get("predictions")
|> Kino.DataTable.new()
```

以下のように推論結果が返ってきます

![スクリーンショット 2022-12-19 0.26.04.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/2f7e8cb6-5236-56ad-5706-9fcd809eb222.png)

#### 並列推論

Flow で並列実行してみましょう

```elixir
0..10
|> Flow.from_enumerable(max_demand: 1)
|> Flow.map(fn _index ->
  client
  |> AWS.SageMakerRuntime.invoke_endpoint(
    "sagemaker-phoenix-endpoint",
    %{
      "Body" => binary,
      "ContentType" => "image/jpeg",
      "Accept" => "application/json"
    },
    recv_timeout: 60_000
  )
end)
|> Enum.to_list()
```

[awslogs](https://github.com/jorgebastida/awslogs) などでログ出力を確認すると、複数インスタンスが起動して並列分散処理されていることが分かります

`variant-1/fac77384aafc5731b17cb634d80f4f3e-c29d16194858454caa496a95f294a13e` などの部分が実際に動作しているインスタンスの ID です

複数 ID が出力されています

```
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/fac77384aafc5731b17cb634d80f4f3e-c29d16194858454caa496a95f294a13e [info] POST /invocations
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/fac77384aafc5731b17cb634d80f4f3e-c29d16194858454caa496a95f294a13e [debug] Processing with ApiWeb.PredictionController.index/2
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/fac77384aafc5731b17cb634d80f4f3e-c29d16194858454caa496a95f294a13e   Parameters: %{}
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/fac77384aafc5731b17cb634d80f4f3e-c29d16194858454caa496a95f294a13e   Pipelines: [:api]
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/94ed1e91bb2a18d06f327889ddb597b3-4c85239a484f4b008e52d3254d8c4d59 [info] Predictions: [%{label: "malamute, malemute, Alaskan malamute", score: 0.9850469827651978}, %{label: "bicycle-built-for-two, tandem bicycle, tandem", score: 0.0035503478720784187}, %{label: "mountain bike, all-terrain bike, off-roader", score: 0.0011567040346562862}, %{label: "Eskimo dog, husky", score: 7.718618726357818e-4}, %{label: "Tibetan mastiff", score: 5.428484291769564e-4}]
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/94ed1e91bb2a18d06f327889ddb597b3-4c85239a484f4b008e52d3254d8c4d59 [info] Sent 200 in 3574ms
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/8635afe3515d0a8fbfaa93dc1bc7358b-264ef0238c7c4fedb750356eb0281700 [info] Predictions: [%{label: "malamute, malemute, Alaskan malamute", score: 0.9850469827651978}, %{label: "bicycle-built-for-two, tandem bicycle, tandem", score: 0.0035503478720784187}, %{label: "mountain bike, all-terrain bike, off-roader", score: 0.0011567040346562862}, %{label: "Eskimo dog, husky", score: 7.718618726357818e-4}, %{label: "Tibetan mastiff", score: 5.428484291769564e-4}]
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/8635afe3515d0a8fbfaa93dc1bc7358b-264ef0238c7c4fedb750356eb0281700 [info] Sent 200 in 5674ms
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/fac77384aafc5731b17cb634d80f4f3e-c29d16194858454caa496a95f294a13e [info] Predictions: [%{label: "malamute, malemute, Alaskan malamute", score: 0.9850469827651978}, %{label: "bicycle-built-for-two, tandem bicycle, tandem", score: 0.0035503478720784187}, %{label: "mountain bike, all-terrain bike, off-roader", score: 0.0011567040346562862}, %{label: "Eskimo dog, husky", score: 7.718618726357818e-4}, %{label: "Tibetan mastiff", score: 5.428484291769564e-4}]
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/fac77384aafc5731b17cb634d80f4f3e-c29d16194858454caa496a95f294a13e [info] Sent 200 in 4758ms
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/d0a505828f0da4b94bc2f61a92e1e4ee-a851e8dbea4046488bf665a6965c7a13 [info] POST /invocations
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/d0a505828f0da4b94bc2f61a92e1e4ee-a851e8dbea4046488bf665a6965c7a13 [debug] Processing with ApiWeb.PredictionController.index/2
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/d0a505828f0da4b94bc2f61a92e1e4ee-a851e8dbea4046488bf665a6965c7a13   Parameters: %{}
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/d0a505828f0da4b94bc2f61a92e1e4ee-a851e8dbea4046488bf665a6965c7a13   Pipelines: [:api]
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/d0a505828f0da4b94bc2f61a92e1e4ee-a851e8dbea4046488bf665a6965c7a13 [info] Predictions: [%{label: "malamute, malemute, Alaskan malamute", score: 0.9850469827651978}, %{label: "bicycle-built-for-two, tandem bicycle, tandem", score: 0.0035503478720784187}, %{label: "mountain bike, all-terrain bike, off-roader", score: 0.0011567040346562862}, %{label: "Eskimo dog, husky", score: 7.718618726357818e-4}, %{label: "Tibetan mastiff", score: 5.428484291769564e-4}]
/aws/sagemaker/Endpoints/sagemaker-phoenix-endpoint variant-1/d0a505828f0da4b94bc2f61a92e1e4ee-a851e8dbea4046488bf665a6965c7a13 [info] Sent 200 in 371ms

```

### 後片付け

作成したリソースを削除します
（余計なコストを発生させないため）

#### SageMaker エンドポイント削除

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

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

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

#### IAM ロール削除

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

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

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

#### ECR リポジトリー削除

全てのイメージを削除してからリポジトリーを削除します

```elixir
image_digest_list =
  client
  |> AWS.ECR.list_images(%{
    "repositoryName" => image
  })
  |> elem(1)
  |> Map.get("imageIds")
  |> Enum.map(&Map.take(&1,["imageDigest"]))
```

```elixir
client
|> AWS.ECR.batch_delete_image(%{
  "repositoryName" => image,
  "imageIds" => image_digest_list
})
```

```elixir
client
|> AWS.ECR.list_images(%{
  "repositoryName" => image
})
|> elem(1)
|> Map.get("imageIds")
|> Kino.DataTable.new()
```

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

#### モデルファイル削除

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

## まとめ

通常の SageMaker とほぼ同じ容量で SageMaker Serverless を使用することができました

特に呼出頻度が低い場合や、極端にアクセス集中する時間がある場合などに有効です

通常 SageMaker で動かしていたコンテナをほぼそのままサーバレス化できるのは素晴らしいですね