diff --git a/Cargo.lock b/Cargo.lock index 4e91ab75..2e4018f3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1567,6 +1567,7 @@ dependencies = [ name = "kernels-python" version = "0.18.0-dev0" dependencies = [ + "eyre", "kernels-common", "pyo3", "serde_json", diff --git a/kernels-common/Cargo.toml b/kernels-common/Cargo.toml index 4a60b12b..dd98dc8f 100644 --- a/kernels-common/Cargo.toml +++ b/kernels-common/Cargo.toml @@ -19,10 +19,9 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" serde-value = "0.7" sha2 = "0.11" +tempfile = "3" thiserror = "1" toml = "0.8" url = { version = "2", features = ["serde"] } walkdir = "2" -[dev-dependencies] -tempfile = "3" diff --git a/kernels-common/src/hf.rs b/kernels-common/src/hf.rs index 099d5baa..8d8cd526 100644 --- a/kernels-common/src/hf.rs +++ b/kernels-common/src/hf.rs @@ -71,15 +71,24 @@ fn expand_tilde_with_home(path: PathBuf, home: Option) -> PathBuf { home.join(rest) } +/// The kernel cache directory could not be determined. +/// +/// This happens when none of `KERNELS_CACHE`, `HF_HUB_CACHE`, +/// `HUGGINGFACE_HUB_CACHE`, `HF_HOME`, and `XDG_CACHE_HOME` are set and the +/// user's home directory cannot be resolved. +#[derive(Clone, Copy, Debug, Error)] +#[error( + "cannot determine the kernel cache directory, set `KERNELS_CACHE`, `HF_HUB_CACHE`, or `HF_HOME`" +)] +pub struct UnknownCacheDir; + /// Error building a Hub client. #[derive(Debug, Error)] #[non_exhaustive] pub enum HFKernelsClientError { - /// The kernel cache directory could not be detepmined. - #[error( - "cannot determine the kernel cache directory, set `KERNELS_CACHE`, `HF_HUB_CACHE`, or `HF_HOME`" - )] - UnknownCacheDir, + /// The kernel cache directory could not be determined. + #[error(transparent)] + UnknownCacheDir(#[from] UnknownCacheDir), /// The underlying `hf-hub` client could not be constructed. #[error("cannot create Hugging Face Hub client")] @@ -138,7 +147,7 @@ impl HFKernelsClientBuilder { fn hf_client_builder(self) -> Result { let cache_dir = match self.cache_dir { Some(cache_dir) => cache_dir, - None => kernels_cache().ok_or(HFKernelsClientError::UnknownCacheDir)?, + None => kernels_cache()?, }; let mut builder = HFClient::builder() @@ -175,8 +184,8 @@ fn hf_hub_cache() -> Option { } /// The kernels cache directory. -fn kernels_cache() -> Option { - resolve_kernels_cache(env_path("KERNELS_CACHE"), hf_hub_cache()) +pub(crate) fn kernels_cache() -> Result { + resolve_kernels_cache(env_path("KERNELS_CACHE"), hf_hub_cache()).ok_or(UnknownCacheDir) } fn resolve_hf_home( diff --git a/kernels-common/src/signing/receipt.rs b/kernels-common/src/signing/receipt.rs index 86875caf..cff9b99c 100644 --- a/kernels-common/src/signing/receipt.rs +++ b/kernels-common/src/signing/receipt.rs @@ -1,6 +1,34 @@ +use std::fs; +use std::io::{self, Write as _}; +use std::path::PathBuf; + use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use thiserror::Error; use crate::git::Oid; +use crate::hf::{UnknownCacheDir, kernels_cache}; + +/// Version of the on-disk receipt format. +pub const CACHE_FORMAT_VERSION: &str = "v1"; + +/// Receipt of a successful kernel verification. +#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] +pub struct VerificationReceipt { + /// The kernel location the receipt applies to. + location: KernelLocation, +} + +impl VerificationReceipt { + pub fn new(location: KernelLocation) -> Self { + VerificationReceipt { location } + } + + /// The kernel location the verification applies to. + pub fn location(&self) -> &KernelLocation { + &self.location + } +} /// Kernel location. /// @@ -27,4 +55,355 @@ impl KernelLocation { variant: variant.into(), } } + + fn receipt_key(&self) -> String { + let mut hasher = Sha256::new(); + match self { + KernelLocation::RemoteKernel { + repo_id, + revision, + variant, + } => { + update_with_string(&mut hasher, "remote_kernel"); + update_with_string(&mut hasher, repo_id); + update_with_string(&mut hasher, revision.as_str()); + update_with_string(&mut hasher, variant); + } + } + hex_encode(&hasher.finalize()) + } +} + +/// Storage for verification receipts. +/// +/// Receipts are stored as JSON files in a cache directory, named by their +/// receipt key. +#[derive(Clone, Debug)] +pub struct ReceiptStore { + dir: PathBuf, +} + +impl ReceiptStore { + /// The receipt store inside the kernels cache. + pub fn in_kernels_cache() -> Result { + let default_dir = kernels_cache()? + .join(".verified-kernels") + .join(CACHE_FORMAT_VERSION); + + Ok(Self::from_path(default_dir)) + } + + /// A receipt store in the given directory. + pub fn from_path(dir: impl Into) -> Self { + ReceiptStore { dir: dir.into() } + } + + /// Load the receipt for the given kernel location. + /// + /// Returns `Ok(None)` when no receipt exists for the location, and an + /// error when a receipt exists but cannot be read, is corrupt, or + /// describes a different kernel. + pub fn load( + &self, + location: &KernelLocation, + ) -> Result, ReceiptStoreError> { + let path = self.dir.join(location.receipt_key()); + let data = match fs::read(&path) { + Ok(data) => data, + Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(None), + Err(source) => return Err(ReceiptStoreError::Read { path, source }), + }; + let receipt: VerificationReceipt = + serde_json::from_slice(&data).map_err(|source| ReceiptStoreError::Corrupt { + path: path.clone(), + source, + })?; + + if receipt.location != *location { + return Err(ReceiptStoreError::LocationMismatch { + path, + found: Box::new(receipt.location), + }); + } + + Ok(Some(receipt)) + } + + /// Store a receipt. + pub fn store(&self, receipt: &VerificationReceipt) -> Result<(), ReceiptStoreError> { + let path = self.dir.join(receipt.location.receipt_key()); + let write_err = |source: io::Error| ReceiptStoreError::Write { + path: path.clone(), + source, + }; + + // This extra ceremony is so that we write the receipt atomically. + let payload = serde_json::to_string(receipt).map_err(|e| write_err(io::Error::other(e)))?; + fs::create_dir_all(&self.dir).map_err(&write_err)?; + let mut tmp_file = tempfile::NamedTempFile::new_in(&self.dir).map_err(&write_err)?; + tmp_file.write_all(payload.as_bytes()).map_err(&write_err)?; + // Without this a crash can leave a zero-length receipt behind, which + // reads back as `Corrupt` rather than as a plain cache miss. + tmp_file.as_file().sync_all().map_err(&write_err)?; + tmp_file.persist(&path).map_err(|e| write_err(e.error))?; + Ok(()) + } +} + +/// Error handling a verification receipt. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum ReceiptStoreError { + /// The receipt file exists but cannot be read. + #[error("cannot read receipt `{path}`")] + Read { + path: PathBuf, + #[source] + source: io::Error, + }, + + /// The receipt cannot be interpreted: invalid JSON, schema mismatch, or + /// bad base64. + #[error("receipt `{path}` is corrupt")] + Corrupt { + path: PathBuf, + #[source] + source: serde_json::Error, + }, + + /// The receipt describes a different kernel than the one it was looked + /// up for. + #[error("receipt `{path}` describes a different kernel: {found:?}")] + LocationMismatch { + path: PathBuf, + /// The location recorded in the receipt. + found: Box, + }, + + /// The receipt cannot be written, e.g. because the cache is not + /// writable. Also covers serialization failures. + #[error("cannot store receipt `{path}`")] + Write { + path: PathBuf, + #[source] + source: io::Error, + }, + + /// The receipt store location could not be determined. + #[error(transparent)] + UnknownCacheDir(#[from] UnknownCacheDir), +} + +/// Length-prefixed string hash. +fn update_with_string(hasher: &mut impl Digest, part: &str) { + hasher.update((part.len() as u64).to_be_bytes()); + hasher.update(part.as_bytes()); +} + +fn hex_encode(bytes: &[u8]) -> String { + use std::fmt::Write as _; + bytes + .iter() + .fold(String::with_capacity(2 * bytes.len()), |mut s, b| { + let _ = write!(s, "{b:02x}"); + s + }) +} + +#[cfg(test)] +mod tests { + use std::str::FromStr; + + use super::*; + use tempfile::TempDir; + + fn oid(hex_digit: &str) -> Oid { + Oid::from_str(&hex_digit.repeat(40)).unwrap() + } + + fn hub_location() -> KernelLocation { + KernelLocation::remote( + "kernels-test/signatures", + oid("a"), + "torch30-cxx11-cu128-x86_64-linux", + ) + } + + #[test] + fn receipt_key_is_deterministic_and_distinguishes_kernels() { + let location = hub_location(); + assert_eq!(location.receipt_key(), location.receipt_key()); + + for other in [ + KernelLocation::remote( + "kernels-test/other", + oid("a"), + "torch30-cxx11-cu128-x86_64-linux", + ), + KernelLocation::remote( + "kernels-test/signatures", + oid("b"), + "torch30-cxx11-cu128-x86_64-linux", + ), + KernelLocation::remote( + "kernels-test/signatures", + oid("a"), + "torch30-cxx11-cpu-x86_64-linux", + ), + ] { + assert_ne!(location.receipt_key(), other.receipt_key()); + } + } + + /// The parts of a location are length-prefixed, so that moving a + /// character from one part to the next cannot produce the same key. + #[test] + fn receipt_key_does_not_confuse_adjacent_parts() { + let split_one = KernelLocation::remote("kernels-test/sig", oid("a"), "natures"); + let split_other = KernelLocation::remote("kernels-test/signatures", oid("a"), ""); + assert_ne!(split_one.receipt_key(), split_other.receipt_key()); + } + + #[test] + fn receipt_roundtrip() -> io::Result<()> { + let dir = TempDir::new()?; + let store = ReceiptStore::from_path(dir.path()); + let location = hub_location(); + let receipt = VerificationReceipt::new(location.clone()); + + store.store(&receipt).expect("receipt should store"); + let loaded = store + .load(&location) + .expect("receipt should load") + .expect("receipt should exist"); + + assert_eq!(loaded.location(), receipt.location()); + Ok(()) + } + + #[test] + fn load_receipt_missing_or_corrupt() { + let dir = TempDir::new().unwrap(); + let store = ReceiptStore::from_path(dir.path()); + let location = hub_location(); + let key = location.receipt_key(); + + assert!(matches!(store.load(&location), Ok(None))); + + fs::write(dir.path().join(&key), "not a receipt").unwrap(); + assert!(matches!( + store.load(&location), + Err(ReceiptStoreError::Corrupt { .. }) + )); + + // Valid JSON, but the location does not match the schema. + fs::write( + dir.path().join(&key), + r#"{"location":{"type":"remote_kernel","repo_id":"kernels-test/signatures"}}"#, + ) + .unwrap(); + assert!(matches!( + store.load(&location), + Err(ReceiptStoreError::Corrupt { .. }) + )); + + // A location variant that this version does not know. A newer + // `kernels` may add one, and reading it must fail cleanly rather + // than be misinterpreted: the caller then re-verifies and overwrites. + fs::write( + dir.path().join(&key), + r#"{"location":{"type":"lunar_kernel","crater":"Tycho"}}"#, + ) + .unwrap(); + assert!(matches!( + store.load(&location), + Err(ReceiptStoreError::Corrupt { .. }) + )); + } + + #[test] + fn load_receipt_rejects_transplanted_receipt() { + let dir = TempDir::new().unwrap(); + let store = ReceiptStore::from_path(dir.path()); + + let signed = hub_location(); + let unsigned = KernelLocation::remote( + "kernels-test/signatures", + oid("b"), + "torch30-cxx11-cu128-x86_64-linux", + ); + + store + .store(&VerificationReceipt::new(signed.clone())) + .expect("receipt should store"); + + // Transplant the receipt onto the other revision's key. + fs::copy( + dir.path().join(signed.receipt_key()), + dir.path().join(unsigned.receipt_key()), + ) + .unwrap(); + + match store.load(&unsigned) { + Err(ReceiptStoreError::LocationMismatch { found, .. }) => { + assert_eq!(*found, signed); + } + other => panic!("expected a location mismatch, got: {other:?}"), + } + + // The receipt is still valid under its own key. + assert!(store.load(&signed).unwrap().is_some()); + } + + #[test] + fn store_receipt_fails_when_cache_not_writable() { + let dir = TempDir::new().unwrap(); + let receipt_dir = dir.path().join("receipts"); + // A regular file where the receipt directory should be. + fs::write(&receipt_dir, "not a directory").unwrap(); + + let receipt = VerificationReceipt::new(hub_location()); + assert!(matches!( + ReceiptStore::from_path(&receipt_dir).store(&receipt), + Err(ReceiptStoreError::Write { .. }) + )); + } + + #[test] + fn receipt_json_format_is_stable() { + let receipt = VerificationReceipt::new(hub_location()); + + let json = serde_json::to_string(&receipt).unwrap(); + let expected = format!( + r#"{{"location":{{"type":"remote_kernel","repo_id":"kernels-test/signatures","revision":"{}","variant":"torch30-cxx11-cu128-x86_64-linux"}}}}"#, + "a".repeat(40) + ); + assert_eq!(json, expected); + + // The pinned format must roundtrip. + let parsed: VerificationReceipt = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed.location(), receipt.location()); + } + + #[test] + fn unknown_receipt_fields_are_ignored() { + let dir = TempDir::new().unwrap(); + let store = ReceiptStore::from_path(dir.path()); + let location = hub_location(); + + fs::write( + dir.path().join(location.receipt_key()), + format!( + r#"{{"location":{{"type":"remote_kernel","repo_id":"kernels-test/signatures","revision":"{}","variant":"torch30-cxx11-cu128-x86_64-linux","from_the_future":[1,2,3]}},"verified_at":"2026-09-10T13:42:47Z"}}"#, + "a".repeat(40) + ), + ) + .unwrap(); + + let loaded = store + .load(&location) + .expect("receipt should load") + .expect("receipt should exist"); + assert_eq!(loaded.location(), &location); + } } diff --git a/kernels/Cargo.toml b/kernels/Cargo.toml index 0b917258..2aef6216 100644 --- a/kernels/Cargo.toml +++ b/kernels/Cargo.toml @@ -13,6 +13,7 @@ path = "rust/lib.rs" crate-type = ["cdylib"] [dependencies] +eyre = "0.6.12" pyo3 = { version = "0.26", features = ["abi3", "abi3-py38"] } serde_json = "1" diff --git a/kernels/rust/lib.rs b/kernels/rust/lib.rs index 4f1e2e16..a502a911 100644 --- a/kernels/rust/lib.rs +++ b/kernels/rust/lib.rs @@ -21,7 +21,7 @@ mod version; use config::{PyBuild, PyGeneral}; use git::PyOid; use lock::{PyKernelLock, PyKernelLocks, PyKernelPaths, PyNixKernelLock, PyNixKernelLocks}; -use signing::PyKernelLocation; +use signing::{PyKernelLocation, PyReceiptStore, PyVerificationReceipt, ReceiptError}; use version::PyVersion; /// A validated kernel name matching `^[a-z][-a-z0-9]*[a-z0-9]$`. @@ -778,10 +778,13 @@ fn data_py(m: &PyBound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add( "DigestValidationError", m.py().get_type::(), )?; + m.add("ReceiptError", m.py().get_type::())?; m.add("__version__", env!("CARGO_PKG_VERSION"))?; Ok(()) diff --git a/kernels/rust/signing.rs b/kernels/rust/signing.rs index 4bf7dafc..8eb2ebe8 100644 --- a/kernels/rust/signing.rs +++ b/kernels/rust/signing.rs @@ -1,8 +1,23 @@ -use kernels_common::signing::receipt::KernelLocation; +use std::path::PathBuf; + +use kernels_common::signing::receipt::{KernelLocation, ReceiptStore, VerificationReceipt}; +use pyo3::exceptions::PyException; use pyo3::prelude::*; use crate::git::PyOid; +pyo3::create_exception!( + _rust, + ReceiptError, + PyException, + "Raised by `ReceiptStore` when a receipt cannot be read, written, or \ + interpreted.\n\n\ + A missing receipt is not an error: `ReceiptStore.load` returns `None` \ + for it. Since a receipt is only a cache of a previous verification, \ + callers can treat this exception as a cache miss and re-verify, at the \ + cost of not noticing a cache that is persistently broken." +); + /// The location of a kernel that a verification applies to. #[pyclass(name = "KernelLocation", frozen, eq, hash)] #[derive(Clone, Debug, Eq, Hash, PartialEq)] @@ -39,3 +54,85 @@ impl PyKernelLocation { } } } + +/// Receipt of a successful kernel verification. +#[pyclass(name = "VerificationReceipt", frozen, eq, hash)] +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub(crate) struct PyVerificationReceipt { + inner: VerificationReceipt, +} + +impl From for PyVerificationReceipt { + fn from(inner: VerificationReceipt) -> Self { + Self { inner } + } +} + +#[pymethods] +impl PyVerificationReceipt { + #[new] + fn new(location: &PyKernelLocation) -> Self { + VerificationReceipt::new(location.inner.clone()).into() + } + + #[getter] + fn location(&self) -> PyKernelLocation { + self.inner.location().clone().into() + } + + fn __repr__(&self) -> String { + format!( + "VerificationReceipt(location={})", + self.location().__repr__() + ) + } +} + +/// Store of kernel verification receipts. +#[pyclass(name = "ReceiptStore", frozen)] +#[derive(Clone, Debug)] +pub(crate) struct PyReceiptStore { + inner: ReceiptStore, +} + +#[pymethods] +impl PyReceiptStore { + /// The receipt store inside the kernels cache. + /// + /// Raises `ReceiptError` when the cache directory cannot be determined, + /// in which case verifications cannot be cached. + #[staticmethod] + fn in_kernels_cache() -> PyResult { + ReceiptStore::in_kernels_cache() + .map(|inner| PyReceiptStore { inner }) + .map_err(|err| ReceiptError::new_err(format!("{:#}", eyre::Report::new(err)))) + } + + /// A receipt store in the given directory. + #[staticmethod] + fn from_path(path: PathBuf) -> Self { + PyReceiptStore { + inner: ReceiptStore::from_path(path), + } + } + + /// The receipt for `location`, or `None` when the kernel has not been + /// verified yet. + /// + /// Raises `ReceiptError` if a receipt exists but cannot be used. + fn load(&self, location: &PyKernelLocation) -> PyResult> { + self.inner + .load(&location.inner) + .map(|receipt| receipt.map(Into::into)) + .map_err(|err| ReceiptError::new_err(format!("{:#}", eyre::Report::new(err)))) + } + + /// Store `receipt`, replacing any existing receipt for its location. + /// + /// Raises `ReceiptError` if the receipt cannot be written. + fn store(&self, receipt: &PyVerificationReceipt) -> PyResult<()> { + self.inner + .store(&receipt.inner) + .map_err(|err| ReceiptError::new_err(format!("{:#}", eyre::Report::new(err)))) + } +} diff --git a/kernels/src/kernels/_versions.py b/kernels/src/kernels/_versions.py index 5ea778a9..65e7cef3 100644 --- a/kernels/src/kernels/_versions.py +++ b/kernels/src/kernels/_versions.py @@ -1,4 +1,6 @@ import logging +import os +import tempfile from pathlib import Path from huggingface_hub import constants @@ -6,16 +8,14 @@ from huggingface_hub.hf_api import GitRefInfo from kernels._rust import KernelVersion, Oid +from kernels.hf_hub import _get_cache_dir logger = logging.getLogger(__name__) def _cached_refs_dir(repo_id: str) -> Path: """The cache directory that holds the refs of a kernel repository.""" - # Lazy import so that we can mock it in tests. - from kernels.hf_hub import CACHE_DIR - - cache_dir = CACHE_DIR or constants.HF_HUB_CACHE + cache_dir = _get_cache_dir() or constants.HF_HUB_CACHE return Path(cache_dir) / repo_folder_name(repo_id=repo_id, repo_type="kernel") / "refs" @@ -131,6 +131,47 @@ def _resolve_ref_from_cache(repo_id: str, ref: str) -> str | None: return None +def _record_ref_in_cache(repo_id: str, ref: str, commit: str) -> None: + """Record that `ref` points at `commit` in the local Hugging Face cache. + + Errors are ignored: not being able to write to the cache must never make + a kernel fail to load. + """ + if ref == commit: + return + + refs_dir = _cached_refs_dir(repo_id) + ref_path = refs_dir / ref + + # Extra guard against path-travesal attacks (in addition to _resolve_ref_from_cache). + try: + ref_path.resolve().relative_to(refs_dir.resolve()) + except (OSError, ValueError): + return + + try: + if ref_path.is_file() and ref_path.read_text() == commit: + return + + ref_path.parent.mkdir(parents=True, exist_ok=True) + + # Write automically to avoid races. Place in the cache directory, since + # os.replace() is not atomic between filesystems. + fd, tmp_name = tempfile.mkstemp(dir=ref_path.parent, prefix=f".{ref_path.name}.") + try: + with os.fdopen(fd, "w") as tmp_file: + tmp_file.write(commit) + os.replace(tmp_name, ref_path) + except BaseException: + try: + os.unlink(tmp_name) + except OSError: + pass + raise + except OSError as e: + logger.warning("Could not record revision '%s' of '%s' in the cache: %s", ref, repo_id, e) + + def _resolve_ref(repo_id: str, ref: str, *, local_files_only: bool) -> Oid: """Resolve a branch, tag, or commit to the commit it points at. @@ -168,13 +209,31 @@ def resolve_kernel_version(repo_id: str, version: KernelVersion, *, local_files_ full Git commit SHA. """ if isinstance(version, KernelVersion.Version): - ref = resolve_version_spec_as_ref(repo_id, version.version, local_files_only=local_files_only) - return Oid.from_str(ref.target_commit) + # `name` rather than `ref`: the cache names its refs `v1`, not + # `refs/heads/v1`. + version_ref = resolve_version_spec_as_ref(repo_id, version.version, local_files_only=local_files_only) + ref, commit = version_ref.name, Oid.from_str(version_ref.target_commit) elif isinstance(version, KernelVersion.Revision): - return _resolve_ref(repo_id, version.revision, local_files_only=local_files_only) + ref, commit = version.revision, _resolve_ref(repo_id, version.revision, local_files_only=local_files_only) else: raise ValueError(f"Invalid version type: {version}") + if not local_files_only: + # Kernels are fetched by commit, since we need the commit hash for receipt + # validation, etc. However, that means that snapshot downloads do not create + # refs in the cache. This causes a kernel fetched by version/ref not to be + # found in offline mode. To work around this problem, create a ref ourselves. + # + # Note that this can create the situation where the ref exists, but no + # snapshot or an incomplete snapshot. However, this is fine for + # huggingface_hub, since it also writes the ref before downloading the + # snapshot: + # + # https://github.com/huggingface/huggingface_hub/blob/5a9cdda63f231a1b57a05eab88dc4357c790ba87/src/huggingface_hub/_snapshot_download.py#L426 + _record_ref_in_cache(repo_id, ref, str(commit)) + + return commit + def resolve_revision_or_version( repo_id: str, diff --git a/kernels/src/kernels/cli/download.py b/kernels/src/kernels/cli/download.py index 785f7e55..f8145cfd 100644 --- a/kernels/src/kernels/cli/download.py +++ b/kernels/src/kernels/cli/download.py @@ -1,7 +1,7 @@ import sys from kernels._rust import KernelLocks -from kernels.hf_hub import CACHE_DIR, _get_hf_api +from kernels.hf_hub import _get_cache_dir, _get_hf_api from kernels.resolver import _BYTECODE_IGNORE_PATTERNS, resolve_hub_kernel @@ -31,7 +31,7 @@ def download_kernels(args): repo_type="kernel", allow_patterns="build/*", ignore_patterns=_BYTECODE_IGNORE_PATTERNS, - cache_dir=CACHE_DIR, + cache_dir=_get_cache_dir(), revision=str(lock.commit), ) else: diff --git a/kernels/src/kernels/cli/info.py b/kernels/src/kernels/cli/info.py index a50a48ac..59cfd68f 100644 --- a/kernels/src/kernels/cli/info.py +++ b/kernels/src/kernels/cli/info.py @@ -7,7 +7,7 @@ from kernels._rust import Metadata from kernels._versions import _get_available_versions, resolve_version_spec_as_ref -from kernels.hf_hub import CACHE_DIR, _get_hf_api +from kernels.hf_hub import _get_cache_dir, _get_hf_api from kernels.variants import ( ArchVariant, Variant, @@ -75,7 +75,7 @@ def _hub_kernel_info( repo_id, repo_type="kernel", filename=f"build/{variants[0].variant_str}/metadata.json", - cache_dir=CACHE_DIR, + cache_dir=_get_cache_dir(), revision=revision, ) metadata = Metadata.read_from_file(metadata_path) diff --git a/kernels/src/kernels/deps.py b/kernels/src/kernels/deps.py index ce2141b9..5dc11e6d 100644 --- a/kernels/src/kernels/deps.py +++ b/kernels/src/kernels/deps.py @@ -13,7 +13,7 @@ RemoteKernel, Resolver, ) -from kernels.validate import MetadataValidator +from kernels.validate import KernelValidator, MetadataValidator # Default state is `None`, to signal that we are not in a kernel # loading context. @@ -68,6 +68,12 @@ def load(self) -> ModuleType: return _import_from_path(self.location.variant_path, repo_info=repo_info, deps=deps) + def validate_kernel(self: "DepTreeNode[LocalKernel]", validator: KernelValidator): + validator.validate_kernel(kernel=self.location) + + for node in self.deps.values(): + node.validate_kernel(validator) + def validate_metadata( self: "DepTreeNode[LocalKernel | RemoteKernel]", validator: MetadataValidator, diff --git a/kernels/src/kernels/hf_hub.py b/kernels/src/kernels/hf_hub.py index 634ea58c..b3040990 100644 --- a/kernels/src/kernels/hf_hub.py +++ b/kernels/src/kernels/hf_hub.py @@ -73,11 +73,10 @@ def _get_hf_api(user_agent: str | dict | None = None) -> HfApi: def _get_cache_dir() -> str | None: - """Returns the kernels cache directory.""" - return os.environ.get("KERNELS_CACHE", None) - + """The kernels cache directory, or `None` to use the Hub's default. -CACHE_DIR: str | None = _get_cache_dir() + This re-reads the envvar on every call, so that we can mock it in tests.""" + return os.environ.get("KERNELS_CACHE", None) @dataclass(frozen=True) diff --git a/kernels/src/kernels/install.py b/kernels/src/kernels/install.py index c95500d5..f726e981 100644 --- a/kernels/src/kernels/install.py +++ b/kernels/src/kernels/install.py @@ -3,7 +3,7 @@ from kernels._rust import KernelDependency from kernels._versions import revision_or_version from kernels.deps import resolve_kernel_tree -from kernels.hf_hub import CACHE_DIR, _get_hf_api +from kernels.hf_hub import _get_cache_dir, _get_hf_api from kernels.locking import extract_dependency_locks from kernels.resolver import _BYTECODE_IGNORE_PATTERNS, HubCacheResolver, HubResolver @@ -91,7 +91,7 @@ def install_kernel_all_variants( repo_type="kernel", allow_patterns="build/*", ignore_patterns=_BYTECODE_IGNORE_PATTERNS, - cache_dir=CACHE_DIR, + cache_dir=_get_cache_dir(), revision=str(lock.commit), ) ) diff --git a/kernels/src/kernels/layer/func.py b/kernels/src/kernels/layer/func.py index e550d610..1a7fb158 100644 --- a/kernels/src/kernels/layer/func.py +++ b/kernels/src/kernels/layer/func.py @@ -18,7 +18,7 @@ get_caller_locked_kernel_revision, get_locked_kernel_revision, ) -from ..validate import AllValidator, default_metadata_validators +from ..validate import AllKernelValidator, AllMetadataValidator, default_kernel_validators, default_metadata_validators from .layer import _create_func_module, use_kernel_forward_from_hub from .repos import RepositoryProtocol @@ -329,7 +329,8 @@ def load(self) -> Type["nn.Module"]: backend=None, kernel=self.kernel_dep, resolver=resolver, - metadata_validator=AllValidator(validators=default_metadata_validators()), + kernel_validator=AllKernelValidator(validators=default_kernel_validators()), + metadata_validator=AllMetadataValidator(validators=default_metadata_validators()), ) return _get_kernel_func(self, kernel) diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index a01fcc7a..c02d5946 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/kernels/src/kernels/layer/layer.py @@ -22,7 +22,7 @@ get_caller_locked_kernel_revision, get_locked_kernel_revision, ) -from ..validate import AllValidator, default_metadata_validators +from ..validate import AllKernelValidator, AllMetadataValidator, default_kernel_validators, default_metadata_validators from .device import Device from .globals import _DISABLE_KERNEL_MAPPING, _KERNEL_MAPPING from .mode import Mode @@ -247,7 +247,8 @@ def load(self) -> Type["nn.Module"]: backend=None, kernel=self.kernel_dep, resolver=resolver, - metadata_validator=AllValidator(validators=default_metadata_validators()), + kernel_validator=AllKernelValidator(validators=default_kernel_validators()), + metadata_validator=AllMetadataValidator(validators=default_metadata_validators()), ) return _get_kernel_layer(self, kernel) diff --git a/kernels/src/kernels/load.py b/kernels/src/kernels/load.py index 768ef82c..52488a03 100644 --- a/kernels/src/kernels/load.py +++ b/kernels/src/kernels/load.py @@ -24,9 +24,12 @@ SequentialResolver, ) from kernels.validate import ( - AllValidator, + AllKernelValidator, + AllMetadataValidator, ArchValidator, + KernelValidator, MetadataValidator, + default_kernel_validators, default_metadata_validators, ) @@ -61,6 +64,7 @@ def get_kernel_with_resolver( backend: str | None, kernel: KernelDependency, resolver: Resolver | None, + kernel_validator: KernelValidator, metadata_validator: MetadataValidator, ) -> ModuleType: """ @@ -76,8 +80,10 @@ def get_kernel_with_resolver( The kernel to load. resolver (`Resolver`, *optional*): The resolver used to resolve the kernel and its (transitive) dependencies. + kernel_validator (`KernelValidator`): + The validator to apply to the kernels in the kernel dependency tree. metadata_validator (`MetadataValidator`): - The validator to apply to the resolved kernel dependency tree. + The validator to apply to the metadata in the kernel dependency tree. Returns: `ModuleType`: The imported kernel module. @@ -90,6 +96,7 @@ def get_kernel_with_resolver( ) tree.validate_metadata(metadata_validator) tree_only_local = tree.install(api=api) + tree_only_local.validate_kernel(kernel_validator) return tree_only_local.load() @@ -170,7 +177,8 @@ def get_kernel( backend=backend, kernel=KernelDependency(repo_id=repo_id, version=kernel_version), resolver=SequentialResolver(resolvers=resolvers), - metadata_validator=AllValidator(validators=validators), + kernel_validator=AllKernelValidator(validators=default_kernel_validators()), + metadata_validator=AllMetadataValidator(validators=validators), ) @@ -223,7 +231,8 @@ def get_local_kernel( # We don't have a name for the kernel, so let's just use the path. kernel=KernelDependency(repo_id=str(repo_path), version=KernelVersion.Version(0)), resolver=SequentialResolver(resolvers), - metadata_validator=AllValidator(validators=default_metadata_validators()), + kernel_validator=AllKernelValidator(validators=default_kernel_validators()), + metadata_validator=AllMetadataValidator(validators=default_metadata_validators()), ) @@ -339,7 +348,8 @@ def load_kernel( backend=backend, kernel=kernel_dep, resolver=resolver, - metadata_validator=AllValidator(validators=default_metadata_validators()), + kernel_validator=AllKernelValidator(validators=default_kernel_validators()), + metadata_validator=AllMetadataValidator(validators=default_metadata_validators()), ) @@ -380,5 +390,6 @@ def get_locked_kernel( backend=None, kernel=kernel_dep, resolver=resolver, - metadata_validator=AllValidator(validators=default_metadata_validators()), + kernel_validator=AllKernelValidator(validators=default_kernel_validators()), + metadata_validator=AllMetadataValidator(validators=default_metadata_validators()), ) diff --git a/kernels/src/kernels/locking.py b/kernels/src/kernels/locking.py index 9f6b3c7f..86651ca1 100644 --- a/kernels/src/kernels/locking.py +++ b/kernels/src/kernels/locking.py @@ -14,7 +14,7 @@ ) from kernels._versions import resolve_kernel_version from kernels.compat import tomllib -from kernels.hf_hub import CACHE_DIR, _check_trust_remote_code +from kernels.hf_hub import _check_trust_remote_code, _get_cache_dir from kernels.variants import get_variants @@ -63,7 +63,7 @@ def lock_kernel_tree( kernel.repo_id, repo_type="kernel", filename=f"build/{variant.variant_str}/metadata.json", - cache_dir=CACHE_DIR, + cache_dir=_get_cache_dir(), revision=str(revision), local_files_only=False, ) diff --git a/kernels/src/kernels/resolver.py b/kernels/src/kernels/resolver.py index ad1d8748..8ebeab55 100644 --- a/kernels/src/kernels/resolver.py +++ b/kernels/src/kernels/resolver.py @@ -7,7 +7,7 @@ from kernels._rust import KernelDependency, KernelLocks, KernelPaths, Metadata, Oid from kernels._versions import _get_available_versions, resolve_kernel_version -from kernels.hf_hub import CACHE_DIR, _check_trust_remote_code +from kernels.hf_hub import _check_trust_remote_code, _get_cache_dir from kernels.variants import ( Variant, get_variants, @@ -68,7 +68,7 @@ def install(self, *, api: HfApi) -> LocalKernel: repo_type="kernel", allow_patterns=allow_patterns, ignore_patterns=_BYTECODE_IGNORE_PATTERNS, - cache_dir=CACHE_DIR, + cache_dir=_get_cache_dir(), revision=str(self.revision), local_files_only=False, ) @@ -138,7 +138,7 @@ def resolve_hub_kernel( repo_id, repo_type="kernel", filename=f"build/{variant.variant_str}/metadata.json", - cache_dir=CACHE_DIR, + cache_dir=_get_cache_dir(), revision=str(revision), local_files_only=False, ) @@ -201,7 +201,7 @@ def resolve_hub_cache_kernel( repo_id, repo_type="kernel", ignore_patterns=_BYTECODE_IGNORE_PATTERNS, - cache_dir=CACHE_DIR, + cache_dir=_get_cache_dir(), revision=str(revision), local_files_only=True, ) diff --git a/kernels/src/kernels/validate.py b/kernels/src/kernels/validate.py index 147cf555..9fb0a3dd 100644 --- a/kernels/src/kernels/validate.py +++ b/kernels/src/kernels/validate.py @@ -1,17 +1,32 @@ import logging +import sys from dataclasses import dataclass -from typing import Protocol +from typing import TYPE_CHECKING, Protocol from packaging.version import InvalidVersion, parse -from kernels._rust import Metadata, Version +if sys.version_info >= (3, 11): + from typing import assert_never +else: + from typing_extensions import assert_never + +from kernels._rust import KernelLocation, Metadata, Version from kernels.archs import _check_arch_incompatibility from kernels.backends import _backend +from kernels.compat import has_sigstore from kernels.python_deps import validate_dependencies +from kernels.resolver import LocalKernel logger = logging.getLogger(__name__) +if TYPE_CHECKING: + from sigstore.verify.policy import VerificationPolicy + + +# Metadata validators. + + class MetadataValidator(Protocol): """Metadata validator for a kernel build variant.""" @@ -111,7 +126,7 @@ def validate_metadata(self, *, metadata: Metadata, variant: str) -> None: @dataclass -class AllValidator: +class AllMetadataValidator: """Apply multiple validators to a kernel dependency tree.""" validators: list[MetadataValidator] @@ -124,3 +139,83 @@ def validate_metadata(self, *, metadata: Metadata, variant: str) -> None: def default_metadata_validators() -> list[MetadataValidator]: """The metadata validators that are applied to every kernel dependency tree.""" return [DependencyValidator(), MinverValidator(), DirtyValidator()] + + +# Kernel validators. + + +class KernelValidator(Protocol): + """Kernel (build variant) validator.""" + + def validate_kernel(self, *, kernel: "LocalKernel") -> None: ... + + +@dataclass +class SignatureValidator: + """Verify the signature of a kernel build variant. + + Only kernels with a known Hub origin are verified, since local kernels + are typically for development and not signed. + + Verification issues are currently reported as warnings. However, an + exception will be raised in future versions.""" + + policy: "VerificationPolicy | None" = None + + def validate_kernel(self, *, kernel: "LocalKernel") -> None: + if not has_sigstore: + return + + if kernel.origin is None: + return + + # sigstore is still an optional dependency, so import lazily. + from kernels.verify import VerificationResult, verify_variant + + location = KernelLocation.remote( + kernel.origin.repo_id, + kernel.origin.revision, + kernel.variant_str, + ) + + result = verify_variant(kernel.variant_path, policy=self.policy, location=location) + + kernel_str = f"Kernel '{kernel.metadata.name}' variant '{kernel.variant_str}'" + + match result: + case VerificationResult.Success(): + logger.debug(f"{kernel_str}: {result}") + case VerificationResult.Failure(): + logger.warning(f"{kernel_str}: {result}", stacklevel=3) + case _ as unreachable: + assert_never(unreachable) + + +def default_kernel_validators() -> list[KernelValidator]: + """The kernel validators that are applied to every kernel dependency tree.""" + return [SignatureValidator()] + + +@dataclass +class AllKernelValidator: + """Apply multiple validators to a kernel dependency tree.""" + + validators: list[KernelValidator] + + def validate_kernel(self, *, kernel: "LocalKernel") -> None: + for validator in self.validators: + validator.validate_kernel(kernel=kernel) + + +# Prototype type checks. + +if TYPE_CHECKING: + # Ensure all validators obey the protocol. + _metadata_validators: tuple[MetadataValidator, ...] = ( + DependencyValidator(), + ArchValidator(), + DirtyValidator(), + MinverValidator(), + AllMetadataValidator([]), + ) + _kernel_validator: tuple[KernelValidator, ...] = (SignatureValidator(), AllKernelValidator([])) diff --git a/kernels/src/kernels/verify.py b/kernels/src/kernels/verify.py index ff51e7ae..ba043e29 100644 --- a/kernels/src/kernels/verify.py +++ b/kernels/src/kernels/verify.py @@ -16,6 +16,9 @@ DigestViolation, KernelLocation, Metadata, + ReceiptError, + ReceiptStore, + VerificationReceipt, ) logger = logging.getLogger(__name__) @@ -192,6 +195,29 @@ def __str__(self) -> str: ) +def _open_receipt_store() -> ReceiptStore | None: + """The receipt store, or `None` when verifications cannot be cached.""" + try: + return ReceiptStore.in_kernels_cache() + except ReceiptError as e: + logger.warning(f"Cannot cache kernel verifications: {e}") + return None + + +def _has_receipt(store: ReceiptStore, location: KernelLocation) -> bool: + """Whether the kernel at `location` was verified before. + + An unusable receipt counts as a cache miss: the kernel is then verified in + full, which overwrites the receipt. A broken cache must never make a kernel + fail to verify. + """ + try: + return store.load(location) is not None + except ReceiptError as e: + logger.warning(f"Ignoring unusable kernel verification receipt: {e}") + return False + + def verify_variant( variant_path: Path, *, @@ -238,6 +264,21 @@ def verify_variant( if not metadata_path.is_file(): return VerificationResult.MetadataMissing() + receipt_store = _open_receipt_store() if cache else None + + if receipt_store is not None and _has_receipt(receipt_store, location): + # The receipt attests that this kernel metadata was verified + # using the signature and the kernel data during the digest + # in the metadata. However, it may have been verified with a + # different policy, so we have to check certificate in the + # bundle against the currently required policy. + try: + verify_policy.verify(signature_bundle.signing_certificate) + except VerificationError as e: + return VerificationResult.SignatureVerificationFailure(reason=str(e)) + + return VerificationResult.Success() + verifier = Verifier.production() metadata_bytes = metadata_path.read_bytes() @@ -271,4 +312,10 @@ def verify_variant( except DigestValidationError as e: return VerificationResult.DigestVerificationFailure(violations=e.violations) + if receipt_store is not None: + try: + receipt_store.store(VerificationReceipt(location)) + except ReceiptError as e: + logger.warning(f"Cannot store kernel verification receipt: {e}") + return VerificationResult.Success() diff --git a/kernels/tests/test_validate.py b/kernels/tests/test_validate.py index 95f470b9..97726fa1 100644 --- a/kernels/tests/test_validate.py +++ b/kernels/tests/test_validate.py @@ -9,13 +9,14 @@ import kernels import kernels.validate as validate_module import kernels.verify as verify_module -from kernels._rust import Metadata, Oid, Version +from kernels._rust import KernelLocation, Metadata, Oid, Version from kernels.deps import DepTreeNode from kernels.resolver import LocalKernel, RemoteKernel from kernels.validate import ( ArchValidator, DirtyValidator, MinverValidator, + SignatureValidator, _installed_version, default_metadata_validators, ) @@ -242,3 +243,71 @@ def fake_verify_variant(variant_path, *, location, policy=None, cache=True): monkeypatch.setattr(verify_module, "verify_variant", fake_verify_variant) return calls, results + + +def test_signature_validator_skips_local_kernels(tmp_path, make_metadata, recorded_verifications): + calls, _ = recorded_verifications + kernel = LocalKernel(variant_path=tmp_path / "torch-cuda", metadata=make_metadata("cuda", None)) + + SignatureValidator().validate_kernel(kernel=kernel) + + assert calls == [] + + +def test_signature_validator_identifies_kernel_by_origin(tmp_path, make_metadata, recorded_verifications): + calls, _ = recorded_verifications + kernel = _hub_kernel(tmp_path, make_metadata("cuda", None)) + + SignatureValidator().validate_kernel(kernel=kernel) + + (call,) = calls + assert call["variant_path"] == kernel.variant_path + assert call["location"] == KernelLocation.remote(_SIGNED_REPO_ID, _SIGNED_REVISION, "torch-cuda") + # Loading a kernel must reuse a previous verification. + assert call["cache"] is True + + +def test_signature_validator_passes_policy(tmp_path, make_metadata, recorded_verifications): + calls, _ = recorded_verifications + kernel = _hub_kernel(tmp_path, make_metadata("cuda", None)) + sentinel = object() + + SignatureValidator(policy=sentinel).validate_kernel(kernel=kernel) + + (call,) = calls + assert call["policy"] is sentinel + + +def test_signature_validator_is_quiet_on_success(tmp_path, make_metadata, recorded_verifications, caplog): + kernel = _hub_kernel(tmp_path, make_metadata("cuda", None)) + + with caplog.at_level(logging.WARNING, logger="kernels.validate"): + SignatureValidator().validate_kernel(kernel=kernel) + + assert caplog.text == "" + + +@pytest.mark.parametrize( + "result", + [ + VerificationResult.SignatureBundleMissing(), + VerificationResult.SignatureBundleInvalid(reason="bad bundle"), + VerificationResult.SignatureVerificationFailure(reason="bad signature"), + VerificationResult.MetadataInvalid(reason="bad metadata"), + VerificationResult.MetadataMissing(), + VerificationResult.DigestMissing(), + VerificationResult.DigestVerificationFailure(violations=[]), + ], +) +def test_signature_validator_warns_but_does_not_raise(tmp_path, make_metadata, recorded_verifications, caplog, result): + _, results = recorded_verifications + results.append(result) + kernel = _hub_kernel(tmp_path, make_metadata("cuda", None)) + + with caplog.at_level(logging.WARNING, logger="kernels.validate"): + SignatureValidator().validate_kernel(kernel=kernel) + + # The message belongs to the result. The validator only says which kernel + # it applies to, so the wording is asserted where it is defined. + assert str(result) in caplog.text + assert "test-kernel" in caplog.text diff --git a/kernels/tests/test_verify.py b/kernels/tests/test_verify.py index 65f9c3a8..a942abcb 100644 --- a/kernels/tests/test_verify.py +++ b/kernels/tests/test_verify.py @@ -1,3 +1,4 @@ +import logging from dataclasses import is_dataclass from pathlib import Path @@ -6,9 +7,9 @@ import kernels.verify as verify_module from kernels import install_kernel -from kernels._rust import DigestViolation, KernelLocation, Oid +from kernels._rust import DigestViolation, KernelLocation, Oid, ReceiptStore from kernels._versions import resolve_revision_or_version -from kernels.hf_hub import CACHE_DIR, _get_hf_api +from kernels.hf_hub import _get_cache_dir, _get_hf_api from kernels.resolver import _BYTECODE_IGNORE_PATTERNS from kernels.verify import VerificationResult, verify_variant @@ -21,6 +22,15 @@ ) +@pytest.fixture +def receipt_store(tmp_path, monkeypatch): + """An isolated receipt store, so that tests do not share verifications.""" + receipt_dir = tmp_path / "receipts" + store = ReceiptStore.from_path(receipt_dir) + monkeypatch.setattr(verify_module, "_open_receipt_store", lambda: store) + return store + + @pytest.fixture def signed_kernel(): """A correctly signed kernel, with the location that identifies it.""" @@ -99,7 +109,7 @@ def test_invalid_metadata_fails(): repo_type="kernel", allow_patterns="build/*", ignore_patterns=_BYTECODE_IGNORE_PATTERNS, - cache_dir=CACHE_DIR, + cache_dir=_get_cache_dir(), revision=str(revision), ) ) @@ -141,7 +151,7 @@ def test_missing_metadata_fails(): repo_type="kernel", allow_patterns="build/*", ignore_patterns=_BYTECODE_IGNORE_PATTERNS, - cache_dir=CACHE_DIR, + cache_dir=_get_cache_dir(), revision=str(revision), ) ) @@ -182,6 +192,67 @@ def test_invalid_signature_fails(): raise RuntimeError(f"Expected SignatureVerificationFailure, was: {other}") +def test_verification_is_cached(receipt_store, signed_kernel, monkeypatch): + variant_path, location = signed_kernel + + assert verify_variant(variant_path, policy=TEST_POLICY, location=location) == VerificationResult.Success() + assert receipt_store.load(location) is not None + + # The second verification must be served from the receipt, without + # rehashing the variant. + _no_hashing(monkeypatch) + assert verify_variant(variant_path, policy=TEST_POLICY, location=location) == VerificationResult.Success() + + +def test_verification_is_not_cached_with_cache_off(receipt_store, signed_kernel, monkeypatch): + variant_path, location = signed_kernel + + result = verify_variant(variant_path, policy=TEST_POLICY, location=location, cache=False) + assert result == VerificationResult.Success() + + # Nothing was recorded, ... + assert receipt_store.load(location) is None + + # ... and a verification with caching off does the full work even when a + # receipt does exist. + assert verify_variant(variant_path, policy=TEST_POLICY, location=location) == VerificationResult.Success() + assert receipt_store.load(location) is not None + + _no_hashing(monkeypatch) + with pytest.raises(AssertionError, match="was rehashed"): + verify_variant(variant_path, policy=TEST_POLICY, location=location, cache=False) + + +def test_cached_verification_still_enforces_policy(receipt_store, signed_kernel, monkeypatch): + variant_path, location = signed_kernel + + # Verify under a policy that accepts this kernel, so a receipt is stored. + assert verify_variant(variant_path, policy=TEST_POLICY, location=location) == VerificationResult.Success() + + # The receipt says the kernel was verified, but not *under which policy*, + # so a policy that does not accept this signer must still reject it. + _no_hashing(monkeypatch) + match verify_variant(variant_path, policy=OTHER_POLICY, location=location): + case VerificationResult.SignatureVerificationFailure(): + pass + case other: + raise RuntimeError(f"Expected SignatureVerificationFailure, was: {other}") + + +def test_unusable_receipt_falls_back_to_verification(receipt_store, signed_kernel, tmp_path, caplog): + variant_path, location = signed_kernel + + assert verify_variant(variant_path, policy=TEST_POLICY, location=location) == VerificationResult.Success() + + (receipt_path,) = list((tmp_path / "receipts").iterdir()) + receipt_path.write_text("not a receipt") + + with caplog.at_level(logging.WARNING, logger="kernels.verify"): + assert verify_variant(variant_path, policy=TEST_POLICY, location=location) == VerificationResult.Success() + + assert "unusable kernel verification receipt" in caplog.text + + ALL_RESULTS = [ VerificationResult.Success(), VerificationResult.SignatureBundleMissing(), diff --git a/kernels/tests/test_versions.py b/kernels/tests/test_versions.py index 8a516a6a..ed722a64 100644 --- a/kernels/tests/test_versions.py +++ b/kernels/tests/test_versions.py @@ -1,7 +1,11 @@ +from pathlib import Path + import pytest from huggingface_hub.file_download import repo_folder_name +import kernels._versions as versions import kernels.hf_hub as hf_hub +from kernels import install_kernel from kernels._rust import KernelVersion, Oid from kernels._versions import _resolve_ref @@ -12,7 +16,7 @@ @pytest.fixture def cached_refs(tmp_path, monkeypatch): """A cache containing a single ref, so offline resolution is hermetic.""" - monkeypatch.setattr(hf_hub, "CACHE_DIR", str(tmp_path)) + monkeypatch.setenv("KERNELS_CACHE", str(tmp_path)) refs = tmp_path / repo_folder_name(repo_id=REPO_ID, repo_type="kernel") / "refs" refs.mkdir(parents=True) (refs / "main").write_text(COMMIT) @@ -79,7 +83,6 @@ def test_a_version_needs_no_ref_resolution(monkeypatch): Guards against reintroducing a resolution step that would suggest a version can name something other than a commit. """ - import kernels._versions as versions def fail(*args, **kwargs): raise AssertionError("a version was resolved as if it were a ref") @@ -93,3 +96,71 @@ def fail(*args, **kwargs): ) assert revision == Oid.from_str(str(revision)) + + +# Kernels are fetched by commit, so snapshot downloads do not create a ref. +# This is not a problem during normal usage, but breaks offline use if the +# kernel is resolved by version or non-commit ref. For this reason, we create +# a ref when resolving a name. The tests below ensure that this behavior is +# correct. + + +def _refs_of(cache_dir: Path, repo_id: str = REPO_ID) -> list[str]: + refs = cache_dir / repo_folder_name(repo_id=repo_id, repo_type="kernel") / "refs" + return sorted(p.name for p in refs.iterdir()) if refs.is_dir() else [] + + +@pytest.mark.parametrize("version", [KernelVersion.Revision("v1"), KernelVersion.Version(1)]) +def test_resolution_records_the_ref(tmp_path, monkeypatch, version): + monkeypatch.setenv("KERNELS_CACHE", str(tmp_path)) + + commit = versions.resolve_kernel_version("kernels-community/relu", version, local_files_only=False) + + assert _refs_of(tmp_path, "kernels-community/relu") == ["v1"] + + # Written verbatim, since huggingface_hub does not strip the file. + ref_path = tmp_path / repo_folder_name(repo_id="kernels-community/relu", repo_type="kernel") / "refs" / "v1" + assert ref_path.read_text() == str(commit) + + +def test_resolving_a_commit_records_nothing(tmp_path, monkeypatch): + monkeypatch.setenv("KERNELS_CACHE", str(tmp_path)) + + versions.resolve_kernel_version(REPO_ID, KernelVersion.Revision(COMMIT), local_files_only=False) + + assert _refs_of(tmp_path) == [] + + +def test_recording_a_ref_tolerates_an_unwritable_cache(tmp_path, monkeypatch, caplog): + monkeypatch.setenv("KERNELS_CACHE", str(tmp_path / "not-a-directory")) + (tmp_path / "not-a-directory").write_text("") + + with caplog.at_level("WARNING", logger="kernels._versions"): + versions._record_ref_in_cache(REPO_ID, "v1", COMMIT) + + assert "Could not record revision" in caplog.text + + +def test_recording_a_ref_stays_inside_the_refs_directory(tmp_path, monkeypatch): + monkeypatch.setenv("KERNELS_CACHE", str(tmp_path)) + + versions._record_ref_in_cache(REPO_ID, "../../../escaped", COMMIT) + + assert not (tmp_path.parent.parent.parent / "escaped").exists() + assert not (tmp_path / "escaped").exists() + + +def test_a_downloaded_revision_resolves_offline(tmp_path, monkeypatch): + monkeypatch.setenv("KERNELS_CACHE", str(tmp_path)) + + install_kernel("kernels-community/relu", revision="v1", backend="cpu") + + assert install_kernel("kernels-community/relu", revision="v1", backend="cpu", local_files_only=True) + + +def test_a_downloaded_version_resolves_offline(tmp_path, monkeypatch): + monkeypatch.setenv("KERNELS_CACHE", str(tmp_path)) + + install_kernel("kernels-community/relu", version=1, backend="cpu") + + assert install_kernel("kernels-community/relu", version=1, backend="cpu", local_files_only=True) diff --git a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py index 06f182c8..62d917f6 100755 --- a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py +++ b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py @@ -4,7 +4,11 @@ from kernels.hf_hub import _get_hf_api from kernels.load import get_kernel_with_resolver from kernels.resolver import KernelPathsResolver, RepoPathsResolver, SequentialResolver -from kernels.validate import AllValidator, default_metadata_validators +from kernels.validate import ( + AllKernelValidator, + AllMetadataValidator, + default_metadata_validators, +) from kernels._rust import KernelDependency, KernelPaths, KernelVersion out = os.getenv("out") @@ -46,5 +50,6 @@ backend=None, kernel=kernel, resolver=SequentialResolver(resolvers=resolvers), - metadata_validator=AllValidator(validators=default_metadata_validators()), + kernel_validator=AllKernelValidator(validators=[]), + metadata_validator=AllMetadataValidator(validators=default_metadata_validators()), )