inference/backend/

onnx.rs

//! ONNX Runtime embedding backend.
//!
//! This is the production default backend.  It wraps the existing session-pool
//! logic that was previously inline in `EmbeddingEngine`.  No functional change —
//! pure extraction to satisfy the [`EmbeddingBackend`] trait.
//!
//! ## GPU mode (DAKERA_USE_GPU=1)
//!
//! When GPU is enabled, `pool_size` is capped to **1** regardless of
//! `DAKERA_ONNX_POOL_SIZE`.  A single session with a `parking_lot::Mutex`
//! naturally serialises CUDA forward passes without a separate semaphore.  The
//! CUDA execution provider is configured with a hard memory limit
//! (`DAKERA_GPU_MEM_LIMIT_GB`, default 15 GB) and
//! `ArenaExtendStrategy::SameAsRequested` to prevent `BFCArena` growth beyond
//! the working set.  Together these guarantee at most one concurrent CUDA call
//! and bounded VRAM usage, replacing the application-level
//! `GPU_INFERENCE_SEMAPHORE` approach from v0.11.79/v0.11.80 (DAK-6134).
//!
//! ## CPU mode
//!
//! A pool of `N` independent ONNX sessions (`N = DAKERA_ONNX_POOL_SIZE`, default
//! 4) serves concurrent callers via round-robin dispatch.  Each session runs
//! inside a `tokio::task::spawn_blocking` call to avoid blocking the async
//! executor.

use crate::backend::{BackendKind, EmbeddingBackend};
use crate::batch::{mean_pooling, normalize_embeddings, BatchProcessor};
use crate::error::{InferenceError, Result};
use crate::models::ModelConfig;
use async_trait::async_trait;
use ort::execution_providers::{ArenaExtendStrategy, CUDAExecutionProvider};
use ort::inputs;
use ort::session::builder::GraphOptimizationLevel;
use ort::session::Session;
use ort::value::Tensor;
use parking_lot::Mutex;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokenizers::Tokenizer;
use tracing::{info, instrument, warn};

/// ONNX Runtime embedding backend with a session pool for concurrent inference.
pub struct OnnxBackend {
    sessions: Vec<Arc<Mutex<Session>>>,
    next_session: AtomicUsize,
    processor: Arc<BatchProcessor>,
    config: ModelConfig,
    dimension: usize,
}

/// Determine the ONNX session pool size based on inference mode.
///
/// GPU: always 1 — a single `parking_lot::Mutex`-guarded session serialises CUDA calls
/// at the allocator level, replacing the application-level `GPU_INFERENCE_SEMAPHORE`.
/// CPU: the configured pool size (minimum 1) for concurrent inference.
fn resolve_pool_size(use_gpu: bool, configured: usize) -> usize {
    if use_gpu {
        1
    } else {
        configured.max(1)
    }
}

impl OnnxBackend {
    /// Build a new `OnnxBackend` by downloading model files and building the session pool.
    #[instrument(skip_all, fields(model = %config.model))]
    pub async fn new(config: &ModelConfig) -> Result<Self> {
        let config = config.clone();
        let use_gpu = std::env::var("DAKERA_USE_GPU")
            .map(|v| v == "1")
            .unwrap_or(config.use_gpu);

        if use_gpu {
            info!("ONNX backend: CUDA execution provider enabled (DAKERA_USE_GPU=1)");
        }
        info!("Initialising ONNX backend: model={}", config.model);

        let (tokenizer_path, onnx_path) = Self::download_model_files(&config, use_gpu).await?;

        info!("Loading tokenizer from {:?}", tokenizer_path);
        let tokenizer = Tokenizer::from_file(&tokenizer_path)
            .map_err(|e| InferenceError::TokenizationError(e.to_string()))?;

        let num_threads = config.num_threads.unwrap_or(4);
        let pool_size = resolve_pool_size(use_gpu, config.session_pool_size);

        // GPU memory limit: default 15 GB on L4 (24 GB total), overrideable via env.
        let gpu_mem_limit_bytes: usize = std::env::var("DAKERA_GPU_MEM_LIMIT_GB")
            .ok()
            .and_then(|v| v.parse::<usize>().ok())
            .unwrap_or(15)
            * 1024
            * 1024
            * 1024;

        if use_gpu {
            info!(
                "ONNX backend: GPU mode — pool_size=1, gpu_mem_limit={}GB",
                gpu_mem_limit_bytes / (1024 * 1024 * 1024)
            );
        }

        let onnx_path_clone = onnx_path.clone();

        let sessions: Vec<Arc<Mutex<Session>>> =
            tokio::task::spawn_blocking(move || -> Result<Vec<Arc<Mutex<Session>>>> {
                (0..pool_size)
                    .map(|_| {
                        let builder = Session::builder()
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                            .with_optimization_level(GraphOptimizationLevel::Level3)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                            .with_intra_threads(num_threads)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?;

                        // DAK-6145: disable memory pattern pre-allocation on CPU.
                        // On memory-constrained servers, ORT's pattern-based pre-allocation
                        // exhausts the BFCArena before inference starts, causing batch=4
                        // (8MB query/Add buffer) to fail even after 3 halvings.  With
                        // memory_pattern=false, allocation is fresh per-run so headroom
                        // remains available for per-layer activations.
                        let mut builder = if use_gpu {
                            builder
                                .with_execution_providers([CUDAExecutionProvider::default()
                                    .with_memory_limit(gpu_mem_limit_bytes)
                                    .with_arena_extend_strategy(
                                        ArenaExtendStrategy::SameAsRequested,
                                    )
                                    .build()])
                                .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                        } else {
                            builder
                                .with_memory_pattern(false)
                                .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                        };

                        let s = builder
                            .commit_from_file(&onnx_path_clone)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?;
                        Ok(Arc::new(Mutex::new(s)))
                    })
                    .collect()
            })
            .await
            .map_err(|e| {
                InferenceError::ModelLoadError(format!("Session pool init panicked: {}", e))
            })??;

        let dimension = config.model.dimension();
        let processor = Arc::new(BatchProcessor::new(
            tokenizer,
            config.model,
            config.max_batch_size,
        ));

        info!(
            "ONNX backend ready: model={}, dimension={}, threads={}, pool={}",
            config.model, dimension, num_threads, pool_size
        );

        Ok(Self {
            sessions,
            next_session: AtomicUsize::new(0),
            processor,
            config,
            dimension,
        })
    }

    /// Number of ONNX sessions in the pool.
    pub fn pool_size(&self) -> usize {
        self.sessions.len()
    }

    // ── File download helpers (shared with CandleBackend) ──────────────────────

    /// Resolve tokenizer and ONNX model files, downloading from HuggingFace if needed.
    #[instrument(skip_all, fields(model = %config.model))]
    pub async fn download_model_files(
        config: &ModelConfig,
        use_gpu: bool,
    ) -> Result<(PathBuf, PathBuf)> {
        let model_id = config.model.model_id();
        let onnx_repo_id = config.model.onnx_repo_id();
        let onnx_filename = if use_gpu {
            config.model.onnx_filename_gpu()
        } else {
            config.model.onnx_filename()
        };

        info!(
            "Resolving model files: tokenizer={}, onnx={}@{}",
            model_id, onnx_filename, onnx_repo_id
        );

        let tokenizer_cache_dir = Self::model_cache_dir(model_id)?;
        let onnx_cache_dir = Self::model_cache_dir(onnx_repo_id)?;

        let onnx_subdir = onnx_cache_dir.join("onnx");
        std::fs::create_dir_all(&onnx_subdir)?;

        let local_tokenizer = tokenizer_cache_dir.join("tokenizer.json");
        let onnx_basename = Path::new(onnx_filename)
            .file_name()
            .and_then(|s| s.to_str())
            .unwrap_or("model_quantized.onnx");
        let local_onnx = onnx_subdir.join(onnx_basename);

        // GPU FP32 model truncation guard (DAK-5976)
        if use_gpu && local_onnx.exists() {
            let cached_size = local_onnx.metadata().map(|m| m.len()).unwrap_or(0);
            if cached_size <= 500_000_000 {
                warn!(
                    "Cached GPU ONNX at {:?} is {} bytes (≤500 MB) — likely truncated. Deleting.",
                    local_onnx, cached_size
                );
                let _ = std::fs::remove_file(&local_onnx);
            }
        }

        if !local_tokenizer.exists() || !local_onnx.exists() {
            let model_id_owned = model_id.to_string();
            let onnx_repo_id_owned = onnx_repo_id.to_string();
            let onnx_filename_owned = onnx_filename.to_string();
            let tokenizer_cache = tokenizer_cache_dir.clone();
            let onnx_cache = onnx_cache_dir.clone();

            tokio::task::spawn_blocking(move || {
                if !tokenizer_cache.join("tokenizer.json").exists() {
                    Self::download_hf_file(&model_id_owned, "tokenizer.json", &tokenizer_cache)
                        .map_err(|e| {
                            InferenceError::HubError(format!("Failed to download tokenizer: {}", e))
                        })?;
                }
                if !onnx_cache.join(&onnx_filename_owned).exists() {
                    Self::download_hf_file(&onnx_repo_id_owned, &onnx_filename_owned, &onnx_cache)
                        .map_err(|e| {
                            InferenceError::HubError(format!(
                                "Failed to download ONNX model: {}",
                                e
                            ))
                        })?;
                }
                Ok::<_, InferenceError>(())
            })
            .await
            .map_err(|e| InferenceError::HubError(format!("Download task panicked: {}", e)))??;
        } else {
            info!("All model files found in local cache");
        }

        let final_onnx = onnx_cache_dir.join(onnx_filename);
        Ok((local_tokenizer, final_onnx))
    }

    /// Get or create the local model cache directory.
    pub fn model_cache_dir(model_id: &str) -> Result<PathBuf> {
        let base = std::env::var("HF_HOME")
            .map(PathBuf::from)
            .unwrap_or_else(|_| {
                let home = std::env::var("HOME").unwrap_or_else(|_| {
                    warn!("HOME environment variable not set, using /tmp for model cache");
                    "/tmp".to_string()
                });
                PathBuf::from(home).join(".cache").join("huggingface")
            });
        let dir = base.join("dakera").join(model_id.replace('/', "--"));
        std::fs::create_dir_all(&dir)?;
        Ok(dir)
    }

    /// Download a single file from HuggingFace using ureq (sync, call inside spawn_blocking).
    pub fn download_hf_file(
        model_id: &str,
        filename: &str,
        cache_dir: &Path,
    ) -> std::result::Result<PathBuf, String> {
        let file_path = cache_dir.join(filename);
        if file_path.exists() {
            info!("Cached: {}/{}", model_id, filename);
            return Ok(file_path);
        }

        if let Some(parent) = file_path.parent() {
            std::fs::create_dir_all(parent)
                .map_err(|e| format!("Failed to create directory {:?}: {}", parent, e))?;
        }

        let url = format!(
            "https://huggingface.co/{}/resolve/main/{}",
            model_id, filename
        );
        info!("Downloading: {}", url);

        let hf_token = std::env::var("HF_TOKEN")
            .or_else(|_| std::env::var("HUGGING_FACE_HUB_TOKEN"))
            .ok();

        let agent = ureq::AgentBuilder::new()
            .redirects(0)
            .timeout(std::time::Duration::from_secs(300))
            .build();

        let mut current_url = url;
        let mut redirects = 0_u32;

        let response = loop {
            let mut req = agent.get(&current_url);
            if let Some(ref token) = hf_token {
                req = req.set("Authorization", &format!("Bearer {}", token));
            }
            let resp = req.call();

            let r = match resp {
                Ok(r) => r,
                Err(ureq::Error::Status(_status, r)) => r,
                Err(e) => return Err(format!("{}: {}", filename, e)),
            };

            let status = r.status();
            if (200..300).contains(&status) {
                break r;
            } else if (300..400).contains(&status) {
                redirects += 1;
                if redirects > 10 {
                    return Err(format!("{}: too many redirects", filename));
                }
                let location = r
                    .header("location")
                    .ok_or_else(|| format!("{}: redirect without Location header", filename))?
                    .to_string();

                current_url = if location.starts_with('/') {
                    let parsed = url::Url::parse(&current_url)
                        .map_err(|e| format!("{}: bad URL: {}", filename, e))?;
                    let host = parsed
                        .host_str()
                        .ok_or_else(|| format!("{}: missing host", filename))?;
                    format!("{}://{}{}", parsed.scheme(), host, location)
                } else {
                    location
                };
            } else {
                return Err(format!("{}: HTTP {}", filename, status));
            }
        };

        let expected_bytes: Option<u64> = response
            .header("x-linked-size")
            .or_else(|| response.header("content-length"))
            .and_then(|v| v.parse::<u64>().ok());

        let mut bytes = Vec::new();
        response
            .into_reader()
            .take(2_147_483_648)
            .read_to_end(&mut bytes)
            .map_err(|e| format!("Failed to read {}: {}", filename, e))?;

        if let Some(expected) = expected_bytes {
            if (bytes.len() as u64) < expected {
                return Err(format!(
                    "{}: download incomplete — received {} of {} bytes",
                    filename,
                    bytes.len(),
                    expected
                ));
            }
        }

        std::fs::write(&file_path, &bytes)
            .map_err(|e| format!("Failed to write {}: {}", filename, e))?;

        info!("Downloaded {} ({} bytes)", filename, bytes.len());
        Ok(file_path)
    }

    /// Public alias for use by other inference modules (e.g. GLiNER NER engine).
    pub fn download_hf_file_pub(
        model_id: &str,
        filename: &str,
        cache_dir: &Path,
    ) -> std::result::Result<PathBuf, String> {
        Self::download_hf_file(model_id, filename, cache_dir)
    }

    // ── Internal embedding logic ───────────────────────────────────────────────

    /// Internal batch embedding: split → distribute across pool → collect.
    ///
    /// On BFCArena / allocator OOM the batch is halved and retried until
    /// `batch_size == 1`.  Starting from the default `max_batch_size=32` this
    /// gives up to 5 halvings (32→16→8→4→2→1) before surfacing the error,
    /// covering the case where even a 4-text batch exceeds available memory
    /// under concurrent load (DAK-6145).
    async fn embed_batch_internal(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
        if texts.is_empty() {
            return Ok(vec![]);
        }

        let pool_len = self.sessions.len();
        let normalize = self.config.model.normalize_embeddings();
        let start_idx = self.next_session.fetch_add(1, Ordering::Relaxed);
        let mut batch_size = self.config.max_batch_size.max(1);

        // DAK-6145: 5 halvings — 32→16→8→4→2→1 — before hard fail.
        // Previous depth of 3 (stopping at batch=4) was insufficient: BGE-Large at
        // batch=4 seq≈503 still needs 8MB for the first query/Add buffer, which the
        // arena cannot provide when concurrent sessions have exhausted system RAM.
        // batch=1 requires only ~2MB and reliably succeeds under memory pressure.
        for attempt in 0_u32..=5 {
            let batches: Vec<Vec<String>> = texts.chunks(batch_size).map(|b| b.to_vec()).collect();

            let mut handles = Vec::with_capacity(batches.len());
            for (i, batch_owned) in batches.into_iter().enumerate() {
                let session = Arc::clone(&self.sessions[(start_idx + i) % pool_len]);
                let processor = Arc::clone(&self.processor);
                // GPU mode: pool_size=1 so all handles point to the same session. The
                // parking_lot::Mutex serialises CUDA forward passes implicitly — no
                // application-level semaphore needed (DAK-6134 deep fix).
                handles.push(tokio::task::spawn_blocking(move || {
                    let mut session_guard = session.lock();
                    Self::process_batch_blocking(
                        &batch_owned,
                        &mut session_guard,
                        &processor,
                        normalize,
                    )
                }));
            }

            let mut all_embeddings: Vec<Vec<f32>> = Vec::with_capacity(texts.len());
            let mut oom: Option<InferenceError> = None;

            for handle in handles {
                match handle.await {
                    Err(panic_err) => {
                        return Err(InferenceError::InferenceError(format!(
                            "Inference task panicked: {panic_err}"
                        )));
                    }
                    Ok(Err(e)) => {
                        if attempt < 5 && Self::is_gpu_oom(&e) {
                            oom = Some(e);
                            break;
                        }
                        return Err(e);
                    }
                    Ok(Ok(batch_embs)) => {
                        all_embeddings.extend(batch_embs);
                    }
                }
            }

            if let Some(_oom_err) = oom {
                let next_batch = (batch_size / 2).max(1);
                warn!(
                    "ONNX allocator OOM (attempt {}/5) — retrying with batch_size {} → {}",
                    attempt + 1,
                    batch_size,
                    next_batch,
                );
                batch_size = next_batch;
                continue;
            }

            return Ok(all_embeddings);
        }

        Err(InferenceError::InferenceError(format!(
            "ONNX inference failed: allocator OOM after 5 batch-halving attempts (batch_size={batch_size})"
        )))
    }

    fn is_gpu_oom(err: &InferenceError) -> bool {
        let msg = err.to_string();
        msg.contains("BFCArena")
            || msg.contains("Failed to allocate memory")
            || msg.contains("CUDA_OUT_OF_MEMORY")
            || msg.contains("CUDA out of memory")
            || (msg.contains("allocate") && msg.contains("buffer of size"))
    }

    fn process_batch_blocking(
        texts: &[String],
        session: &mut Session,
        processor: &BatchProcessor,
        normalize: bool,
    ) -> Result<Vec<Vec<f32>>> {
        let prepared = processor.tokenize_batch(texts)?;
        let batch_size = prepared.batch_size;
        let seq_len = prepared.seq_len;
        let attention_mask_flat = prepared.attention_mask.clone();

        let input_ids_tensor =
            Tensor::<i64>::from_array(([batch_size, seq_len], prepared.input_ids))
                .map_err(|e| InferenceError::InferenceError(e.to_string()))?;
        let attention_mask_tensor =
            Tensor::<i64>::from_array(([batch_size, seq_len], prepared.attention_mask))
                .map_err(|e| InferenceError::InferenceError(e.to_string()))?;
        let token_type_ids_tensor =
            Tensor::<i64>::from_array(([batch_size, seq_len], prepared.token_type_ids))
                .map_err(|e| InferenceError::InferenceError(e.to_string()))?;

        let outputs = session
            .run(inputs![
                "input_ids" => input_ids_tensor,
                "attention_mask" => attention_mask_tensor,
                "token_type_ids" => token_type_ids_tensor
            ])
            .map_err(|e: ort::Error| InferenceError::InferenceError(e.to_string()))?;

        let (ort_shape, lhs_slice) = outputs[0]
            .try_extract_tensor::<f32>()
            .map_err(|e| InferenceError::InferenceError(e.to_string()))?;

        if ort_shape.len() != 3 {
            return Err(InferenceError::InferenceError(format!(
                "Expected 3D last_hidden_state, got {} dims",
                ort_shape.len()
            )));
        }
        let hidden_size = ort_shape[2] as usize;

        let mut embeddings = mean_pooling(
            lhs_slice,
            batch_size,
            seq_len,
            hidden_size,
            &attention_mask_flat,
        );

        if normalize {
            normalize_embeddings(&mut embeddings);
        }

        Ok(embeddings)
    }
}

#[async_trait]
impl EmbeddingBackend for OnnxBackend {
    async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
        self.embed_batch_internal(texts).await
    }

    fn dimension(&self) -> usize {
        self.dimension
    }

    fn backend_kind(&self) -> BackendKind {
        BackendKind::Onnx
    }
}

#[cfg(test)]
mod tests {
    use super::{resolve_pool_size, OnnxBackend};
    use crate::error::InferenceError;

    #[test]
    fn gpu_mode_always_pool_size_one() {
        assert_eq!(resolve_pool_size(true, 1), 1);
        assert_eq!(
            resolve_pool_size(true, 4),
            1,
            "GPU overrides configured pool_size=4 → 1"
        );
        assert_eq!(resolve_pool_size(true, 0), 1, "GPU overrides zero → 1");
    }

    #[test]
    fn cpu_mode_respects_configured_pool_size() {
        assert_eq!(resolve_pool_size(false, 4), 4);
        assert_eq!(resolve_pool_size(false, 1), 1);
        assert_eq!(
            resolve_pool_size(false, 0),
            1,
            "CPU clamps zero to minimum 1"
        );
    }

    // ── is_gpu_oom detection (DAK-6145) ─────────────────────────────────────

    fn oom_err(msg: &str) -> InferenceError {
        InferenceError::InferenceError(msg.to_string())
    }

    #[test]
    fn detects_bfcarena_oom() {
        let e = oom_err("Non-zero status code returned while running Add node. \
            Status Message: bfc_arena.cc:358 void *onnxruntime::BFCArena::\
            AllocateRawInternal(size_t, bool, Stream *) Failed to allocate memory \
            for requested buffer of size 8241152");
        assert!(OnnxBackend::is_gpu_oom(&e), "BFCArena OOM must be detected");
    }

    #[test]
    fn detects_cuda_out_of_memory() {
        let e = oom_err("CUDA_OUT_OF_MEMORY: out of memory on device 0");
        assert!(OnnxBackend::is_gpu_oom(&e));
    }

    #[test]
    fn detects_allocate_buffer_pattern() {
        let e = oom_err("Failed to allocate memory for requested buffer of size 1234");
        assert!(OnnxBackend::is_gpu_oom(&e));
    }

    #[test]
    fn non_oom_error_not_detected() {
        let e = oom_err("Shape mismatch: expected [4, 512] got [4, 256]");
        assert!(!OnnxBackend::is_gpu_oom(&e), "shape error must not trigger OOM retry");
    }

    /// Verify that halving from max_batch_size=32 reaches batch_size=1 in ≤5 steps.
    #[test]
    fn batch_halving_reaches_one_in_five_steps() {
        let mut batch_size = 32_usize;
        let mut halvings = 0_u32;
        while batch_size > 1 {
            batch_size = (batch_size / 2).max(1);
            halvings += 1;
        }
        assert_eq!(batch_size, 1);
        assert!(halvings <= 5, "expected ≤5 halvings, got {halvings}");
    }
}