kernel/discovery/

hf_scanner.rs

//! Scans a Hugging Face hub cache: each `models--<org>--<repo>` directory's
//! current snapshot is inspected (`config.json` / `model_index.json` / a bare
//! GGUF) for a modality hint, its blobs summed for the footprint, and its
//! shard/blob completeness checked to flag a still-downloading model.
//!
//! The diffusers `model_index.json` path yields only a generic job hint until the
//! pipeline-family registry it needs is ported.

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

use crate::discovery::gguf_models::is_mmproj_name;
use crate::discovery::gguf_shards::group;
use crate::discovery::modality_hints::{self, Hint};
use crate::discovery::scanner::{DiscoveredModel, ScanResult, StoreScanner};
use crate::discovery::weights::{gguf_tree, primary_of};
use crate::records::{ExecutionMode, JsonValue, Modality, ModelSource, SourceKind};
use crate::resolution::has_ggml_magic;

/// Filenames that, alongside a text/unknown model, mark a sentence-transformers
/// embedding model.
const SENTENCE_TRANSFORMERS_MARKERS: [&str; 2] = ["config_sentence_transformers.json", "1_Pooling"];

/// A scanner over one or more Hugging Face hub cache roots.
pub struct HFCacheScanner {
    roots: Vec<PathBuf>,
    user_roots: Vec<PathBuf>,
}

impl HFCacheScanner {
    /// A scanner over the given standard cache `roots` (a missing root is skipped).
    pub fn new(roots: Vec<PathBuf>) -> Self {
        Self {
            roots,
            user_roots: Vec::new(),
        }
    }

    /// A scanner over a single root.
    pub fn single(root: impl Into<PathBuf>) -> Self {
        Self::new(vec![root.into()])
    }

    /// A scanner over standard `roots` (missing → skipped) plus `user_roots`
    /// (missing → a scan failure, since the user pointed at them explicitly).
    pub fn with_user_roots(roots: Vec<PathBuf>, user_roots: Vec<PathBuf>) -> Self {
        Self { roots, user_roots }
    }

    fn scan_root(&self, root: &Path, required: bool, result: &mut ScanResult) {
        if !root.exists() {
            if required {
                mark_failed(result);
            }
            return;
        }
        let Ok(entries) = std::fs::read_dir(root) else {
            mark_failed(result);
            return;
        };

        for entry in entries.flatten() {
            let dir = entry.path();
            let Some(dir_name) = dir.file_name().and_then(|name| name.to_str()) else {
                continue;
            };
            let Some(rest) = dir_name.strip_prefix("models--") else {
                continue;
            };
            let repo = rest.replace("--", "/");

            let Some((snapshot, revision)) = current_snapshot(&dir) else {
                result
                    .issues
                    .push(format!("hf-cache: {repo} has no usable snapshot"));
                continue;
            };

            let names = snapshot_file_names(&snapshot);
            // Walked once and shared: a repo that keeps a directory per
            // quantization has no weights at the snapshot root, so the hint,
            // the completeness check and the primary weight all have to look
            // below it.
            let ggufs = gguf_tree(&snapshot).weights;
            let mut diagnostics = Vec::new();
            let hint = resolve_hint(&snapshot, &names, &ggufs, &mut diagnostics);

            let downloading = has_incomplete_blobs(&dir.join("blobs"))
                || index_references_missing_shard(&snapshot, &names)
                || gguf_shards_incomplete(&ggufs);

            // Last non-empty path segment (empty segments are skipped, so a
            // trailing slash doesn't yield an empty name).
            let name = repo
                .rsplit('/')
                .find(|segment| !segment.is_empty())
                .unwrap_or(&repo)
                .to_owned();
            let mut source = ModelSource::new(SourceKind::huggingface_cache(), &display(&dir));
            source.repo = Some(repo);
            source.reference = Some(revision);

            let mut discovered = DiscoveredModel::new(name, source);
            discovered.modality_hint = hint.modality;
            discovered.capabilities_hint = hint.capabilities;
            discovered.execution_hint = hint.execution;
            discovered.footprint_bytes = directory_bytes(&dir.join("blobs"));
            discovered.primary_weight_path = largest_weight(&snapshot, &ggufs);
            discovered.diagnostics = diagnostics;
            discovered.context_length_hint = hint.context_length;
            discovered.downloading = downloading;
            result.discovered.push(discovered);
        }
    }
}

impl StoreScanner for HFCacheScanner {
    fn kinds(&self) -> Vec<SourceKind> {
        vec![SourceKind::huggingface_cache()]
    }

    fn scan(&self) -> ScanResult {
        let mut result = ScanResult::default();
        for root in &self.roots {
            self.scan_root(root, false, &mut result);
        }
        for root in &self.user_roots {
            self.scan_root(root, true, &mut result);
        }
        result
    }
}

fn mark_failed(result: &mut ScanResult) {
    let kind = SourceKind::huggingface_cache();
    if !result.failed_kinds.contains(&kind) {
        result.failed_kinds.push(kind);
    }
}

/// Pick the snapshot to represent a repo: `refs/main` if it points at a present
/// snapshot, else the most-recently-modified snapshot directory.
fn current_snapshot(repo_dir: &Path) -> Option<(PathBuf, String)> {
    let snapshots = repo_dir.join("snapshots");

    if let Ok(revision) = std::fs::read_to_string(repo_dir.join("refs/main")) {
        let trimmed = revision.trim();
        if !trimmed.is_empty() {
            let snapshot = snapshots.join(trimmed);
            if snapshot.exists() {
                return Some((snapshot, trimmed.to_owned()));
            }
        }
    }

    let mut newest: Option<(PathBuf, std::time::SystemTime)> = None;
    for entry in std::fs::read_dir(&snapshots)
        .into_iter()
        .flatten()
        .flatten()
    {
        if entry
            .file_name()
            .to_str()
            .is_some_and(|name| name.starts_with('.'))
        {
            continue;
        }
        let modified = entry
            .metadata()
            .and_then(|meta| meta.modified())
            .unwrap_or(std::time::UNIX_EPOCH);
        // Strict `>` keeps the first of equal-mtime snapshots.
        if newest.as_ref().is_none_or(|(_, best)| modified > *best) {
            newest = Some((entry.path(), modified));
        }
    }
    newest.map(|(path, _)| {
        let revision = path
            .file_name()
            .and_then(|name| name.to_str())
            .unwrap_or_default()
            .to_owned();
        (path, revision)
    })
}

fn snapshot_file_names(snapshot: &Path) -> BTreeSet<String> {
    let mut names = BTreeSet::new();
    for entry in std::fs::read_dir(snapshot).into_iter().flatten().flatten() {
        if let Some(name) = entry.file_name().to_str()
            && !name.starts_with('.')
        {
            names.insert(name.to_owned());
        }
    }
    names
}

/// Determine the modality hint for a snapshot from its files, then apply the
/// sentence-transformers and missing-tokenizer refinements.
fn resolve_hint(
    snapshot: &Path,
    names: &BTreeSet<String>,
    ggufs: &[(PathBuf, u64)],
    diagnostics: &mut Vec<String>,
) -> Hint {
    let mut hint = if names.contains("model_index.json") {
        modality_hints::from_model_index(&snapshot.join("model_index.json"))
    } else if names.contains("config.json") {
        modality_hints::from_config_json(&snapshot.join("config.json"))
            .unwrap_or_else(|| Hint::unknown(ExecutionMode::Sync))
    } else if !ggufs.is_empty() {
        modality_hints::gguf_hint()
    } else {
        diagnostics.push("no config.json or model_index.json in snapshot".to_owned());
        Hint::unknown(ExecutionMode::Sync)
    };

    let text = Some(Modality::text());
    if (hint.modality.is_none() || hint.modality == text)
        && names
            .iter()
            .any(|name| SENTENCE_TRANSFORMERS_MARKERS.contains(&name.as_str()))
    {
        let mut embedding = modality_hints::embedding_hint();
        embedding.context_length = hint.context_length;
        hint = embedding;
    }

    if hint.modality == text
        && !names
            .iter()
            .any(|name| name.starts_with("tokenizer") || name == "vocab.json")
    {
        diagnostics.push("no tokenizer found".to_owned());
    }

    hint
}

fn has_incomplete_blobs(blobs: &Path) -> bool {
    for entry in std::fs::read_dir(blobs).into_iter().flatten().flatten() {
        let is_incomplete = entry
            .file_name()
            .to_str()
            .is_some_and(|name| name.ends_with(".incomplete"));
        if is_incomplete
            && entry
                .file_type()
                .map(|kind| kind.is_file())
                .unwrap_or(false)
        {
            return true;
        }
    }
    false
}

fn gguf_shards_incomplete(ggufs: &[(PathBuf, u64)]) -> bool {
    // Sizes are irrelevant to completeness, which counts members against the
    // total each name declares.
    let sized: Vec<(PathBuf, i64)> = ggufs.iter().map(|(path, _)| (path.clone(), 0)).collect();
    let (groups, _) = group(&sized);
    groups.iter().any(|shard_group| !shard_group.complete())
}

/// Whether a `*.safetensors.index.json` weight map references a shard whose file
/// (resolving symlinks into the blob store) is missing — the sign of a partial
/// safetensors download.
fn index_references_missing_shard(snapshot: &Path, names: &BTreeSet<String>) -> bool {
    for index_name in names
        .iter()
        .filter(|name| name.ends_with(".safetensors.index.json"))
    {
        let Ok(bytes) = std::fs::read(snapshot.join(index_name)) else {
            continue;
        };
        let Ok(JsonValue::Object(json)) = serde_json::from_slice::<JsonValue>(&bytes) else {
            continue;
        };
        let Some(JsonValue::Object(weight_map)) = json.get("weight_map") else {
            continue;
        };
        let shards: BTreeSet<&str> = weight_map.values().filter_map(JsonValue::as_str).collect();
        // `exists()` follows the snapshot's symlink into `blobs/`, so a dangling
        // link (a shard not yet downloaded) reads as missing.
        if shards.iter().any(|shard| !snapshot.join(shard).exists()) {
            return true;
        }
    }
    false
}

fn directory_bytes(dir: &Path) -> i64 {
    let mut total = 0;
    walk_bytes(dir, &mut total);
    total
}

fn walk_bytes(dir: &Path, total: &mut i64) {
    for entry in std::fs::read_dir(dir).into_iter().flatten().flatten() {
        let path = entry.path();
        match entry.file_type() {
            Ok(kind) if kind.is_dir() => walk_bytes(&path, total),
            // Follow symlinks for the size (a snapshot's weights are links into
            // `blobs/`); count only what resolves to a regular file.
            _ => {
                if let Ok(meta) = std::fs::metadata(&path)
                    && meta.is_file()
                {
                    *total += meta.len() as i64;
                }
            }
        }
    }
}

/// The weight file a runtime would load, resolved through its symlink: the
/// largest weight in the snapshot, where a GGUF is chosen by the rule
/// identification reads the header by, and every other format is measured at
/// the snapshot root, which is where those formats put their weights.
///
/// A shard set is weighed by its largest member but answers with its first
/// file, which is what a server loads the rest through.
fn largest_weight(snapshot: &Path, ggufs: &[(PathBuf, u64)]) -> Option<String> {
    let gguf = primary_of(ggufs).map(|path| {
        let largest = ggufs.iter().map(|(_, size)| *size).max().unwrap_or(0);
        (path, largest)
    });
    let best = match (gguf, largest_other_weight(snapshot)) {
        (Some(gguf), Some(other)) if other.1 > gguf.1 => other,
        (Some(gguf), _) => gguf,
        (None, other) => other?,
    };
    Some(resolve(&best.0))
}

/// The largest safetensors or GGML weight at the snapshot root, with a tie
/// going to the earlier path so one snapshot always reads the same way.
fn largest_other_weight(snapshot: &Path) -> Option<(PathBuf, u64)> {
    let mut best: Option<(PathBuf, u64)> = None;
    for entry in std::fs::read_dir(snapshot).into_iter().flatten().flatten() {
        let path = entry.path();
        let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
            continue;
        };
        if name.starts_with('.') || !is_other_weight_file(&path) {
            continue;
        }
        let Ok(metadata) = std::fs::metadata(&path) else {
            continue;
        };
        if !metadata.is_file() {
            continue;
        }
        let size = metadata.len();
        let better = best.as_ref().is_none_or(|(best_path, best_size)| {
            size > *best_size || (size == *best_size && path < *best_path)
        });
        if better {
            best = Some((path, size));
        }
    }
    best
}

/// Whether `path` is a weight file in a format other than GGUF, which
/// [`primary_of`] answers for.
fn is_other_weight_file(path: &Path) -> bool {
    let name = path
        .file_name()
        .and_then(|name| name.to_str())
        .unwrap_or_default();
    if is_mmproj_name(name) {
        return false;
    }
    match path
        .extension()
        .and_then(|ext| ext.to_str())
        .map(str::to_ascii_lowercase)
        .as_deref()
    {
        Some("safetensors") => true,
        Some("bin") => has_ggml_magic(path),
        _ => false,
    }
}

/// A path resolved through symlinks (falling back to itself), as a string.
fn resolve(path: &Path) -> String {
    std::fs::canonicalize(path)
        .unwrap_or_else(|_| path.to_path_buf())
        .to_string_lossy()
        .into_owned()
}

fn display(path: &Path) -> String {
    path.to_string_lossy().into_owned()
}