DarkRoom/core/dr-inference-engine/src/api.rs

//! The API table `ort` runs on, chosen once (docs/dev/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, unless `DARKROOM_ORT_DIR`
/// names a runtime: the same variable the desktop honours, so an example
/// can be pointed at the runtime the app uses without learning `init`.
pub fn ensure_installed() {
    if RUNTIME.get().is_none() {
        let dirs: Vec<PathBuf> = std::env::var_os("DARKROOM_ORT_DIR")
            .map(PathBuf::from)
            .into_iter()
            .collect();
        install(&dirs);
    }
}

/// 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 {
        find_library(dir, name).ok_or("not present")?
    };

    // 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 })
    }
}

/// `libonnxruntime.so` in `dir`, or a versioned spelling of it —
/// `libonnxruntime.so.1.30.0` is what the Python wheel ships, and a package
/// that installs only the versioned file is not wrong.
#[cfg(feature = "native")]
fn find_library(dir: &std::path::Path, name: &str) -> Option<PathBuf> {
    let exact = dir.join(name);
    if exact.is_file() {
        return Some(exact);
    }
    let prefix = format!("{name}.");
    let mut versioned: Vec<PathBuf> = std::fs::read_dir(dir)
        .ok()?
        .filter_map(|e| e.ok())
        .map(|e| e.path())
        .filter(|p| {
            p.is_file()
                && p.file_name()
                    .and_then(|n| n.to_str())
                    .is_some_and(|n| n.starts_with(&prefix))
        })
        .collect();
    versioned.sort();
    versioned.pop()
}