2台の端末で起動した Livebook を接続し、 Distributed な AI を体感する

## はじめに

Elixir は分散処理が得意な言語です

と、言いつつ、今までの記事では同一の端末上で起動した Livebook を別ノードとして論理的な分散状態を作っていました

それでは分散処理の本領が発揮できないため、物理的に個別の端末から Livebook を起動し、ノードを接続してみたいと思います

実装するのは Jose がライブコーディングしてくれた、 Bumblebee による分散AIです

https://news.livebook.dev/distributed2-machine-learning-notebooks-with-elixir-and-livebook---launch-week-1---day-2-1aIlaw

サーバー側のノートブックでAIモデルを読み込み、クライアント側のノートブックから呼び出します

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

- サーバー側

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/bumblebee/distributed/distributed_ai_server.livemd

- クライアント側

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/bumblebee/distributed/distributed_ai_client.livemd

## 実行環境

実行端末

- サーバー側(AIモデルが動く端末):
  - MacBook Pro 2019
  - macOS Ventura 13.4
  - Elixir 1.14.4
  - Eralng 25.2.3
  - IP: 192.168.8.157

- クライアント側(サーバー側に接続し、AIを呼ぶ端末):
  - MacBook Air 2017
  - macOS Catalina 10.15.7
  - Elixir 1.14.4
  - Eralng 25.2.3
  - IP: 192.168.8.171

共通してインストールしているもの

- [asdf](https://asdf-vm.com/)

ネットワーク

- 2台の端末は同一ローカルネットワークに接続している
- サブネット: 192.168.8.0/24

## 準備

いつもはコンテナで Livebook を起動しているのですが、コンテナだと hostname や IP アドレスなど、様々な障壁があるため、今回は Livebook をローカルで起動します

2台の端末でそれぞれ、以下のコマンドを実行します
（環境構築方法は他の方法でもいいので、参考まで）

### asdf plugin 追加

```bash
asdf plugin add elixir
asdf plugin add erlang
```

### Elixir インストール

```bash
asdf install elixir 1.14.4
asdf global elixir 1.14.4
```

### Erlang インストール

```bash
asdf install erlang 25.2.3
asdf global erlang 25.2.3
```

## Livebook 起動

Livebook ソースコードの取得

```bash
git clone https://github.com/livebook-dev/livebook.git
cd livebook
```

依存モジュールのインストール

```bash
mix deps.get --only prod
```

ここが肝です

ローカルネットワーク内からアクセスできるように、ノード名を IP アドレス付にするよう指定します

また、外部からのアクセスを受け付けるため、LIVEBOOK_IPは `0.0.0.0` を指定します

```bash
LIVEBOOK_DISTRIBUTION=name \
  LIVEBOOK_NODE="livebook-server@<各端末のIPアドレス>" \
  LIVEBOOK_IP=0.0.0.0 \
  MIX_ENV=prod mix phx.server
```

起動したそれぞれの Livebook の URL `http://<IPアドレス>:8080` をブラウザで開きます

それぞれ認証トークンを入力し、ホーム画面に遷移します

![スクリーンショット 2023-06-06 19.32.51.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/ed72b650-0098-2976-3e9c-84cb54c12869.png)

## サーバー側処理

サーバー側で新しいノートブックを開き、以下のコードを実行します

### サーバー側セットアップ

Bumblebee や EXLA をインストールします

高速化のため、 Nx のバックエンドを EXLA に指定します

```elixir
Mix.install(
  [
    {:kino_bumblebee, "~> 0.3.0"},
    {:exla, "~> 0.5.1"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)
```

### サーバーの開始

今回は Vision Transformer による画像分類を実行します

Bumblebee でモデルをロードした後、 `Kino.start_child` で子プロセスを起動します

これにより、接続した別ノードから `ViT` の名前で画像分類処理を呼び出すことができます

```elixir
{:ok, model_info} = Bumblebee.load_model({:hf, "google/vit-base-patch16-224"})
{:ok, featurizer} = Bumblebee.load_featurizer({:hf, "google/vit-base-patch16-224"})

serving =
  Bumblebee.Vision.image_classification(model_info, featurizer,
    compile: [batch_size: 1],
    defn_options: [compiler: EXLA]
  )

Kino.start_child({Nx.Serving, name: ViT, serving: serving})
```

### 接続情報の取得

外部から接続するためのノード名、クッキーを取得します

```elixir
{node(), Node.get_cookie()}
```

実行結果は以下のようになります

ノード名やクッキーはノートブックをセットアップする度に変更されます

```elixir
{:"3k4lpr4q-livebook-server@192.168.8.157",
 :"c_vDI22GjTvH5x-kPdObzLkDYgOzk2hVfQGRck2vfh-lQV1-xirDoN"}
```

## クライアント側処理

クライアント側で新しいノートブックを開き、以下のコードを実行します

### クライアント側セットアップ

サーバー側と同じです

```elixir
Mix.install(
  [
    {:kino_bumblebee, "~> 0.3.0"},
    {:exla, "~> 0.5.1"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)
```

### サーバー側との接続

サーバー側で取得したノード名、クッキーを設定し、サーバー側のノートブックに接続します

```elixir
{node, coockie} =
  {:"3k4lpr4q-livebook-server@192.168.8.157",
   :"c_vDI22GjTvH5x-kPdObzLkDYgOzk2hVfQGRck2vfh-lQV1-xirDoN"}

Node.set_cookie(node, coockie)

Node.connect(node)
```

### サーバー側 AI の呼び出し

`Nx.Serving.batched_run` にサーバー側で用意した `ViT` を指定し、画像分類処理を呼び出しています

```elixir
image_input = Kino.Input.image("Image", size: {224, 224})
form = Kino.Control.form([image: image_input], submit: "Run")
frame = Kino.Frame.new()

Kino.listen(form, fn %{data: %{image: image}} ->
  if image do
    Kino.Frame.render(frame, Kino.Text.new("Running..."))

    image = image.data |> Nx.from_binary(:u8) |> Nx.reshape({image.height, image.width, 3})

    output = Nx.Serving.batched_run(ViT, image)

    output.predictions
    |> Enum.map(&{&1.label, &1.score})
    |> Kino.Bumblebee.ScoredList.new()
    |> then(&Kino.Frame.render(frame, &1))
  end
end)

Kino.Layout.grid([form, frame], boxed: true, gap: 16)
```

セルを実行して表示される画像入力から画像を選択し、 Run ボタンをクリックすると、画像分類が実行できます

![スクリーンショット 2023-06-06 19.52.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/f7648a02-aa5e-c803-5fd3-5f45249d9c1a.png)

## まとめ

ノードを接続するだけで、簡単に別ノードの処理を呼び出すことができました

以下の記事では、サーバー側を複数ノードにして画像分類処理を分類させています

https://qiita.com/RyoWakabayashi/items/1b8379920c72895637b6