candding 0.1.0

//! Native downloads: hf-hub lists a revision's files and streams their bytes, and candding writes the cache itself,
//! naming each blob by the hash it computes, so no value from a response header becomes a path.

mod blob;
mod failure;
mod listing;
mod lock;
#[cfg(test)]
mod tests;
mod verify;

use std::fs;
use std::path::PathBuf;
use std::thread;

use failure::Step;
use futures_util::StreamExt;
use hf_hub::HFClient;
use hf_hub::HFRepository;
use hf_hub::HFResult;
use hf_hub::RepoTypeModel;
use hf_hub::repository::RepoTreeEntry;
use hf_hub::split_id;

use crate::error::CanddingError;
use crate::error::Result;
use crate::hub::HubConfig;
use crate::hub::cache;
use crate::hub::include_patterns;
use crate::hub::pattern;
use crate::registry::WeightsFormat;

/// Fetches the files `format` needs at `revision` that the snapshot lacks, and returns the snapshot directory.
pub fn fetch(
  config: &HubConfig,
  repo_id: &str,
  revision: &str,
  format: WeightsFormat,
) -> Result<PathBuf> {
  if config.offline {
    return cache::snapshot_dir(&config.cache, repo_id, revision).ok_or_else(|| {
      CanddingError::Offline {
        id: repo_id.to_string(),
      }
    });
  }
  if !valid_repo_id(repo_id) {
    return Err(CanddingError::RepositoryNotFound {
      id: repo_id.to_string(),
    });
  }
  let _lock = lock::repository(&config.cache, repo_id)?;
  apart(|| {
    let runtime = tokio::runtime::Builder::new_current_thread()
      .enable_all()
      .build()?;
    let repo = repository(config, repo_id)
      .map_err(|error| failure::map(Step::Listing, error, repo_id, revision))?;
    runtime.block_on(online(&repo, config, repo_id, revision, format))
  })
}

/// hf-hub is async: a thread of its own lets a caller that is already inside a runtime block on the fetch.
fn apart<T: Send>(work: impl FnOnce() -> T + Send) -> T {
  thread::scope(|scope| {
    let worker = scope.spawn(work);
    worker
      .join()
      .unwrap_or_else(|panic| std::panic::resume_unwind(panic))
  })
}

async fn online(
  repo: &HFRepository<RepoTypeModel>,
  config: &HubConfig,
  repo_id: &str,
  revision: &str,
  format: WeightsFormat,
) -> Result<PathBuf> {
  let failed = |step| move |error| failure::map(step, error, repo_id, revision);
  let commit = repo
    .info()
    .revision(revision)
    .send()
    .await
    .map_err(failed(Step::Listing))?
    .sha
    .ok_or_else(|| CanddingError::DownloadFailed {
      id: repo_id.to_string(),
      reason: format!("the hub named no commit for revision `{revision}`"),
    })?;
  listing::check_commit(repo_id, &commit)?;
  let files = listing::files(tree(repo, &commit).await.map_err(failed(Step::Listing))?);
  listing::check_paths(repo_id, &files)?;
  let paths: Vec<String> = files.iter().map(|file| file.path.clone()).collect();
  let selected = pattern::select(include_patterns(format), &paths);
  let wanted = files
    .iter()
    .filter(|file| selected.contains(&file.path.as_str()));
  tracing::info!(
    repo_id,
    revision,
    commit = commit.as_str(),
    files = selected.len(),
    "downloading"
  );
  let repo_dir = config.cache.join(cache::repo_dir_name(repo_id));
  let snapshot = repo_dir.join("snapshots").join(&commit);
  // Marked before the first write and cleared after the ref, so an interrupted fetch never passes for complete.
  let marker = cache::in_progress_marker(&config.cache, repo_id, &commit);
  if let Some(parent) = marker.parent() {
    fs::create_dir_all(parent)?;
  }
  fs::write(&marker, "")?;
  for file in wanted {
    let pointer = snapshot.join(&file.path);
    if verify::intact(&pointer, file)? {
      continue;
    }
    tracing::debug!(file = file.path, "fetching");
    verify::remove(&pointer)?;
    blob::store(repo, file, &commit, &repo_dir, &pointer, failed(Step::File)).await?;
  }
  fs::create_dir_all(&snapshot)?;
  if revision != commit {
    let reference = repo_dir.join("refs").join(revision);
    if let Some(parent) = reference.parent() {
      fs::create_dir_all(parent)?;
    }
    fs::write(reference, &commit)?;
  }
  fs::remove_file(marker)?;
  Ok(snapshot)
}

async fn tree(repo: &HFRepository<RepoTypeModel>, commit: &str) -> HFResult<Vec<RepoTreeEntry>> {
  let stream = repo.list_tree().revision(commit).recursive(true).send()?;
  let mut stream = std::pin::pin!(stream);
  let mut entries = Vec::new();
  while let Some(entry) = stream.next().await {
    entries.push(entry?);
  }
  Ok(entries)
}

/// huggingface_hub's `validate_repo_id`, so a mistyped path such as `./name` is refused before the URL parser
/// drops its `.` and asks the hub for `name`.
pub(crate) fn valid_repo_id(id: &str) -> bool {
  let word = |c: char| c.is_alphanumeric() || c == '_';
  let parts: Vec<&str> = id.split('/').collect();
  parts.len() <= 2
    && parts.last().is_some_and(|name| name.chars().count() <= 96)
    && !id.contains("--")
    && !id.contains("..")
    && !id.ends_with(".git")
    && parts.iter().all(|part| {
      part.starts_with(word)
        && part.ends_with(word)
        && part.chars().all(|c| word(c) || c == '-' || c == '.')
    })
}

/// hf-hub writes nothing to the cache: candding does, and a disabled cache makes any hf-hub write an error.
fn repository(config: &HubConfig, repo_id: &str) -> HFResult<HFRepository<RepoTypeModel>> {
  let mut builder = HFClient::builder().cache_enabled(false);
  if let Some(endpoint) = &config.endpoint {
    builder = builder.endpoint(endpoint.clone());
  }
  if let Some(token) = &config.token {
    builder = builder.token(token.clone());
  }
  let (owner, name) = split_id(repo_id);
  Ok(builder.build()?.model(owner, name))
}