feat: add WASI-NN bindings to Wasm Workers Server

Cargo.toml

@@ -13,7 +13,7 @@ repository = { workspace = true }
 [workspace.package]
 version = "1.4.0"
 edition = "2021"
-authors = [ "Wasm Labs <https://wasmlabs.dev>" ]
+authors = ["Wasm Labs <https://wasmlabs.dev>"]
 license = "Apache-2.0"
 repository = "https://github.com/vmware-labs/wasm-workers-server/"
 
@@ -60,15 +60,16 @@ members = [
   "crates/worker",
   "kits/rust",
   "kits/rust/worker",
-  "kits/javascript"
+  "kits/javascript",
 ]
 # Exclude examples
 exclude = [
   "examples/pdf-create",
   "examples/rust-basic",
   "examples/rust-fetch",
   "examples/rust-kv",
-  "examples/rust-params"
+  "examples/rust-params",
+  "examples/rust-wasi-nn",
 ]
 
 [workspace.dependencies]
@@ -93,6 +94,7 @@ wws-api-manage = { path = "./crates/api-manage" }
 wws-api-manage-openapi = { path = "./crates/api-manage-openapi" }
 wasmtime = "10.0.1"
 wasmtime-wasi = "10.0.1"
+wasmtime-wasi-nn = "10.0.1"
 wasi-common = "10.0.1"
 path-slash = "0.2.1"
 openssl = { version = "=0.10.55" }

crates/worker/Cargo.toml

@@ -19,6 +19,7 @@ tokio = { workspace = true }
 toml = { workspace = true }
 wasmtime = { workspace = true }
 wasmtime-wasi = { workspace = true }
+wasmtime-wasi-nn = { workspace = true }
 wasi-common = { workspace = true }
 wws-config = { workspace = true }
 wws-data-kv = { workspace = true }

crates/worker/src/config.rs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::features::http_requests::HttpRequestsConfig;
+use crate::features::wasi_nn::WasiNnConfig;
 use crate::features::{data::ConfigData, folders::Folder};
 use anyhow::{anyhow, Result};
 use serde::{Deserialize, Deserializer};
@@ -13,9 +14,12 @@ use wws_data_kv::KVConfigData;
 
 /// List all available features for a worker
 #[derive(Deserialize, Clone, Default)]
+#[serde(default)]
 pub struct Features {
     /// Allow to perform http requests from a worker
     pub http_requests: HttpRequestsConfig,
+    /// Enables WASI-NN bindings for Machine Learning inference
+    pub wasi_nn: WasiNnConfig,
 }
 
 /// Workers configuration. These files are optional when no configuration change is required.

crates/worker/src/features/mod.rs

@@ -4,3 +4,4 @@
 pub mod data;
 pub mod folders;
 pub mod http_requests;
+pub mod wasi_nn;

crates/worker/src/features/wasi_nn.rs

@@ -0,0 +1,13 @@
+// Copyright 2023 VMware, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use serde::Deserialize;
+
+pub const WASI_NN_BACKEND_OPENVINO: &str = "openvino";
+
+#[derive(Deserialize, Clone, Default)]
+#[serde(default)]
+pub struct WasiNnConfig {
+    /// List of Machine Learning backends. For now, only "openvino" option is supported
+    pub allowed_backends: Vec<String>,
+}

crates/worker/src/lib.rs

@@ -11,15 +11,18 @@ use actix_web::HttpRequest;
 use anyhow::{anyhow, Result};
 use bindings::http::{add_to_linker as http_add_to_linker, HttpBindings};
 use config::Config;
+use features::wasi_nn::WASI_NN_BACKEND_OPENVINO;
 use io::{WasmInput, WasmOutput};
 use sha256::digest as sha256_digest;
 use std::fs::{self, File};
 use std::path::PathBuf;
+use std::sync::Arc;
 use std::{collections::HashMap, path::Path};
 use stdio::Stdio;
 use wasi_common::WasiCtx;
 use wasmtime::{Engine, Linker, Module, Store};
 use wasmtime_wasi::{ambient_authority, Dir, WasiCtxBuilder};
+use wasmtime_wasi_nn::WasiNnCtx;
 use wws_config::Config as ProjectConfig;
 use wws_runtimes::{init_runtime, Runtime};
 
@@ -43,6 +46,7 @@ pub struct Worker {
 
 struct WorkerState {
     pub wasi: WasiCtx,
+    pub wasi_nn: Option<Arc<WasiNnCtx>>,
     pub http: HttpBindings,
 }
 
@@ -134,12 +138,36 @@ impl Worker {
             }
         }
 
+        // WASI-NN
+        let wasi_nn;
+        let allowed_backends = &self.config.features.wasi_nn.allowed_backends;
+
+        if !allowed_backends.is_empty() {
+            // For now, we only support OpenVINO
+            if allowed_backends.len() != 1
+                || !allowed_backends.contains(&WASI_NN_BACKEND_OPENVINO.to_string())
+            {
+                eprintln!("❌ The only WASI-NN supported backend name is \"{WASI_NN_BACKEND_OPENVINO}\". Please, update your config.");
+                wasi_nn = None;
+            } else {
+                wasmtime_wasi_nn::add_to_linker(&mut linker, |s: &mut WorkerState| {
+                    Arc::get_mut(s.wasi_nn.as_mut().unwrap())
+                        .expect("wasi-nn is not implemented with multi-threading support")
+                })?;
+
+                wasi_nn = Some(Arc::new(WasiNnCtx::new()?));
+            }
+        } else {
+            wasi_nn = None;
+        }
+
         // Pass to the runtime to add any WASI specific requirement
         wasi_builder = self.runtime.prepare_wasi_ctx(wasi_builder)?;
 
         let wasi = wasi_builder.build();
         let state = WorkerState {
             wasi,
+            wasi_nn,
             http: HttpBindings {
                 http_config: self.config.features.http_requests.clone(),
             },

examples/Makefile

@@ -20,9 +20,14 @@ rust-params:
 	cargo build --target wasm32-wasi --release && \
 	mv target/wasm32-wasi/release/rust-params.wasm "./[id].wasm"
 
+rust-wasi-nn:
+	cd rust-wasi-nn && \
+	cargo build --target wasm32-wasi --release && \
+	mv target/wasm32-wasi/release/rust-wasi-nn.wasm "./inference.wasm"
+
 rust-pdf-create:
 	cd rust-pdf-create && \
 	cargo build --target wasm32-wasi --release && \
 	mv target/wasm32-wasi/release/rust-pdf-create.wasm ./index.wasm
 
-all: rust-basic rust-fetch rust-kv rust-params
+all: rust-basic rust-fetch rust-kv rust-params rust-wasi-nn

examples/rust-wasi-nn/.gitignore

@@ -0,0 +1,3 @@
+_images/image.jpg
+_models/*.bin
+_models/*.xml