feat(wasi-nn): process-wide graph registry, single-flight loads, in-place model updates

crates/http-service/src/executor/http.rs

@@ -351,6 +351,10 @@ mod tests {
             &self.engine
         }
 
+        fn graph_registry(&self) -> runtime::CachedGraphRegistry {
+            runtime::CachedGraphRegistry::new()
+        }
+
         fn make_secret_store(&self, _secrets: &Vec<SecretOption>) -> anyhow::Result<SecretStore> {
             todo!()
         }

crates/http-service/src/lib.rs

@@ -223,7 +223,12 @@ where
         linker.allow_shadowing(true);
         wasmtime_wasi_http::p2::add_to_linker_async(linker)?;
         wasmtime_wasi_nn::wit::add_to_linker(linker, |data: &mut runtime::Data<_>| {
-            WasiNnView::new(&mut data.table, &mut data.wasi_nn)
+            // Executors preload (`StoreBuilder::preload_wasi_nn`), so this is set;
+            // a store built without it (tests) gets an empty context on first use.
+            let wasi_nn = data
+                .wasi_nn
+                .get_or_insert_with(|| runtime::CachedGraphRegistry::ctx_for(&[]));
+            WasiNnView::new(&mut data.table, wasi_nn)
         })?;
         Ok(())
     }

crates/runtime/Cargo.toml

@@ -38,6 +38,7 @@ chrono = { version = "0.4", features = ["serde"] }
 [dev-dependencies]
 claims = "0.8"
 wat = "1"
+tempfile = "3"
 
 [lints]
 workspace = true

crates/runtime/src/lib.rs

@@ -40,6 +40,7 @@ pub mod stub;
 pub mod trace;
 pub mod util;
 
+pub use registry::CachedGraphRegistry;
 pub use trace::new_traceparent;
 
 use crate::access_log::{AccessLogSender, NoopAccessLogSender};
@@ -147,7 +148,10 @@ impl<T> HttpHooks<T> {
 pub struct Data<T: 'static> {
     hooks: HttpHooks<T>,
     wasi: Wasi,
-    pub wasi_nn: WasiNnCtx,
+    /// `None` for a store whose builder was never `preload_wasi_nn`ed — the
+    /// proxy-wasm engine does not link wasi-nn, so its stores skip the context
+    /// (and its allocations) entirely.
+    pub wasi_nn: Option<WasiNnCtx>,
     // memory usage limiter
     store_limits: ProxyLimiter,
     pub timeout: u64,
@@ -583,6 +587,9 @@ pub struct WasmEngine<T: 'static> {
     inner: Engine,
     component_linker: ComponentLinker<T>,
     module_linker: ModuleLinker<T>,
+    /// One wasi-nn graph cache for all stores of this engine, so a model is
+    /// loaded once per process rather than once per executor (see `registry`).
+    graph_registry: crate::registry::CachedGraphRegistry,
 }
 
 // Manual impl: `derive(Clone)` would incorrectly require `T: Clone`, but the
@@ -595,6 +602,7 @@ impl<T> Clone for WasmEngine<T> {
             inner: self.inner.clone(),
             component_linker: self.component_linker.clone(),
             module_linker: self.module_linker.clone(),
+            graph_registry: self.graph_registry.clone(),
         }
     }
 }
@@ -606,6 +614,7 @@ pub struct WasmEngineBuilder<T: 'static> {
     engine: Engine,
     component_linker: ComponentLinker<T>,
     module_linker: ModuleLinker<T>,
+    graph_registry: crate::registry::CachedGraphRegistry,
 }
 
 impl<T: Send + Sync> WasmEngine<T> {
@@ -615,7 +624,7 @@ impl<T: Send + Sync> WasmEngine<T> {
     }
 
     pub fn store_builder(&self, version: WasiVersion) -> StoreBuilder {
-        StoreBuilder::new(self.inner.clone(), version)
+        StoreBuilder::new(self.inner.clone(), version, self.graph_registry.clone())
     }
 
     /// Creates a new [`InstancePre`] for the given [`Component`].
@@ -641,9 +650,18 @@ impl<T: Send + Sync> WasmEngineBuilder<T> {
             engine: Engine::clone(engine),
             component_linker,
             module_linker,
+            graph_registry: crate::registry::CachedGraphRegistry::new(),
         })
     }
 
+    /// Shares `registry` with this engine instead of the private one it starts
+    /// with. Every engine of a process must get the same registry, or each one
+    /// loads its own copy of a model (see [`ContextT::graph_registry`]).
+    pub fn graph_registry(&mut self, registry: crate::registry::CachedGraphRegistry) -> &mut Self {
+        self.graph_registry = registry;
+        self
+    }
+
     pub fn component_linker_ref(&mut self) -> &mut ComponentLinker<T> {
         &mut self.component_linker
     }
@@ -658,6 +676,7 @@ impl<T: Send + Sync> WasmEngineBuilder<T> {
             inner: self.engine,
             component_linker: self.component_linker,
             module_linker: self.module_linker,
+            graph_registry: self.graph_registry,
         }
     }
 }
@@ -686,6 +705,12 @@ pub trait ContextT {
 
     fn engine_ref(&self) -> &Engine;
 
+    /// The process-wide wasi-nn graph cache. Every engine built from this
+    /// context (one per service) shares it, so a model is loaded once per
+    /// process rather than once per service. Return a clone of one long-lived
+    /// instance, never a fresh registry per call.
+    fn graph_registry(&self) -> crate::registry::CachedGraphRegistry;
+
     fn make_secret_store(&self, secrets: &Vec<SecretOption>) -> anyhow::Result<SecretStore>;
 
     fn make_key_value_store(&self, stores: &Vec<KvStoreOption>) -> key_value_store::Builder;