From 955014305e1a97385e6549a1123e618ed6a8fcb3 Mon Sep 17 00:00:00 2001 From: Ruslan Pislari Date: Mon, 12 Oct 2026 14:26:01 +0300 Subject: [PATCH] feat(wasi-nn): process-wide graph registry, single-flight loads, in-place model updates --- Cargo.lock | 1 + crates/http-service/src/executor/http.rs | 4 + crates/http-service/src/lib.rs | 7 +- crates/runtime/Cargo.toml | 1 + crates/runtime/src/lib.rs | 29 +- crates/runtime/src/registry.rs | 592 ++++++++++++++++++++--- crates/runtime/src/service.rs | 1 + crates/runtime/src/store.rs | 147 ++++-- src/context.rs | 40 +- src/main.rs | 5 +- 10 files changed, 724 insertions(+), 103 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fa5f7b2..aafea3f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2844,6 +2844,7 @@ dependencies = [ "serde", "serde_json", "smol_str", + "tempfile", "tokio", "tokio-util", "tracing", diff --git a/crates/http-service/src/executor/http.rs b/crates/http-service/src/executor/http.rs index 9f94ab2..6997297 100644 --- a/crates/http-service/src/executor/http.rs +++ b/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) -> anyhow::Result { todo!() } diff --git a/crates/http-service/src/lib.rs b/crates/http-service/src/lib.rs index e6093c4..e717d21 100644 --- a/crates/http-service/src/lib.rs +++ b/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(()) } diff --git a/crates/runtime/Cargo.toml b/crates/runtime/Cargo.toml index 6bb6bb2..89d404f 100644 --- a/crates/runtime/Cargo.toml +++ b/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 diff --git a/crates/runtime/src/lib.rs b/crates/runtime/src/lib.rs index 8e7ff6e..0abd979 100644 --- a/crates/runtime/src/lib.rs +++ b/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 HttpHooks { pub struct Data { hooks: HttpHooks, 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, // memory usage limiter store_limits: ProxyLimiter, pub timeout: u64, @@ -583,6 +587,9 @@ pub struct WasmEngine { inner: Engine, component_linker: ComponentLinker, module_linker: ModuleLinker, + /// 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 Clone for WasmEngine { 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 { engine: Engine, component_linker: ComponentLinker, module_linker: ModuleLinker, + graph_registry: crate::registry::CachedGraphRegistry, } impl WasmEngine { @@ -615,7 +624,7 @@ impl WasmEngine { } 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 WasmEngineBuilder { 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 { &mut self.component_linker } @@ -658,6 +676,7 @@ impl WasmEngineBuilder { 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) -> anyhow::Result; fn make_key_value_store(&self, stores: &Vec) -> key_value_store::Builder; diff --git a/crates/runtime/src/registry.rs b/crates/runtime/src/registry.rs index 38c9c81..6ebf770 100644 --- a/crates/runtime/src/registry.rs +++ b/crates/runtime/src/registry.rs @@ -1,22 +1,168 @@ +//! wasi-nn graph (ML model) loading and caching. +//! +//! Graphs are big — a production model weighs ~2.3 GiB resident — and slow to +//! load (seconds: the OpenVINO plugin reads and compiles the network). So they +//! are loaded once per process in [`CachedGraphRegistry`], owned by the +//! embedder's context (`ContextT::graph_registry`) and shared by every +//! `WasmEngine` built from it and every `StoreBuilder` of those, and an executor resolves its app's graphs +//! when it is created (`StoreBuilder::preload_wasi_nn`, on the blocking pool), +//! not on the request path. Each `Store` then only clones `Graph` handles +//! (`Arc`s) into its `WasiNnCtx`. +//! +//! Which graphs an app gets is driven by its env: `FASTEDGE_NN_BACKENDS` lists +//! the backends (`openvino`, `candle`), and every key starting with +//! `FASTEDGE_NN_` names a directory the backend can load a graph from. +//! The guest refers to the graph by that directory's last path component. + use std::collections::HashMap; -use std::ops::Deref; -use std::path::Path; -use std::time::Duration; +use std::hash::{Hash, Hasher}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::{Duration, Instant}; use anyhow::{Context, Result, anyhow, bail}; +use moka::ops::compute::Op; use moka::sync::Cache; use wasmtime_wasi_nn::backend::candle::CandleBackend; +use wasmtime_wasi_nn::wit::WasiNnCtx; use wasmtime_wasi_nn::wit::types::GraphEncoding; use wasmtime_wasi_nn::{ GraphRegistry, Registry, - backend::{BackendFromDir, openvino::OpenvinoBackend}, + backend::openvino::OpenvinoBackend, wit::types::ExecutionTarget, {Backend, Graph}, }; +/// How long a graph stays loaded with no `Store` using it. +/// +/// Long on purpose: dropping a graph only pays off for a model nobody asks for +/// any more (a replaced one), while expiring one that is merely quiet costs a +/// multi-second reload and, briefly, a second resident copy next to the handles +/// idle executors still hold. Hours, not minutes. +pub const GRAPH_TIME_TO_IDLE: Duration = Duration::from_secs(4 * 60 * 60); + +/// `(backend kind, graph directory)` pairs an app asks to have loaded. +pub type PreloadSpec = Vec<(String, String)>; + +/// What the process-wide cache is keyed by: the backend and the graph's +/// canonical directory. Two apps pointing at `/tenant-a/model` and +/// `/tenant-b/model` are two graphs; the name the guest sees (`model`) is a +/// per-store matter ([`CachedGraphRegistry::registry_for`]), never the key. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct GraphKey { + kind: String, + path: PathBuf, +} + +impl GraphKey { + fn new(kind: &str, path: &str) -> Result { + // Canonical, so `/models/current -> /models/v2` and `/models/v2` are one + // graph, and a re-pointed symlink is a new one. + let path = std::fs::canonicalize(path) + .map_err(|_| anyhow!("preload directory is not a valid directory: {}", path))?; + Ok(Self { + kind: kind.to_ascii_lowercase(), + path, + }) + } +} + +/// Identity of a graph directory's *contents* at preload time: name, size and +/// mtime of every regular file in it. A model replaced in place (rsync, cp) +/// changes it, so the next executor build loads the new files instead of +/// serving the cached graph until the process restarts. One `stat` per file; +/// no hashing of multi-GiB weights. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct Revision(u64); + +/// What one scan of a graph directory tells: its [`Revision`] and the bytes +/// on disk (for the load log; the resident size is larger, see `rss_bytes`). +#[derive(Clone, Copy, Debug)] +struct DirFiles { + revision: Revision, + bytes: u64, +} + +impl DirFiles { + fn scan(dir: &Path) -> Result { + let mut entries = std::fs::read_dir(dir) + .and_then(|it| it.collect::>>()) + .with_context(|| format!("read graph directory {}", dir.display()))?; + entries.sort_by_key(|e| e.file_name()); + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + let mut bytes = 0; + for entry in entries { + let meta = entry + .metadata() + .with_context(|| format!("stat {}", entry.path().display()))?; + if meta.is_file() { + entry.file_name().hash(&mut hasher); + meta.len().hash(&mut hasher); + meta.modified().ok().hash(&mut hasher); + bytes += meta.len(); + } + } + Ok(Self { + revision: Revision(hasher.finish()), + bytes, + }) + } +} + +/// This process's resident set, for the before/after of a model load. Linux +/// only (`/proc/self/statm`); `None` elsewhere or on any read error. +fn rss_bytes() -> Option { + #[cfg(target_os = "linux")] + { + let statm = std::fs::read_to_string("/proc/self/statm").ok()?; + let resident_pages: u64 = statm.split_whitespace().nth(1)?.parse().ok()?; + Some(resident_pages * 4096) + } + #[cfg(not(target_os = "linux"))] + { + None + } +} + +/// A cache entry: the graph plus the revision of the files it was loaded from. +#[derive(Clone)] +struct CachedGraph { + revision: Revision, + graph: Graph, +} + +/// One graph as resolved for a store builder: the shared cache entry it came +/// from, the name the guest loads it by, and the handle. +#[derive(Clone)] +pub struct LoadedGraph { + pub key: GraphKey, + pub name: String, + pub graph: Graph, +} + +/// Process-wide cache of loaded graphs, keyed by [`GraphKey`]. +/// +/// Cheap to clone (the cache is reference-counted); one instance lives on the +/// embedder's context, is shared by every `WasmEngine` built from it, and is +/// handed to every `StoreBuilder` those engines create. +/// +/// Entries are strong: an executor rebuild (config change, debug toggle) finds +/// the graph here and never reloads it. A graph is dropped after +/// [`GRAPH_TIME_TO_IDLE`] without *traffic* — executors only read the cache at +/// preload, so every `Store` creation refreshes its graphs' idle clocks through +/// [`CachedGraphRegistry::mark_used`]. Without that, a model under constant +/// load would expire on schedule and the next executor rebuild would load a +/// second ~2.3 GiB copy next to the one live executors still hold. #[derive(Clone)] -pub struct CachedGraphRegistry(Cache); +pub struct CachedGraphRegistry(Cache); +impl Default for CachedGraphRegistry { + fn default() -> Self { + Self::new() + } +} + +/// The graphs one `Store` can see, by name. pub struct StoreRegistry(HashMap); pub fn backends() -> Vec { @@ -36,67 +182,403 @@ impl GraphRegistry for StoreRegistry { } } +/// Reads an app's graph preload requests from its env. +/// +/// `FASTEDGE_NN_BACKENDS=openvino,candle` selects the backends; for each, every +/// env key starting with `FASTEDGE_NN_` (backend upper-cased) holds a +/// graph directory. Order follows the backend list, then the env order. +pub fn preload_spec, V: AsRef>(env: &[(K, V)]) -> PreloadSpec { + let Some((_, backends)) = env + .iter() + .find(|(k, _)| k.as_ref() == "FASTEDGE_NN_BACKENDS") + else { + return Vec::new(); + }; + + backends + .as_ref() + .split(',') + .map(str::trim) + .filter(|b| !b.is_empty()) + .flat_map(|backend| { + let prefix = format!("FASTEDGE_NN_{}", backend.to_ascii_uppercase()); + env.iter().filter_map(move |(k, v)| { + k.as_ref() + .starts_with(&prefix) + .then(|| (backend.to_string(), v.as_ref().to_string())) + }) + }) + .collect() +} + +/// The name a graph is registered under: its directory's last path component. +fn graph_name(path: &str) -> Result { + Path::new(path) + .file_name() + .map(|s| s.to_string_lossy().into_owned()) + .ok_or_else(|| anyhow!("no file name in graph path: {}", path)) +} + +/// Loads one graph from `path` with the backend `kind`. Slow and memory-hungry; +/// only ever called through the cache. +fn load_graph(kind: &str, path: &str) -> Result { + let dir = Path::new(path); + if !dir.is_dir() { + bail!("preload directory is not a valid directory: {}", path); + } + let encoding = kind.parse().unwrap_or(GraphEncoding::Autodetect); + let mut backends = backends(); + let backend = backends + .iter_mut() + .find(|b| b.encoding() == encoding) + .ok_or_else(|| anyhow!("unsupported backend: {}", kind))? + .as_dir_loadable() + .ok_or_else(|| anyhow!("{} does not support directory loading", kind))?; + Ok(backend.load_from_dir(dir, ExecutionTarget::Cpu)?) +} + impl CachedGraphRegistry { pub fn new() -> Self { + Self::with_time_to_idle(GRAPH_TIME_TO_IDLE) + } + + fn with_time_to_idle(time_to_idle: Duration) -> Self { let builder = Cache::builder() // Max 100 entries .max_capacity(100) - // Time to idle (TTI): 30 minutes - .time_to_idle(Duration::from_secs(30 * 60)); + .time_to_idle(time_to_idle) + // Models are few and each is gigabytes: say when one leaves the + // cache and why. Memory is only freed once executors built on it + // drop their handles. + .eviction_listener(|key: Arc, _: CachedGraph, cause| { + tracing::info!( + backend = %key.kind, + path = %key.path.display(), + ?cause, + "wasi-nn graph dropped from the process cache" + ); + }); Self(builder.build()) } - /// Load a graph from the files contained in the `path` directory. + /// Records that `graphs` are in use, so their entries do not expire while + /// requests keep coming. Called on every `Store` creation; a Moka read is + /// a hash lookup plus a buffered access note, nothing on the order of the + /// store itself. + pub fn mark_used(&self, graphs: &[LoadedGraph]) { + for g in graphs { + let _ = self.0.get(&g.key); + } + } + + /// Resolves every graph in `spec`, loading the ones not cached yet — or + /// cached from files that have since changed ([`Revision`]). /// - /// This expects the backend to know how to load graphs (i.e., ML model) - /// from a directory. The name used in the registry is the directory's last - /// suffix: if the backend can find the files it expects in `/my/model/foo`, - /// the registry will contain a new graph named `foo`. - fn load(&self, backend: &mut dyn BackendFromDir, path: impl AsRef) -> Result { - let path = path.as_ref(); - if !path.is_dir() { - bail!( - "preload directory is not a valid directory: {}", - path.display() - ); + /// The per-key compute is atomic: concurrent requests for the same graph + /// wait for a single load instead of each running one (a cold graph used + /// to be loaded once per concurrent miss, each copy costing the full model + /// footprint), and a stale entry is replaced in place, so the old graph is + /// freed as soon as the executors built on it are gone. The handles + /// returned keep a graph alive for the executor's lifetime; the cache + /// entry's own lifetime is a matter of traffic (see [`Self::mark_used`]). + pub fn load_graphs(&self, spec: &[(String, String)]) -> Result> { + let mut graphs: Vec = Vec::with_capacity(spec.len()); + for (kind, path) in spec { + let key = GraphKey::new(kind, path)?; + let name = graph_name(path)?; + // The guest addresses graphs by name, so within one app two + // directories with the same basename would shadow each other. + if let Some(dup) = graphs.iter().find(|g| g.name == name) { + bail!( + "wasi-nn graph name {:?} used by both {} and {}", + name, + dup.key.path.display(), + path + ); + } + let files = DirFiles::scan(&key.path)?; + let revision = files.revision; + let graph = self + .0 + .entry(key.clone()) + .and_try_compute_with(|current| match current { + Some(entry) if entry.value().revision == revision => { + tracing::debug!(backend = %kind, path = %path, name = %name, + "wasi-nn graph served from the process cache"); + Ok(Op::Nop) + } + current => { + let replacing = current.is_some(); + let rss_before = rss_bytes(); + let started = Instant::now(); + let graph = load_graph(kind, path)?; + tracing::info!( + backend = %kind, + path = %path, + name = %name, + file_bytes = files.bytes, + rss_delta_bytes = rss_before + .zip(rss_bytes()) + .map(|(before, after)| after.saturating_sub(before)), + elapsed_ms = started.elapsed().as_millis() as u64, + replacing, + "wasi-nn graph loaded into the process cache" + ); + Ok(Op::Put(CachedGraph { revision, graph })) + } + }) + .map_err(|error: anyhow::Error| { + anyhow!("wasi-nn backend load {}: {:#}", path, error) + })? + .into_entry() + .expect("Nop on a present entry or Put: the entry is there") + .into_value() + .graph; + graphs.push(LoadedGraph { key, name, graph }); } - let name = path - .file_name() - .map(|s| s.to_string_lossy()) - .ok_or(anyhow!("no file name in path"))?; - - let graph = backend.load_from_dir(path, ExecutionTarget::Cpu)?; - self.0.insert(name.into_owned(), graph.clone()); - Ok(graph) - } - - pub fn preload_graphs( - &self, - preload: Vec<(String, String)>, - ) -> Result<(impl IntoIterator, Registry)> { - let mut backends = backends(); - let mut registry = StoreRegistry(HashMap::new()); - for (kind, path) in preload { - let name = Path::new(&path) - .file_name() - .map(|s| s.to_string_lossy()) - .ok_or(anyhow!("no file name in path"))?; - let graph = if let Some(graph) = self.0.get(name.deref()) { - graph - } else { - let kind_ = kind.parse().unwrap_or(GraphEncoding::Autodetect); - let backend = backends - .iter_mut() - .find(|b| b.encoding() == kind_) - .ok_or(anyhow!("unsupported backend: {}", kind))? - .as_dir_loadable() - .ok_or(anyhow!("{} does not support directory loading", kind))?; - self.load(backend, &path) - .with_context(|| format!("wasi-nn backend load {}", path))? - }; - registry.0.insert(name.into_owned(), graph); + Ok(graphs) + } + + /// Builds a `Store`'s view of already-resolved graphs. Only `Arc` clones. + pub fn registry_for(graphs: &[LoadedGraph]) -> Registry { + Registry::from(StoreRegistry( + graphs + .iter() + .map(|g| (g.name.clone(), g.graph.clone())) + .collect(), + )) + } + + /// A `Store`'s wasi-nn context over already-resolved graphs: the backends + /// plus [`Self::registry_for`]. Only `Arc` clones and a few small allocs. + pub fn ctx_for(graphs: &[LoadedGraph]) -> WasiNnCtx { + WasiNnCtx::new(backends(), Self::registry_for(graphs)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn env(pairs: &[(&str, &str)]) -> Vec<(String, String)> { + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() + } + + #[test] + fn preload_spec_is_empty_without_backends() { + let e = env(&[("FASTEDGE_NN_OPENVINO_MODEL", "/models/a"), ("X", "y")]); + assert!(preload_spec(&e).is_empty()); + } + + #[test] + fn preload_spec_collects_dirs_per_backend_in_order() { + let e = env(&[ + ("FASTEDGE_NN_OPENVINO_MODEL_B", "/models/b"), + ("FASTEDGE_NN_BACKENDS", "openvino, candle"), + ("FASTEDGE_NN_CANDLE_MODEL", "/models/c"), + ("FASTEDGE_NN_OPENVINO_MODEL_A", "/models/a"), + ("UNRELATED", "/models/z"), + ]); + assert_eq!( + preload_spec(&e), + vec![ + ("openvino".to_string(), "/models/b".to_string()), + ("openvino".to_string(), "/models/a".to_string()), + ("candle".to_string(), "/models/c".to_string()), + ] + ); + } + + #[test] + fn preload_spec_matches_backend_prefix_case_insensitively() { + let e = env(&[ + ("FASTEDGE_NN_BACKENDS", "OpenVINO"), + ("FASTEDGE_NN_OPENVINO", "/models/a"), + ]); + assert_eq!( + preload_spec(&e), + vec![("OpenVINO".to_string(), "/models/a".to_string())] + ); + } + + #[test] + fn load_graphs_rejects_a_missing_directory_naming_it() { + let registry = CachedGraphRegistry::new(); + let spec = vec![("openvino".to_string(), "/definitely/not/here".to_string())]; + let error = registry + .load_graphs(&spec) + .err() + .expect("missing dir must fail") + .to_string(); + assert!(error.contains("/definitely/not/here"), "{error}"); + } + + #[test] + fn load_graphs_rejects_an_unsupported_backend() { + // A valid encoding no configured backend provides. (An unparseable kind + // falls back to `Autodetect`, which a backend does accept.) + let dir = std::env::temp_dir(); + let registry = CachedGraphRegistry::new(); + let spec = vec![("onnx".to_string(), dir.to_string_lossy().into_owned())]; + let error = registry + .load_graphs(&spec) + .err() + .expect("unknown backend must fail") + .to_string(); + assert!(error.contains("unsupported backend"), "{error}"); + } + + /// A graph that cannot run; enough to occupy a cache slot. + struct StubGraph; + + impl wasmtime_wasi_nn::backend::BackendGraph for StubGraph { + fn init_execution_context( + &self, + ) -> std::result::Result< + wasmtime_wasi_nn::ExecutionContext, + wasmtime_wasi_nn::backend::BackendError, + > { + Err(wasmtime_wasi_nn::backend::BackendError::NoInputs) } - Ok((backends, Registry::from(registry))) + } + + fn stub_graph() -> Graph { + Graph::from(Box::new(StubGraph) as Box) + } + + fn cached(graph: Graph) -> CachedGraph { + CachedGraph { + revision: Revision(0), + graph, + } + } + + /// Replacing a model's files in place must be seen by the next executor + /// build; touching nothing must not (or every rebuild would reload). + #[test] + fn revision_follows_the_files_contents_not_the_directory() { + let dir = tempfile::tempdir().unwrap(); + let bin = dir.path().join("model.bin"); + std::fs::write(&bin, b"v1").unwrap(); + std::fs::write(dir.path().join("model.xml"), b"").unwrap(); + let first = DirFiles::scan(dir.path()).unwrap().revision; + assert_eq!( + first, + DirFiles::scan(dir.path()).unwrap().revision, + "stable when untouched" + ); + + std::fs::write(&bin, b"v2-longer").unwrap(); + assert_ne!( + first, + DirFiles::scan(dir.path()).unwrap().revision, + "size change seen" + ); + + std::fs::write(dir.path().join("extra.bin"), b"x").unwrap(); + let with_extra = DirFiles::scan(dir.path()).unwrap().revision; + std::fs::remove_file(dir.path().join("extra.bin")).unwrap(); + assert_ne!( + with_extra, + DirFiles::scan(dir.path()).unwrap().revision, + "added/removed file seen" + ); + } + + /// The idle clock must follow store creations, not preloads: a model that + /// serves requests continuously stays cached past its idle timeout, and + /// one that stops being used is dropped. + #[test] + fn mark_used_keeps_a_busy_graph_cached_until_traffic_stops() { + let tti = Duration::from_millis(300); + let registry = CachedGraphRegistry::with_time_to_idle(tti); + let key = GraphKey { + kind: "stub".into(), + path: "/models/m".into(), + }; + registry.0.insert(key.clone(), cached(stub_graph())); + let held = vec![LoadedGraph { + key: key.clone(), + name: "m".into(), + graph: registry.0.get(&key).unwrap().graph, + }]; + + // Busy for 3x the idle timeout, touched every third of it. + for _ in 0..9 { + std::thread::sleep(tti / 3); + registry.mark_used(&held); + } + assert!(registry.0.get(&key).is_some(), "busy graph expired"); + + // Idle for 2x the timeout; still held by `held`, like an idle executor + // would, which must not keep the cache entry alive. + std::thread::sleep(tti * 2); + assert!(registry.0.get(&key).is_none(), "idle graph still cached"); + } + + /// Same basename, different directories: two cache entries, each app gets + /// its own graph under the name `model` — and a different backend for the + /// same directory is a third. + #[test] + fn graph_key_separates_same_basename_under_different_paths_and_backends() { + let root = tempfile::tempdir().unwrap(); + let a = root.path().join("tenant-a").join("model"); + let b = root.path().join("tenant-b").join("model"); + std::fs::create_dir_all(&a).unwrap(); + std::fs::create_dir_all(&b).unwrap(); + let (a, b) = (a.to_str().unwrap(), b.to_str().unwrap()); + + let ka = GraphKey::new("openvino", a).unwrap(); + let kb = GraphKey::new("openvino", b).unwrap(); + let kc = GraphKey::new("candle", a).unwrap(); + assert_ne!(ka, kb); + assert_ne!(ka, kc); + assert_eq!( + ka, + GraphKey::new("OpenVINO", a).unwrap(), + "kind is case-insensitive" + ); + assert_eq!(graph_name(a).unwrap(), "model"); + assert_eq!(graph_name(b).unwrap(), "model"); + } + + /// Within one app two directories with the same basename cannot both be + /// addressed by the guest, so that is a config error, not a silent shadow. + #[test] + fn load_graphs_rejects_duplicate_names_within_an_app() { + let root = tempfile::tempdir().unwrap(); + let a = root.path().join("a").join("model"); + let b = root.path().join("b").join("model"); + std::fs::create_dir_all(&a).unwrap(); + std::fs::create_dir_all(&b).unwrap(); + let registry = CachedGraphRegistry::with_time_to_idle(Duration::from_secs(60)); + // Cached at the directory's current revision, so `a` hits and the + // failure is the duplicate name, not a load. + registry.0.insert( + GraphKey::new("openvino", a.to_str().unwrap()).unwrap(), + CachedGraph { + revision: DirFiles::scan(&a).unwrap().revision, + graph: stub_graph(), + }, + ); + let error = registry + .load_graphs(&[ + ("openvino".into(), a.to_str().unwrap().into()), + ("openvino".into(), b.to_str().unwrap().into()), + ]) + .err() + .expect("duplicate name must fail") + .to_string(); + assert!(error.contains("\"model\""), "{error}"); + } + + #[test] + fn registry_for_exposes_graphs_by_name_and_empty_is_fine() { + let registry = CachedGraphRegistry::registry_for(&[]); + assert!(registry.get("anything").is_none()); } } diff --git a/crates/runtime/src/service.rs b/crates/runtime/src/service.rs index 9a3772f..5b0fb74 100644 --- a/crates/runtime/src/service.rs +++ b/crates/runtime/src/service.rs @@ -43,6 +43,7 @@ where pub fn build(self) -> anyhow::Result { let engine = { let mut builder = WasmEngine::builder(self.context.engine_ref())?; + builder.graph_registry(self.context.graph_registry()); S::configure_engine(&mut builder)?; builder.build() }; diff --git a/crates/runtime/src/store.rs b/crates/runtime/src/store.rs index 162ae55..0c87536 100644 --- a/crates/runtime/src/store.rs +++ b/crates/runtime/src/store.rs @@ -1,6 +1,6 @@ use crate::limiter::ProxyLimiter; use crate::logger::Logger; -use crate::registry::CachedGraphRegistry; +use crate::registry::{CachedGraphRegistry, LoadedGraph}; use crate::util::stats::StatsVisitor; use crate::{DEFAULT_EPOCH_TICK_INTERVAL, Data, HttpHooks, Wasi, WasiVersion}; use anyhow::Result; @@ -126,7 +126,13 @@ pub struct StoreBuilder { logger: Option, version: WasiVersion, properties: HashMap, + /// Process-wide graph cache, shared with every other builder of the engine. registry: CachedGraphRegistry, + /// The app's wasi-nn graphs, resolved once by [`Self::preload_wasi_nn`] so + /// building a `Store` only clones handles. `None` until preloaded, and a + /// store built from a never-preloaded builder gets no wasi-nn context at + /// all: preloading is the opt-in (proxy-wasm does not link wasi-nn). + nn_graphs: Option>, secret_store: SecretStore, key_value_store: key_value_store::Builder, dictionary: Dictionary, @@ -138,7 +144,11 @@ pub struct StoreBuilder { impl StoreBuilder { // Called by Engine::store_builder. - pub(crate) fn new(engine: wasmtime::Engine, version: WasiVersion) -> Self { + pub(crate) fn new( + engine: wasmtime::Engine, + version: WasiVersion, + registry: CachedGraphRegistry, + ) -> Self { Self { engine, max_duration: 100, @@ -147,7 +157,8 @@ impl StoreBuilder { logger: None, version, properties: Default::default(), - registry: CachedGraphRegistry::new(), + registry, + nn_graphs: None, secret_store: Default::default(), key_value_store: key_value_store::Builder::default(), dictionary: Default::default(), @@ -268,35 +279,40 @@ impl StoreBuilder { } } - pub fn make_wasi_nn(&self) -> Result { - // initialize application specific graph - let backends: Vec<&str> = self - .env - .iter() - .find(|(k, _)| k == "FASTEDGE_NN_BACKENDS") - .map(|(_, v)| v.split(",").collect()) - .unwrap_or_default(); - - debug!("wasi-nn backends: {:?}", backends); - - let preload_graphs: Vec<(String, String)> = backends - .into_iter() - .flat_map(|b| { - let entry = self.env.iter().filter_map(|(k, v)| { - if k.starts_with(&format!("FASTEDGE_NN_{}", b.to_ascii_uppercase())) { - Some((b.to_string(), v.to_string())) - } else { - None - } - }); - entry - }) - .collect(); - debug!("preload graphs: {:?}", preload_graphs); + /// Resolves the app's wasi-nn graphs (from `FASTEDGE_NN_*` env, see + /// [`crate::registry`]) now, so that every `Store` built later gets them + /// for the price of an `Arc` clone. + /// + /// Call this where the executor is created — a cache miss loads a model, + /// which takes seconds and gigabytes, and that belongs on the blocking + /// pool, not on the request path. Loads are coalesced process-wide, so + /// concurrent executor creations for the same app share one load. + pub fn preload_wasi_nn(self) -> Result { + let spec = crate::registry::preload_spec(&self.env); + debug!("preload graphs: {:?}", spec); + let nn_graphs = self.registry.load_graphs(&spec)?; + Ok(Self { + nn_graphs: Some(nn_graphs), + ..self + }) + } - let (backends, registry) = self.registry.preload_graphs(preload_graphs)?; + /// Whether the app's env asks for wasi-nn graphs (`FASTEDGE_NN_*`). + /// Lets an executor factory whose engine lacks wasi-nn flag the misconfig + /// instead of silently ignoring it. + pub fn has_wasi_nn_env(&self) -> bool { + !crate::registry::preload_spec(&self.env).is_empty() + } - Ok(WasiNnCtx::new(backends, registry)) + /// The store's wasi-nn context, or `None` for a builder that was never + /// [`Self::preload_wasi_nn`]ed — that is the opt-in, and a store without it + /// pays nothing here (no backends, no registry, no allocations). + pub fn make_wasi_nn(&self) -> Option { + let graphs = self.nn_graphs.as_ref()?; + // Keep the process-wide entries alive while this app has traffic; + // nothing else reads them between executor rebuilds. + self.registry.mark_used(graphs); + Some(CachedGraphRegistry::ctx_for(graphs)) } /// Builds a [`Store`] from this builder with `Default` WasiCtxBuilder @@ -321,7 +337,7 @@ impl StoreBuilder { let table = ResourceTable::new(); wasi_ctx_builder.envs(&self.env); - let wasi_nn = self.make_wasi_nn()?; + let wasi_nn = self.make_wasi_nn(); let logger = if let Some(mut logger) = self.logger { logger.extend(self.properties); @@ -651,16 +667,20 @@ mod tests { let engine = make_engine(); let before = instances::live(); - let store = StoreBuilder::new(engine.clone(), WasiVersion::Preview1) - .build(NoStats) - .expect("build store"); + let store = StoreBuilder::new( + engine.clone(), + WasiVersion::Preview1, + CachedGraphRegistry::new(), + ) + .build(NoStats) + .expect("build store"); assert_eq!( instances::live(), before + 1, "building a store must count a live instance" ); - let second = StoreBuilder::new(engine, WasiVersion::Preview1) + let second = StoreBuilder::new(engine, WasiVersion::Preview1, CachedGraphRegistry::new()) .build(NoStats) .expect("build store"); assert_eq!(instances::live(), before + 2); @@ -675,4 +695,61 @@ mod tests { "dropping the store must release the count" ); } + + /// An app without wasi-nn env preloads nothing, and a store's wasi-nn + /// context then comes from the preloaded (empty) set without touching the + /// loader. (`make_wasi_nn` is called directly: building a whole `Store` + /// here would race the live-instance counts asserted by the test above.) + #[test] + fn preload_wasi_nn_without_nn_env_is_a_noop() { + let builder = StoreBuilder::new( + make_engine(), + WasiVersion::Preview1, + CachedGraphRegistry::new(), + ) + .set_env(&[("UNRELATED", "1")]) + .preload_wasi_nn() + .expect("preload"); + assert_eq!(builder.nn_graphs.as_ref().map(Vec::len), Some(0)); + builder + .make_wasi_nn() + .expect("wasi-nn ctx from preloaded graphs"); + } + + /// A builder that never preloaded (proxy-wasm) builds no wasi-nn context, + /// whatever the env says: the engine would not link it anyway. + #[test] + fn make_wasi_nn_is_none_without_preload() { + let builder = StoreBuilder::new( + make_engine(), + WasiVersion::Preview1, + CachedGraphRegistry::new(), + ) + .set_env(&[ + ("FASTEDGE_NN_BACKENDS", "openvino"), + ("FASTEDGE_NN_OPENVINO_MODEL", "/no/such/model-dir"), + ]); + assert!(builder.has_wasi_nn_env()); + assert!(builder.make_wasi_nn().is_none()); + } + + /// A graph the executor cannot load fails executor creation — with the + /// offending directory in the error — instead of every request later on. + #[test] + fn preload_wasi_nn_reports_an_unloadable_graph_at_executor_creation() { + let error = StoreBuilder::new( + make_engine(), + WasiVersion::Preview1, + CachedGraphRegistry::new(), + ) + .set_env(&[ + ("FASTEDGE_NN_BACKENDS", "openvino"), + ("FASTEDGE_NN_OPENVINO_MODEL", "/no/such/model-dir"), + ]) + .preload_wasi_nn() + .err() + .expect("missing model dir must fail") + .to_string(); + assert!(error.contains("/no/such/model-dir"), "{error}"); + } } diff --git a/src/context.rs b/src/context.rs index 1bc6481..3b0dae5 100644 --- a/src/context.rs +++ b/src/context.rs @@ -14,8 +14,8 @@ use runtime::app::{KvStoreOption, SecretOption}; use runtime::logger::{Console, Logger}; use runtime::util::stats::{CdnPhase, StatsVisitor}; use runtime::{ - App, ContextT, ExecutorCache, PreCompiledLoader, Router, WasiVersion, WasmEngine, - componentize_if_necessary, + App, CachedGraphRegistry, ContextT, ExecutorCache, PreCompiledLoader, Router, WasiVersion, + WasmEngine, componentize_if_necessary, }; use secret::SecretStore; use smol_str::SmolStr; @@ -34,9 +34,10 @@ pub struct Context { pub(crate) engine: Engine, pub(crate) app: Option, pub(crate) backend: Backend>, - pub(crate) wasm_bytes: Vec, + pub(crate) wasm_bytes: Arc<[u8]>, pub(crate) wasi_http: bool, pub(crate) cache_backend: Arc, + pub(crate) graph_registry: CachedGraphRegistry, } impl PreCompiledLoader for Context { @@ -69,6 +70,10 @@ impl ContextT for Context { &self.engine } + fn graph_registry(&self) -> CachedGraphRegistry { + self.graph_registry.clone() + } + fn make_secret_store(&self, secrets: &Vec) -> anyhow::Result { let mut secret_impl = SecretImpl::default(); for s in secrets { @@ -101,15 +106,15 @@ impl ContextT for Context { } } -impl ExecutorFactory>> for Context { - type Executor = RunExecutor; - - async fn get_executor( +impl Context { + /// Builds an executor for `app`. Blocking and CPU-bound (component compilation, + /// wasi-nn graph preload), so it must run on the blocking pool. + fn make_executor( &self, name: SmolStr, app: &App, engine: &WasmEngine>>, - ) -> anyhow::Result { + ) -> anyhow::Result { let mut dictionary = Dictionary::new(); for (k, v) in app.env.iter() { dictionary.insert(k.to_string(), v.to_string()); @@ -125,6 +130,7 @@ impl ExecutorFactory>> for Context { let store_builder = engine .store_builder(version) .set_env(&env) + .preload_wasi_nn()? .max_memory_size(app.mem_limit) .max_epoch_ticks(app.max_duration) .logger(logger) @@ -155,6 +161,24 @@ impl ExecutorFactory>> for Context { } } +impl ExecutorFactory>> for Context { + type Executor = RunExecutor; + + async fn get_executor( + &self, + name: SmolStr, + app: &App, + engine: &WasmEngine>>, + ) -> anyhow::Result { + // Building an executor compiles the component and may load a multi-GiB wasi-nn + // graph, so keep it off the async worker (as the server's `get_executor` does). + let this = self.clone(); + let app = app.clone(); + let engine = engine.clone(); + tokio::task::spawn_blocking(move || this.make_executor(name, &app, &engine)).await? + } +} + impl ExecutorCache for Context { async fn remove(&self, _name: &str) { unreachable!() diff --git a/src/main.rs b/src/main.rs index 09a0bfd..cfe49de 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,7 +17,7 @@ use hyper_tls::HttpsConnector; use hyper_util::client::legacy::connect::HttpConnector; use runtime::app::{CacheMode, KvStoreOption, Status}; use runtime::service::{Service, ServiceBuilder}; -use runtime::{App, SecretValue, WasmConfig}; +use runtime::{App, CachedGraphRegistry, SecretValue, WasmConfig}; use smol_str::{SmolStr, ToSmolStr}; use std::collections::HashMap; use std::path::PathBuf; @@ -192,9 +192,10 @@ async fn main() -> anyhow::Result<()> { engine, app: Some(cli_app), backend, - wasm_bytes, + wasm_bytes: wasm_bytes.into(), wasi_http: run.wasi_http.unwrap_or_default(), cache_backend: Arc::new(MemoryCacheBackend::new()), + graph_registry: CachedGraphRegistry::new(), }; let http: HttpService = ServiceBuilder::new(context).build()?;