diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 2db5779b2806..26d8ebfb20ce 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -422,6 +422,10 @@ 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@v8 + if: runner.arch == 'X64' + # 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 +526,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@v4 - 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 +753,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 d2726d6872c1..d7ed5667e31f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2638,6 +2638,7 @@ dependencies = [ "sha2", "url", "wasi", + "wasi-nn", "wit-bindgen", ] @@ -3047,6 +3048,15 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "wasi-nn" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7031683cc05a71515d9200fb159b28d717ded3c40dbb979d1602cf46f3a68f40" +dependencies = [ + "thiserror", +] + [[package]] name = "wasi-preview1-component-adapter" version = "17.0.0" @@ -3835,11 +3845,14 @@ name = "wasmtime-wasi-nn" version = "17.0.0" dependencies = [ "anyhow", + "cap-std", "openvino", + "test-programs-artifacts", "thiserror", "tracing", "walkdir", "wasmtime", + "wasmtime-wasi", "wiggle", ] diff --git a/ci/run-wasi-nn-example.sh b/ci/run-wasi-nn-example.sh deleted file mode 100755 index 6f433bc3eff2..000000000000 --- a/ci/run-wasi-nn-example.sh +++ /dev/null @@ -1,57 +0,0 @@ -#!/bin/bash - -# The following script demonstrates how to execute a machine learning inference -# using the wasi-nn module optionally compiled into Wasmtime. Calling it will -# download the necessary model and tensor files stored separately in $FIXTURE -# into $TMP_DIR (optionally pass a directory with existing files as the first -# argument to re-try the script). Then, it will compile and run several examples -# in the Wasmtime CLI. -set -e -WASMTIME_DIR=$(dirname "$0" | xargs dirname) -FIXTURE=https://github.com/intel/openvino-rs/raw/main/crates/openvino/tests/fixtures/mobilenet -if [ -z "${1+x}" ]; then - # If no temporary directory is specified, create one. - TMP_DIR=$(mktemp -d -t ci-XXXXXXXXXX) - REMOVE_TMP_DIR=1 -else - # If a directory was specified, use it and avoid removing it. - TMP_DIR=$(realpath $1) - REMOVE_TMP_DIR=0 -fi - -# One of the examples expects to be in a specifically-named directory. -mkdir -p $TMP_DIR/mobilenet -TMP_DIR=$TMP_DIR/mobilenet - -# Build Wasmtime with wasi-nn enabled; we attempt this first to avoid extra work -# if the build fails. -cargo build -p wasmtime-cli --features wasi-nn - -# Download all necessary test fixtures to the temporary directory. -wget --no-clobber $FIXTURE/mobilenet.bin --output-document=$TMP_DIR/model.bin -wget --no-clobber $FIXTURE/mobilenet.xml --output-document=$TMP_DIR/model.xml -wget --no-clobber $FIXTURE/tensor-1x224x224x3-f32.bgr --output-document=$TMP_DIR/tensor.bgr - -# Now build an example that uses the wasi-nn API. Run the example in Wasmtime -# (note that the example uses `fixture` as the expected location of the -# model/tensor files). -pushd $WASMTIME_DIR/crates/wasi-nn/examples/classification-example -cargo build --release --target=wasm32-wasi -cp target/wasm32-wasi/release/wasi-nn-example.wasm $TMP_DIR -popd -cargo run -- run --dir $TMP_DIR::fixture -S nn $TMP_DIR/wasi-nn-example.wasm - -# Build and run another example, this time using Wasmtime's graph flag to -# preload the model. -pushd $WASMTIME_DIR/crates/wasi-nn/examples/classification-example-named -cargo build --release --target=wasm32-wasi -cp target/wasm32-wasi/release/wasi-nn-example-named.wasm $TMP_DIR -popd -cargo run -- run --dir $TMP_DIR::fixture -S nn,nn-graph=openvino::$TMP_DIR \ - $TMP_DIR/wasi-nn-example-named.wasm - -# Clean up the temporary directory only if it was not specified (users may want -# to keep the directory around). -if [[ $REMOVE_TMP_DIR -eq 1 ]]; then - rm -rf $TMP_DIR -fi diff --git a/crates/test-programs/Cargo.toml b/crates/test-programs/Cargo.toml index a17b5631af8f..86854364ff43 100644 --- a/crates/test-programs/Cargo.toml +++ b/crates/test-programs/Cargo.toml @@ -12,6 +12,7 @@ workspace = true [dependencies] anyhow = { workspace = true } wasi = "0.11.0" +wasi-nn = "0.6.0" wit-bindgen = { workspace = true, features = ['default'] } libc = { workspace = true } getrandom = "0.2.9" diff --git a/crates/test-programs/artifacts/build.rs b/crates/test-programs/artifacts/build.rs index b97422d3b953..193c2b3134dd 100644 --- a/crates/test-programs/artifacts/build.rs +++ b/crates/test-programs/artifacts/build.rs @@ -65,14 +65,6 @@ fn build_and_generate_tests() { generated_code += &format!("pub const {camel}: &'static str = {wasm:?};\n"); - let adapter = match target.as_str() { - "reactor" => &reactor_adapter, - s if s.starts_with("api_proxy") => &proxy_adapter, - _ => &command_adapter, - }; - let path = compile_component(&wasm, adapter); - generated_code += &format!("pub const {camel}_COMPONENT: &'static str = {path:?};\n"); - // Bucket, based on the name of the test, into a "kind" which generates // a `foreach_*` macro below. let kind = match target.as_str() { @@ -81,6 +73,7 @@ fn build_and_generate_tests() { s if s.starts_with("preview2_") => "preview2", s if s.starts_with("cli_") => "cli", s if s.starts_with("api_") => "api", + s if s.starts_with("nn_") => "nn", // If you're reading this because you hit this panic, either add it // to a test suite above or add a new "suite". The purpose of the // categorization above is to have a static assertion that tests @@ -93,6 +86,18 @@ fn build_and_generate_tests() { if !kind.is_empty() { kinds.entry(kind).or_insert(Vec::new()).push(target); } + + // Generate a component from each test. + if kind == "nn" { + continue; + } + let adapter = match target.as_str() { + "reactor" => &reactor_adapter, + s if s.starts_with("api_proxy") => &proxy_adapter, + _ => &command_adapter, + }; + let path = compile_component(&wasm, adapter); + generated_code += &format!("pub const {camel}_COMPONENT: &'static str = {path:?};\n"); } for (kind, targets) in kinds { diff --git a/crates/test-programs/src/bin/nn_image_classification.rs b/crates/test-programs/src/bin/nn_image_classification.rs new file mode 100644 index 000000000000..f81b89154ed1 --- /dev/null +++ b/crates/test-programs/src/bin/nn_image_classification.rs @@ -0,0 +1,59 @@ +use anyhow::Result; +use std::fs; +use wasi_nn::*; + +pub fn main() -> Result<()> { + let xml = fs::read_to_string("fixture/model.xml").unwrap(); + println!("Read graph XML, first 50 characters: {}", &xml[..50]); + + let weights = fs::read("fixture/model.bin").unwrap(); + println!("Read graph weights, size in bytes: {}", weights.len()); + + let graph = GraphBuilder::new(GraphEncoding::Openvino, ExecutionTarget::CPU) + .build_from_bytes([&xml.into_bytes(), &weights])?; + println!("Loaded graph into wasi-nn with ID: {}", graph); + + let mut context = graph.init_execution_context()?; + println!("Created wasi-nn execution context with ID: {}", context); + + // Load a tensor that precisely matches the graph input tensor (see + // `fixture/frozen_inference_graph.xml`). + let data = fs::read("fixture/tensor.bgr").unwrap(); + println!("Read input tensor, size in bytes: {}", data.len()); + context.set_input(0, wasi_nn::TensorType::F32, &[1, 3, 224, 224], &data)?; + + // Execute the inference. + context.compute()?; + println!("Executed graph inference"); + + // Retrieve the output. + let mut output_buffer = vec![0f32; 1001]; + context.get_output(0, &mut output_buffer[..])?; + println!( + "Found results, sorted top 5: {:?}", + &sort_results(&output_buffer)[..5] + ); + + Ok(()) +} + +// Sort the buffer of probabilities. The graph places the match probability for +// each class at the index for that class (e.g. the probability of class 42 is +// placed at buffer[42]). Here we convert to a wrapping InferenceResult and sort +// the results. It is unclear why the MobileNet output indices are "off by one" +// but the `.skip(1)` below seems necessary to get results that make sense (e.g. +// 763 = "revolver" vs 762 = "restaurant"). +fn sort_results(buffer: &[f32]) -> Vec { + let mut results: Vec = buffer + .iter() + .skip(1) + .enumerate() + .map(|(c, p)| InferenceResult(c, *p)) + .collect(); + results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + results +} + +// A wrapper for class ID and match probabilities. +#[derive(Debug, PartialEq)] +struct InferenceResult(usize, f32); diff --git a/crates/test-programs/src/bin/nn_image_classification_named.rs b/crates/test-programs/src/bin/nn_image_classification_named.rs new file mode 100644 index 000000000000..9e70770efc10 --- /dev/null +++ b/crates/test-programs/src/bin/nn_image_classification_named.rs @@ -0,0 +1,53 @@ +use anyhow::Result; +use std::fs; +use wasi_nn::*; + +pub fn main() -> Result<()> { + let graph = GraphBuilder::new(GraphEncoding::Openvino, ExecutionTarget::CPU) + .build_from_cache("mobilenet")?; + println!("Loaded a graph: {:?}", graph); + + let mut context = graph.init_execution_context()?; + println!("Created an execution context: {:?}", context); + + // Load a tensor that precisely matches the graph input tensor (see + // `fixture/frozen_inference_graph.xml`). + let tensor_data = fs::read("fixture/tensor.bgr")?; + println!("Read input tensor, size in bytes: {}", tensor_data.len()); + context.set_input(0, TensorType::F32, &[1, 3, 224, 224], &tensor_data)?; + + // Execute the inference. + context.compute()?; + println!("Executed graph inference"); + + // Retrieve the output. + let mut output_buffer = vec![0f32; 1001]; + context.get_output(0, &mut output_buffer[..])?; + + println!( + "Found results, sorted top 5: {:?}", + &sort_results(&output_buffer)[..5] + ); + Ok(()) +} + +// Sort the buffer of probabilities. The graph places the match probability for +// each class at the index for that class (e.g. the probability of class 42 is +// placed at buffer[42]). Here we convert to a wrapping InferenceResult and sort +// the results. It is unclear why the MobileNet output indices are "off by one" +// but the `.skip(1)` below seems necessary to get results that make sense (e.g. +// 763 = "revolver" vs 762 = "restaurant"). +fn sort_results(buffer: &[f32]) -> Vec { + let mut results: Vec = buffer + .iter() + .skip(1) + .enumerate() + .map(|(c, p)| InferenceResult(c, *p)) + .collect(); + results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + results +} + +// A wrapper for class ID and match probabilities. +#[derive(Debug, PartialEq)] +struct InferenceResult(usize, f32); diff --git a/crates/wasi-nn/Cargo.toml b/crates/wasi-nn/Cargo.toml index 9b42aa00fb1a..fbc941f5cda6 100644 --- a/crates/wasi-nn/Cargo.toml +++ b/crates/wasi-nn/Cargo.toml @@ -29,3 +29,9 @@ thiserror = { workspace = true } [build-dependencies] walkdir = { workspace = true } + +[dev-dependencies] +cap-std = { workspace = true } +test-programs-artifacts = { workspace = true } +wasmtime-wasi = { workspace = true, features = ["sync"] } +wasmtime = { workspace = true, features = ["cranelift"] } diff --git a/crates/wasi-nn/README.md b/crates/wasi-nn/README.md index 4933fcba339c..c726c7c90cd6 100644 --- a/crates/wasi-nn/README.md +++ b/crates/wasi-nn/README.md @@ -37,9 +37,3 @@ An end-to-end example demonstrating ML classification is included in [examples]: `examples/classification-example` contains a standalone Rust project that uses the [wasi-nn] APIs and is compiled to the `wasm32-wasi` target using the high-level `wasi-nn` [bindings]. - -Run the example from the Wasmtime project directory: - -```sh -$ ci/run-wasi-nn-example.sh -``` diff --git a/crates/wasi-nn/src/lib.rs b/crates/wasi-nn/src/lib.rs index 71d089d07489..e66c59fb2a05 100644 --- a/crates/wasi-nn/src/lib.rs +++ b/crates/wasi-nn/src/lib.rs @@ -4,6 +4,7 @@ mod registry; pub mod backend; pub use ctx::{preload, WasiNnCtx}; pub use registry::{GraphRegistry, InMemoryRegistry}; +pub mod testing; pub mod wit; pub mod witx; diff --git a/crates/wasi-nn/src/testing.rs b/crates/wasi-nn/src/testing.rs new file mode 100644 index 000000000000..b960335dc500 --- /dev/null +++ b/crates/wasi-nn/src/testing.rs @@ -0,0 +1,97 @@ +//! This is testing-specific code--it is public only so that it can be +//! accessible both in unit and integration tests. +//! +//! This module checks: +//! - that OpenVINO can be found in the environment +//! - that some ML model artifacts can be downloaded and cached. + +use anyhow::{anyhow, Context, Result}; +use std::{env, fs, path::Path, path::PathBuf, process::Command, sync::Mutex}; + +/// 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! check_test { + () => { + if let Err(e) = $crate::testing::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)), + } +} + +/// Protect `check_openvino_artifacts_are_available` from concurrent access; +/// when running tests in parallel, we want to avoid two threads attempting to +/// create the same directory or download the same file. +static ARTIFACTS: Mutex<()> = Mutex::new(()); + +/// Return `Ok` if we find the cached MobileNet test artifacts; this will +/// download the artifacts if necessary. +fn check_openvino_artifacts_are_available() -> Result<()> { + let _exclusively_retrieve_artifacts = ARTIFACTS.lock().unwrap(); + 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(()) +} diff --git a/crates/wasi-nn/tests/all.rs b/crates/wasi-nn/tests/all.rs new file mode 100644 index 000000000000..de1f9f507f17 --- /dev/null +++ b/crates/wasi-nn/tests/all.rs @@ -0,0 +1,92 @@ +//! Run the wasi-nn tests in `crates/test-programs`. + +use anyhow::Result; +use std::path::Path; +use test_programs_artifacts::*; +use wasmtime::{Config, Engine, Linker, Module, Store}; +use wasmtime_wasi::sync::{Dir, WasiCtxBuilder}; +use wasmtime_wasi::WasiCtx; +use wasmtime_wasi_nn::{backend, testing, InMemoryRegistry, WasiNnCtx}; + +const PREOPENED_DIR_NAME: &str = "fixture"; + +/// Run a wasi-nn test program. This is modeled after +/// `crates/wasi/tests/all/main.rs` but still uses the older preview1 API for +/// file reads. +fn run(path: &str, preload_model: bool) -> Result<()> { + wasmtime_wasi_nn::check_test!(); + let path = Path::new(path); + let config = Config::new(); + let engine = Engine::new(&config)?; + let mut linker = Linker::new(&engine); + wasmtime_wasi_nn::witx::add_to_linker(&mut linker, |s: &mut Ctx| &mut s.wasi_nn)?; + wasmtime_wasi::add_to_linker(&mut linker, |s: &mut Ctx| &mut s.wasi)?; + let module = Module::from_file(&engine, path)?; + let mut store = Store::new(&engine, Ctx::new(&testing::artifacts_dir(), preload_model)?); + let instance = linker.instantiate(&mut store, &module)?; + let start = instance.get_typed_func::<(), ()>(&mut store, "_start")?; + start.call(&mut store, ())?; + Ok(()) +} + +/// The host state for running wasi-nn tests. +struct Ctx { + wasi: WasiCtx, + wasi_nn: WasiNnCtx, +} +impl Ctx { + fn new(preopen_dir: &Path, preload_model: bool) -> Result { + // Create the WASI context. + let preopen_dir = Dir::open_ambient_dir(preopen_dir, cap_std::ambient_authority())?; + let mut builder = WasiCtxBuilder::new(); + builder + .inherit_stdio() + .preopened_dir(preopen_dir, PREOPENED_DIR_NAME)?; + let wasi = builder.build(); + + // Create the wasi-nn context. + let mut openvino = backend::openvino::OpenvinoBackend::default(); + let mut registry = InMemoryRegistry::new(); + let mobilenet_dir = testing::artifacts_dir(); + if preload_model { + registry.load(&mut openvino, &mobilenet_dir)?; + } + let wasi_nn = WasiNnCtx::new([openvino.into()], registry.into()); + + Ok(Self { wasi, wasi_nn }) + } +} + +// Check that every wasi-nn test in `crates/test-programs` has its +// manually-added `#[test]` function. +macro_rules! assert_test_exists { + ($name:ident) => { + #[allow(unused_imports)] + use self::$name as _; + }; +} +foreach_nn!(assert_test_exists); + +#[cfg_attr( + not(all( + target_arch = "x86_64", + any(target_os = "linux", target_os = "windows") + )), + ignore +)] +#[test] +fn nn_image_classification() { + run(NN_IMAGE_CLASSIFICATION, false).unwrap() +} + +#[cfg_attr( + not(all( + target_arch = "x86_64", + any(target_os = "linux", target_os = "windows") + )), + ignore +)] +#[test] +fn nn_image_classification_named() { + run(NN_IMAGE_CLASSIFICATION_NAMED, true).unwrap() +} diff --git a/supply-chain/audits.toml b/supply-chain/audits.toml index 064868b341ba..073eb4560779 100644 --- a/supply-chain/audits.toml +++ b/supply-chain/audits.toml @@ -2291,6 +2291,12 @@ who = "Pat Hickey " criteria = "safe-to-deploy" version = "0.3.0" +[[audits.wasi-nn]] +who = "Andrew Brown " +criteria = "safe-to-deploy" +version = "0.6.0" +notes = "This crate contains `unsafe` code due to its purpose: it wraps up `witx-bindgen`-generated code that calls the raw wasi-nn API." + [[audits.wasm-bindgen-shared]] who = "Pat Hickey " criteria = "safe-to-deploy"