nornir 0.4.42

//! Make the ort CUDA execution provider "just work" without a manual
//! `LD_LIBRARY_PATH`.
//!
//! The onnxruntime CUDA provider (`libonnxruntime_providers_cuda.so`) is
//! `dlopen`ed lazily when the session is built, and it `NEEDS` the CUDA
//! runtime libs (`libcudart`, `libcublas`, `libcudnn`, …). glibc fixes the
//! loader search path at process start, so setting `LD_LIBRARY_PATH` from
//! inside the process is ignored. The robust fix is to **`dlopen` the CUDA
//! libs ourselves with `RTLD_GLOBAL`** before the provider loads — then its
//! `NEEDED` sonames resolve against the already-loaded globals.
//!
//! Discovery order (first dir that contains a given lib wins):
//!   1. `NORNIR_CUDA_LIBS` — explicit colon-separated dirs (highest priority).
//!   2. Active Python env: `$VIRTUAL_ENV` / `$CONDA_PREFIX` →
//!      `…/site-packages/nvidia/*/lib` (pip `nvidia-*-cu12` wheels).
//!   3. `NORNIR_CUDA_SCAN_ROOTS` — colon-separated roots scanned for
//!      `**/site-packages/nvidia/*/lib` (one level of venvs).
//!   4. Known system dirs: ollama bundles, `/usr/local/cuda*/lib64`, etc.
//!
//! Best-effort: anything missing just means ort falls back to CPU (no panic).
//!
//! Cargo feature: `embed-ort`.

use std::path::{Path, PathBuf};
use std::sync::OnceLock;

/// Leaf-first load order: a lib's own `NEEDED` CUDA deps must already be
/// global when it loads. `cudnn` (and its split sub-libs) come last.
const ORDERED_SONAMES: &[&str] = &[
    "libcudart.so.12",
    "libnvJitLink.so.12",
    "libcublasLt.so.12",
    "libcublas.so.12",
    "libcufft.so.11",
    "libcurand.so.10",
    "libcusparse.so.12",
    "libcusolver.so.11",
];

/// Outcome of a preload attempt (for logging / diagnostics).
#[derive(Debug, Default, Clone)]
pub struct CudaPreload {
    /// Sonames successfully `dlopen`ed (kept resident for the process life).
    pub loaded: Vec<String>,
    /// Directories that were searched.
    pub dirs: Vec<PathBuf>,
    /// True if a cuDNN lib was found and loaded.
    pub cudnn: bool,
}

// Loaded libraries are leaked on purpose: they must stay resident for the
// whole process so the provider can resolve against them.
static PRELOAD: OnceLock<CudaPreload> = OnceLock::new();

/// Discover + `dlopen` CUDA libs once. Idempotent; safe to call before every
/// session build.
pub fn ensure() -> &'static CudaPreload {
    PRELOAD.get_or_init(run)
}

/// Find `libonnxruntime.so` for the **dynamic-load** ort backend (#9). With
/// `ort/load-dynamic` the binary links no onnxruntime; ort `dlopen`s it at
/// runtime from `$ORT_DYLIB_PATH` (or the bare soname). The runtime selector
/// calls this to decide whether the ort/GPU path is even *usable* on this box:
/// `Some(path)` ⇒ onnxruntime is present and loadable, `None` ⇒ fall back to the
/// pure-Rust tract CPU embedder. Search order:
///   1. `$ORT_DYLIB_PATH` (explicit; honoured as-is if it loads).
///   2. The same dirs CUDA discovery uses ([`candidate_dirs`]) — so dropping a
///      `libonnxruntime.so` next to the CUDA libs in `/opt/nornir/cuda` is enough.
///   3. The bare soname `libonnxruntime.so` (system loader path).
/// Whatever is found is verified to actually `dlopen` (a stale/broken file
/// returns `None`, so the selector degrades to CPU rather than letting ort panic).
pub fn onnxruntime_dylib() -> Option<PathBuf> {
    // 1. Explicit override — trust it only if it loads.
    if let Some(p) = std::env::var_os("ORT_DYLIB_PATH") {
        let path = PathBuf::from(p);
        if !path.as_os_str().is_empty() && loadable(&path) {
            return Some(path);
        }
    }
    // 2. The CUDA discovery dirs (incl. /opt/nornir/cuda) — colocated runtime.
    if let Some(path) = find_lib(&candidate_dirs(), "libonnxruntime.so") {
        if loadable(&path) {
            return Some(path);
        }
    }
    // 3. Bare soname via the system loader.
    let bare = PathBuf::from("libonnxruntime.so");
    if loadable(&bare) {
        return Some(bare);
    }
    None
}

/// True if `path` can be `dlopen`ed *and* exports `OrtGetApiBase` (i.e. it is a
/// real onnxruntime, not a same-named decoy). The handle is dropped immediately;
/// ort re-loads (and keeps) its own copy. A bare soname is resolved by the loader.
fn loadable(path: &Path) -> bool {
    use libloading::os::unix::{Library, RTLD_NOW};
    let arg: &std::ffi::OsStr = path.as_os_str();
    // SAFETY: probing a shared object; handle dropped at scope end.
    let Ok(lib) = (unsafe { Library::open(Some(arg), RTLD_NOW) }) else {
        return false;
    };
    // SAFETY: looking up a symbol on a live handle; we only test existence.
    let ok = unsafe { lib.get::<unsafe extern "C" fn() -> *const std::ffi::c_void>(b"OrtGetApiBase") }.is_ok();
    ok
}

/// Probe + arm the dynamic-load ort path (#9). If [`onnxruntime_dylib`] finds a
/// usable onnxruntime, set `ORT_DYLIB_PATH` to its absolute location (so ort's
/// loader picks exactly that file) and return `true`. Otherwise leave the env
/// untouched and return `false` — the caller must then use the tract CPU
/// embedder, never touch ort (whose `api()` would panic if the dylib is absent).
/// Idempotent and cheap to call from the selector.
pub fn arm_onnxruntime() -> bool {
    match onnxruntime_dylib() {
        Some(path) => {
            if path.is_absolute() {
                // SAFETY: single-threaded selection point, before any ort use.
                std::env::set_var("ORT_DYLIB_PATH", &path);
            }
            true
        }
        None => false,
    }
}

fn run() -> CudaPreload {
    let dirs = candidate_dirs();
    let mut out = CudaPreload {
        dirs: dirs.clone(),
        ..Default::default()
    };
    for soname in ORDERED_SONAMES {
        if let Some(path) = find_lib(&dirs, soname) {
            if dlopen_global(&path) {
                out.loaded.push(soname.to_string());
            }
        }
    }
    // cuDNN 9 is split: load its sub-libs (precompiled engines, ops, …) before
    // the umbrella `libcudnn.so.9`, all from the same dir.
    if let Some(dir) = dirs.iter().find(|d| d.join("libcudnn.so.9").exists()) {
        let mut subs: Vec<PathBuf> = std::fs::read_dir(dir)
            .into_iter()
            .flatten()
            .flatten()
            .map(|e| e.path())
            .filter(|p| {
                p.file_name()
                    .and_then(|n| n.to_str())
                    .is_some_and(|n| n.starts_with("libcudnn") && n.contains(".so.") && !n.ends_with("libcudnn.so.9"))
            })
            .collect();
        subs.sort();
        for p in subs {
            dlopen_global(&p);
        }
        if dlopen_global(&dir.join("libcudnn.so.9")) {
            out.loaded.push("libcudnn.so.9".to_string());
            out.cudnn = true;
        }
    }
    out
}

/// True if the NVIDIA **userspace driver** (`libcuda.so.1`) is loadable — the lib
/// the GPU driver package installs alongside the kernel module. nornir can
/// **detect** this but **cannot install or embed it**: the kernel driver needs
/// root, a module matching the running kernel, and (usually) a reboot — it is not
/// a redistributable `.so` you can bundle in a binary. Only the *runtime* libs
/// (cudart/cublas/cuDNN) are redistributable and discoverable via [`ensure`].
pub fn driver_present() -> bool {
    use libloading::os::unix::{Library, RTLD_NOW};
    // SAFETY: probing a well-known system soname; handle dropped immediately.
    unsafe { Library::open(Some(std::ffi::OsStr::new("libcuda.so.1")), RTLD_NOW) }.is_ok()
}

fn yn(b: bool) -> &'static str {
    if b {
        "found"
    } else {
        "MISSING"
    }
}

/// Preflight the CUDA GPU path: `(gpu_ready, human_report)`. Checks the NVIDIA
/// driver, the runtime libs [`ensure`] resolved, and cuDNN, then appends concrete
/// advice for whatever is missing. `gpu_ready` is true iff embed-ort should run
/// on the GPU (driver + cudart + cuDNN all present). Drives `nornir vector doctor`.
pub fn preflight() -> (bool, String) {
    let p = ensure();
    let driver = driver_present();
    let cudart = p.loaded.iter().any(|s| s.starts_with("libcudart"));
    let gpu_ready = driver && cudart && p.cudnn;

    let mut s = String::from("nornir CUDA preflight (embed-ort GPU path)\n");
    s.push_str(&format!("  NVIDIA driver libcuda.so.1 : {}\n", yn(driver)));
    s.push_str(&format!("  libcudart (CUDA runtime)   : {}\n", yn(cudart)));
    s.push_str(&format!("  cuDNN 9                    : {}\n", yn(p.cudnn)));
    s.push_str(&format!(
        "  CUDA libs loaded           : {}\n",
        if p.loaded.is_empty() { "(none)".into() } else { p.loaded.join(", ") }
    ));
    s.push_str(&format!(
        "  dirs searched              : {}\n",
        p.dirs.iter().map(|d| d.display().to_string()).collect::<Vec<_>>().join(", ")
    ));
    let missing: Vec<&str> = ORDERED_SONAMES
        .iter()
        .copied()
        .filter(|n| !p.loaded.iter().any(|l| l == n))
        .collect();
    if !missing.is_empty() {
        s.push_str(&format!("  runtime libs not found     : {}\n", missing.join(", ")));
    }
    s.push_str(&format!(
        "\n  verdict: GPU embedding {}\n",
        if gpu_ready { "READY ✓" } else { "unavailable → CPU fallback" }
    ));
    if !driver {
        s.push_str(
            "  → install the NVIDIA GPU driver for your distro (e.g. `sudo apt install \
             nvidia-driver-XXX`, `sudo dnf install akmod-nvidia`, or NVIDIA's .run). nornir \
             can't install the kernel driver — it needs root + a matching kernel module + reboot.\n",
        );
    }
    if driver && (!p.cudnn || !cudart) {
        s.push_str(
            "  → CUDA runtime libs incomplete. Put a matched CUDA-12 / cuDNN-9 .so set \
             (libcudart.so.12, libcublas*, libcudnn.so.9 + its sub-libs) in ONE dir and set \
             NORNIR_CUDA_LIBS to it (e.g. /opt/nornir/cuda). See `.nornir/vector.md`.\n",
        );
    }
    if gpu_ready {
        s.push_str("  → all set; embed-ort runs on the GPU.\n");
    }
    (gpu_ready, s)
}

/// `nornir vector setup-cuda` — pin a complete CUDA runtime set into
/// `/opt/nornir/cuda` (the built-in search dir) by **copying** the libs from
/// wherever discovery finds them on this box (pip `nvidia-*` wheels in a venv,
/// an ollama bundle, a system toolkit). After this, the GPU works for every
/// nornir process with no env — and survives the source venv being deleted.
/// Returns `(copied, missing)` soname lists. Downloading from NVIDIA's redist
/// CDN (for boxes with no local source) is a planned follow-up.
pub fn setup(target: &Path) -> anyhow::Result<(Vec<String>, Vec<String>)> {
    use anyhow::Context;
    let dirs = candidate_dirs();
    std::fs::create_dir_all(target)
        .with_context(|| format!("create {} (need root? try sudo)", target.display()))?;
    let mut copied = Vec::new();
    let mut missing = Vec::new();

    // The ordered runtime set + cuDNN (umbrella + sub-libs from the same dir).
    for soname in ORDERED_SONAMES {
        match find_lib(&dirs, soname) {
            Some(src) => {
                let dst = target.join(src.file_name().unwrap_or_default());
                std::fs::copy(&src, &dst)
                    .with_context(|| format!("copy {} -> {}", src.display(), dst.display()))?;
                // The plain `.so.N` name the loader asks for, alongside any
                // fully-versioned file name we copied.
                let alias = target.join(soname);
                if !alias.exists() {
                    std::fs::copy(&src, &alias).ok();
                }
                copied.push(soname.to_string());
            }
            None => missing.push(soname.to_string()),
        }
    }
    if let Some(dir) = dirs.iter().find(|d| d.join("libcudnn.so.9").exists()) {
        let mut n = 0usize;
        for e in std::fs::read_dir(dir).into_iter().flatten().flatten() {
            let name = e.file_name();
            let name = name.to_string_lossy();
            if name.starts_with("libcudnn") && name.contains(".so") {
                std::fs::copy(e.path(), target.join(e.file_name()))
                    .with_context(|| format!("copy {}", e.path().display()))?;
                n += 1;
            }
        }
        copied.push(format!("libcudnn.so.9 (+{n} files)"));
    } else {
        missing.push("libcudnn.so.9".to_string());
    }
    Ok((copied, missing))
}

fn candidate_dirs() -> Vec<PathBuf> {
    let mut dirs: Vec<PathBuf> = Vec::new();
    let push = |p: PathBuf, dirs: &mut Vec<PathBuf>| {
        if p.is_dir() && !dirs.contains(&p) {
            dirs.push(p);
        }
    };

    // 1. Explicit override.
    if let Some(v) = std::env::var_os("NORNIR_CUDA_LIBS") {
        for p in std::env::split_paths(&v) {
            push(p, &mut dirs);
        }
    }
    // 2. Active Python env → pip nvidia wheels.
    for key in ["VIRTUAL_ENV", "CONDA_PREFIX"] {
        if let Some(root) = std::env::var_os(key) {
            for d in nvidia_pkg_dirs(Path::new(&root)) {
                push(d, &mut dirs);
            }
        }
    }
    // 3. Extra scan roots (each scanned for venvs one level deep).
    if let Some(v) = std::env::var_os("NORNIR_CUDA_SCAN_ROOTS") {
        for root in std::env::split_paths(&v) {
            for d in scan_root_for_nvidia(&root) {
                push(d, &mut dirs);
            }
        }
    }
    // 4. Known system locations: nornir's own conventional CUDA dir first (drop a
    //    matched CUDA-12 / cuDNN-9 .so set here and the GPU "just works", no env),
    //    then ollama bundles + a system CUDA toolkit.
    for sys in [
        "/opt/nornir/cuda",
        "/usr/local/lib/ollama/cuda_v12",
        "/usr/local/lib/ollama/cuda_v13",
        "/usr/local/cuda/lib64",
        "/usr/local/cuda-12/lib64",
        "/opt/cuda/lib64",
        "/usr/lib/x86_64-linux-gnu",
    ] {
        push(PathBuf::from(sys), &mut dirs);
    }
    dirs
}

/// `<root>/lib*/python*/site-packages/nvidia/*/lib` for a venv/conda root.
fn nvidia_pkg_dirs(root: &Path) -> Vec<PathBuf> {
    let mut out = Vec::new();
    for libname in ["lib", "lib64"] {
        let pyroot = root.join(libname);
        let Ok(entries) = std::fs::read_dir(&pyroot) else {
            continue;
        };
        for e in entries.flatten() {
            let nvidia = e.path().join("site-packages/nvidia");
            if let Ok(pkgs) = std::fs::read_dir(&nvidia) {
                for p in pkgs.flatten() {
                    let lib = p.path().join("lib");
                    if lib.is_dir() {
                        out.push(lib);
                    }
                }
            }
        }
    }
    out
}

/// Scan one level of subdirs of `root` as candidate venv roots.
fn scan_root_for_nvidia(root: &Path) -> Vec<PathBuf> {
    let mut out = nvidia_pkg_dirs(root);
    if let Ok(entries) = std::fs::read_dir(root) {
        for e in entries.flatten() {
            if e.path().is_dir() {
                out.extend(nvidia_pkg_dirs(&e.path()));
            }
        }
    }
    out
}

/// First dir containing `soname` (exact) or a `soname.<minor>` variant.
fn find_lib(dirs: &[PathBuf], soname: &str) -> Option<PathBuf> {
    for d in dirs {
        let exact = d.join(soname);
        if exact.exists() {
            return Some(exact);
        }
        // e.g. libcudart.so.12 → match libcudart.so.12.8.90
        if let Ok(entries) = std::fs::read_dir(d) {
            for e in entries.flatten() {
                if let Some(name) = e.file_name().to_str() {
                    if name.starts_with(soname) {
                        return Some(e.path());
                    }
                }
            }
        }
    }
    None
}

fn dlopen_global(path: &Path) -> bool {
    use libloading::os::unix::{Library, RTLD_GLOBAL, RTLD_NOW};
    // SAFETY: loading a CUDA shared lib by absolute path; the handle is leaked
    // so it stays resident for the provider to resolve against.
    match unsafe { Library::open(Some(path), RTLD_NOW | RTLD_GLOBAL) } {
        Ok(lib) => {
            std::mem::forget(lib);
            true
        }
        Err(_) => false,
    }
}