wasi-nn: switch testing infrastructure to Rust

.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

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

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 \
     $@

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"]

crates/wasi-nn/examples/classification-example-named/Cargo.toml → 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"

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.

crates/wasi-nn/examples/classification-example/Cargo.toml → 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"

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.

crates/wasi-nn/src/backend/openvino.rs

@@ -178,3 +178,56 @@ fn read(path: &Path) -> anyhow::Result<Vec<u8>> {
     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::<f32>(),
+            )
+        };
+        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(())
+    }
+}

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 {

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<Path>) -> 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`"
+);