From 105ac9855f6cae37ce0b96089553daf52329c7ad Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Tue, 18 Aug 2026 09:10:34 +0000 Subject: [PATCH 1/2] kernel-builder: move `build.toml` parsing to `kernels-data` This change moves over `parse_and_validate` to `kernels-data` and exposes it as `Build::open`. We also make this function available through the Python API as well as a small number of methods from `Build`. The `Build` methods that are made available through the Python API are guided by what is needed for dependency support. In general, we do not want to do any `build.toml` parsing in `kernels`, since it's a build-side file format. However, to be able to build a kernel with dependencies in Nix, we need to lock the dependencies and this is the same locking code as used by `kernels lock`. The only difference is that we take the kernels from `build.toml` and not `pyproject.toml`. For this reason, exposing `build.toml` to the client in a minor way seems justified for this specific purpose. --- kernel-builder/src/card.rs | 7 +- kernel-builder/src/main.rs | 11 +-- kernel-builder/src/pyproject/mod.rs | 6 +- kernel-builder/src/upload.rs | 5 +- kernel-builder/src/util.rs | 43 +---------- kernels-data/Cargo.toml | 4 +- kernels-data/bindings/python/kernels_data.pyi | 35 +++++++++ kernels-data/bindings/python/src/config.rs | 74 +++++++++++++++++++ kernels-data/bindings/python/src/lib.rs | 10 ++- .../bindings/python/tests/test_config.py | 40 ++++++++++ kernels-data/src/config/compat.rs | 8 ++ kernels-data/src/config/mod.rs | 37 ++++++---- kernels-data/src/config/parse.rs | 40 ++++++++++ 13 files changed, 243 insertions(+), 77 deletions(-) create mode 100644 kernels-data/bindings/python/src/config.rs create mode 100644 kernels-data/bindings/python/tests/test_config.py create mode 100644 kernels-data/src/config/parse.rs diff --git a/kernel-builder/src/card.rs b/kernel-builder/src/card.rs index f486b23f0..6874b870b 100644 --- a/kernel-builder/src/card.rs +++ b/kernel-builder/src/card.rs @@ -5,12 +5,11 @@ use std::{ }; use eyre::{bail, Context, Result}; +use kernels_data::config::Build; use minijinja::{context, Environment}; use rustpython_parser::{ast, Parse}; -use kernels_data::config::Build; - -use crate::util::{check_or_infer_kernel_dir, parse_build}; +use crate::util::check_or_infer_kernel_dir; fn extract_all(kernel_dir: &Path, module_name: &str) -> Option> { let init_path = ["torch-ext", "tvm-ffi-ext"] @@ -147,7 +146,7 @@ fn render_card(build: &Build, kernel_dir: &Path) -> Result { pub fn fill_card(kernel_dir: Option, output: Option) -> Result<()> { let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?; - let build = parse_build(&kernel_dir)?; + let build = Build::open(&kernel_dir)?; let content = render_card(&build, &kernel_dir)?; match output { diff --git a/kernel-builder/src/main.rs b/kernel-builder/src/main.rs index 73a8380ac..6fa9db982 100644 --- a/kernel-builder/src/main.rs +++ b/kernel-builder/src/main.rs @@ -5,6 +5,7 @@ use std::path::PathBuf; use clap::{Args, CommandFactory, Parser, Subcommand}; use clap_complete::Shell; use eyre::{Context, Result}; +use kernels_data::config::{v5, Build, BuildCompat}; use kernels_data::git::Oid; mod card; @@ -37,20 +38,16 @@ use upload::{run_upload, RepoTypeArg, UploadArgs}; mod pyproject; use pyproject::{clean_pyproject, create_pyproject}; -use kernels_data::config::{v5, Build, BuildCompat}; - mod nix; mod skills; mod util; -use util::{check_or_infer_kernel_dir, parse_and_validate}; +use util::check_or_infer_kernel_dir; mod validate_builds; use validate_builds::check_builds; -use crate::util::parse_and_validate_compat; - /// Build-time metadata gathered by the [`built`](https://crates.io/crates/built) /// crate (see `build.rs`), including the `kernel-builder` git provenance that is /// recorded in the build metadata of the kernels it builds. @@ -470,13 +467,13 @@ fn main() -> Result<()> { fn check_config(kernel_dir: Option) -> Result<()> { let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?; - parse_and_validate(kernel_dir)?; + Build::open(kernel_dir)?; Ok(()) } fn update_build(kernel_dir: Option) -> Result<()> { let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?; - let build_compat: BuildCompat = parse_and_validate_compat(&kernel_dir)?; + let build_compat = BuildCompat::open(&kernel_dir)?; if matches!(build_compat, BuildCompat::V5(_)) { return Ok(()); diff --git a/kernel-builder/src/pyproject/mod.rs b/kernel-builder/src/pyproject/mod.rs index fb911377b..cc8cb25f4 100644 --- a/kernel-builder/src/pyproject/mod.rs +++ b/kernel-builder/src/pyproject/mod.rs @@ -12,7 +12,7 @@ use minijinja::Environment; use crate::{ pyproject::ops_identifier::KernelIdentifier, - util::{check_or_infer_kernel_dir, check_or_infer_target_dir, parse_build}, + util::{check_or_infer_kernel_dir, check_or_infer_target_dir}, }; pub(crate) mod common; @@ -55,7 +55,7 @@ pub fn create_pyproject( ) -> Result<()> { let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?; let target_dir = check_or_infer_target_dir(&kernel_dir, target_dir)?; - let build = parse_build(&kernel_dir)?; + let build = Build::open(&kernel_dir)?; // Assemble build provenance. Prefer an explicitly provided kernel git // provenance (e.g. passed by Nix builds, where the source tree has no @@ -89,7 +89,7 @@ pub fn clean_pyproject( ) -> Result<()> { let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?; let target_dir = check_or_infer_target_dir(&kernel_dir, target_dir)?; - let build = parse_build(&kernel_dir)?; + let build = Build::open(&kernel_dir)?; // Provenance is irrelevant when computing the set of files to clean. let kernel_id = KernelIdentifier::new(&kernel_dir, build.general.name.python_name(), unique_id); diff --git a/kernel-builder/src/upload.rs b/kernel-builder/src/upload.rs index 3daef03cc..823a78b34 100644 --- a/kernel-builder/src/upload.rs +++ b/kernel-builder/src/upload.rs @@ -13,13 +13,14 @@ use hf_hub::{ HFError, HFRepositorySync, RepoType, RepoTypeKernel, RepoTypeModel, }; use indicatif::{ProgressBar, ProgressStyle}; +use kernels_data::config::Build; use kernels_data::metadata::Metadata; use serde::Serialize; use walkdir::WalkDir; use crate::{ hf::{self, repo_handle}, - util::{check_or_infer_kernel_dir, discover_variants, parse_build}, + util::{check_or_infer_kernel_dir, discover_variants}, }; /// Bridges `ProgressHandler` events to an `indicatif::ProgressBar`. @@ -146,7 +147,7 @@ fn get_repo_and_branch( branch: Option, variants: &[PathBuf], ) -> Result<(String, Option)> { - let build = parse_build(kernel_dir); + let build = Build::open(kernel_dir); let build_branch = build .as_ref() diff --git a/kernel-builder/src/util.rs b/kernel-builder/src/util.rs index 6bfc80e55..c474fdb36 100644 --- a/kernel-builder/src/util.rs +++ b/kernel-builder/src/util.rs @@ -1,17 +1,9 @@ use std::env::current_dir; -use std::fs::{self, File}; -use std::io::Read; +use std::fs; use std::path::{Path, PathBuf}; use eyre::{bail, ensure, Context, Result}; -use kernels_data::config::{Build, BuildCompat, CurrentConfig}; - -pub(crate) fn parse_build(kernel_dir: impl AsRef) -> Result { - let build_compat = parse_and_validate(kernel_dir)?; - Ok(build_compat.into()) -} - pub(crate) fn check_or_infer_kernel_dir(kernel_dir: Option>) -> Result { match kernel_dir { Some(kernel_dir) => { @@ -45,39 +37,6 @@ pub(crate) fn check_or_infer_target_dir( } } -pub(crate) fn parse_and_validate(kernel_dir: impl AsRef) -> Result { - // Only v4 is auto-upgraded to v5 on load; older editions must be migrated - // explicitly with `update-build`. - match parse_and_validate_compat(kernel_dir)? { - BuildCompat::V5(build) => Ok(build), - BuildCompat::V4(build) => { - eprintln!( - "⚠️ build.toml uses the legacy v4 format; upgrading to edition 5 in memory. \ - Run `kernel-builder update-build` to persist the upgrade." - ); - Ok(Build::from(build).into()) - } - BuildCompat::V3(_) => bail!( - "build.toml uses an unsupported legacy format; migrate it with \ - `kernel-builder update-build`" - ), - } -} - -pub(crate) fn parse_and_validate_compat(kernel_dir: impl AsRef) -> Result { - let build_toml = kernel_dir.as_ref().join("build.toml"); - let mut toml_data = String::new(); - File::open(&build_toml) - .wrap_err_with(|| format!("Cannot open {} for reading", build_toml.to_string_lossy()))? - .read_to_string(&mut toml_data) - .wrap_err_with(|| format!("Cannot read from {}", build_toml.to_string_lossy()))?; - - let config: BuildCompat = toml::from_str(&toml_data) - .wrap_err_with(|| format!("Cannot parse TOML in {}", build_toml.to_string_lossy()))?; - - Ok(config) -} - /// Discover build variant directories (contain `metadata.json`). /// Checks `result` symlink (Nix store output) first, then falls back to `build/`. pub(crate) fn discover_variants(kernel_dir: &Path) -> Result<(PathBuf, Vec)> { diff --git a/kernels-data/Cargo.toml b/kernels-data/Cargo.toml index 3306d4c95..0205c5b89 100644 --- a/kernels-data/Cargo.toml +++ b/kernels-data/Cargo.toml @@ -18,10 +18,10 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" serde-value = "0.7" sha2 = "0.11" -walkdir = "2" thiserror = "1" +toml = "0.8" url = { version = "2", features = ["serde"] } +walkdir = "2" [dev-dependencies] tempfile = "3" -toml = "0.8" diff --git a/kernels-data/bindings/python/kernels_data.pyi b/kernels-data/bindings/python/kernels_data.pyi index a294f8915..0e3ea2cbd 100644 --- a/kernels-data/bindings/python/kernels_data.pyi +++ b/kernels-data/bindings/python/kernels_data.pyi @@ -7,6 +7,8 @@ from typing import Optional, final __all__ = [ "Backend", "BackendInfo", + "Build", + "General", "Provenance", "DigestAlgorithm", "GitStatus", @@ -324,6 +326,39 @@ class KernelDependency: def __eq__(self, value: object, /) -> bool: ... def __hash__(self) -> int: ... +@final +class General: + """General kernel configuration common to all backends.""" + + @property + def backends(self) -> list[Backend]: + """Backends the kernel supports.""" + ... + + def __repr__(self) -> str: ... + +@final +class Build: + """Parsed and validated `build.toml` configuration for a kernel.""" + + @staticmethod + def open(kernel_dir: os.PathLike[str] | str) -> "Build": + """Parse and validate the `build.toml` in `kernel_dir`. + + Raises: + ValueError: If the build configuration cannot be parsed or validated. + """ + ... + + @property + def general(self) -> General: + """General kernel configuration.""" + ... + + def all_kernel_depends(self, backend: Backend) -> list[KernelDependency]: + """Get the general + backend-specific kernel dependencies for `backend`.""" + ... + @final class Metadata: """Parsed `metadata.json` for a kernel build variant.""" diff --git a/kernels-data/bindings/python/src/config.rs b/kernels-data/bindings/python/src/config.rs new file mode 100644 index 000000000..e0d23041c --- /dev/null +++ b/kernels-data/bindings/python/src/config.rs @@ -0,0 +1,74 @@ +use std::path::PathBuf; + +use kernels_data::config::{Build, General}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; + +use crate::{PyBackend, PyKernelDependency}; + +/// General kernel configuration common to all backends. +#[pyclass(name = "General", frozen)] +#[derive(Clone, Debug)] +pub(crate) struct PyGeneral { + backends: Vec, +} + +impl From<&General> for PyGeneral { + fn from(general: &General) -> Self { + Self { + backends: general.backends.iter().copied().map(Into::into).collect(), + } + } +} + +#[pymethods] +impl PyGeneral { + #[getter] + fn backends(&self) -> Vec { + self.backends.clone() + } + + fn __repr__(&self) -> String { + format!("General(backends={:?})", self.backends) + } +} + +/// Parsed and validated `build.toml` configuration for a kernel. +#[pyclass(name = "Build", frozen)] +pub(crate) struct PyBuild { + inner: Build, +} + +#[pymethods] +impl PyBuild { + /// Parse and validate the `build.toml` in `kernel_dir`. + /// + /// Raises `ValueError` if the build configuration cannot be parsed or + /// validated. + #[staticmethod] + fn open(kernel_dir: PathBuf) -> PyResult { + Build::open(&kernel_dir) + .map(|inner| PyBuild { inner }) + .map_err(|err| { + PyValueError::new_err(format!( + "Cannot parse build configuration in `{}`: {err:#}", + kernel_dir.display() + )) + }) + } + + #[getter] + fn general(&self) -> PyGeneral { + PyGeneral::from(&self.inner.general) + } + + /// Get the general + backend-specific kernel dependencies for `backend`. + fn all_kernel_depends(&self, backend: PyBackend) -> Vec { + self.inner + .general + .all_kernel_depends(backend.into()) + .into_iter() + .map(Into::into) + .collect() + } +} diff --git a/kernels-data/bindings/python/src/lib.rs b/kernels-data/bindings/python/src/lib.rs index 32088ed4e..bce6e81fb 100644 --- a/kernels-data/bindings/python/src/lib.rs +++ b/kernels-data/bindings/python/src/lib.rs @@ -13,6 +13,10 @@ use pyo3::Bound as PyBound; use pyo3::exceptions::{PyException, PyOSError, PyRuntimeError, PyValueError}; use pyo3::prelude::*; +mod config; + +use config::{PyBuild, PyGeneral}; + /// A dotted numeric version (e.g. `12.8.0`). Trailing zeros are stripped /// during normalization. #[pyclass(name = "Version", frozen, eq, ord, hash)] @@ -76,7 +80,7 @@ impl PyKernelName { /// Kernel backend (hardware target). #[pyclass(name = "Backend", eq, frozen, hash)] #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] -enum PyBackend { +pub(crate) enum PyBackend { #[pyo3(name = "CANN")] Cann, #[pyo3(name = "CPU")] @@ -349,7 +353,7 @@ impl PyKernelVersion { /// A dependency on another kernel. #[pyclass(name = "KernelDependency", frozen, eq, hash)] #[derive(Clone, Debug, Eq, Hash, PartialEq)] -struct PyKernelDependency { +pub(crate) struct PyKernelDependency { repo_id: String, version: PyKernelVersion, } @@ -755,6 +759,8 @@ fn kernels_data_py(m: &PyBound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/kernels-data/bindings/python/tests/test_config.py b/kernels-data/bindings/python/tests/test_config.py new file mode 100644 index 000000000..14154d0d9 --- /dev/null +++ b/kernels-data/bindings/python/tests/test_config.py @@ -0,0 +1,40 @@ +import pytest + +from kernels_data import Backend, Build + + +def _write_build_toml(path, backends): + backends_toml = ", ".join(f'"{b}"' for b in backends) + path.write_text( + f"""\ +[general] +name = "my-kernel" +version = 1 +edition = 5 +license = "Apache-2.0" +backends = [{backends_toml}] + +[torch-noarch] +""" + ) + return path + + +def test_build_open_and_general_backends(tmp_path): + _write_build_toml(tmp_path / "build.toml", backends=["cpu", "cuda"]) + + build = Build.open(tmp_path) + assert isinstance(build, Build) + assert build.general.backends == [Backend.CPU, Backend.CUDA] + + +def test_build_open_missing_file(tmp_path): + with pytest.raises(ValueError): + Build.open(tmp_path) + + +def test_build_all_kernel_depends_empty(tmp_path): + _write_build_toml(tmp_path / "build.toml", backends=["cpu"]) + + build = Build.open(tmp_path) + assert build.all_kernel_depends(Backend.CPU) == [] diff --git a/kernels-data/src/config/compat.rs b/kernels-data/src/config/compat.rs index d66c44a8f..69d070306 100644 --- a/kernels-data/src/config/compat.rs +++ b/kernels-data/src/config/compat.rs @@ -1,3 +1,5 @@ +use std::path::Path; + use eyre::Result; use serde::{Deserialize, de}; use serde_value::Value; @@ -14,6 +16,12 @@ pub enum BuildCompat { V5(v5::Build), } +impl BuildCompat { + pub fn open(kernel_dir: impl AsRef) -> Result { + super::parse::parse_and_validate_compat(kernel_dir) + } +} + #[derive(Deserialize)] struct EditionProbe { #[serde(default)] diff --git a/kernels-data/src/config/mod.rs b/kernels-data/src/config/mod.rs index 861c79390..18505ce2a 100644 --- a/kernels-data/src/config/mod.rs +++ b/kernels-data/src/config/mod.rs @@ -1,7 +1,7 @@ use std::{ collections::{BTreeMap, HashMap}, fmt::Display, - path::PathBuf, + path::{Path, PathBuf}, str::FromStr, }; @@ -24,6 +24,8 @@ pub use kernel_deps::{KernelDependency, KernelVersion}; mod name; pub use name::KernelName; +mod parse; + pub mod v3; pub mod v4; pub mod v5; @@ -41,6 +43,25 @@ pub struct Build { pub framework: Framework, } +impl Build { + pub fn open(kernel_dir: impl AsRef) -> Result { + let build_compat = parse::parse_and_validate(kernel_dir)?; + Ok(build_compat.into()) + } + + pub fn is_noarch(&self) -> bool { + matches!(self.framework, Framework::TorchNoarch(_)) + } + + pub fn branch(&self) -> Option<&str> { + self.general.hub.as_ref().and_then(|h| h.branch.as_deref()) + } + + pub fn repo_id(&self) -> Option<&str> { + self.general.hub.as_ref().and_then(|h| h.repo_id.as_deref()) + } +} + pub enum Framework { Torch(Torch), TorchNoarch(TorchNoarch), @@ -81,20 +102,6 @@ impl Framework { } } -impl Build { - pub fn is_noarch(&self) -> bool { - matches!(self.framework, Framework::TorchNoarch(_)) - } - - pub fn branch(&self) -> Option<&str> { - self.general.hub.as_ref().and_then(|h| h.branch.as_deref()) - } - - pub fn repo_id(&self) -> Option<&str> { - self.general.hub.as_ref().and_then(|h| h.repo_id.as_deref()) - } -} - pub struct General { pub name: KernelName, diff --git a/kernels-data/src/config/parse.rs b/kernels-data/src/config/parse.rs new file mode 100644 index 000000000..44dbd65c5 --- /dev/null +++ b/kernels-data/src/config/parse.rs @@ -0,0 +1,40 @@ +use std::fs::File; +use std::io::Read; +use std::path::Path; + +use eyre::{Context, Result, bail}; + +use super::{Build, BuildCompat, CurrentConfig}; + +pub(crate) fn parse_and_validate(kernel_dir: impl AsRef) -> Result { + // Only v4 is auto-upgraded to v5 on load; older editions must be migrated + // explicitly with `update-build`. + match parse_and_validate_compat(kernel_dir)? { + BuildCompat::V5(build) => Ok(build), + BuildCompat::V4(build) => { + eprintln!( + "⚠️ build.toml uses the legacy v4 format; upgrading to edition 5 in memory. \ + Run `kernel-builder update-build` to persist the upgrade." + ); + Ok(Build::from(build).into()) + } + BuildCompat::V3(_) => bail!( + "build.toml uses an unsupported legacy format; migrate it with \ + `kernel-builder update-build`" + ), + } +} + +pub(crate) fn parse_and_validate_compat(kernel_dir: impl AsRef) -> Result { + let build_toml = kernel_dir.as_ref().join("build.toml"); + let mut toml_data = String::new(); + File::open(&build_toml) + .wrap_err_with(|| format!("Cannot open {} for reading", build_toml.to_string_lossy()))? + .read_to_string(&mut toml_data) + .wrap_err_with(|| format!("Cannot read from {}", build_toml.to_string_lossy()))?; + + let config: BuildCompat = toml::from_str(&toml_data) + .wrap_err_with(|| format!("Cannot parse TOML in {}", build_toml.to_string_lossy()))?; + + Ok(config) +} From e38cc32c4206fa2e69e159902dbecae7a9e99fc9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Tue, 18 Aug 2026 09:42:33 +0000 Subject: [PATCH 2/2] Test Python binding for `Build` with kernel dependencies --- .../bindings/python/tests/test_config.py | 38 ++++++++++++++++++- 1 file changed, 36 insertions(+), 2 deletions(-) diff --git a/kernels-data/bindings/python/tests/test_config.py b/kernels-data/bindings/python/tests/test_config.py index 14154d0d9..5d072557f 100644 --- a/kernels-data/bindings/python/tests/test_config.py +++ b/kernels-data/bindings/python/tests/test_config.py @@ -1,6 +1,6 @@ import pytest -from kernels_data import Backend, Build +from kernels_data import Backend, Build, KernelDependency, KernelVersion def _write_build_toml(path, backends): @@ -13,6 +13,10 @@ def _write_build_toml(path, backends): edition = 5 license = "Apache-2.0" backends = [{backends_toml}] +kernel-depends = [{{ repo-id = "kernels-community/activation", version = 1 }}] + +[general.cuda] +kernel-depends = [{{ repo-id = "kernels-community/cuda-helper", version = 2 }}] [torch-noarch] """ @@ -33,8 +37,38 @@ def test_build_open_missing_file(tmp_path): Build.open(tmp_path) +def test_build_all_kernel_depends(tmp_path): + _write_build_toml(tmp_path / "build.toml", backends=["cpu", "cuda"]) + + build = Build.open(tmp_path) + assert build.all_kernel_depends(Backend.CUDA) == [ + KernelDependency( + repo_id="kernels-community/activation", version=KernelVersion.Version(1) + ), + KernelDependency( + repo_id="kernels-community/cuda-helper", version=KernelVersion.Version(2) + ), + ] + assert build.all_kernel_depends(Backend.CPU) == [ + KernelDependency( + repo_id="kernels-community/activation", version=KernelVersion.Version(1) + ) + ] + + def test_build_all_kernel_depends_empty(tmp_path): - _write_build_toml(tmp_path / "build.toml", backends=["cpu"]) + (tmp_path / "build.toml").write_text( + """\ +[general] +name = "my-kernel" +version = 1 +edition = 5 +license = "Apache-2.0" +backends = ["cpu"] + +[torch-noarch] +""" + ) build = Build.open(tmp_path) assert build.all_kernel_depends(Backend.CPU) == []