feat: add WASI-NN bindings to Wasm Workers Server (#201)

* feat: add WASI-NN bindings to Wasm Workers Server

* fix: remove unnecessary var initialization

* fix: typo error when closing a conditional assignment

* feat: add new documentation about ML inference and WASI-NN

* docs: improve docs and remove typos

* fix: remove support for tensor files and clarifies why we need to skip 1

---------

Co-authored-by: Rafael Fernández López <[email protected]>

Author: Angelmmiguel

Cargo.lock

@@ -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" }

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/Cargo.toml

@@ -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/config.rs

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

crates/worker/src/features/mod.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/features/wasi_nn.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,35 @@ impl Worker {
             }
         }
 
+        // WASI-NN
+        let allowed_backends = &self.config.features.wasi_nn.allowed_backends;
+
+        let wasi_nn = 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.");
+                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")
+                })?;
+
+                Some(Arc::new(WasiNnCtx::new()?))
+            }
+        } else {
+            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(),
             },

crates/worker/src/lib.rs

@@ -1,8 +1,13 @@
 ---
+title: Dynamic Routes
 sidebar_position: 4
 ---
 
-# Dynamic routes
+:::info
+
+[Available since v1.0](https://github.com/vmware-labs/wasm-workers-server/releases/tag/v1.0.0)
+
+:::
 
 Defining static routes may not be enough for some applications. You may need a worker to process URLs that includes identifiers. **To create a worker associated with a dynamic route, include the route parameter in brackets when setting the worker filename**.
 

docs/docs/features/dynamic-routes.md

@@ -1,8 +1,13 @@
 ---
+title: HTTP Requests (fetch)
 sidebar_position: 3
 ---
 
-# HTTP Requests (fetch)
+:::info
+
+[Available since v1.4](https://github.com/vmware-labs/wasm-workers-server/releases/tag/v1.4.0)
+
+:::
 
 Often times, workers require to access data from an external resource like a website or an API. This feature allows workers to perform HTTP requests to external resources. It follows the capability-based model, so workers cannot perform any HTTP request until you configure the allowed hosts and HTTP methods.