Powered by AppSignal & Oban Pro
Would you like to see your link here? Contact us

Nx.Serving Preprocess Postprocess remote

preprocess_postprocess_remote.livemd

Nx.Serving Preprocess Postprocess remote

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

Connect to server

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

Call serving

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

Nx.Serving.batched_run(DefaultServing, batch)
Node.self()
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)