wasmedge_sdk/
plugin.rs

1//! Defines plugin related structs.
2
3use crate::{instance::Instance, WasmEdgeResult};
4use wasmedge_sys::{self as sys};
5
6#[cfg(feature = "wasi_nn")]
7use wasmedge_types::error::WasmEdgeError;
8
9/// Defines low-level types used in Plugin development.
10pub mod ffi {
11    pub use wasmedge_sys::ffi::{
12        WasmEdge_ModuleDescriptor, WasmEdge_ModuleInstanceContext, WasmEdge_PluginDescriptor,
13    };
14}
15
16/// Preload config for initializing the wasi_nn plug-in.
17#[cfg(feature = "wasi_nn")]
18#[cfg_attr(docsrs, doc(cfg(feature = "wasi_nn")))]
19#[derive(Debug)]
20pub struct NNPreload {
21    /// The alias of the model in the WASI-NN environment.
22    alias: String,
23    /// The inference backend.
24    backend: GraphEncoding,
25    /// The execution target, on which the inference runs.
26    target: ExecutionTarget,
27    /// The path to the model file. Note that the path is the guest path instead of the host path.
28    path: std::path::PathBuf,
29}
30#[cfg(feature = "wasi_nn")]
31#[cfg_attr(docsrs, doc(cfg(feature = "wasi_nn")))]
32impl NNPreload {
33    /// Creates a new preload config.
34    ///
35    /// # Arguments
36    ///
37    /// * `alias` - The alias of the model in the WASI-NN environment.
38    ///
39    /// * `backend` - The inference backend.
40    ///
41    /// * `target` - The execution target, on which the inference runs.
42    ///
43    /// * `path` - The path to the model file. Note that the path is the guest path instead of the host path.
44    ///
45    pub fn new(
46        alias: impl AsRef<str>,
47        backend: GraphEncoding,
48        target: ExecutionTarget,
49        path: impl AsRef<std::path::Path>,
50    ) -> Self {
51        Self {
52            alias: alias.as_ref().to_owned(),
53            backend,
54            target,
55            path: path.as_ref().to_owned(),
56        }
57    }
58}
59#[cfg(feature = "wasi_nn")]
60#[cfg_attr(docsrs, doc(cfg(feature = "wasi_nn")))]
61impl std::fmt::Display for NNPreload {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        write!(
64            f,
65            "{alias}:{backend}:{target}:{path}",
66            alias = self.alias,
67            backend = self.backend,
68            target = self.target,
69            path = self.path.to_string_lossy().into_owned()
70        )
71    }
72}
73#[cfg(feature = "wasi_nn")]
74#[cfg_attr(docsrs, doc(cfg(feature = "wasi_nn")))]
75impl std::str::FromStr for NNPreload {
76    type Err = WasmEdgeError;
77
78    fn from_str(preload: &str) -> std::result::Result<Self, Self::Err> {
79        let nn_preload: Vec<&str> = preload.split(':').collect();
80        if nn_preload.len() != 4 {
81            return Err(WasmEdgeError::Operation(format!(
82                "Failed to convert to NNPreload value. Invalid preload string: {}. The correct format is: 'alias:backend:target:path'",
83                preload
84            )));
85        }
86        let (alias, backend, target, path) = (
87            nn_preload[0].to_string(),
88            nn_preload[1]
89                .parse::<GraphEncoding>()
90                .map_err(|err| WasmEdgeError::Operation(err.to_string()))?,
91            nn_preload[2]
92                .parse::<ExecutionTarget>()
93                .map_err(|err| WasmEdgeError::Operation(err.to_string()))?,
94            std::path::PathBuf::from(nn_preload[3]),
95        );
96
97        Ok(Self::new(alias, backend, target, path))
98    }
99}
100
101#[cfg(feature = "wasi_nn")]
102#[test]
103fn test_generate_nnpreload_from_str() {
104    use std::str::FromStr;
105
106    // valid preload string
107    let preload = "default:GGML:CPU:llama-2-7b-chat.Q5_K_M.gguf";
108    let result = NNPreload::from_str(preload);
109    assert!(result.is_ok());
110    let nnpreload = result.unwrap();
111    assert_eq!(nnpreload.alias, "default");
112    assert_eq!(nnpreload.backend, GraphEncoding::GGML);
113    assert_eq!(nnpreload.target, ExecutionTarget::CPU);
114    assert_eq!(
115        nnpreload.path,
116        std::path::PathBuf::from("llama-2-7b-chat.Q5_K_M.gguf")
117    );
118
119    // invalid preload string
120    let preload = "default:CPU:GGML:llama-2-7b-chat.Q5_K_M.gguf";
121    let result = NNPreload::from_str(preload);
122    assert!(result.is_err());
123    let err = result.unwrap_err();
124    assert_eq!(
125        WasmEdgeError::Operation(
126            "Failed to convert to NNBackend value. Unknown NNBackend type: CPU".to_string()
127        ),
128        err
129    );
130
131    // invalid preload string: unsupported target
132    let preload = "default:GGML:NPU:llama-2-7b-chat.Q5_K_M.gguf";
133    let result = NNPreload::from_str(preload);
134    assert!(result.is_err());
135    let err = result.unwrap_err();
136    assert_eq!(
137        WasmEdgeError::Operation(
138            "Failed to convert to ExecutionTarget value. Unknown ExecutionTarget type: NPU"
139                .to_string()
140        ),
141        err
142    );
143
144    // invalid preload string: invalid format
145    let preload = "default:GGML:CPU";
146    let result = NNPreload::from_str(preload);
147    assert!(result.is_err());
148    let err = result.unwrap_err();
149    assert_eq!(
150        WasmEdgeError::Operation(
151            "Failed to convert to NNPreload value. Invalid preload string: default:GGML:CPU. The correct format is: 'alias:backend:target:path'"
152                .to_string()
153        ),
154        err
155    );
156}
157
158/// Describes the encoding of the graph.
159#[cfg(feature = "wasi_nn")]
160#[cfg_attr(docsrs, doc(cfg(feature = "wasi_nn")))]
161#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
162#[allow(non_camel_case_types)]
163pub enum GraphEncoding {
164    OpenVINO,
165    ONNX,
166    TensorFlow,
167    PyTorch,
168    TensorFlowLite,
169    Autodetect,
170    GGML,
171}
172#[cfg(feature = "wasi_nn")]
173#[cfg_attr(docsrs, doc(cfg(feature = "wasi_nn")))]
174impl std::fmt::Display for GraphEncoding {
175    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
176        match self {
177            GraphEncoding::PyTorch => write!(f, "PyTorch"),
178            GraphEncoding::TensorFlowLite => write!(f, "TensorflowLite"),
179            GraphEncoding::TensorFlow => write!(f, "Tensorflow"),
180            GraphEncoding::OpenVINO => write!(f, "OpenVINO"),
181            GraphEncoding::GGML => write!(f, "GGML"),
182            GraphEncoding::ONNX => write!(f, "ONNX"),
183            GraphEncoding::Autodetect => write!(f, "Autodetect"),
184        }
185    }
186}
187#[cfg(feature = "wasi_nn")]
188#[cfg_attr(docsrs, doc(cfg(feature = "wasi_nn")))]
189impl std::str::FromStr for GraphEncoding {
190    type Err = WasmEdgeError;
191
192    fn from_str(s: &str) -> Result<Self, Self::Err> {
193        match s.to_lowercase().as_str() {
194            "openvino" => Ok(GraphEncoding::OpenVINO),
195            "onnx" => Ok(GraphEncoding::ONNX),
196            "tensorflow" => Ok(GraphEncoding::TensorFlow),
197            "pytorch" => Ok(GraphEncoding::PyTorch),
198            "tensorflowlite" => Ok(GraphEncoding::TensorFlowLite),
199            "autodetect" => Ok(GraphEncoding::Autodetect),
200            "ggml" => Ok(GraphEncoding::GGML),
201            _ => Err(WasmEdgeError::Operation(format!(
202                "Failed to convert to NNBackend value. Unknown NNBackend type: {}",
203                s
204            ))),
205        }
206    }
207}
208
209/// Define where the graph should be executed.
210#[cfg(feature = "wasi_nn")]
211#[cfg_attr(docsrs, doc(cfg(feature = "wasi_nn")))]
212#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
213#[allow(non_camel_case_types)]
214pub enum ExecutionTarget {
215    CPU,
216    GPU,
217    TPU,
218    AUTO,
219}
220#[cfg(feature = "wasi_nn")]
221#[cfg_attr(docsrs, doc(cfg(feature = "wasi_nn")))]
222impl std::fmt::Display for ExecutionTarget {
223    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
224        match self {
225            ExecutionTarget::CPU => write!(f, "CPU"),
226            ExecutionTarget::GPU => write!(f, "GPU"),
227            ExecutionTarget::TPU => write!(f, "TPU"),
228            ExecutionTarget::AUTO => write!(f, "AUTO"),
229        }
230    }
231}
232#[cfg(feature = "wasi_nn")]
233#[cfg_attr(docsrs, doc(cfg(feature = "wasi_nn")))]
234impl std::str::FromStr for ExecutionTarget {
235    type Err = WasmEdgeError;
236
237    fn from_str(s: &str) -> Result<Self, Self::Err> {
238        match s.to_uppercase().as_str() {
239            "CPU" => Ok(ExecutionTarget::CPU),
240            "GPU" => Ok(ExecutionTarget::GPU),
241            "TPU" => Ok(ExecutionTarget::TPU),
242            "AUTO" => Ok(ExecutionTarget::AUTO),
243            _ => Err(WasmEdgeError::Operation(format!(
244                "Failed to convert to ExecutionTarget value. Unknown ExecutionTarget type: {}",
245                s
246            ))),
247        }
248    }
249}
250
251/// Defines the API to manage plugins.
252#[derive(Debug)]
253pub struct PluginManager {}
254impl PluginManager {
255    /// Load plugins from the given path.
256    ///
257    /// * If the path is not given, then the default plugin paths will be used. The default plugin paths are
258    ///
259    ///     * The environment variable "WASMEDGE_PLUGIN_PATH".
260    ///
261    ///     * The `../plugin/` directory related to the WasmEdge installation path.
262    ///
263    ///     * The `wasmedge/` directory under the library path if the WasmEdge is installed under the "/usr".
264    ///
265    /// * If the path is given, then
266    ///
267    ///     * If the path is pointing at a file , then it indicates that a single plugin will be loaded from the file.
268    ///
269    ///     * If the path is pointing at a directory, then the method will load plugins from the files in the directory.
270    ///
271    /// # Argument
272    ///
273    /// * `path` - A path to a plugin file or a directory holding plugin files. If `None`, then the default plugin path will be used.
274    pub fn load(path: Option<&std::path::Path>) -> WasmEdgeResult<()> {
275        match path {
276            Some(p) => sys::plugin::PluginManager::load_plugins(p),
277            None => {
278                sys::plugin::PluginManager::load_plugins_from_default_paths();
279                Ok(())
280            }
281        }
282    }
283
284    /// Initialize the wasi_nn plug-in with the preloads.
285    ///
286    /// Note that this function is only available after loading the wasi_nn plug-in and before creating, and before creating the module instance from the plug-in.
287    ///
288    /// # Argument
289    ///
290    /// * `preloads` - The preload list.
291    ///
292    /// # Example
293    ///
294    /// ```ignore
295    /// // load wasinn-pytorch-plugin from the default plugin directory: /usr/local/lib/wasmedge
296    /// PluginManager::load(None)?;
297    /// // preload named model
298    /// PluginManager::nn_preload(vec![NNPreload::new(
299    ///     "default",
300    ///     GraphEncoding::GGML,
301    ///     ExecutionTarget::CPU,
302    ///     "llama-2-7b-chat.Q5_K_M.gguf",
303    /// )]);
304    /// ```
305    ///
306    /// If a preload is string, then use `NNPreload::from_str` to create a `NNPreload` instance:
307    ///
308    /// ```ignore
309    /// use std::str::FromStr;
310    ///
311    /// // load wasinn-pytorch-plugin from the default plugin directory: /usr/local/lib/wasmedge
312    /// PluginManager::load(None)?;
313    /// // preload named model
314    /// PluginManager::nn_preload(vec![NNPreload::from_str("default:GGML:CPU:llama-2-7b-chat.Q5_K_M.gguf")?]);
315    ///
316    /// ```
317    #[cfg(feature = "wasi_nn")]
318    #[cfg_attr(docsrs, doc(cfg(feature = "wasi_nn")))]
319    pub fn nn_preload(preloads: Vec<NNPreload>) {
320        let mut nn_preloads = Vec::new();
321        for preload in preloads {
322            nn_preloads.push(preload.to_string());
323        }
324
325        let nn_preloads_str: Vec<&str> = nn_preloads.iter().map(|s| s.as_str()).collect();
326
327        sys::plugin::PluginManager::nn_preload(nn_preloads_str);
328    }
329
330    /// Returns the count of loaded plugins.
331    pub fn count() -> u32 {
332        sys::plugin::PluginManager::count()
333    }
334
335    /// Returns the names of all loaded plugins.
336    pub fn names() -> Vec<String> {
337        sys::plugin::PluginManager::names()
338    }
339
340    /// Returns the target plugin by its name.
341    ///
342    /// # Argument
343    ///
344    /// * `name` - The name of the target plugin.
345    ///
346    /// # Error
347    ///
348    /// If failed to return the plugin module instance, then return [PluginError::NotFound](wasmedge_types::error::PluginError::NotFound) error.
349    pub fn find(name: impl AsRef<str>) -> WasmEdgeResult<Plugin> {
350        sys::plugin::PluginManager::find(name.as_ref()).map(|p| Plugin { inner: p })
351    }
352
353    pub fn create_plugin_instance(
354        pname: impl AsRef<str>,
355        mname: impl AsRef<str>,
356    ) -> WasmEdgeResult<PluginInstance> {
357        let plugin = sys::plugin::PluginManager::create_plugin_instance(pname, mname)?;
358        Ok(plugin)
359    }
360
361    /// Initializes the `wasmedge_process` plugin module instance with the parameters.
362    ///
363    /// # Arguments
364    ///
365    /// * `allowed_cmds` - A white list of commands.
366    ///
367    /// * `allowed` - Determines if wasmedge_process is allowed to execute all commands on the white list.
368    #[cfg(all(
369        target_os = "linux",
370        feature = "wasmedge_process",
371        not(feature = "static")
372    ))]
373    #[cfg_attr(
374        docsrs,
375        doc(cfg(all(
376            target_os = "linux",
377            feature = "wasmedge_process",
378            not(feature = "static")
379        )))
380    )]
381    pub fn init_wasmedge_process(allowed_cmds: Option<Vec<&str>>, allowed: bool) {
382        sys::plugin::PluginManager::init_wasmedge_process(allowed_cmds, allowed);
383    }
384
385    pub fn auto_detect_plugins() -> WasmEdgeResult<Vec<Instance>> {
386        let mut plugin_mods = vec![];
387        for plugin_name in PluginManager::names().iter() {
388            if let Ok(plugin) = PluginManager::find(plugin_name) {
389                for mod_name in plugin.mod_names().iter() {
390                    if let Ok(mod_instance) = plugin.mod_instance(mod_name) {
391                        plugin_mods.push(mod_instance)
392                    }
393                }
394            }
395        }
396        Ok(plugin_mods)
397    }
398}
399
400impl PluginManager {
401    pub fn load_plugin_wasi_nn() -> WasmEdgeResult<Instance> {
402        Self::create_plugin_instance("wasi_nn", "wasi_nn")
403    }
404
405    pub fn load_wasi_crypto_common() -> WasmEdgeResult<Instance> {
406        Self::create_plugin_instance("wasi_crypto", "wasi_crypto_common")
407    }
408    pub fn load_wasi_crypto_asymmetric_common() -> WasmEdgeResult<Instance> {
409        Self::create_plugin_instance("wasi_crypto", "wasi_crypto_asymmetric_common")
410    }
411    pub fn load_wasi_crypto_kx() -> WasmEdgeResult<Instance> {
412        Self::create_plugin_instance("wasi_crypto", "wasi_crypto_kx")
413    }
414    pub fn load_wasi_crypto_signatures() -> WasmEdgeResult<Instance> {
415        Self::create_plugin_instance("wasi_crypto", "wasi_crypto_signatures")
416    }
417    pub fn load_wasi_crypto_symmetric() -> WasmEdgeResult<Instance> {
418        Self::create_plugin_instance("wasi_crypto", "wasi_crypto_symmetric")
419    }
420}
421
422/// Represents a loaded plugin. It provides the APIs for accessing the plugin.
423#[derive(Debug)]
424pub struct Plugin {
425    inner: sys::plugin::Plugin,
426}
427impl Plugin {
428    /// Returns the name of this plugin.
429    pub fn name(&self) -> String {
430        self.inner.name()
431    }
432
433    /// Returns the count of the module instances in this plugin.
434    pub fn mod_count(&self) -> u32 {
435        self.inner.mod_count()
436    }
437
438    /// Returns the names of all module instances in this plugin.
439    pub fn mod_names(&self) -> Vec<String> {
440        self.inner.mod_names()
441    }
442
443    /// Returns a module instance that is generated from the module with the given name in this plugin.
444    ///
445    /// # Argument
446    ///
447    /// * `name` - The name of the target module.
448    ///
449    /// # Error
450    ///
451    /// If failed to return the plugin module instance, then return [PluginError::Create](wasmedge_types::error::PluginError::Create) error.
452    pub fn mod_instance(&self, name: impl AsRef<str>) -> WasmEdgeResult<PluginInstance> {
453        self.inner.mod_instance(name.as_ref())
454    }
455}
456
457pub type PluginInstance = Instance;
wasmedge_sdk/plugin.rs

wasmedge_sdk/
plugin.rs
