Add dr-inference-engine and route every model session through it
One crate names the runtime, the providers and the devices; dr-face and dr-segment ask it for a session by role. It hands ort an API table once per process — from a libonnxruntime it dlopens when the app names a directory holding one, otherwise from tract — so the Rust build stays free of C on every target and a package can install the runtime as a file (docs/inference.md §3). Sessions live in a registry behind a Model handle that holds the bytes, not the session: every use refreshes a timestamp and a reaper unloads whatever sat idle past the decay. A scan that runs the detector on each image never lets it go idle; a click in the develop view lets the segmenter go after thirty seconds; a handle used after that reloads, and reloads on a higher rung if a compiled engine has landed meanwhile. The probe walks the platform's ladder by building strict sessions and timing them against the CPU provider, caches the choice against a fingerprint of the runtime, driver, hardware and models, and compiles engines for the selected rung in the background, smallest model first. Nothing in this commit turns the native path on: the apps still run on tract until they call init with a runtime directory.
This commit is contained in:
@@ -0,0 +1,142 @@
|
||||
//! The API table `ort` runs on, chosen once (docs/inference.md §3).
|
||||
//!
|
||||
//! `ort` with `alternative-backend` links no runtime and asks, on first use,
|
||||
//! for an `OrtApi` — a struct of function pointers. Two things can fill it:
|
||||
//! a `libonnxruntime` this module `dlopen`s, or `ort-tract`. The Rust build
|
||||
//! is identical either way; the difference is whether a file was found.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::OnceLock;
|
||||
|
||||
/// What supplied the table.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Runtime {
|
||||
/// Pure Rust, one core, every operator these graphs use. The floor.
|
||||
Tract,
|
||||
/// The C++ ONNX Runtime, loaded from `path`.
|
||||
OnnxRuntime { path: PathBuf, version: String },
|
||||
}
|
||||
|
||||
impl Runtime {
|
||||
pub fn label(&self) -> String {
|
||||
match self {
|
||||
Runtime::Tract => "tract".into(),
|
||||
Runtime::OnnxRuntime { version, .. } => format!("ONNX Runtime {version}"),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_native(&self) -> bool {
|
||||
matches!(self, Runtime::OnnxRuntime { .. })
|
||||
}
|
||||
}
|
||||
|
||||
static RUNTIME: OnceLock<Runtime> = OnceLock::new();
|
||||
|
||||
/// The runtime in use; tract until something installs another.
|
||||
pub fn runtime() -> Runtime {
|
||||
RUNTIME.get().cloned().unwrap_or(Runtime::Tract)
|
||||
}
|
||||
|
||||
/// Install a table if none is installed yet — tract, since no directories
|
||||
/// were named. What a test or an example gets.
|
||||
pub fn ensure_installed() {
|
||||
if RUNTIME.get().is_none() {
|
||||
install(&[]);
|
||||
}
|
||||
}
|
||||
|
||||
/// Look for `libonnxruntime` in `dirs`, in order, and hand `ort` the first
|
||||
/// table that loads; otherwise tract. Once per process.
|
||||
pub fn install(dirs: &[PathBuf]) -> Runtime {
|
||||
RUNTIME
|
||||
.get_or_init(|| {
|
||||
#[cfg(feature = "native")]
|
||||
for dir in dirs {
|
||||
match load_native(dir) {
|
||||
Ok(rt) => return rt,
|
||||
Err(e) => log::info!("inference: no runtime in {}: {e}", dir.display()),
|
||||
}
|
||||
}
|
||||
#[cfg(not(feature = "native"))]
|
||||
let _ = dirs;
|
||||
install_tract()
|
||||
})
|
||||
.clone()
|
||||
}
|
||||
|
||||
#[cfg(feature = "tract")]
|
||||
fn install_tract() -> Runtime {
|
||||
let _ = ort::set_api(ort_tract::api());
|
||||
Runtime::Tract
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "tract"))]
|
||||
fn install_tract() -> Runtime {
|
||||
// A build with neither tract nor a runtime file has nothing to run
|
||||
// models on; every `open` will report the un-set API rather than panic
|
||||
// somewhere deeper.
|
||||
log::error!("inference: no ONNX Runtime found and tract is not compiled in");
|
||||
Runtime::Tract
|
||||
}
|
||||
|
||||
#[cfg(feature = "native")]
|
||||
fn load_native(dir: &std::path::Path) -> Result<Runtime, String> {
|
||||
let name = if cfg!(target_os = "windows") {
|
||||
"onnxruntime.dll"
|
||||
} else if cfg!(any(target_os = "macos", target_os = "ios")) {
|
||||
"libonnxruntime.dylib"
|
||||
} else {
|
||||
"libonnxruntime.so"
|
||||
};
|
||||
// An empty dir means the bare name: the system loader's search, which on
|
||||
// Android includes the APK's own native libraries.
|
||||
let path = if dir.as_os_str().is_empty() {
|
||||
PathBuf::from(name)
|
||||
} else {
|
||||
let p = dir.join(name);
|
||||
if !p.is_file() {
|
||||
return Err("not present".into());
|
||||
}
|
||||
p
|
||||
};
|
||||
|
||||
// SAFETY: the library's initialisers are ONNX Runtime's own; the symbol
|
||||
// is the documented entry point with the documented signature; the table
|
||||
// is copied out and the library handle is leaked, so every pointer in
|
||||
// the copy stays valid for the life of the process.
|
||||
unsafe {
|
||||
let lib = libloading::Library::new(&path).map_err(|e| e.to_string())?;
|
||||
let get_base: libloading::Symbol<
|
||||
unsafe extern "system" fn() -> *const ort_sys::OrtApiBase,
|
||||
> = lib.get(b"OrtGetApiBase\0").map_err(|e| e.to_string())?;
|
||||
let base = get_base();
|
||||
if base.is_null() {
|
||||
return Err("OrtGetApiBase returned null".into());
|
||||
}
|
||||
let version = std::ffi::CStr::from_ptr(((*base).GetVersionString)())
|
||||
.to_string_lossy()
|
||||
.into_owned();
|
||||
let api = ((*base).GetApi)(ort_sys::ORT_API_VERSION);
|
||||
if api.is_null() {
|
||||
return Err(format!(
|
||||
"ONNX Runtime {version} is older than API version {}",
|
||||
ort_sys::ORT_API_VERSION
|
||||
));
|
||||
}
|
||||
if !ort::set_api((*api).clone()) {
|
||||
return Err("an API table was already installed".into());
|
||||
}
|
||||
std::mem::forget(lib);
|
||||
|
||||
// Qualcomm's DSP loader finds the Hexagon skel through this variable,
|
||||
// and only through it; the runtime's own directory is where the APK
|
||||
// put it. Harmless anywhere else.
|
||||
#[cfg(target_os = "android")]
|
||||
if !dir.as_os_str().is_empty() {
|
||||
std::env::set_var("ADSP_LIBRARY_PATH", dir);
|
||||
}
|
||||
|
||||
log::info!("inference: ONNX Runtime {version} from {}", path.display());
|
||||
Ok(Runtime::OnnxRuntime { path, version })
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user