[WASI-NN] Support RPC mode

RPC mode allows using another Wasi-NN instance that is running on a remote
WasmEdge instance, via `ssh -R remote.sock:local.sock`.

An example usecase is to allow Linux VM (e.g., Lima) guest to use the host GPU.

The gRPC proto can be repurposed for non-WASM applications as well.

- - -

Build
=====

Set `WASMEDGE_BUILD_WASI_NN_RPC` to `ON`.
Enabled by default when gRPC (libgrpc++-dev) is installed.

Usage
=====

Host 1 (rpc server / ssh client, e.g., Lima host with physical GPU)
-----

```
wasi_nn_rpcserver \
  --nn-rpc-uri unix:///$HOME/nn_server.sock \
  --nn-preload default:GGML:AUTO:llama-2-7b-chat.Q5_K_M.gguf

ssh \
  -R /tmp/nn_client.sock:$HOME/nn_server.sock \
  host2
```

Host 2 (rpc client / ssh server, e.g., Lima guest)
-----

```
wasmedge \
  --nn-rpc-uri unix:///tmp/nn_client.sock \
  wasmedge-ggml-llama-interactive.wasm \
  default "1 + 1 = ?"
```

See <https://github.com/second-state/WasmEdge-WASINN-examples/tree/master/wasmedge-ggml-llama-interactive>
for how to obtain `llama-2-7b-chat.Q5_K_M.gguf` and `wasmedge-ggml-llama-interactive.wasm`.

Signed-off-by: Akihiro Suda <[email protected]>

.github/workflows/build-extensions.yml

@@ -100,7 +100,7 @@ jobs:
         shell: bash
         run: |
           apt update
-          apt install -y unzip libopenblas-dev pkg-config
+          apt install -y unzip libopenblas-dev pkg-config protobuf-compiler-grpc libgrpc-dev libgrpc++-dev
           bash utils/wasi-nn/install-pytorch.sh
       - name: Build and test WASI-NN using ${{ matrix.compiler }} with ${{ matrix.build_type }} mode
         shell: bash
@@ -125,6 +125,21 @@ jobs:
             cp -f ${output_dir}/${output_bin} ${output_bin}
             tar -zcvf plugin_${plugin_array[$i]}.tar.gz ${output_bin}
           done
+      - name: Test WASI-NN RPC mode with GGML
+        shell: bash
+        run: |
+          set -eux
+          # wasi_nn_rpcserver is built in a clean "build2" dir
+          cmake -Bbuild2 -GNinja
+          cmake --build build2
+          export WASI_NN_RPC_TEST_URI=unix:///tmp/wasi_nn_rpc.sock
+          WASMEDGE_PLUGIN_PATH=build/plugins/wasi_nn build2/tools/wasmedge/wasi_nn_rpcserver \
+            --nn-rpc-uri=$WASI_NN_RPC_TEST_URI \
+            --nn-preload=default:GGML:AUTO:build/test/plugins/wasi_nn/wasinn_ggml_fixtures/orca_mini.gguf &
+          sleep 3
+          cd ${test_dir}
+          # The test binary consumes $WASI_NN_RPC_TEST_URI
+          ./${test_bin} --gtest_filter=WasiNNTest.GGMLBackendWithRPC
       - name: Upload artifact - wasi_nn-pytorch
         uses: actions/upload-artifact@v3
         with:

.github/workflows/release.yml

@@ -145,7 +145,7 @@ jobs:
         shell: bash
         run: |
           apt update
-          apt install -y unzip libopenblas-dev pkg-config
+          apt install -y unzip libopenblas-dev pkg-config protobuf-compiler-grpc libgrpc-dev libgrpc++-dev
           bash utils/wasi-nn/install-pytorch.sh
       - name: Build WASI-NN plugin
         shell: bash

.github/workflows/reusable-build-on-fedora.yml

@@ -21,7 +21,8 @@ jobs:
       - name: Install dependency
         run: |
           dnf update -y
-          dnf install -y cmake ninja-build llvm llvm-devel lld-devel clang git file rpm-build dpkg-dev spdlog-devel
+          dnf install -y cmake ninja-build llvm llvm-devel lld-devel clang git file rpm-build dpkg-dev spdlog-devel \
+            pkgconf-pkg-config protobuf-c-compiler grpc-cpp grpc-plugins grpc-devel
       - uses: actions/checkout@v3
         with:
           fetch-depth: 0

.github/workflows/reusable-build-on-macos.yml

@@ -33,7 +33,7 @@ jobs:
       - name: Setup build environment
         run: |
           eval $(/opt/homebrew/bin/brew shellenv)
-          brew install llvm@16 ninja cmake wabt
+          brew install llvm@16 ninja cmake wabt grpc
       - name: Set environment variables for release
         if: ${{ inputs.release }}
         run: |

.github/workflows/reusable-build-on-ubuntu.yml

@@ -24,6 +24,10 @@ jobs:
       BUILD_TESTS: OFF
     container: wasmedge/wasmedge:${{ matrix.docker_tag }}
     steps:
+      - name: Install dependencies
+        run: |
+          apt-get update
+          apt-get install -y pkg-config protobuf-compiler-grpc libgrpc-dev libgrpc++-dev
       - uses: actions/checkout@v3
         with:
           fetch-depth: 0

CMakeLists.txt

@@ -13,6 +13,8 @@ set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreadedDLL")
 
 set(WASMEDGE_CAPI_VERSION "0.1.0" CACHE STRING "WasmEdge C API library version")
 set(WASMEDGE_CAPI_SOVERSION "0" CACHE STRING "WasmEdge C API library soversion")
+set(WASMEDGE_WASI_NN_VERSION "0.1.0" CACHE STRING "WasmEdge WASI-NN library version")
+set(WASMEDGE_WASI_NN_SOVERSION "0" CACHE STRING "WasmEdge WASI-NN library soversion")
 
 find_program(GIT_CMD git)
 execute_process(COMMAND
@@ -56,6 +58,7 @@ option(WASMEDGE_BUILD_TOOLS "Generate wasmedge and wasmedgec tools. Depend on an
 option(WASMEDGE_BUILD_FUZZING "Generate fuzzing test tools. Couldn't build with wasmedge tools and unit tests." OFF)
 option(WASMEDGE_BUILD_PLUGINS "Generate plugins." ON)
 option(WASMEDGE_BUILD_EXAMPLE "Generate examples." OFF)
+option(WASMEDGE_BUILD_WASI_NN_RPC "Generate Wasi-NN RPC." OFF)
 option(WASMEDGE_FORCE_DISABLE_LTO "Forcefully disable link time optimization when linking even in Release/RelWithDeb build." OFF)
 option(WASMEDGE_LINK_LLVM_STATIC "Statically link the LLVM library into the WasmEdge tools and libraries." OFF)
 option(WASMEDGE_LINK_TOOLS_STATIC "Statically link the wasmedge and wasmedgec tools. Will forcefully link the LLVM library statically." OFF)
@@ -108,6 +111,24 @@ if(WASMEDGE_BUILD_PLUGINS)
   endif()
 endif()
 
+# * Homebrew: grpc
+# * Debian, Ubuntu: libgrpc-dev, libgrpc++-dev
+find_package(PkgConfig)
+if(PkgConfig_FOUND)
+  pkg_check_modules(gRPCPP grpc++)
+endif()
+# Do not check find_package(gRPC), because libgrpc-dev for Ubuntu 22.04 does not contain cmake files.
+# https://packages.ubuntu.com/search?keywords=libgrpc-dev
+# Do not check find_package(protobuf), because libprotobuf-dev for Ubuntu does not contain cmake files.
+# https://bugs.debian.org/cgi-bin/bugreport.cgi?bug=1027876
+if(gRPCPP_FOUND AND WASMEDGE_BUILD_SHARED_LIB)
+  message(STATUS "Setting WASMEDGE_BUILD_WASI_NN_RPC to ON. If you see an error related to gRPC or protobuf, try setting this to OFF.")
+  set(WASMEDGE_BUILD_WASI_NN_RPC ON)
+endif()
+if(WASMEDGE_BUILD_WASI_NN_RPC AND NOT WASMEDGE_BUILD_SHARED_LIB)
+  message(FATAL_ERROR "WASMEDGE_BUILD_WASI_NN_RPC depends on WASMEDGE_BUILD_SHARED_LIB")
+endif()
+
 set(WASMEDGE_BUILD_PACKAGE "DEB;RPM" CACHE STRING "Package generate types")
 set(CPACK_PROJECT_CONFIG_FILE ${CMAKE_CURRENT_SOURCE_DIR}/cmake/cpack_config.cmake)
 

cmake/Helper.cmake

@@ -164,6 +164,13 @@ endfunction()
 function(wasmedge_add_library target)
   add_library(${target} ${ARGN})
   wasmedge_setup_target(${target})
+  # Linux needs an explicit INSTALL_RPATH to allow libwasmedge.so to find libwasiNNRPC.so
+  # in the same directory
+  if(CMAKE_SYSTEM_NAME MATCHES "Linux")
+    set_target_properties(${target} PROPERTIES
+      INSTALL_RPATH "$ORIGIN"
+    )
+  endif()
 endfunction()
 
 function(wasmedge_add_executable target)

include/api/wasmedge/wasmedge.h

@@ -3873,6 +3873,20 @@ WASMEDGE_CAPI_EXPORT extern int WasmEdge_Driver_Compiler(int Argc,
 WASMEDGE_CAPI_EXPORT extern int WasmEdge_Driver_Tool(int Argc,
                                                      const char *Argv[]);
 
+#ifdef WASMEDGE_BUILD_WASI_NN_RPC
+/// Entrypoint for the Wasi-NN RPC server tool.
+///
+/// This function provides an entrypoint to the WasmEdge Wasi-NN RPC server tool
+/// with the command line arguments.
+///
+/// \param Argc the argument count.
+/// \param Argv the argument vector.
+///
+/// \returns the execution status.
+WASMEDGE_CAPI_EXPORT extern int
+WasmEdge_Driver_WasiNNRPCServer(int Argc, const char *Argv[]);
+#endif
+
 /// Entrypoint for the unified tool.
 ///
 /// This function provides an entrypoint to the WasmEdge unified tool with the

include/driver/wasiNNRPCServerTool.h

@@ -0,0 +1,9 @@
+#pragma once
+
+namespace WasmEdge {
+namespace Driver {
+
+int WasiNNRPCServer(int Argc, const char *Argv[]) noexcept;
+
+} // namespace Driver
+} // namespace WasmEdge