mecha10_models/services/model_service/

mod.rs

//! Model management service for downloading and caching AI models from HuggingFace Hub
//!
//! This service provides a unified interface for:
//! - Downloading models from HuggingFace Hub using the model catalog
//! - Caching models to the project's models/ directory
//! - Listing available and installed models
//! - Managing model lifecycle (pull, remove, info)
//!
//! # Design
//!
//! - Models are defined in `model_catalog.toml` embedded in the CLI binary
//! - Models are downloaded from HuggingFace Hub on first use (lazy loading)
//! - Models are cached to `<project_root>/models/<name>.onnx`
//! - HF cache is also used (`~/.cache/huggingface/hub/`) for offline support
//!
//! # Usage
//!
//! ```ignore
//! use mecha10_models::services::ModelService;
//!
//! let service = ModelService::new()?;
//!
//! // List catalog models
//! let catalog = service.list_catalog()?;
//!
//! // Download a model
//! service.pull("yolov8n", None)?;
//!
//! // List installed models
//! let installed = service.list_installed()?;
//! ```
//!
//! # Module layout
//!
//! - [`catalog`]: catalog/HuggingFace config schema types (re-exported)
//! - [`download`]: downloading the `.onnx` model artifact
//! - [`labels`]: downloading/generating the `labels.txt` file
//! - [`preprocessing`]: preprocessing auto-detection and `config.json` generation
//! - [`quantize`]: INT8 quantization
//! - [`registry`]: downloading a versioned artifact from the private per-project
//!   model registry (`LAB-1802`)
//!
//! # Storage layout (`LAB-1802`)
//!
//! Catalog and registry artifacts used to both land at the same
//! `<models_dir>/<name>/model.onnx` path, so a registry model sharing a name with
//! a catalog entry (or vice versa) would silently overwrite the other's file. Each
//! source now gets its own namespaced subtree so a name collision can never
//! clobber the wrong artifact:
//!
//! - Catalog: `<models_dir>/catalog/<name>/model.onnx` (versionless - see
//!   [`ModelService::catalog_model_dir`])
//! - Registry: `<models_dir>/registry/<name>/<version>/model.onnx` (see
//!   [`ModelService::registry_model_dir`])

mod catalog;
#[cfg(feature = "model-download")]
mod download;
#[cfg(feature = "model-download")]
mod labels;
mod preprocessing;
#[cfg(feature = "model-download")]
mod quantize;
mod registry;

pub use catalog::{ModelCatalogEntry, PreprocessingPreset, QuantizeConfig};

use anyhow::{Context, Result};
use catalog::ModelCatalog;
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use tokio::fs;

/// Which system an installed model artifact came from (`LAB-1802`) - determines
/// which namespaced subtree of `models_dir` it lives under. See the module docs
/// above for the storage layout.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ModelOrigin {
    /// The curated, unversioned, HuggingFace-backed public catalog.
    Catalog,
    /// The versioned, per-project model registry (`packages/mecha10-api/src/models/registry`).
    Registry,
}

/// Installed model metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InstalledModel {
    /// Model name
    pub name: String,
    /// Local file path
    pub path: PathBuf,
    /// File size in bytes
    pub size: u64,
    /// Catalog entry if from catalog
    pub catalog_entry: Option<ModelCatalogEntry>,
    /// Which system this artifact was downloaded from (`LAB-1802`).
    pub source: ModelOrigin,
    /// Registry version, when `source == ModelOrigin::Registry`. Always `None`
    /// for catalog entries (the public catalog is versionless).
    pub version: Option<i32>,
}

/// Model management service
pub struct ModelService {
    /// Model catalog
    catalog: Vec<ModelCatalogEntry>,
    /// Models directory (defaults to ./models)
    models_dir: PathBuf,
}

impl ModelService {
    /// Create a new ModelService with default settings
    #[allow(dead_code)]
    pub fn new() -> Result<Self> {
        Self::with_models_dir(PathBuf::from("models"))
    }

    /// Create a ModelService with custom models directory
    pub fn with_models_dir(models_dir: PathBuf) -> Result<Self> {
        // Load embedded model catalog
        let catalog_toml = include_str!("../../../model_catalog.toml");
        let catalog: ModelCatalog = toml::from_str(catalog_toml).context("Failed to parse model_catalog.toml")?;

        Ok(Self {
            catalog: catalog.models,
            models_dir,
        })
    }

    /// List all models in the catalog
    pub fn list_catalog(&self) -> Result<Vec<ModelCatalogEntry>> {
        Ok(self.catalog.clone())
    }

    /// Get catalog entry for a model by name
    pub fn get_catalog_entry(&self, name: &str) -> Option<&ModelCatalogEntry> {
        self.catalog.iter().find(|m| m.name == name)
    }

    /// `models_dir/catalog/<name>/` - the public catalog's namespaced
    /// subdirectory. See the module docs above (`LAB-1802`).
    pub(super) fn catalog_model_dir(&self, name: &str) -> PathBuf {
        self.models_dir.join("catalog").join(name)
    }

    /// `models_dir/registry/` - the root under which every registry model's
    /// per-version subdirectories live. See the module docs above (`LAB-1802`).
    fn registry_root_dir(&self) -> PathBuf {
        self.models_dir.join("registry")
    }

    /// List all installed models, from both the catalog and registry namespaced
    /// subtrees (`LAB-1802`) - this is what makes a registry model (e.g. one
    /// auto-registered from a completed training job and then locally fetched via
    /// `mecha10 model set`) show up in `mecha10 model list`/`installed`
    /// alongside catalog models, instead of only being visible through the
    /// separate `mecha10 model` registry verb.
    pub async fn list_installed(&self) -> Result<Vec<InstalledModel>> {
        let mut installed = self.list_installed_catalog().await?;
        installed.extend(self.list_installed_registry().await?);
        Ok(installed)
    }

    async fn list_installed_catalog(&self) -> Result<Vec<InstalledModel>> {
        let catalog_dir = self.models_dir.join("catalog");
        if !catalog_dir.exists() {
            return Ok(Vec::new());
        }

        let mut installed = Vec::new();
        let mut entries = fs::read_dir(&catalog_dir).await?;

        while let Some(entry) = entries.next_entry().await? {
            let path = entry.path();

            // Only look at directories (each model has its own directory)
            if !path.is_dir() {
                continue;
            }

            // Check if model.onnx exists in this directory
            let model_path = path.join("model.onnx");
            if !model_path.exists() {
                continue;
            }

            let metadata = fs::metadata(&model_path).await?;
            let size = metadata.len();

            // Extract model name from directory name
            let name = path
                .file_name()
                .and_then(|s| s.to_str())
                .unwrap_or("unknown")
                .to_string();

            // Try to find catalog entry
            let catalog_entry = self.get_catalog_entry(&name).cloned();

            installed.push(InstalledModel {
                name,
                path: model_path,
                size,
                catalog_entry,
                source: ModelOrigin::Catalog,
                version: None,
            });
        }

        Ok(installed)
    }

    async fn list_installed_registry(&self) -> Result<Vec<InstalledModel>> {
        let registry_dir = self.registry_root_dir();
        if !registry_dir.exists() {
            return Ok(Vec::new());
        }

        let mut installed = Vec::new();
        let mut model_entries = fs::read_dir(&registry_dir).await?;

        while let Some(model_entry) = model_entries.next_entry().await? {
            let model_dir = model_entry.path();
            if !model_dir.is_dir() {
                continue;
            }
            let name = model_dir
                .file_name()
                .and_then(|s| s.to_str())
                .unwrap_or("unknown")
                .to_string();

            let mut version_entries = fs::read_dir(&model_dir).await?;
            while let Some(version_entry) = version_entries.next_entry().await? {
                let version_dir = version_entry.path();
                if !version_dir.is_dir() {
                    continue;
                }

                let model_path = version_dir.join("model.onnx");
                if !model_path.exists() {
                    continue;
                }

                let metadata = fs::metadata(&model_path).await?;
                let version = version_dir
                    .file_name()
                    .and_then(|s| s.to_str())
                    .and_then(|s| s.parse().ok());

                installed.push(InstalledModel {
                    name: name.clone(),
                    path: model_path,
                    size: metadata.len(),
                    catalog_entry: None,
                    source: ModelOrigin::Registry,
                    version,
                });
            }
        }

        Ok(installed)
    }

    /// Remove an installed catalog model. Registry models are versioned (a name
    /// alone doesn't identify a unique artifact - see the module docs), so this
    /// only ever removes from the catalog namespace; if `name` only exists under
    /// the registry namespace, this returns a clear error pointing at the
    /// versioned directory instead of silently doing nothing or guessing a
    /// version to delete (`LAB-1802`).
    pub async fn remove(&self, name: &str) -> Result<()> {
        let catalog_dir = self.catalog_model_dir(name);

        if catalog_dir.exists() {
            fs::remove_dir_all(&catalog_dir)
                .await
                .context(format!("Failed to remove model '{}'", name))?;
            return Ok(());
        }

        let registry_model_dir = self.registry_root_dir().join(name);
        if registry_model_dir.exists() {
            anyhow::bail!(
                "Model '{}' is a registry model (versioned) - remove a specific version's directory under '{}' directly",
                name,
                registry_model_dir.display()
            );
        }

        anyhow::bail!("Model '{}' is not installed", name);
    }

    /// Get the path to a catalog model's ONNX file
    ///
    /// Note: This method is no longer used by CLI (replaced by node-runner in Phase 2).
    /// Kept for potential future use.
    #[allow(dead_code)]
    pub fn get_model_path(&self, name: &str) -> PathBuf {
        self.catalog_model_dir(name).join("model.onnx")
    }

    /// Check if a catalog model is installed
    ///
    /// Note: This method is no longer used by CLI (replaced by node-runner in Phase 2).
    /// Kept for potential future use.
    #[allow(dead_code)]
    pub async fn is_installed(&self, name: &str) -> bool {
        let model_path = self.get_model_path(name);
        model_path.exists()
    }

    /// Get info about a model (catalog or installed), from either origin
    /// (`LAB-1808`) - a Registry-origin model is just as valid a target here as
    /// a Catalog one, since `list_installed` already merges both namespaced
    /// subtrees.
    ///
    /// `version` disambiguates which installed entry to report when a
    /// Registry model has more than one version cached locally (a Catalog
    /// entry is versionless/always-latest, so `version` is ignored for it -
    /// there is only ever one candidate). When `version` is `None`, the
    /// highest installed Registry version is reported (Catalog has no
    /// versions to pick between).
    pub async fn info(&self, name: &str, version: Option<i32>) -> Result<ModelInfo> {
        let catalog_entry = self.get_catalog_entry(name).cloned();
        let installed = self.list_installed().await?;

        let installed_info = match version {
            Some(requested) => installed
                .into_iter()
                .find(|m| m.name == name && m.version == Some(requested)),
            None => {
                let mut candidates: Vec<InstalledModel> = installed.into_iter().filter(|m| m.name == name).collect();
                candidates.sort_by_key(|m| m.version.unwrap_or(0));
                candidates.pop()
            }
        };

        Ok(ModelInfo {
            name: name.to_string(),
            catalog_entry,
            installed_info,
        })
    }

    /// Validate a model file is a valid ONNX file
    #[allow(dead_code)]
    pub async fn validate(&self, path: &Path) -> Result<bool> {
        // Basic validation: check file exists and has .onnx extension
        if !path.exists() {
            return Ok(false);
        }

        if path.extension().and_then(|s| s.to_str()) != Some("onnx") {
            return Ok(false);
        }

        // TODO: Add ONNX format validation using ort crate
        // For now, just check the magic bytes
        let bytes = fs::read(path).await?;

        // ONNX files are Protocol Buffers, check for protobuf signature
        // This is a very basic check - full validation would require parsing
        Ok(bytes.len() > 4)
    }
}

/// Combined model information (catalog + installed)
#[derive(Debug, Clone, Serialize)]
pub struct ModelInfo {
    pub name: String,
    pub catalog_entry: Option<ModelCatalogEntry>,
    pub installed_info: Option<InstalledModel>,
}

impl ModelInfo {
    /// Check if model is installed
    #[allow(dead_code)]
    pub fn is_installed(&self) -> bool {
        self.installed_info.is_some()
    }
}

#[cfg(test)]
mod info_tests {
    use super::*;

    fn service(models_dir: PathBuf) -> ModelService {
        ModelService::with_models_dir(models_dir).unwrap()
    }

    async fn write_model(dir: PathBuf) {
        fs::create_dir_all(&dir).await.unwrap();
        fs::write(dir.join("model.onnx"), b"fake-onnx-bytes").await.unwrap();
    }

    #[tokio::test]
    async fn info_finds_a_registry_origin_model_not_just_catalog() {
        // LAB-1808: `info` must work for a Registry-origin model too, not just
        // Catalog - `models info` shouldn't require the model to also exist in
        // the (versionless) public catalog.
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        write_model(models_dir.join("registry").join("rover-detector").join("3")).await;

        let service = service(models_dir);
        let info = service.info("rover-detector", None).await.unwrap();

        assert!(info.catalog_entry.is_none());
        let installed = info.installed_info.expect("registry model should be found");
        assert_eq!(installed.source, ModelOrigin::Registry);
        assert_eq!(installed.version, Some(3));
    }

    #[tokio::test]
    async fn info_with_no_version_reports_the_highest_installed_registry_version() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        write_model(models_dir.join("registry").join("rover-detector").join("1")).await;
        write_model(models_dir.join("registry").join("rover-detector").join("3")).await;
        write_model(models_dir.join("registry").join("rover-detector").join("2")).await;

        let service = service(models_dir);
        let info = service.info("rover-detector", None).await.unwrap();

        assert_eq!(info.installed_info.unwrap().version, Some(3));
    }

    #[tokio::test]
    async fn info_with_an_explicit_version_reports_that_version() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        write_model(models_dir.join("registry").join("rover-detector").join("1")).await;
        write_model(models_dir.join("registry").join("rover-detector").join("2")).await;

        let service = service(models_dir);
        let info = service.info("rover-detector", Some(1)).await.unwrap();

        assert_eq!(info.installed_info.unwrap().version, Some(1));
    }

    #[tokio::test]
    async fn info_with_an_unknown_explicit_version_reports_not_installed() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        write_model(models_dir.join("registry").join("rover-detector").join("1")).await;

        let service = service(models_dir);
        let info = service.info("rover-detector", Some(99)).await.unwrap();

        assert!(info.installed_info.is_none());
    }

    #[tokio::test]
    async fn info_for_a_catalog_model_ignores_version_and_reports_it() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        write_model(models_dir.join("catalog").join("yolov8n")).await;

        let service = service(models_dir);
        let info = service.info("yolov8n", None).await.unwrap();

        let installed = info.installed_info.expect("catalog model should be found");
        assert_eq!(installed.source, ModelOrigin::Catalog);
        assert_eq!(installed.version, None);
    }

    #[tokio::test]
    async fn info_for_a_completely_unknown_model_reports_not_installed_and_no_catalog_entry() {
        let tmp = tempfile::tempdir().unwrap();
        let service = service(tmp.path().to_path_buf());

        let info = service.info("nonexistent-model", None).await.unwrap();

        assert!(info.catalog_entry.is_none());
        assert!(info.installed_info.is_none());
    }
}