candding 0.1.0

Pure candle embeddings for Rust with no ONNX: dense, sparse, multi-vector and rerank models today, image models on the roadmap
Documentation
use std::fs;
use std::path::Path;
use std::path::PathBuf;

use tempfile::tempdir;

use super::ID;
use super::REPO;
use super::SHA;
use super::body;
use super::config;
use super::hub;
use super::layout::blob_name;
use super::server::Reply;
use super::server::Server;
use super::server::closed_url;
use crate::error::CanddingError;
use crate::hub::cache;
use crate::hub::download::fetch;
use crate::hub::resolve_hub;
use crate::registry::WeightsFormat;

/// `hub`, except that the HEAD of `file` is refused.
fn refusing(file: &'static str) -> impl Fn(&str, &str) -> Reply + Send + Sync + 'static {
  move |method, path| match method == "HEAD" && path.ends_with(&format!("/resolve/{SHA}/{file}")) {
    true => Reply::new(403, ""),
    false => hub(method, path),
  }
}

/// Written before the downloads, the way huggingface_hub writes it, the ref would make this snapshot resolve.
#[test]
fn an_interrupted_branch_download_writes_no_ref_and_resolves_again() {
  let server = Server::start(refusing("config.json"));
  let cache = tempdir().unwrap();
  let config = config(cache.path(), &server.url);
  assert!(fetch(&config, ID, "main", WeightsFormat::Safetensors).is_err());
  let repo = cache.path().join(REPO);
  let first = repo
    .join("snapshots")
    .join(SHA)
    .join("1_Pooling/config.json");
  assert!(first.is_file(), "the first file arrived");
  assert!(!repo.join("refs/main").exists());
  let before = server.requests().len();
  assert!(resolve_hub(&config, ID, "main", WeightsFormat::Safetensors).is_err());
  assert!(server.requests().len() > before, "taken as complete");
}

/// A commit gets no ref, so the marker is what keeps a pinned resolve from taking the snapshot.
#[test]
fn an_interrupted_pinned_download_is_not_taken_as_complete() {
  let server = Server::start(refusing("config.json"));
  let cache = tempdir().unwrap();
  let config = config(cache.path(), &server.url);
  assert!(fetch(&config, ID, SHA, WeightsFormat::Safetensors).is_err());
  assert!(cache::in_progress_marker(cache.path(), ID, SHA).is_file());
  let before = server.requests().len();
  assert!(resolve_hub(&config, ID, SHA, WeightsFormat::Safetensors).is_err());
  assert!(server.requests().len() > before, "taken as complete");
}

#[test]
fn an_interrupted_repair_keeps_a_branch_from_resolving_offline() {
  let cache = tempdir().unwrap();
  let complete = Server::start(hub);
  let mut config = config(cache.path(), &complete.url);
  let snapshot = fetch(&config, ID, "main", WeightsFormat::Safetensors).unwrap();
  fs::remove_file(snapshot.join("model.safetensors")).unwrap();
  config.endpoint = Some(Server::start(refusing("model.safetensors")).url);
  assert!(fetch(&config, ID, "main", WeightsFormat::Safetensors).is_err());
  config.offline = true;
  match resolve_hub(&config, ID, "main", WeightsFormat::Safetensors) {
    Err(CanddingError::Offline { id }) => assert_eq!(id, ID),
    other => panic!("{other:?}"),
  }
}

#[test]
fn a_completed_pinned_download_resolves_without_a_request() {
  let server = Server::start(hub);
  let cache = tempdir().unwrap();
  let config = config(cache.path(), &server.url);
  let snapshot = fetch(&config, ID, SHA, WeightsFormat::Safetensors).unwrap();
  assert!(!cache::in_progress_marker(cache.path(), ID, SHA).exists());
  assert!(!cache.path().join(REPO).join("refs").exists());
  let before = server.requests().len();
  let resolved = resolve_hub(&config, ID, SHA, WeightsFormat::Safetensors).unwrap();
  assert_eq!(resolved, snapshot);
  assert_eq!(server.requests().len(), before);
}

/// What `snapshot_download(revision=<commit>)` leaves: blobs, relative links to them, no ref and no marker.
#[cfg(unix)]
fn another_tool_s_pinned_snapshot(cache: &Path) -> PathBuf {
  let repo = cache.join(REPO);
  let snapshot = repo.join("snapshots").join(SHA);
  for file in ["1_Pooling/config.json", "config.json", "model.safetensors"] {
    let name = blob_name(file);
    fs::create_dir_all(repo.join("blobs")).unwrap();
    fs::write(repo.join("blobs").join(&name), body(file)).unwrap();
    let pointer = snapshot.join(file);
    fs::create_dir_all(pointer.parent().unwrap()).unwrap();
    let up = "../".repeat(file.matches('/').count() + 2);
    std::os::unix::fs::symlink(format!("{up}blobs/{name}"), pointer).unwrap();
  }
  fs::create_dir_all(cache.join(".locks").join(REPO)).unwrap();
  snapshot
}

#[cfg(unix)]
#[test]
fn a_pinned_snapshot_another_tool_wrote_resolves_offline_unless_marked_in_progress() {
  let cache = tempdir().unwrap();
  let snapshot = another_tool_s_pinned_snapshot(cache.path());
  let mut config = config(cache.path(), &closed_url());
  config.offline = true;
  let resolved = resolve_hub(&config, ID, SHA, WeightsFormat::Safetensors).unwrap();
  assert_eq!(resolved, snapshot);
  let marker = cache::in_progress_marker(cache.path(), ID, SHA);
  fs::create_dir_all(marker.parent().unwrap()).unwrap();
  fs::write(&marker, "").unwrap();
  match resolve_hub(&config, ID, SHA, WeightsFormat::Safetensors) {
    Err(CanddingError::Offline { id }) => assert_eq!(id, ID),
    other => panic!("{other:?}"),
  }
}