From 7fbf0fd8f8a96ea3a592d19cbca57106ebd21501 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Tue, 15 Sep 2026 09:46:14 +0000 Subject: [PATCH 1/3] kernels: verify the signature of a kernel while loading The basics of kernel signature verification were already in `kernels`, but verification would only be performed when explicitly running the `kernels verify-signature` command. This change enables kernel signature verification during kernel loading. However, as part of the step-wise roll-out, signature verification will now only warn if the metadata is not signed, not correctly signed, or the data's digest is incorrect. Signature verification during kernel loading is implemented as follows. First, we add the `KernelValidator` protocol, this is the counterpart to `MetadataValidator` that can only run when the kernel is fully downloaded. `get_kernel_with_resolver` now requires an additional `kernel_validator` argument and runs the given validators on the kernel dependency tree after the kernels are downloaded ('installed'). We add `SignatureValidator` as the first validator to implement the `KernelValidator` protocol and enable it in the relevant `get_.*` functions. This validator calls the existing verification structure. Validation is only performed on remote kernels, since local kernels are used in development and are typically not signed. There was one more wrinkle in that validating the hash of all kernel files on each load is fairly expensive. So upon the first successful verification, we write a receipt that attests that an earlier verification is done. The receipt is content-addressed by the repo id and the resolved Git commit SHA, so it does not attest for newly-downloaded versions. The receipts allow us to skip the most expensive parts of the verification. However, the user might use a different policy after the verification that produced the receipt, so we will always test the kernel's signing certificate against the policy. This ensures that the kernel signature will be rejected if the kernel violates the policy used during a kernel load. --- Cargo.lock | 1 + kernels-common/Cargo.toml | 3 +- kernels-common/src/hf.rs | 25 +- kernels-common/src/signing/receipt.rs | 379 ++++++++++++++++++++++++++ kernels/Cargo.toml | 1 + kernels/rust/lib.rs | 5 +- kernels/rust/signing.rs | 99 ++++++- kernels/src/kernels/deps.py | 8 +- kernels/src/kernels/layer/func.py | 5 +- kernels/src/kernels/layer/layer.py | 5 +- kernels/src/kernels/load.py | 23 +- kernels/src/kernels/validate.py | 101 ++++++- kernels/src/kernels/verify.py | 47 ++++ kernels/tests/test_validate.py | 71 ++++- kernels/tests/test_verify.py | 73 ++++- 15 files changed, 818 insertions(+), 28 deletions(-) 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/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/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/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..dd2b6480 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,7 +7,7 @@ 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.resolver import _BYTECODE_IGNORE_PATTERNS @@ -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.""" @@ -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(), From 607996fdd5ff1a9e5680ea190b57da790c082043 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Tue, 15 Sep 2026 13:01:28 +0000 Subject: [PATCH 2/3] kernels: write ref to the cache when needed We resolve kernels refs to full commit SHAs as early as possible (since we need immutable identifiers for e.g. receipt caching). However, that also means that we are doing snapshot downloads by Git SHA. As a result, refs such as `v1` do not get cached, breaking non-locked offline use. --- kernels/src/kernels/_versions.py | 73 +++++++++++++++++++++++++--- kernels/src/kernels/cli/download.py | 4 +- kernels/src/kernels/cli/info.py | 4 +- kernels/src/kernels/hf_hub.py | 7 ++- kernels/src/kernels/install.py | 4 +- kernels/src/kernels/locking.py | 4 +- kernels/src/kernels/resolver.py | 8 +-- kernels/tests/test_verify.py | 6 +-- kernels/tests/test_versions.py | 75 ++++++++++++++++++++++++++++- 9 files changed, 157 insertions(+), 28 deletions(-) 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/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/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/tests/test_verify.py b/kernels/tests/test_verify.py index dd2b6480..a942abcb 100644 --- a/kernels/tests/test_verify.py +++ b/kernels/tests/test_verify.py @@ -9,7 +9,7 @@ from kernels import install_kernel 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 @@ -109,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), ) ) @@ -151,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), ) ) 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) From 4bb1309648d3dd26b6106dca4b6643ea2c3d6a8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Tue, 15 Sep 2026 13:49:35 +0000 Subject: [PATCH 3/3] nix-builder: update get-kernel-check-hook --- .../pkgs/get-kernel-check/get-kernel-check-hook.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) 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()), )