tenferro_xla/pjrt/
plugin.rs

1use std::fmt;
2use std::path::{Path, PathBuf};
3
4use libloading::Library;
5
6use crate::{Error, Result};
7
8use super::sys::{GetPjrtApiFn, PJRT_Api};
9
10/// Dynamically loaded PJRT plugin.
11///
12/// # Examples
13///
14/// ```
15/// use tenferro_xla::{Error, PjrtPlugin};
16///
17/// let err = PjrtPlugin::load_from_env("__TENFERRO_XLA_DOCS_UNSET").unwrap_err();
18/// assert!(matches!(err, Error::MissingEnv { .. }));
19/// ```
20pub struct PjrtPlugin {
21    path: PathBuf,
22    _api: *const PJRT_Api,
23    _library: Library,
24}
25
26impl fmt::Debug for PjrtPlugin {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        f.debug_struct("PjrtPlugin")
29            .field("path", &self.path)
30            .finish_non_exhaustive()
31    }
32}
33
34impl PjrtPlugin {
35    /// Load a PJRT plugin path from an environment variable.
36    ///
37    /// # Examples
38    ///
39    /// ```
40    /// use tenferro_xla::{Error, PjrtPlugin};
41    ///
42    /// let err = PjrtPlugin::load_from_env("__TENFERRO_XLA_DOCS_UNSET").unwrap_err();
43    /// assert!(matches!(err, Error::MissingEnv { .. }));
44    /// ```
45    ///
46    /// # Errors
47    ///
48    /// Returns `Error::MissingEnv` when `var` is unset, or `Error::PluginLoad`
49    /// with the typed dynamic-library source when loading or symbol lookup
50    /// fails.
51    pub fn load_from_env(var: &'static str) -> Result<Self> {
52        let path = super::plugin_path_from_env(var)?;
53        Self::load_path(path)
54    }
55
56    /// Load a PJRT plugin from an explicit dynamic-library path.
57    ///
58    /// # Examples
59    ///
60    /// ```
61    /// use tenferro_xla::{Error, PjrtPlugin};
62    ///
63    /// let err = PjrtPlugin::load_path("/definitely/missing/pjrt.so").unwrap_err();
64    /// assert!(matches!(err, Error::PluginLoad { .. }));
65    /// ```
66    ///
67    /// # Errors
68    ///
69    /// Returns `Error::PluginLoad` with the original dynamic-library error as
70    /// its source when the file cannot be opened or `GetPjrtApi` is missing,
71    /// or `Error::PjrtCall` when the plugin returns a null API table.
72    pub fn load_path(path: impl Into<PathBuf>) -> Result<Self> {
73        let path = path.into();
74        // SAFETY: Loading a dynamic library is inherently unsafe. The library is
75        // retained in `Self` for at least as long as the API pointer is exposed.
76        let library = unsafe { Library::new(&path) }.map_err(|source| Error::PluginLoad {
77            path: path.clone(),
78            source: Box::new(source),
79        })?;
80        // SAFETY: OpenXLA PJRT plugins export `GetPjrtApi` with this signature.
81        // The returned table is owned by the plugin and remains valid while the
82        // dynamic library is loaded.
83        let api = unsafe {
84            let symbol: libloading::Symbol<'_, GetPjrtApiFn> =
85                library
86                    .get(b"GetPjrtApi")
87                    .map_err(|source| Error::PluginLoad {
88                        path: path.clone(),
89                        source: Box::new(source),
90                    })?;
91            symbol()
92        };
93        if api.is_null() {
94            return Err(Error::PjrtCall {
95                call: "GetPjrtApi",
96                message: format!("plugin at {path:?} returned a null API table"),
97            });
98        }
99        Ok(Self {
100            path,
101            _api: api,
102            _library: library,
103        })
104    }
105
106    /// Return the path used to load the plugin.
107    ///
108    /// # Examples
109    ///
110    /// ```
111    /// use std::path::Path;
112    /// use tenferro_xla::PjrtPlugin;
113    ///
114    /// let _path_type = std::any::type_name::<&Path>();
115    /// let _method: fn(&PjrtPlugin) -> &Path = PjrtPlugin::path;
116    /// ```
117    pub fn path(&self) -> &Path {
118        &self.path
119    }
120
121    pub(crate) fn api(&self) -> *const PJRT_Api {
122        self._api
123    }
124}
tenferro_xla/pjrt/plugin.rs

tenferro_xla/pjrt/
plugin.rs
