Nx.Serving の前処理、後処理でノード間通信のエラーを防ぐ

## はじめに

`Nx.Serving` を使うと、行列演算の一連処理を簡単に他ノードから実行できます

https://hexdocs.pm/nx/Nx.Serving.html

ただし、ノード間通信をする場合、データの前処理、後処理を適切にしないとエラーが発生するケースがあります

`Nx.Serving` には前処理、後処理の仕組みが備えられているので、それらを使ってエラーを回避します

今回も Livebook 上で実装してみます

https://livebook.dev/

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

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/nx/serving/preprocess_postprocess.livemd

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/nx/serving/preprocess_postprocess_remote.livemd

## サーバー側

新しいノートブックを開き、行列演算を実行する側のコードを実行します

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

[Nx](https://github.com/elixir-nx/nx/tree/main/nx) と [EXLA](https://github.com/elixir-nx/nx/tree/main/exla) と [Kino](https://github.com/livebook-dev/kino) をインストールします

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

```elixir
Mix.install(
  [
    {:nx, "~> 0.5"},
    {:exla, "~> 0.5"},
    {:kino, "~> 0.9"}
  ],
  config: [
    nx: [
      default_backend: EXLA.Backend
    ]
  ]
)
```

### 前処理、後処理をしない場合

今回は単に2倍する行列演算 `&Nx.multiply(&1, 2)` を提供します

また、 `Kino.start_child` で他ノードからもアクセスできるように子プロセスを開始します

```elixir
default_serving =
  fn opts -> Nx.Defn.jit(&Nx.multiply(&1, 2), opts) end
  |> Nx.Serving.new(compiler: EXLA)

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

この `Nx.Serving` に対して入力を与える場合、データを `Nx.Batch` にする必要があります

この変換はいちいち手間に感じます

```elixir
batch = Nx.Batch.stack([Nx.tensor([1])])

Nx.Serving.run(default_serving, batch)
```

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

```elixir
#Nx.Tensor<
  s64[1][1]
  EXLA.Backend<host:0, 0.511354218.2511208458.206349>
  [
    [2]
  ]
>
```

つまり、 `Nx.Batch.stack([<バッチ1>])` が `Nx.tensor([<バッチ1の演算結果>])` になっています

少し直観的ではないですね

子プロセスを呼び出す場合も同様です

```elixir
Nx.Serving.batched_run(DefaultServing, batch)
```

実行結果は同じようになります

```elixir
#Nx.Tensor<
  s64[1][1]
  EXLA.Backend<host:0, 0.511354218.2511208458.206349>
  [
    [2]
  ]
>
```

### 前処理、後処理をする場合

前処理、後処理を設定します

`IO.inspect` で「どの関数なのか」と「どのノードで実行したのか」を表示します

```elixir
pre_post_serving =
  fn opts -> Nx.Defn.jit(&Nx.multiply(&1, 2), opts) end
  |> Nx.Serving.new(compiler: EXLA)
  |> Nx.Serving.client_preprocessing(fn input ->
    # 前処理
    IO.inspect("client_preprocessing")
    IO.inspect(Node.self())
    # テンソルで与えられた入力をバッチに変換する
    {Nx.Batch.stack([input]), :client_info}
  end)
  |> Nx.Serving.client_postprocessing(fn output, _metadata, _multi? ->
    # 後処理
    IO.inspect("client_postprocessing")
    IO.inspect(Node.self())
    # 出力テンソルの次元を減らす
    Nx.squeeze(output, axes: [0])
  end)
  |> Nx.Serving.distributed_postprocessing(fn output ->
    # 分散している場合の後処理
    IO.inspect("distributed_postprocessing")
    IO.inspect(Node.self())
    # バイナリバックエンドに変換する
    Nx.backend_transfer(output, Nx.BinaryBackend)
  end)

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

それぞれ、以下の関数で処理を追加できます

- 前処理: [`Nx.Serving.client_preprocessing`](https://hexdocs.pm/nx/Nx.Serving.html#client_preprocessing/2)
- 後処理: [`Nx.Serving.client_postprocessing`](https://hexdocs.pm/nx/Nx.Serving.html#client_postprocessing/2)
- 分散している場合の後処理: [`Nx.Serving.distributed_postprocessing`](https://hexdocs.pm/nx/Nx.Serving.html#distributed_postprocessing/2)

現在のノードを確認します

```elixir
Node.self()
```

結果は以下のようになりました

```elixir
:"dog6ujac-livebook_server@eb526776ce9e"
```

### 直接実行

前処理内でバッチ化するため、テンソルを直接入力にすることができます

```elixir
Nx.Serving.run(pre_post_serving, Nx.tensor([1]))
```

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

前処理、後処理に仕込んだ `IO.inspect` が動いています

（今回はローカル実行なので、 `distributed_postprocessing` は動いていません）

また、次元が減って入力のテンソルと同じ形になっています

```elixir
"client_preprocessing"
:"dog6ujac-livebook_server@eb526776ce9e"
"client_postprocessing"
:"dog6ujac-livebook_server@eb526776ce9e"

#Nx.Tensor<
  s64[1]
  EXLA.Backend<host:0, 0.511354218.2511208458.206383>
  [2]
>
```

### ローカル子プロセス実行

ローカル（同じノード内）で子プロセスを呼び出します

```elixir
Nx.Serving.batched_run(PrePostServing, Nx.tensor([1]))
```

ローカル内にあるため、結果は直接実行の場合と同じになります

```elixir
"client_preprocessing"
:"dog6ujac-livebook_server@eb526776ce9e"
"client_postprocessing"
:"dog6ujac-livebook_server@eb526776ce9e"

#Nx.Tensor<
  s64[1]
  EXLA.Backend<host:0, 0.511354218.2511208458.206394>
  [2]
>
```

### 擬似分散実行

`Nx.Serving.batched_run` の第1引数を `{:distributed, <子プロセス名>}` のようにすると、別ノードから実行したときと同じ挙動をします（分散実行を明示しています）

```elixir
distributed_preprocessing = fn input ->
  # 分散している場合の前処理
  IO.inspect("distributed_preprocessing")
  IO.inspect(Node.self())
  # バイナリバックエンドに変換する
  Nx.backend_transfer(input, Nx.BinaryBackend)
end

Nx.Serving.batched_run({:distributed, PrePostServing}, Nx.tensor([1]), distributed_preprocessing)
```

実行すると、以下のようになります

```elixir
"distributed_preprocessing"
:"dog6ujac-livebook_server@eb526776ce9e"
"client_preprocessing"
:"dog6ujac-livebook_server@eb526776ce9e"
"client_postprocessing"
:"dog6ujac-livebook_server@eb526776ce9e"
"distributed_postprocessing"
:"dog6ujac-livebook_server@eb526776ce9e"

#Nx.Tensor<
  s64[1]
  [2]
>
```

実行される順番が以下のようになっていることが分かります

- distributed_preprocessing
- client_preprocessing
- client_postprocessing
- distributed_postprocessing

### クッキーの取得

次に別ノードから接続するため、クッキーを表示しておきます

（ノード名は先程取得済です）

```elixir
Node.get_cookie()
```

## クライアント側

別のノートブックを開き、以下のコードを実行します

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

サーバー側と同じセットアップを実行します

```elixir
Mix.install(
  [
    {:nx, "~> 0.5"},
    {:exla, "~> 0.5"},
    {:kino, "~> 0.9"}
  ],
  config: [
    nx: [
      default_backend: EXLA.Backend
    ]
  ]
)
```

### サーバーへの接続

ノード名とクッキーのテキスト入力を作ります

```elixir
node_name_input = Kino.Input.text("SERVER_NODE_NAME")
cookie_input = Kino.Input.text("SERVER_COOKIE")

[node_name_input, cookie_input]
|> Kino.Layout.grid(columns: 2)
```

表示されたテキスト入力にサーバー側のノード名、クッキーを入力します

![スクリーンショット 2023-06-13 11.24.40.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/0b09cc13-a092-1c0a-5b79-b3dbeb17f94e.png)

サーバーに接続します

```elixir
node_name =
  node_name_input
  |> Kino.Input.read()
  |> String.to_atom()

cookie =
  cookie_input
  |> Kino.Input.read()
  |> String.to_atom()

Node.set_cookie(node_name, cookie)

Node.connect(node_name)
```

`true` が返ってくれば接続できています

### 分散実行

では、実際に別ノードから実行してみましょう

まず、前処理、後処理のない子プロセスを呼んでみます

```elixir
batch = Nx.Batch.stack([Nx.tensor([1])])

Nx.Serving.batched_run(DefaultServing, batch)
```

すると、以下のようなエラーが発生します

（セル自体は正常終了しますが、エラーメッセージが表示され、演算処理は実行されません）

```elixir
#Inspect.Error<
  got RuntimeError with message:

      """
      Unable to get buffer.
      """

  while inspecting:

      %{
        __struct__: Nx.Tensor,
        data: %EXLA.Backend{
          buffer: %EXLA.DeviceBuffer{
            ref: #Reference<17018.511354218.2511208458.206405>,
            client_name: :host,
            device_id: 0,
            shape: %EXLA.Shape{
              ref: #Reference<17018.511354218.2511470593.75873>,
              dims: {1, 1},
              dtype: {:s, 64}
            }
          }
        },
        names: [nil, nil],
        shape: {1, 1},
        type: {:s, 64}
      }

  Stacktrace:

    (exla 0.5.3) lib/exla/device_buffer.ex:55: EXLA.DeviceBuffer.unwrap!/1
    (exla 0.5.3) lib/exla/backend.ex:133: EXLA.Backend.inspect/2
    (nx 0.5.3) lib/nx/tensor.ex:165: Inspect.Nx.Tensor.inspect/2
    (elixir 1.14.2) lib/inspect/algebra.ex:341: Inspect.Algebra.to_doc/2
    (elixir 1.14.2) lib/kernel.ex:2254: Kernel.inspect/2
    (kino 0.9.4) lib/kino/output.ex:468: Kino.Output.inspect/2
    lib/livebook/runtime/evaluator/default_formatter.ex:41: Livebook.Runtime.Evaluator.DefaultFormatter.to_output/1
    lib/livebook/runtime/evaluator.ex:464: Livebook.Runtime.Evaluator.continue_do_evaluate_code/5

>
```

これは、 Nx のドキュメントにも書いてあることです

https://hexdocs.pm/nx/0.5.3/Nx.Serving.html

> batched_run/3 receives an optional distributed_preprocessing callback as third argument for preprocessing the input for distributed requests. When using libraries like EXLA or Torchx, the tensor is often allocated in memory inside a third-party library so it is necessary to either transfer or copy the tensor to the binary backend before sending it to another node. This can be done by passing either Nx.backend_transfer/1 or Nx.backend_copy/1 as third argument:
>
> ```elixir
> Nx.Serving.batched_run(MyDistributedServing, input, &Nx.backend_copy/1)
> ```

和訳（DeepL）
> batched_run/3 は、分散リクエストのための入力の前処理のために、オプションの distributed_preprocessing コールバックを第 3 引数として受け取ります。 EXLA や Torchx のようなライブラリを使用する場合、テンソルはサードパーティライブラリ内のメモリに割り当てられることが多いので、他のノードに送信する前にテンソルをバイナリバックエンドに転送またはコピーすることが必要です。これは、第 3 引数に Nx.backend_transfer/1 または Nx.backend_copy/1 を渡すことで実行できます。

というわけなので、前処理、後処理付の子プロセスを呼び出します

```elixir
distributed_preprocessing = fn input ->
  # 分散している場合の前処理
  IO.inspect("distributed_preprocessing")
  IO.inspect(Node.self())
  # バイナリバックエンドに変換する
  Nx.backend_transfer(input, Nx.BinaryBackend)
end

Nx.Serving.batched_run(PrePostServing, Nx.tensor([1]), distributed_preprocessing)
```

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

```elixir
"distributed_preprocessing"
:"hscwsmdh-livebook_server@eb526776ce9e"
"client_preprocessing"
:"dog6ujac-livebook_server@eb526776ce9e"
"client_postprocessing"
:"dog6ujac-livebook_server@eb526776ce9e"
"distributed_postprocessing"
:"dog6ujac-livebook_server@eb526776ce9e"

#Nx.Tensor<
  s64[1]
  [2]
>
```

`distributed_preprocessing` だけが自分のノードで実行され、後はサーバー側で実行されました

## まとめ

EXLA などのバックエンドを使っている場合、分散処理をするには前処理、後処理が必須になります

また、前処理、後処理で入出力の形式を変更することができます

（Bumblebee は内部でこれを利用しています）