wasmtime_wasi_nn/registry/
in_memory.rs

1//! Implement a [`GraphRegistry`] with a hash map.
2
3use super::{Graph, GraphRegistry};
4use crate::backend::BackendFromDir;
5use crate::wit::ExecutionTarget;
6use anyhow::{anyhow, bail};
7use std::{collections::HashMap, path::Path};
8
9pub struct InMemoryRegistry(HashMap<String, Graph>);
10impl InMemoryRegistry {
11    pub fn new() -> Self {
12        Self(HashMap::new())
13    }
14
15    /// Load a graph from the files contained in the `path` directory.
16    ///
17    /// This expects the backend to know how to load graphs (i.e., ML model)
18    /// from a directory. The name used in the registry is the directory's last
19    /// suffix: if the backend can find the files it expects in `/my/model/foo`,
20    /// the registry will contain a new graph named `foo`.
21    pub fn load(&mut self, backend: &mut dyn BackendFromDir, path: &Path) -> anyhow::Result<()> {
22        if !path.is_dir() {
23            bail!(
24                "preload directory is not a valid directory: {}",
25                path.display()
26            );
27        }
28        let name = path
29            .file_name()
30            .map(|s| s.to_string_lossy())
31            .ok_or(anyhow!("no file name in path"))?;
32
33        let graph = backend.load_from_dir(path, ExecutionTarget::Cpu)?;
34        self.0.insert(name.into_owned(), graph);
35        Ok(())
36    }
37}
38
39impl GraphRegistry for InMemoryRegistry {
40    fn get(&self, name: &str) -> Option<&Graph> {
41        self.0.get(name)
42    }
43    fn get_mut(&mut self, name: &str) -> Option<&mut Graph> {
44        self.0.get_mut(name)
45    }
46}
wasmtime_wasi_nn/registry/in_memory.rs

wasmtime_wasi_nn/registry/
in_memory.rs
