car_inference/

models.rs

//! Model registry — tracks available Qwen3 models, handles download-on-first-use.

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

use serde::{Deserialize, Serialize};
use tracing::info;

use crate::InferenceError;

/// Role a model is suited for.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelRole {
    /// Fast classification, embedding, routing
    Small,
    /// Code reasoning, skill repair, policy eval
    Medium,
    /// Full reasoning, complex generation
    Large,
    /// Maximum quality via MoE (3B active / 30B total)
    Expert,
    /// Dedicated embedding model (semantic similarity, retrieval)
    Embedding,
}

/// Metadata about a model in the registry.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelInfo {
    pub name: String,
    pub hf_repo: String,
    pub hf_filename: String,
    pub tokenizer_repo: String,
    pub role: ModelRole,
    pub param_count: &'static str,
    pub quantized_size_mb: u64,
    pub downloaded: bool,
}

/// Registry of available models and their local paths.
pub struct ModelRegistry {
    models_dir: PathBuf,
    catalog: Vec<ModelSpec>,
}

struct ModelSpec {
    name: &'static str,
    hf_repo: &'static str,
    hf_filename: &'static str,
    tokenizer_repo: &'static str,
    role: ModelRole,
    param_count: &'static str,
    quantized_size_mb: u64,
}

impl ModelRegistry {
    pub fn new(models_dir: PathBuf) -> Self {
        Self {
            models_dir,
            catalog: builtin_catalog(),
        }
    }

    /// List all models with their download status.
    pub fn list_models(&self) -> Vec<ModelInfo> {
        self.catalog
            .iter()
            .map(|spec| {
                let local_path = self.models_dir.join(spec.name).join("model.gguf");
                ModelInfo {
                    name: spec.name.to_string(),
                    hf_repo: spec.hf_repo.to_string(),
                    hf_filename: spec.hf_filename.to_string(),
                    tokenizer_repo: spec.tokenizer_repo.to_string(),
                    role: spec.role,
                    param_count: spec.param_count,
                    quantized_size_mb: spec.quantized_size_mb,
                    downloaded: local_path.exists(),
                }
            })
            .collect()
    }

    /// Find a catalog entry by name (case-insensitive).
    fn find_spec(&self, name: &str) -> Option<&ModelSpec> {
        self.catalog
            .iter()
            .find(|s| s.name.eq_ignore_ascii_case(name))
    }

    /// Ensure a model is downloaded, returning its local directory path.
    pub async fn ensure_model(&self, name: &str) -> Result<PathBuf, InferenceError> {
        let spec = self
            .find_spec(name)
            .ok_or_else(|| InferenceError::ModelNotFound(name.to_string()))?;

        let model_dir = self.models_dir.join(spec.name);
        let model_path = model_dir.join("model.gguf");
        let tokenizer_path = model_dir.join("tokenizer.json");

        // Presence is not integrity: a zero-length partial write or a dangling
        // symlink (pruned blob) satisfies `exists()` but is unusable. Gate on
        // `cache_file_usable` so a broken dest re-downloads instead of being
        // returned as ready.
        if crate::download::cache_file_usable(&model_path)
            && crate::download::cache_file_usable(&tokenizer_path)
        {
            return Ok(model_dir);
        }

        // Serialize concurrent ensure/pull/remove of the same model — without
        // this, two callers can race the dest relink and truncate a shared HF
        // blob. Mirrors the registry pull path's locking.
        let _guard = crate::download::acquire_model_lock(spec.name).await;

        std::fs::create_dir_all(&model_dir)?;

        // Download model weights
        if !crate::download::cache_file_usable(&model_path) {
            info!(
                model = spec.name,
                repo = spec.hf_repo,
                "downloading model weights"
            );
            download_file(spec.hf_repo, spec.hf_filename, &model_path).await?;
        }

        // Download tokenizer
        if !crate::download::cache_file_usable(&tokenizer_path) {
            info!(
                model = spec.name,
                repo = spec.tokenizer_repo,
                "downloading tokenizer"
            );
            download_file(spec.tokenizer_repo, "tokenizer.json", &tokenizer_path).await?;
        }

        Ok(model_dir)
    }

    /// Legacy removal is deliberately disabled.
    ///
    /// Model removal must flow through [`crate::InferenceEngine::remove_model_from_car`]
    /// so CAR can verify its ownership receipt, drain live allocations, and
    /// preserve shared Hugging Face artifacts.
    #[deprecated(note = "use InferenceEngine::remove_model_from_car")]
    pub fn remove_model(&self, name: &str) -> Result<(), InferenceError> {
        Err(InferenceError::InferenceFailed(format!(
            "legacy removal for {name} is disabled; use receipt-backed models.remove"
        )))
    }
}

/// Download a single file from a HuggingFace repo.
async fn download_file(repo: &str, filename: &str, dest: &Path) -> Result<(), InferenceError> {
    let api = hf_hub::api::tokio::Api::new()
        .map_err(|e| InferenceError::DownloadFailed(e.to_string()))?;

    let repo = api.model(repo.to_string());
    let path = repo
        .get(filename)
        .await
        .map_err(|e| InferenceError::DownloadFailed(format!("{filename}: {e}")))?;

    // `get()` returns a cached file on mere presence, with no integrity check.
    // If that cached pointer is dangling/empty (the shared cache was pruned or
    // a prior download was interrupted), force a real re-download before we
    // link it into our models dir.
    let path = if crate::download::cache_file_usable(&path) {
        path
    } else {
        let fresh = repo
            .download(filename)
            .await
            .map_err(|e| InferenceError::DownloadFailed(format!("{filename}: {e}")))?;
        // The re-download must yield a usable pointer; if it's still
        // dangling/empty, fail loudly rather than link a broken target.
        if !crate::download::cache_file_usable(&fresh) {
            return Err(InferenceError::DownloadFailed(format!(
                "{filename}: re-download produced an unusable file at {}",
                fresh.display()
            )));
        }
        fresh
    };

    // hf-hub caches to its own dir; symlink or copy to our location. Re-link
    // when our dest is missing OR stale (dangling symlink / zero-length).
    if crate::download::cache_file_usable(dest) {
        return Ok(());
    }
    // A stale dest (e.g. dangling symlink to a pruned blob) must be cleared
    // before we can re-create the link/copy.
    if dest.exists() || std::fs::symlink_metadata(dest).is_ok() {
        let _ = std::fs::remove_file(dest);
    }

    // Try symlink first, fall back to copy.
    #[cfg(unix)]
    {
        if std::os::unix::fs::symlink(&path, dest).is_ok() {
            return Ok(());
        }
    }

    // Copy fallback: write to a temp file in the dest dir, then atomically
    // rename into place. A plain `copy` onto `dest` would, if `dest` were a
    // surviving symlink, follow it and truncate the shared HF blob; rename
    // replaces the path itself and never writes through a link.
    let tmp = dest.with_extension("download.partial");
    std::fs::copy(&path, &tmp)
        .map_err(|e| InferenceError::DownloadFailed(format!("copy to {}: {e}", tmp.display())))?;
    std::fs::rename(&tmp, dest).map_err(|e| {
        let _ = std::fs::remove_file(&tmp);
        InferenceError::DownloadFailed(format!("install to {}: {e}", dest.display()))
    })?;
    Ok(())
}

/// Built-in catalog of Qwen3 models.
fn builtin_catalog() -> Vec<ModelSpec> {
    vec![
        ModelSpec {
            name: "Qwen3-Embedding-0.6B",
            hf_repo: "Qwen/Qwen3-Embedding-0.6B-GGUF",
            hf_filename: "Qwen3-Embedding-0.6B-Q8_0.gguf",
            tokenizer_repo: "Qwen/Qwen3-Embedding-0.6B",
            role: ModelRole::Embedding,
            param_count: "0.6B",
            quantized_size_mb: 639,
        },
        ModelSpec {
            name: "Qwen3-0.6B",
            hf_repo: "Qwen/Qwen3-0.6B-GGUF",
            hf_filename: "Qwen3-0.6B-Q8_0.gguf",
            tokenizer_repo: "Qwen/Qwen3-0.6B",
            role: ModelRole::Small,
            param_count: "0.6B",
            quantized_size_mb: 650,
        },
        ModelSpec {
            name: "Qwen3-1.7B",
            hf_repo: "Qwen/Qwen3-1.7B-GGUF",
            hf_filename: "Qwen3-1.7B-Q8_0.gguf",
            tokenizer_repo: "Qwen/Qwen3-1.7B",
            role: ModelRole::Medium,
            param_count: "1.7B",
            quantized_size_mb: 1800,
        },
        ModelSpec {
            name: "Qwen3-4B",
            hf_repo: "Qwen/Qwen3-4B-GGUF",
            hf_filename: "Qwen3-4B-Q4_K_M.gguf",
            tokenizer_repo: "Qwen/Qwen3-4B",
            role: ModelRole::Medium,
            param_count: "4B",
            quantized_size_mb: 2500,
        },
        ModelSpec {
            name: "Qwen3-8B",
            hf_repo: "Qwen/Qwen3-8B-GGUF",
            hf_filename: "Qwen3-8B-Q4_K_M.gguf",
            tokenizer_repo: "Qwen/Qwen3-8B",
            role: ModelRole::Large,
            param_count: "8B",
            quantized_size_mb: 4900,
        },
        ModelSpec {
            name: "Qwen3-30B-A3B",
            hf_repo: "Qwen/Qwen3-30B-A3B-GGUF",
            hf_filename: "Qwen3-30B-A3B-Q4_K_M.gguf",
            tokenizer_repo: "Qwen/Qwen3-30B-A3B",
            role: ModelRole::Expert,
            param_count: "30B (3B active)",
            quantized_size_mb: 17000,
        },
    ]
}