Elixir Bumblebee を AWS Lambda で動かす（ローカル実行編）

はじめに

前回の続きで、今回はまずローカル（Dockerコンテナ上）実行します

前回を読んでいない方はまずこちらをどうぞ

実装の全文はこちら

参考にさせていただいた Elixir on Lambda の先駆者様はこちら

先駆者様が実装してくださったモジュール FaasBase はこちら

実行環境

• macOS 13.1
• Rancher Desktop 1.7.0

プロジェクトの実装

まずプロジェクトを作成します

mix new resnet

mix.exs

生成された mix.exs に追記します

...
  def application do
    [
+     mod: {ResnetApplication, []},
      extra_applications: [:logger]
    ]
  end
...

...
  defp deps do
-   []
+   [
+     {:faas_base, "~> 1.1.0"},
+     {:bumblebee, "~> 0.1"},
+     {:stb_image, "~> 0.6"},
+     {:exla, "~> 0.4"}
+   ]
  end
end

resnet.ex

lib/resnet.ex を以下の内容で上書きします

defmodule Resnet do
  use FaasBase, service: :aws
  alias FaasBase.Logger
  alias FaasBase.Aws.Request
  alias FaasBase.Aws.Response

  @impl FaasBase
  def init(context) do
    {:ok, context}
  end

  @impl FaasBase
  def handle(_, %{"Payload" => base64}, _) do
    predictions =
      base64
      |> Base.decode64!()
      |> predict()

    {:ok, Response.to_response(%{"predictions" => predictions}, %{}, 200)}
  end

  defp predict(binary) do
    tensor =
      binary
      |> StbImage.read_binary!()
      |> StbImage.to_nx()

    resnet = Serving.get(:resnet)

    resnet
    |> Nx.Serving.run(tensor)
    |> then(& &1.predictions)
  end
end

resnet.ex が処理の本体です

init と handle は FaasBase で定義されたビヘイビアを実装しています

handle　が Lambda に来たリクエストを処理することになります

handle の第2引数に Lambda へのリクエストボディーが入ってきます

今回は BASE64 文字列で来るようにして、デコードしてから推論しています

serving.ex

続いて lib/serving.ex を追加します

Bumblebee で Nx.Serving による画像識別サービスを生成し、保持します

defmodule Serving do
  use Agent

  @resnet_id "microsoft/resnet-50"
  @cache_dir "/opt/ml/model"

  def start_link(_opts) do
    {model, featurizer} = load_model()

    resnet =
      Bumblebee.Vision.image_classification(
        model,
        featurizer,
        defn_options: [compiler: EXLA]
      )

    Agent.start_link(fn ->
      %{resnet: resnet}
    end, name: __MODULE__)
  end

  def load_model() do
    {:ok, model} =
      Bumblebee.load_model({:hf, @resnet_id, cache_dir: @cache_dir})

    {:ok, featurizer} =
      Bumblebee.load_featurizer({:hf, @resnet_id, cache_dir: @cache_dir})

    {model, featurizer}
  end

  # 使用時に Agent から取り出す
  def get(key) do
    Agent.get(__MODULE__, &Map.get(&1, key))
  end
end

SageMaker のときの使い回しですが、モデルのロードを load_model 関数として分離しました

また、 cache_dir を指定して、モデルをダウンロード済のファイルから読み込みます

resnet_application.ex

lib/resnet_application.ex に FaasBase.Aws.Application を改造したものを実装します

defmodule ResnetApplication do
  use Application

  alias FaasBase.Aws.Logger
  alias FaasBase.Aws.BaseTask

  @doc """
  Start Application.
  """
  def start(_type, _args) do
    context = System.get_env()

    children = [
      Serving,
      {Logger, context |> Map.get("LOG_LEVEL", "INFO") |> String.downcase() |> String.to_atom()},
      {FaasBase.Logger, Logger},
      {BaseTask, context}
    ]

    Supervisor.start_link(children, strategy: :one_for_all)
  end
end

元の FaasBase.Aws.Application はこちら

Serving を children に加えただけです

これによって serving.ex から起動時に Nx.Serving が生成され、 reset.ex から呼び出せるようになります

download_models.ex

このまま何も対処しないでこの関数を Lambda で動かすと、関数呼び出し時にモデルファイルのダウンロードが始まり、かなり時間がかかってしまいます

これでは無駄に遅くなる上にタイムアウトになる危険性も高いです

まして Lambda はインスタンスが生成されるたびにストレージがリセットされるので（コンテナなので当然ですが）、これは回避しなければ行けません

というわけで、先に一回 load_model 関数を動かして、 cache_dir の中身をコンテナ定義に含んでしまいましょう

load_model 関数を呼ぶ mix コマンド download_models を実装します

lib/mix/tasks/download_models.ex

defmodule Mix.Tasks.DownloadModels do
  @moduledoc "Download ResNet models"
  use Mix.Task

  @shortdoc "Download ResNet models."
  def run(_) do
    Serving.load_model()
  end
end

こちらの記事を参考にしました

config

環境毎の設定ファイルを作ります

これは Nx バックエンドを制御するためです

• 本番環境: EXLA.Backend で動かす
• 開発環境: Nx.BinaryBackend で動かす

本番環境は当然 EXLA.Backend で動かさないといけません

そうでないと ResNet の推論はいつまで経っても終わりません

しかし、開発環境（コンテナビルド時）で EXLA.Backend を使うと、 mix download_models 実行時に以下のようなエラーが発生します

no process: the process is not alive or there’s no process currently associated with the given name, possibly because its application isn’t started

mix download_models はモデルのダウンロードだけが目的なので Nx.BinaryBackend を指定すれば問題ありません

• config/config.exs

import Config

import_config "#{config_env()}.exs"

• config/dev.exs

import Config

config :nx, default_backend: Nx.BinaryBackend

• config/prod.exs

import Config

config :nx, default_backend: EXLA.Backend

• config/runtime.exs

import Config

config :nx, default_backend: EXLA.Backend

コンテナの実装

Dockerfile

以下の内容で Dockerfile を作成します

FROM amazonlinux:2022.0.20221207.4 as build

RUN set -e \
    && yum -y update \
    && yum -y groupinstall "Development Tools" \
    && yum -y install \
        ncurses-devel \
        openssl \
        openssl-devel \
        gcc-c++ \
    && yum clean all

ARG OTP_VERSION="25.2"

RUN set -e \
    && OTP_DOWNLOAD_URL="https://github.com/erlang/otp/archive/OTP-${OTP_VERSION}.tar.gz" \
    && curl -fSL -o otp-src.tar.gz "$OTP_DOWNLOAD_URL" \
    && export ERL_TOP="/usr/src/otp_src_${OTP_VERSION%%@*}" \
    && mkdir -vp $ERL_TOP \
    && tar -xzf otp-src.tar.gz -C $ERL_TOP --strip-components=1 \
    && rm otp-src.tar.gz \
    && ( cd $ERL_TOP \
        && ./otp_build autoconf \
        && ./configure \
        && make -j$(nproc) \
        && make install ) \
    && find /usr/local -name examples | xargs rm -rf

ARG REBAR3_VERSION="3.20.0"

RUN set -xe \
    && REBAR3_DOWNLOAD_URL="https://github.com/erlang/rebar3/archive/${REBAR3_VERSION}.tar.gz" \
    && mkdir -p /usr/src/rebar3-src \
    && curl -fSL -o rebar3-src.tar.gz "$REBAR3_DOWNLOAD_URL" \
    && tar -xzf rebar3-src.tar.gz -C /usr/src/rebar3-src --strip-components=1 \
    && rm rebar3-src.tar.gz \
    && cd /usr/src/rebar3-src \
    && HOME=$PWD ./bootstrap \
    && install -v ./rebar3 /usr/local/bin/ \
    && rm -rf /usr/src/rebar3-src

ARG ELIXIR_VERSION=1.14.2

ENV LANG=en_US.UTF-8

ENV MIX_REBAR3=/usr/local/bin/rebar3

RUN set -xe \
    && ELIXIR_DOWNLOAD_URL="https://github.com/elixir-lang/elixir/archive/v${ELIXIR_VERSION}.tar.gz" \
    && curl -fsSL $ELIXIR_DOWNLOAD_URL -o elixir-src.tar.gz \
    && mkdir -p /usr/src/elixir-src \
    && tar -xzf elixir-src.tar.gz -C /usr/src/elixir-src --strip-components=1 \
    && rm elixir-src.tar.gz \
    && cd /usr/src/elixir-src \
    && make -j$(nproc) \
    && make install \
    && rm -rf /usr/src/elixir-src \
    && cd $HOME \
    && mix local.hex --force \
    && mix hex.info

WORKDIR /tmp

COPY mix.exs /tmp/
COPY lib /tmp/lib
COPY config /tmp/config

RUN mix deps.get \
    && mix compile

RUN mix download_models

RUN MIX_ENV=prod mix aws.release Resnet

FROM amazonlinux:2022.0.20221207.4

ADD https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/latest/download/aws-lambda-rie /usr/local/bin/aws-lambda-rie

COPY entry.sh "/entry.sh"

RUN chmod +x /usr/local/bin/aws-lambda-rie
RUN chmod +x /entry.sh

ENV LANG=en_US.UTF-8
ENV TZ=:/etc/localtime
ENV PATH=/var/lang/bin:/usr/local/bin:/usr/bin/:/bin:/opt/bin
ENV LD_LIBRARY_PATH=/var/lang/lib:/lib64:/usr/lib64:/var/runtime:/var/runtime/lib:/var/task:/var/task/lib:/opt/lib
ENV LAMBDA_TASK_ROOT=/var/task
ENV LAMBDA_RUNTIME_DIR=/var/runtime

WORKDIR /var/task

COPY --from=build /tmp/_aws/docker/bootstrap /var/runtime/
COPY --from=build /tmp/_aws/docker/ /var/task/
COPY --from=build /opt/ml/model /opt/ml/model

ENTRYPOINT [ "/bin/bash", "/entry.sh" ]

CMD [ "Resnet" ]

ここは先駆者様の実装方法から変更しています

まず、マルチステージビルドを使うことで、 Dockerfile 一つでコンテナを定義しています

マルチステージビルドとは、　FROM を複数回使うコンテナイメージの実装方法です

今回はビルド用ステージで mix release を実行し、生成された実行可能ファイルやモデルファイルを実行用ステージにコピーしています

• ビルド用ステージでやっていること

  • erlang や Elixir など、ビルドに必要なものをインストールする
  • ローカルで実装した Elixir プロジェクトのコード（mix.exs や lib 、 config ディレクトリー）をコンテナ内にコピーする
  • 依存モジュールを取得し、コンパイルする
  • mix download_models でモデルファイルをダウンロードする
  • MIX_ENV=prod mix aws.release Resnet で実行可能ファイルを生成する

• 実行用ステージでやっていること

  • ローカルで Lambda の動作をエミュレートするための AWS Lambda Runtime Interface Emulator を取得する
  • コンテナ起動時に実行する entry.sh (後述)をコンテナ内にコピーする
  • 以下の aws-lambda-base の Dockerfile を参考に、必要な環境変数を設定する
  • ビルド用ステージからビルド結果、モデルファイルをコピーする
  • コンテナ起動時に entry.sh を実行する
  • entry.sh の引数に Resnet を渡す

entry.sh

以下の内容で entry.sh を追加します

内容はカスタムランタイムで実行されるものと同じです

#!/bin/bash

if [ $# -ne 1 ]; then
  echo "entrypoint requires the handler name to be the first argument" 1>&2
  exit 142
fi

export _HANDLER="$1"

RUNTIME_ENTRYPOINT=/var/runtime/bootstrap
if [ -z "${AWS_LAMBDA_RUNTIME_API}" ]; then
    exec /usr/local/bin/aws-lambda-rie $RUNTIME_ENTRYPOINT
else
    exec $RUNTIME_ENTRYPOINT
fi

/var/runtime/bootstrap は FaasBase が生成してくれるファイルで、内容は以下のようになっています

#!/bin/sh
set -euo pipefail
export HOME=/
$(bin/resnet start)

bin/resnet は ビルド結果のファイルで、 resnet start によって起動しています

ローカル実行の場合は /usr/local/bin/aws-lambda-rie を使って Lambda の動作をエミュレートします

docker-compose.yml

ローカル実行を簡単にするために docker-compose.yml を作ります

---

version: '3.2'
services:
  faas:
    container_name: faas_resnet
    build: .
    command: Resnet
    tty: true
    ports:
      - '9000:8080'
    environment:
      - LOG_LEVEL=debug

コンテナ内のポート番号 8080 をホストのポート番号 9000 に割り当てます

環境変数 LOG_LEVEL を debug に設定して、デバッグログを見えるようにしています

コンテナのビルド、実行

以下のコマンドでコンテナをビルド、実行します

$ docker-compose up --build
...
[+] Running 1/1
 ⠿ Container faas_resnet  Recreated                                                                              0.1s
Attaching to faas_resnet
faas_resnet  | 22 Dec 2022 05:44:36,694 [INFO] (rapid) exec '/var/runtime/bootstrap' (cwd=/var/task, handler=)
faas_resnet  | 22 Dec 2022 05:46:17,751 [INFO] (rapid) extensionsDisabledByLayer(/opt/disable-extensions-jwigqn8j) -> stat /opt/disable-extensions-jwigqn8j: no such file or directory
faas_resnet  | 22 Dec 2022 05:46:17,752 [WARNING] (rapid) Cannot list external agents error=open /opt/extensions: no such file or directory

しばらくすると bootstrap が実行され、 WARNING などが表示されます（特に問題ありません）

この状態で別のターミナルからリクエストを送信してみます

jq がない場合は入れた方が結果が見やすいのでインストールしてください

$ echo {\"Payload\":\"$(base64 -i sample.jpg)\"}  | \                            
    curl -X POST -H "Content-Type: application/json" --data @- \
    "http://localhost:9000/2015-03-31/functions/function/invocations" \
    | jq
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100  216k  100   388  100  215k    528   293k --:--:-- --:--:-- --:--:--  295k
{
  "body": {
    "predictions": [
      {
        "label": "notebook, notebook computer",
        "score": 0.6154259443283081
      },
      {
        "label": "laptop, laptop computer",
        "score": 0.1152949333190918
      },
      {
        "label": "bow tie, bow-tie, bowtie",
        "score": 0.039064694195985794
      },
      {
        "label": "projector",
        "score": 0.026744747534394264
      },
      {
        "label": "dining table, board",
        "score": 0.024911964312195778
      }
    ]
  },
  "headers": {},
  "isBase64Encoded": false,
  "statusCode": 200
}

Resnet.handle の第2引数で待ち受けている通り、 {"Payload": "<BASE64エンコードした画像ファイル>"} をデータとして、 http://localhost:9000/2015-03-31/functions/function/invocations に POST しています

推論結果が predictions として返ってきます

コンテナを実行した方のターミナルを見ると、ログが出力されています

faas_resnet  | START RequestId: f1c22be6-30aa-4438-967e-430d71d57c8d Version: $LATEST
faas_resnet  | [DEBUG] %{"Payload" => "/9j/4AAQSkZJRgABAQAASABIAAD/4QgkRXhpZgAA..." <> ...}
faas_resnet  | [DEBUG] %{"RELEASE_BOOT_SCRIPT" => "start", ... "_HANDLER" => "Resnet"}
faas_resnet  | [DEBUG] f1c22be6-30aa-4438-967e-430d71d57c8d
faas_resnet  | [DEBUG] %FaasBase.Aws.Response{body: %{"predictions" => [%{label: "notebook, notebook computer", score: 0.6154259443283081}, ...]}, headers: %{}, status_code: 200, is_base64_encoded: false}
faas_resnet  | END RequestId: f1c22be6-30aa-4438-967e-430d71d57c8d
faas_resnet  | REPORT RequestId: f1c22be6-30aa-4438-967e-430d71d57c8d   Duration: 672.73 ms     Billed Duration: 673 ms Memory Size: 3008 MB     Max Memory Used: 3008 MB

Lambda 実行時と同じように、 処理の開始と終了、実行にかかった時間、使われたメモリ量などが分かります

まとめ

これで Lambda にデプロイする準備ができました

あとはこれを SAM や Terraform などに組み込めば OK です

が、次回の記事ではあえて Livebook からデプロイし、関数を呼び出してみます