1 use super::PREOPENED_DIR_NAME;
2 use crate::check::artifacts_dir;
3 use std::path::Path;
4 use wasmtime::component::{Component, Linker, ResourceTable};
5 use wasmtime::{Config, Engine, Result, Store, format_err};
6 use wasmtime_wasi::p2::bindings::sync::Command;
7 use wasmtime_wasi::{DirPerms, FilePerms, WasiCtx, WasiCtxView};
8 use wasmtime_wasi_nn::wit::WasiNnView;
9 use wasmtime_wasi_nn::{Backend, InMemoryRegistry, wit::WasiNnCtx};
10
11 /// Run a wasi-nn test program. This is modeled after
12 /// `crates/wasi/tests/all/main.rs` but still uses the older p1 API for
13 /// file reads.
run(path: &str, backend: Backend, preload_model: bool) -> Result<()>14 pub fn run(path: &str, backend: Backend, preload_model: bool) -> Result<()> {
15 let path = Path::new(path);
16 let engine = Engine::new(&Config::new())?;
17 let mut linker = Linker::new(&engine);
18 wasmtime_wasi_nn::wit::add_to_linker(&mut linker, |c: &mut Ctx| {
19 WasiNnView::new(&mut c.table, &mut c.wasi_nn)
20 })?;
21 wasmtime_wasi::p2::add_to_linker_sync(&mut linker)?;
22 let module = Component::from_file(&engine, path)?;
23 let mut store = Store::new(&engine, Ctx::new(&artifacts_dir(), preload_model, backend)?);
24 let command = Command::instantiate(&mut store, &module, &linker)?;
25 let result = command.wasi_cli_run().call_run(&mut store)?;
26 result.map_err(|_| format_err!("failed to run command"))
27 }
28
29 /// The host state for running wasi-nn component tests.
30 struct Ctx {
31 wasi: WasiCtx,
32 wasi_nn: WasiNnCtx,
33 table: ResourceTable,
34 }
35
36 impl Ctx {
new(preopen_dir: &Path, preload_model: bool, mut backend: Backend) -> Result<Self>37 fn new(preopen_dir: &Path, preload_model: bool, mut backend: Backend) -> Result<Self> {
38 let mut builder = WasiCtx::builder();
39 builder.inherit_stdio().preopened_dir(
40 preopen_dir,
41 PREOPENED_DIR_NAME,
42 DirPerms::READ,
43 FilePerms::READ,
44 )?;
45 let wasi = builder.build();
46
47 let mut registry = InMemoryRegistry::new();
48 let mobilenet_dir = artifacts_dir();
49 if preload_model {
50 registry.load((backend).as_dir_loadable().unwrap(), &mobilenet_dir)?;
51 }
52 let wasi_nn = WasiNnCtx::new([backend], registry.into());
53
54 let table = ResourceTable::new();
55
56 Ok(Self {
57 wasi,
58 wasi_nn,
59 table,
60 })
61 }
62 }
63
64 impl wasmtime_wasi::WasiView for Ctx {
ctx(&mut self) -> WasiCtxView<'_>65 fn ctx(&mut self) -> WasiCtxView<'_> {
66 WasiCtxView {
67 ctx: &mut self.wasi,
68 table: &mut self.table,
69 }
70 }
71 }
72