diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index b4cb5cfd7d80..f1677b9cdb93 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -422,6 +422,12 @@ jobs: - run: echo CARGO_BUILD_TARGET=${{ matrix.target }} >> $GITHUB_ENV if: matrix.target != '' + # Install OpenVINO for testing wasmtime-wasi-nn. + - uses: abrown/install-openvino-action@v7 + with: + version: 2022.3.0 + - run: echo OPENVINO_BUILD_DIR="${{ env.OPENVINO_INSTALL_DIR }}" >> $GITHUB_ENV + # Fix an ICE for now in gcc when compiling zstd with debuginfo (??) - run: echo CFLAGS=-g0 >> $GITHUB_ENV if: matrix.target == 'x86_64-pc-windows-gnu' @@ -522,32 +528,6 @@ jobs: # Windows fails GitHub Actions will confusingly mark the failed Windows job # as cancelled instead of failed. - # Build and test the wasi-nn module. - test_wasi_nn: - needs: determine - if: needs.determine.outputs.run-full - name: Test wasi-nn module - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v3 - with: - submodules: true - - uses: ./.github/actions/install-rust - - run: rustup target add wasm32-wasi - - uses: abrown/install-openvino-action@v7 - with: - version: 2022.3.0 - apt: true - - run: ./ci/run-wasi-nn-example.sh - env: - RUST_BACKTRACE: 1 - - # common logic to cancel the entire run if this job fails - - run: gh run cancel ${{ github.run_id }} - if: failure() && github.event_name != 'pull_request' - env: - GH_TOKEN: ${{ github.token }} - build-preview1-component-adapter: name: Build wasi-preview1-component-adapter needs: determine @@ -775,7 +755,6 @@ jobs: - checks - checks_winarm64 - fuzz_targets - - test_wasi_nn - bench - meta_deterministic_check - verify-publish diff --git a/Cargo.lock b/Cargo.lock index 7de50af65fd8..57ff5398e8da 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3725,6 +3725,7 @@ dependencies = [ "tracing", "walkdir", "wasmtime", + "wasmtime-wasi", "wiggle", ] diff --git a/Cargo.toml b/Cargo.toml index d89162d30eba..f57e079b2ca5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -167,7 +167,7 @@ wasmtime-runtime = { path = "crates/runtime", version = "=17.0.0" } wasmtime-wast = { path = "crates/wast", version = "=17.0.0" } wasmtime-wasi = { path = "crates/wasi", version = "17.0.0", default-features = false } wasmtime-wasi-http = { path = "crates/wasi-http", version = "=17.0.0", default-features = false } -wasmtime-wasi-nn = { path = "crates/wasi-nn", version = "17.0.0" } +wasmtime-wasi-nn = { path = "crates/wasi-nn", version = "17.0.0", default-features = false } wasmtime-wasi-threads = { path = "crates/wasi-threads", version = "17.0.0" } wasmtime-component-util = { path = "crates/component-util", version = "=17.0.0" } wasmtime-component-macro = { path = "crates/component-macro", version = "=17.0.0" } diff --git a/ci/run-tests.sh b/ci/run-tests.sh index db338dc2409e..738332b736ad 100755 --- a/ci/run-tests.sh +++ b/ci/run-tests.sh @@ -5,6 +5,7 @@ cargo test \ --features wasi-http \ --features component-model \ --features serve \ + --features wasmtime-wasi-nn/test-check \ --workspace \ --exclude test-programs \ $@ diff --git a/crates/wasi-nn/Cargo.toml b/crates/wasi-nn/Cargo.toml index 9b42aa00fb1a..3e04b56c4c6e 100644 --- a/crates/wasi-nn/Cargo.toml +++ b/crates/wasi-nn/Cargo.toml @@ -20,7 +20,7 @@ anyhow = { workspace = true } wiggle = { workspace = true } # This dependency is necessary for the WIT-generation macros to work: -wasmtime = { workspace = true, features = ["component-model"] } +wasmtime = { workspace = true, features = ["component-model", "cranelift"] } # These dependencies are necessary for the wasi-nn implementation: tracing = { workspace = true } @@ -29,3 +29,17 @@ thiserror = { workspace = true } [build-dependencies] walkdir = { workspace = true } + +[dev-dependencies] +wasmtime-wasi = { workspace = true } + +[features] +# When turned on, checks that the system has all the necessary ML frameworks and +# artifacts to run this crate's tests. If these are not available, the tests +# will be skipped; set `FORCE_WASINN_TEST_CHECK=1` or `CI=true` to fail if the +# checks fail. +test-check = [] +# We'll enable the `test-check` feature by default so `cargo test` will "just +# work"; this means we need to carefully turn off `default-features` at the use +# sites to avoid unwanted dependencies. +default = ["test-check"] diff --git a/crates/wasi-nn/examples/classification-example-named/README.md b/crates/wasi-nn/examples/classification-example-named/README.md deleted file mode 100644 index aa56ad0cbaf7..000000000000 --- a/crates/wasi-nn/examples/classification-example-named/README.md +++ /dev/null @@ -1,2 +0,0 @@ -This example project demonstrates using the `wasi-nn` API to perform ML inference. It consists of Rust code that is -built using the `wasm32-wasi` target. See `ci/run-wasi-nn-example.sh` for how this is used. diff --git a/crates/wasi-nn/examples/classification-example/README.md b/crates/wasi-nn/examples/classification-example/README.md deleted file mode 100644 index aa56ad0cbaf7..000000000000 --- a/crates/wasi-nn/examples/classification-example/README.md +++ /dev/null @@ -1,2 +0,0 @@ -This example project demonstrates using the `wasi-nn` API to perform ML inference. It consists of Rust code that is -built using the `wasm32-wasi` target. See `ci/run-wasi-nn-example.sh` for how this is used. diff --git a/crates/wasi-nn/examples/classification-example-named/Cargo.lock b/crates/wasi-nn/examples/image-classification-named/Cargo.lock similarity index 98% rename from crates/wasi-nn/examples/classification-example-named/Cargo.lock rename to crates/wasi-nn/examples/image-classification-named/Cargo.lock index 6a1c9eb0f62f..55e7015b4c7a 100644 --- a/crates/wasi-nn/examples/classification-example-named/Cargo.lock +++ b/crates/wasi-nn/examples/image-classification-named/Cargo.lock @@ -2,6 +2,13 @@ # It is not intended for manual editing. version = 3 +[[package]] +name = "image-classification-named" +version = "0.0.0" +dependencies = [ + "wasi-nn", +] + [[package]] name = "proc-macro2" version = "1.0.66" @@ -65,10 +72,3 @@ checksum = "03d01b90f0cca3f19682e90e1bc3f5e3e441031e19e56ce7dbf034f3b3597552" dependencies = [ "thiserror", ] - -[[package]] -name = "wasi-nn-example-named" -version = "0.0.0" -dependencies = [ - "wasi-nn", -] diff --git a/crates/wasi-nn/examples/classification-example-named/Cargo.toml b/crates/wasi-nn/examples/image-classification-named/Cargo.toml similarity index 90% rename from crates/wasi-nn/examples/classification-example-named/Cargo.toml rename to crates/wasi-nn/examples/image-classification-named/Cargo.toml index b4653659bd3d..85e44e0e4ed5 100644 --- a/crates/wasi-nn/examples/classification-example-named/Cargo.toml +++ b/crates/wasi-nn/examples/image-classification-named/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "wasi-nn-example-named" +name = "image-classification-named" version = "0.0.0" authors = ["The Wasmtime Project Developers"] readme = "README.md" diff --git a/crates/wasi-nn/examples/image-classification-named/README.md b/crates/wasi-nn/examples/image-classification-named/README.md new file mode 100644 index 000000000000..623f99d7b15a --- /dev/null +++ b/crates/wasi-nn/examples/image-classification-named/README.md @@ -0,0 +1,2 @@ +This example project is similar to +[`image-classification`](../image-classification/) but uses named models. diff --git a/crates/wasi-nn/examples/classification-example-named/src/main.rs b/crates/wasi-nn/examples/image-classification-named/src/main.rs similarity index 100% rename from crates/wasi-nn/examples/classification-example-named/src/main.rs rename to crates/wasi-nn/examples/image-classification-named/src/main.rs diff --git a/crates/wasi-nn/examples/classification-example/Cargo.lock b/crates/wasi-nn/examples/image-classification/Cargo.lock similarity index 92% rename from crates/wasi-nn/examples/classification-example/Cargo.lock rename to crates/wasi-nn/examples/image-classification/Cargo.lock index a649a0429289..fcc72d489e28 100644 --- a/crates/wasi-nn/examples/classification-example/Cargo.lock +++ b/crates/wasi-nn/examples/image-classification/Cargo.lock @@ -3,14 +3,14 @@ version = 3 [[package]] -name = "wasi-nn" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c909acded993dc129e02f64a7646eb7b53079f522a814024a88772f41558996" - -[[package]] -name = "wasi-nn-example" +name = "image-classification" version = "0.0.0" dependencies = [ "wasi-nn", ] + +[[package]] +name = "wasi-nn" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c909acded993dc129e02f64a7646eb7b53079f522a814024a88772f41558996" diff --git a/crates/wasi-nn/examples/classification-example/Cargo.toml b/crates/wasi-nn/examples/image-classification/Cargo.toml similarity index 92% rename from crates/wasi-nn/examples/classification-example/Cargo.toml rename to crates/wasi-nn/examples/image-classification/Cargo.toml index 7cedd6252c43..917a9ed27376 100644 --- a/crates/wasi-nn/examples/classification-example/Cargo.toml +++ b/crates/wasi-nn/examples/image-classification/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "wasi-nn-example" +name = "image-classification" version = "0.0.0" authors = ["The Wasmtime Project Developers"] readme = "README.md" diff --git a/crates/wasi-nn/examples/image-classification/README.md b/crates/wasi-nn/examples/image-classification/README.md new file mode 100644 index 000000000000..9d00ca1b9a1d --- /dev/null +++ b/crates/wasi-nn/examples/image-classification/README.md @@ -0,0 +1,3 @@ +This example project demonstrates using the `wasi-nn` API to perform ML +inference. It consists of Rust code that is built using the `wasm32-wasi` +target. It is used by `wasi-nn`'s tests. diff --git a/crates/wasi-nn/examples/classification-example/src/main.rs b/crates/wasi-nn/examples/image-classification/src/main.rs similarity index 100% rename from crates/wasi-nn/examples/classification-example/src/main.rs rename to crates/wasi-nn/examples/image-classification/src/main.rs diff --git a/crates/wasi-nn/src/backend/openvino.rs b/crates/wasi-nn/src/backend/openvino.rs index be9c0250dbb7..f8eaeb5c93c2 100644 --- a/crates/wasi-nn/src/backend/openvino.rs +++ b/crates/wasi-nn/src/backend/openvino.rs @@ -178,3 +178,56 @@ fn read(path: &Path) -> anyhow::Result> { file.read_to_end(&mut buffer)?; Ok(buffer) } + +#[cfg(test)] +mod tests { + use super::*; + use anyhow::Result; + use std::{mem, slice}; + + #[test] + fn image_classification_directly_on_backend() -> Result<()> { + crate::test_check!(); + + // Compute a MobileNet classification using the test artifacts. + let mut backend = OpenvinoBackend::default(); + let graph = backend.load_from_dir( + crate::test_check::artifacts_dir().as_path(), + ExecutionTarget::Cpu, + )?; + let mut context = graph.init_execution_context()?; + let data = read( + crate::test_check::artifacts_dir() + .join("tensor.bgr") + .as_path(), + )?; + let tensor = Tensor { + dimensions: vec![1, 3, 224, 224], + tensor_type: TensorType::Fp32, + data, + }; + context.set_input(0, &tensor)?; + context.compute()?; + let mut destination = vec![0f32; 1001]; + let destination_ = unsafe { + slice::from_raw_parts_mut( + destination.as_mut_ptr().cast(), + destination.len() * mem::size_of::(), + ) + }; + context.get_output(0, destination_)?; + + // Find the top score which should be the entry for "pizza" (see + // https://github.com/leferrad/tensorflow-mobilenet/blob/master/imagenet/labels.txt, + // e.g.) + let (id, score) = destination + .iter() + .enumerate() + .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap()) + .unwrap(); + println!("> top match: label #{} = {}", id, score); + assert_eq!(id, 964); + + Ok(()) + } +} diff --git a/crates/wasi-nn/src/ctx.rs b/crates/wasi-nn/src/ctx.rs index 40cfb9d53d2d..0863bf9a7ede 100644 --- a/crates/wasi-nn/src/ctx.rs +++ b/crates/wasi-nn/src/ctx.rs @@ -56,6 +56,12 @@ impl WasiNnCtx { } } +impl Default for WasiNnCtx { + fn default() -> Self { + Self::new(backend::list(), Registry::from(InMemoryRegistry::new())) + } +} + /// Possible errors while interacting with [WasiNnCtx]. #[derive(Debug, Error)] pub enum WasiNnError { diff --git a/crates/wasi-nn/src/lib.rs b/crates/wasi-nn/src/lib.rs index 71d089d07489..01b0a74a90cd 100644 --- a/crates/wasi-nn/src/lib.rs +++ b/crates/wasi-nn/src/lib.rs @@ -83,3 +83,127 @@ where Self(Box::new(value)) } } + +/// For testing, this module checks: +/// - that OpenVINO can be found in the environment +/// - that some ML model artifacts can be downloaded and cached. +#[cfg(feature = "test-check")] +pub mod test_check { + use anyhow::{anyhow, Context, Result}; + use std::{env, fs, path::Path, path::PathBuf, process::Command}; + + /// Return the directory in which the test artifacts are stored. + pub fn artifacts_dir() -> PathBuf { + PathBuf::from(env!("OUT_DIR")).join("mobilenet") + } + + /// Early-return from a test if the test environment is not met. If the `CI` + /// or `FORCE_WASINN_TEST_CHECK` environment variables are set, though, this + /// will return an error instead. + #[macro_export] + macro_rules! test_check { + () => { + if let Err(e) = $crate::test_check::check() { + if std::env::var_os("CI").is_some() + || std::env::var_os("FORCE_WASINN_TEST_CHECK").is_some() + { + return Err(e); + } else { + println!("> ignoring test: {}", e); + return Ok(()); + } + } + }; + } + + /// Return `Ok` if all checks pass. + pub fn check() -> Result<()> { + check_openvino_is_installed()?; + check_openvino_artifacts_are_available()?; + Ok(()) + } + + /// Return `Ok` if we find a working OpenVINO installation. + fn check_openvino_is_installed() -> Result<()> { + match std::panic::catch_unwind(|| { + println!("> found openvino version: {}", openvino::version()) + }) { + Ok(_) => Ok(()), + Err(e) => Err(anyhow!("unable to find an OpenVINO installation: {:?}", e)), + } + } + + /// Return `Ok` if we find the cached MobileNet test artifacts; this will + /// download the artifacts if necessary. + fn check_openvino_artifacts_are_available() -> Result<()> { + const BASE_URL: &str = "https://github.com/intel/openvino-rs/raw/main/crates/openvino/tests/fixtures/mobilenet"; + let artifacts_dir = artifacts_dir(); + if !artifacts_dir.is_dir() { + fs::create_dir(&artifacts_dir)?; + } + for (from, to) in [ + ("mobilenet.bin", "model.bin"), + ("mobilenet.xml", "model.xml"), + ("tensor-1x224x224x3-f32.bgr", "tensor.bgr"), + ] { + let remote_url = [BASE_URL, from].join("/"); + let local_path = artifacts_dir.join(to); + if !local_path.is_file() { + download(&remote_url, &local_path) + .with_context(|| "unable to retrieve test artifact")?; + } else { + println!("> using cached artifact: {}", local_path.display()) + } + } + Ok(()) + } + + /// Retrieve the bytes at the `from` URL and place them in the `to` file. + fn download(from: &str, to: &Path) -> anyhow::Result<()> { + let mut curl = Command::new("curl"); + curl.arg("--location").arg(from).arg("--output").arg(to); + println!("> downloading: {:?}", &curl); + let result = curl.output().unwrap(); + if !result.status.success() { + panic!( + "curl failed: {}\n{}", + result.status, + String::from_utf8_lossy(&result.stderr) + ); + } + Ok(()) + } + + /// Build the given crate as `wasm32-wasi` and return the path to the built + /// module. + pub fn cargo_build(crate_dir: impl AsRef) -> PathBuf { + let crate_dir = crate_dir.as_ref(); + let crate_name = crate_dir.file_name().unwrap().to_str().unwrap(); + let cargo_toml = crate_dir.join("Cargo.toml"); + let wasm = crate_dir.join(format!("target/wasm32-wasi/release/{}.wasm", crate_name)); + let result = Command::new("cargo") + .arg("build") + .arg("--release") + .arg("--target=wasm32-wasi") + .arg("--manifest-path") + .arg(cargo_toml) + .output() + .unwrap(); + if !wasm.is_file() { + panic!("no file found at: {}", wasm.display()); + } + if !result.status.success() { + panic!( + "cargo build failed: {}\n{}", + result.status, + String::from_utf8_lossy(&result.stderr) + ); + } + wasm + } +} + +#[cfg(all(test, not(feature = "test-check")))] +compile_error!( + "to run wasi-nn tests we need to enable a feature: `cargo test --features test-check`" +); diff --git a/crates/wasi-nn/tests/e2e.rs b/crates/wasi-nn/tests/e2e.rs new file mode 100644 index 000000000000..755625adb6f2 --- /dev/null +++ b/crates/wasi-nn/tests/e2e.rs @@ -0,0 +1,79 @@ +//! Embed wasi-nn in Wasmtime and compute an inference by: +//! - downloading any necessary model artifacts (`test_check!`) +//! - setting up a wasi + wasi-nn environment +//! - build an `example` crate into a `*.wasm` file +//! - run the `*.wasm` file. + +use anyhow::Result; +use wasmtime::{Engine, Linker, Module, Store}; +use wasmtime_wasi::{ambient_authority, Dir, WasiCtx, WasiCtxBuilder}; +use wasmtime_wasi_nn::{backend, InMemoryRegistry, WasiNnCtx}; + +#[test] +fn image_classification() -> Result<()> { + wasmtime_wasi_nn::test_check!(); + + // Set up a WASI environment that includes wasi-nn and opens the MobileNet + // artifacts directory as `fixture` in the guest. + let engine = Engine::default(); + let (mut store, mut linker) = embed_wasi_nn(&engine, WasiNnCtx::default())?; + + // Build and run the example crate. + let wasm_file = wasmtime_wasi_nn::test_check::cargo_build("examples/image-classification"); + let module = Module::from_file(&engine, wasm_file)?; + linker.module(&mut store, "", &module)?; + linker + .get_default(&mut store, "")? + .typed::<(), ()>(&store)? + .call(&mut store, ())?; + + Ok(()) +} + +#[test] +fn image_classification_with_names() -> Result<()> { + wasmtime_wasi_nn::test_check!(); + + // Set up a WASI environment that includes wasi-nn and uses a registry with + // the "mobilenet" name populated. + let engine = Engine::default(); + let mut openvino = backend::openvino::OpenvinoBackend::default(); + let mut registry = InMemoryRegistry::new(); + let mobilenet_dir = wasmtime_wasi_nn::test_check::artifacts_dir(); + registry.load(&mut openvino, &mobilenet_dir)?; + let wasi_nn = WasiNnCtx::new([openvino.into()], registry.into()); + let (mut store, mut linker) = embed_wasi_nn(&engine, wasi_nn)?; + + // Build and run the example crate. + let wasm_file = + wasmtime_wasi_nn::test_check::cargo_build("examples/image-classification-named"); + let module = Module::from_file(&engine, wasm_file)?; + linker.module(&mut store, "", &module)?; + linker + .get_default(&mut store, "")? + .typed::<(), ()>(&store)? + .call(&mut store, ())?; + + Ok(()) +} + +struct Host { + wasi: WasiCtx, + wasi_nn: WasiNnCtx, +} + +fn embed_wasi_nn(engine: &Engine, wasi_nn: WasiNnCtx) -> Result<(Store, Linker)> { + let mut linker = Linker::new(&engine); + let host_dir = Dir::open_ambient_dir( + wasmtime_wasi_nn::test_check::artifacts_dir(), + ambient_authority(), + )?; + let wasi = WasiCtxBuilder::new() + .inherit_stdio() + .preopened_dir(host_dir, "fixture")? + .build(); + let store = Store::::new(&engine, Host { wasi, wasi_nn }); + wasmtime_wasi_nn::witx::add_to_linker(&mut linker, |s: &mut Host| &mut s.wasi_nn)?; + wasmtime_wasi::add_to_linker(&mut linker, |s: &mut Host| &mut s.wasi)?; + Ok((store, linker)) +} diff --git a/tests/all/cli_tests.rs b/tests/all/cli_tests.rs index 03783844b7b3..d7cda6716e6a 100644 --- a/tests/all/cli_tests.rs +++ b/tests/all/cli_tests.rs @@ -609,6 +609,44 @@ fn run_simple_with_wasi_threads() -> Result<()> { Ok(()) } +#[cfg(feature = "wasi-nn")] +#[test] +fn image_classification_with_wasi_nn() -> Result<()> { + wasmtime_wasi_nn::test_check!(); + let wasm = + wasmtime_wasi_nn::test_check::cargo_build("crates/wasi-nn/examples/image-classification"); + let artifacts_dir = wasmtime_wasi_nn::test_check::artifacts_dir(); + let artifacts_dir = artifacts_dir.display(); + let stdout = run_wasmtime(&[ + "run", + "--wasi-modules=experimental-wasi-nn", + &format!("--mapdir=fixture::{artifacts_dir}"), + &format!("{}", wasm.display()), + ])?; + assert!(stdout.contains("InferenceResult(963")); + Ok(()) +} + +#[cfg(feature = "wasi-nn")] +#[test] +fn image_classification_with_wasi_nn_and_named_models() -> Result<()> { + wasmtime_wasi_nn::test_check!(); + let wasm = wasmtime_wasi_nn::test_check::cargo_build( + "crates/wasi-nn/examples/image-classification-named", + ); + let artifacts_dir = wasmtime_wasi_nn::test_check::artifacts_dir(); + let artifacts_dir = artifacts_dir.display(); + let stdout = run_wasmtime(&[ + "run", + "--wasi-modules=experimental-wasi-nn", + &format!("--mapdir=fixture::{artifacts_dir}"), + &format!("--wasi-nn-graph=openvino::{artifacts_dir}"), + &format!("{}", wasm.display()), + ])?; + assert!(stdout.contains("InferenceResult(963")); + Ok(()) +} + #[test] fn wasm_flags() -> Result<()> { // Any argument after the wasm module should be interpreted as for the