wasi-nn: add named models #6854
Conversation
abrown
requested review from
fitzgen and
pchickey
and removed request for
a team
| cp target/wasm32-wasi/release/wasi-nn-example-named.wasm $TMP_DIR | ||
| popd | ||
| cargo run -- run --mapdir fixture::$TMP_DIR --graph openvino::$TMP_DIR \ | ||
| --wasi-modules=experimental-wasi-nn $TMP_DIR/wasi-nn-example-named.wasm |
There was a problem hiding this comment.
I'm not quite happy with using this script for testing. It was good initially to get things working but I would prefer to be using cargo test instead. The problem is that there is a dependency expectation: this can only run if the system has OpenVINO installed, which will not always be the case for developers running cargo test.
I can think of two workarounds:
#[ignore]the OpenVINO-specific tests and add some comments telling developers to turn them on to fully test the system- add some checks at the beginning of each test to early-exit if OpenVINO is not available; the issue with this is that users may see "false success" if something goes wrong in the environment, so perhaps a
FORCE_RUN_OPENVINO_TESTSenv variable is necessary?
Interested in some feedback on whether one of those approaches is better than this script!
There was a problem hiding this comment.
Personally I like your second suggestion and matches what I was going to suggest as well. We do run a higher risk of not actually running the tests in CI but I think running something by default as part of cargo test --workspace is a good mitigating factor for that.
There was a problem hiding this comment.
If the tests are for wasi-nn, we should just have a test backend that does something basic like an affine layer.
If this is testing integration we should think about how we handle multiple backends. This same issue will come up with pytorch, onnxruntime, tf, etc.
There was a problem hiding this comment.
@alexcrichton, I have a whole other set of commits ready for switching to that kind of testing. Just waiting on a review of this one because the new testing also covers the load_by_name logic added here.
There was a problem hiding this comment.
Ok, check out #6895 if you're interested in where I went with this.
| pub use in_memory::InMemoryRegistry; | ||
|
|
||
| pub trait GraphRegistry: Send + Sync { | ||
| fn get_mut(&mut self, name: &str) -> Option<&mut Graph>; |
There was a problem hiding this comment.
Another thing I'm not quite sure of is whether this should return Option<&mut Graph>, Option<&Graph>, or even Option<Graph> (seeing as how Graph is just an Arc-wrapper). We immediately clone the graph and I removed the mutability requirement for BackendGraph::init_execution_context... @geekbeast, any opinions from implementing other kinds of GraphRegistry?
There was a problem hiding this comment.
I can't review the wasi-nn stuff specifically really, it all seems reasonable though. The src/commands/run.rs bits look good to me, although I might recommend renaming --graph to something like --wasi-nn-graph or --nn-graph or something like that perhaps to make it a bit more clear it's for wasi-nn as opposed to something else graph-related (not that we have much else like that right now)
This change adds a way to retrieve preloaded ML models (i.e., "graphs" in wasi-nn terms) from a registry. The wasi-nn specification includes a new function, `load_by_name`, that can be used to access these models more efficiently than before; previously, a user's only option was to read/download/etc. all of the bytes of an ML model and pass them to the `load` function. [named models]: WebAssembly/wasi-nn#36 In Wasmtime's implementation of wasi-nn, we call the registry that holds the models a `GraphRegistry`. We include a simplistic `InMemoryRegistry` for use in the Wasmtime CLI (more on this later) but the idea is that production use will involve some more complex caching and thus a new implementation of a registry--a `Box<dyn GraphRegistry>`--passed into the wasi-nn context. Note that, because we now must be able to `clone` a graph out of the registry and into the "used graphs" table, the OpenVINO `BackendGraph` is updated to be easier to copy around. To allow experimentation with this "preload a named model" functionality, this change also adds a new Wasmtime CLI flag: `--graph <encoding>:<host dir>`. Wasmtime CLI users can now preload a model from a directory; the directory `basename` is used as the model name. Loading models from a directory is probably not desired in Wasmtime embeddings so it is cordoned off into a separate `BackendFromDir` extension trait.
Add a new example crate which loads a model by name and performs image classification. It uses the same MobileNet model as the existing test but a new version of the Rust bindings. The new crate is built and run with the new CLI flag in the `ci/run-wasi-nn-example.sh` script. prtest:full
* wasi-nn: add [named models] This change adds a way to retrieve preloaded ML models (i.e., "graphs" in wasi-nn terms) from a registry. The wasi-nn specification includes a new function, `load_by_name`, that can be used to access these models more efficiently than before; previously, a user's only option was to read/download/etc. all of the bytes of an ML model and pass them to the `load` function. [named models]: WebAssembly/wasi-nn#36 In Wasmtime's implementation of wasi-nn, we call the registry that holds the models a `GraphRegistry`. We include a simplistic `InMemoryRegistry` for use in the Wasmtime CLI (more on this later) but the idea is that production use will involve some more complex caching and thus a new implementation of a registry--a `Box<dyn GraphRegistry>`--passed into the wasi-nn context. Note that, because we now must be able to `clone` a graph out of the registry and into the "used graphs" table, the OpenVINO `BackendGraph` is updated to be easier to copy around. To allow experimentation with this "preload a named model" functionality, this change also adds a new Wasmtime CLI flag: `--graph <encoding>:<host dir>`. Wasmtime CLI users can now preload a model from a directory; the directory `basename` is used as the model name. Loading models from a directory is probably not desired in Wasmtime embeddings so it is cordoned off into a separate `BackendFromDir` extension trait. * wasi-nn: add "named model" test Add a new example crate which loads a model by name and performs image classification. It uses the same MobileNet model as the existing test but a new version of the Rust bindings. The new crate is built and run with the new CLI flag in the `ci/run-wasi-nn-example.sh` script. prtest:full * review: rename `--graph` to `--wasi-nn-graph`
This implements named models in Wasmtime; see the commit messages for more details.