Revamp metadata writing and add kernel id (#471)

* kernel-builder: revamp metadata writing

* Strongly type the kernel identifier within kernel-builder.
* Write the kernel identifier to the kernel metadata.
* Make noarch kernels include the backend in the identifier as well.
* Fix tvm-ffi build-time metadata update to include the arches.
* Rename build-time metadata update script to be more general.
* Run the build-time metadata update script for non-CUDA/ROCm backends.
  This is not strictly necessary now, but allows us to add other
  metadata bits in the future.
* Add a build hook that validates the metadata after build using
  the Rust-side data structures. This ensures that build-time changes
  are valid and schema-conforming.

* kernels: support kernel id metadata and use as kernel name

This change adds support for kernel id metadata and uses the kernel id
(when present) as the identifier in the Python module table in place of
the path hash.

* Clippy fix

* Update CLI docs

* Add additional metadata fields to the docs

* Add neuron to noarch _ops

* metadata warning: tell user how to fix it

* Make id optional

kernels-data is going to be used by `kernels`, so we need to make sure
that we can also read kernels that do not have an id yet.

Author: danieldk

docs/source/builder-cli.md

@@ -11,12 +11,13 @@ This document contains the help content for the `kernel-builder` command-line pr
 * [`kernel-builder build-and-copy`↴](#kernel-builder-build-and-copy)
 * [`kernel-builder build-and-upload`↴](#kernel-builder-build-and-upload)
 * [`kernel-builder upload`↴](#kernel-builder-upload)
+* [`kernel-builder check-config`↴](#kernel-builder-check-config)
+* [`kernel-builder check-builds`↴](#kernel-builder-check-builds)
 * [`kernel-builder create-pyproject`↴](#kernel-builder-create-pyproject)
 * [`kernel-builder devshell`↴](#kernel-builder-devshell)
 * [`kernel-builder list-variants`↴](#kernel-builder-list-variants)
 * [`kernel-builder testshell`↴](#kernel-builder-testshell)
 * [`kernel-builder update-build`↴](#kernel-builder-update-build)
-* [`kernel-builder validate`↴](#kernel-builder-validate)
 * [`kernel-builder skills`↴](#kernel-builder-skills)
 * [`kernel-builder skills add`↴](#kernel-builder-skills-add)
 * [`kernel-builder clean-pyproject`↴](#kernel-builder-clean-pyproject)
@@ -35,12 +36,13 @@ Build Hugging Face Hub kernels
 * `build-and-copy` — Build the kernel and copy artifacts locally
 * `build-and-upload` — Build the kernel and upload to Hugging Face Hub
 * `upload` — Upload kernel build artifacts to the Hugging Face Hub
+* `check-config` — Validate the build.toml file
+* `check-builds` — Validate kernel builds
 * `create-pyproject` — Generate CMake files for a kernel extension build
 * `devshell` — Spawn a kernel development shell
 * `list-variants` — List build variants
 * `testshell` — Spawn a kernel test shell
 * `update-build` — Update a `build.toml` to the current format
-* `validate` — Validate the build.toml file
 * `skills` — Install skills for AI coding assistants (Claude, Codex, OpenCode)
 * `clean-pyproject` — Clean generated artifacts
 
@@ -169,6 +171,30 @@ Upload kernel build artifacts to the Hugging Face Hub
 
 
 
+## `kernel-builder check-config`
+
+Validate the build.toml file
+
+**Usage:** `kernel-builder check-config [KERNEL_DIR]`
+
+###### **Arguments:**
+
+* `<KERNEL_DIR>`
+
+
+
+## `kernel-builder check-builds`
+
+Validate kernel builds
+
+**Usage:** `kernel-builder check-builds [KERNEL_DIR]`
+
+###### **Arguments:**
+
+* `<KERNEL_DIR>`
+
+
+
 ## `kernel-builder create-pyproject`
 
 Generate CMake files for a kernel extension build
@@ -183,7 +209,7 @@ Generate CMake files for a kernel extension build
 ###### **Options:**
 
 * `-f`, `--force` — Force-overwrite existing files
-* `--ops-id <OPS_ID>` — This is an optional unique identifier that is suffixed to the kernel name to avoid name collisions. (e.g. Git SHA)
+* `--unique-id <UNIQUE_ID>` — This is an optional unique identifier that is suffixed to the kernel name to avoid name collisions. (e.g. Git SHA)
 
 
 
@@ -253,18 +279,6 @@ Update a `build.toml` to the current format
 
 
 
-## `kernel-builder validate`
-
-Validate the build.toml file
-
-**Usage:** `kernel-builder validate [KERNEL_DIR]`
-
-###### **Arguments:**
-
-* `<KERNEL_DIR>`
-
-
-
 ## `kernel-builder skills`
 
 Install skills for AI coding assistants (Claude, Codex, OpenCode)

docs/source/builder/writing-kernels.md

@@ -184,7 +184,7 @@ The following sections enumerate all supported options for `build.toml`.
 
 - `name` (required): the name of the kernel. The Python code for a Torch
   extension must be stored in `torch-ext/<name>`.
-- `version` (int, **experimental**): the major version of the kernel.
+- `version` (int): the major version of the kernel.
   The version is written to the kernel's `metadata.json` and is used
   by the `kernels upload` command to upload the kernel to a version
   branch named `v<version>`.

docs/source/kernel-requirements.md

@@ -36,16 +36,52 @@ must be available for that combination.
 
 ## Kernel metadata
 
-The build variant directory can optionally contain a `metadata.json` file.
-Currently the metadata specifies the kernel's version and Python dependencies,
-for example:
+The build variant directory must contain a `metadata.json` file with kernel
+metadata. Currently the following top-level keys are supported:
+
+- `id` (`str`, required): a unique identifier for the kernel. This
+  identifier must also be a valid Python module name. If the kernel
+  registers Torch ops, they must be registered as `torch.ops.<id>`
+- `version` (`int`, required): the kernel version number.
+- `backend` (`dict`, required): information about the compute backend that
+  this build variant supports.
+- `python-depends` (`list[str]`, optional): list of Python dependencies
+  from a curated set of Python dependencies.
+
+Example `metadata.json`:
 
 ```json
 {
+  "id": "_mykernel_cuda_be238e4",
   "python-depends": ["einops"],
-  "version": 1
+  "version": 1,
+  "backend": {
+    "type": "cuda",
+    "archs": ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0+PTX"]
+  }
+}
+```
+
+The `metadata.json` file is generated automatically by `kernel-builder`.
+
+## Backend
+
+The `backend` specifies a dictionary of the following form:
+
+```json
+{
+  # ...
+  "backend": {
+    "type": "cuda",
+    "archs": ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0+PTX"]
+  }
+}
 ```
 
+The backend `type` must be one of `cann`, `cpu`, `cuda`, `metal`, `neuron`,
+`rocm`, or `xpu`. For CUDA and ROCm, the supported architectures must
+be specified in the `archs` field.
+
 ### Python dependencies
 
 You can specify Python dependencies that your kernel requires. Dependencies can be either general (required for all backends) or backend-specific (required only for certain compute backends like CUDA, ROCm, XPU, Metal, or CPU).

kernel-builder/src/main.rs

@@ -40,6 +40,9 @@ mod skills;
 mod util;
 use util::{check_or_infer_kernel_dir, parse_and_validate};
 
+mod validate_builds;
+use validate_builds::check_builds;
+
 #[derive(Args, Debug)]
 struct NixArgs {
     /// Maximum number of parallel Nix build jobs.
@@ -127,6 +130,18 @@ enum Commands {
     /// Upload kernel build artifacts to the Hugging Face Hub.
     Upload(UploadArgs),
 
+    /// Validate the build.toml file.
+    CheckConfig {
+        #[arg(name = "KERNEL_DIR")]
+        kernel_dir: Option<PathBuf>,
+    },
+
+    /// Validate kernel builds.
+    CheckBuilds {
+        #[arg(name = "KERNEL_DIR")]
+        kernel_dir: Option<PathBuf>,
+    },
+
     /// Generate CMake files for a kernel extension build.
     CreatePyproject {
         #[arg(name = "KERNEL_DIR")]
@@ -144,7 +159,7 @@ enum Commands {
         /// This is an optional unique identifier that is suffixed to the
         /// kernel name to avoid name collisions. (e.g. Git SHA)
         #[arg(long)]
-        ops_id: Option<String>,
+        unique_id: Option<String>,
     },
 
     /// Spawn a kernel development shell.
@@ -201,12 +216,6 @@ enum Commands {
         kernel_dir: Option<PathBuf>,
     },
 
-    /// Validate the build.toml file.
-    Validate {
-        #[arg(name = "KERNEL_DIR")]
-        kernel_dir: Option<PathBuf>,
-    },
-
     /// Install skills for AI coding assistants (Claude, Codex, OpenCode).
     Skills {
         #[command(subcommand)]
@@ -333,8 +342,8 @@ fn main() -> Result<()> {
             kernel_dir,
             force,
             target_dir,
-            ops_id,
-        } => create_pyproject(kernel_dir, target_dir, force, ops_id),
+            unique_id,
+        } => create_pyproject(kernel_dir, target_dir, force, unique_id),
         Commands::Devshell {
             kernel_dir,
             variant,
@@ -359,8 +368,12 @@ fn main() -> Result<()> {
             variant,
         ),
         Commands::UpdateBuild { kernel_dir } => update_build(kernel_dir),
-        Commands::Validate { kernel_dir } => {
-            validate(kernel_dir)?;
+        Commands::CheckConfig { kernel_dir } => {
+            check_config(kernel_dir)?;
+            Ok(())
+        }
+        Commands::CheckBuilds { kernel_dir } => {
+            check_builds(kernel_dir)?;
             Ok(())
         }
         Commands::GenerateDocs => {
@@ -390,7 +403,7 @@ fn main() -> Result<()> {
     }
 }
 
-fn validate(kernel_dir: Option<PathBuf>) -> Result<()> {
+fn check_config(kernel_dir: Option<PathBuf>) -> Result<()> {
     let kernel_dir = check_or_infer_kernel_dir(kernel_dir)?;
     parse_and_validate(kernel_dir)?;
     Ok(())

kernel-builder/src/pyproject/common.rs

@@ -6,9 +6,11 @@ use itertools::Itertools;
 use kernels_data::config::{Backend, General};
 use kernels_data::metadata::{BackendInfo, Metadata};
 
+use crate::pyproject::ops_identifier::KernelIdentifier;
 use crate::pyproject::FileSet;
 
 static COMPAT_PY: &str = include_str!("templates/compat.py");
+static ADD_BUILD_METADATA_PY: &str = include_str!("templates/torch/add_build_metadata.py");
 
 pub fn write_compat_py(file_set: &mut FileSet) -> Result<()> {
     let mut path = PathBuf::new();
@@ -18,7 +20,11 @@ pub fn write_compat_py(file_set: &mut FileSet) -> Result<()> {
     Ok(())
 }
 
-pub fn write_metadata(general: &General, file_set: &mut FileSet) -> Result<()> {
+pub fn write_metadata(
+    general: &General,
+    kernel_id: &KernelIdentifier,
+    file_set: &mut FileSet,
+) -> Result<()> {
     for backend in &Backend::all() {
         let writer = file_set.entry(format!("metadata-{backend}.json"));
 
@@ -33,6 +39,7 @@ pub fn write_metadata(general: &General, file_set: &mut FileSet) -> Result<()> {
             .collect::<Result<Vec<_>>>()?;
 
         let metadata = Metadata {
+            id: Some(kernel_id.to_string_for_backend(*backend)),
             version: general.version,
             license: general.license.clone(),
             upstream: general.upstream.clone(),
@@ -61,6 +68,14 @@ where
         .join(";")
 }
 
+pub fn write_add_build_metadata_py(file_set: &mut FileSet) {
+    write_cmake_file(
+        file_set,
+        "add_build_metadata.py",
+        ADD_BUILD_METADATA_PY.as_bytes(),
+    );
+}
+
 /// Helper function to write a file to the cmake subdirectory
 pub fn write_cmake_file(file_set: &mut FileSet, filename: &str, content: &[u8]) {
     let mut path = PathBuf::new();

kernel-builder/src/pyproject/mod.rs

@@ -8,7 +8,10 @@ use eyre::{bail, Result};
 use kernels_data::config::{Build, Framework};
 use minijinja::Environment;
 
-use crate::util::{check_or_infer_kernel_dir, check_or_infer_target_dir, parse_build};
+use crate::{
+    pyproject::ops_identifier::KernelIdentifier,
+    util::{check_or_infer_kernel_dir, check_or_infer_target_dir, parse_build},
+};
 
 pub(crate) mod common;
 pub mod deps;
@@ -21,21 +24,17 @@ mod tvm_ffi;
 pub use fileset::FileSet;
 pub use kernels_data::metadata::parse_metadata;
 
-pub fn create_pyproject_file_set(
-    build: Build,
-    target_dir: impl AsRef<Path>,
-    ops_id: Option<String>,
-) -> Result<FileSet> {
+pub fn create_pyproject_file_set(build: Build, kernel_id: &KernelIdentifier) -> Result<FileSet> {
     let mut env = Environment::new();
     env.set_trim_blocks(true);
     minijinja_embed::load_templates!(&mut env);
 
     let file_set = if matches!(build.framework, Framework::TvmFfi(_)) {
-        tvm_ffi::write_tvm_ffi_ext(&env, &build, target_dir, ops_id)?
+        tvm_ffi::write_tvm_ffi_ext(&env, &build, kernel_id)?
     } else if build.is_noarch() {
-        torch::write_torch_ext_noarch(&env, &build, target_dir, ops_id)?
+        torch::write_torch_ext_noarch(&env, &build, kernel_id)?
     } else {
-        torch::write_torch_ext(&env, &build, target_dir, ops_id)?
+        torch::write_torch_ext(&env, &build, kernel_id)?
     };
 
     Ok(file_set)
@@ -45,12 +44,13 @@ pub fn create_pyproject(
     kernel_dir: Option<PathBuf>,
     target_dir: Option<PathBuf>,
     force: bool,
-    ops_id: Option<String>,
+    unique_id: Option<String>,
 ) -> 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 file_set = create_pyproject_file_set(build, &target_dir, ops_id)?;
+    let kernel_id = KernelIdentifier::new(&kernel_dir, build.general.name.python_name(), unique_id);
+    let file_set = create_pyproject_file_set(build, &kernel_id)?;
     file_set.write(&target_dir, force)?;
 
     Ok(())
@@ -61,14 +61,14 @@ pub fn clean_pyproject(
     target_dir: Option<PathBuf>,
     dry_run: bool,
     force: bool,
-    ops_id: Option<String>,
+    unique_id: Option<String>,
 ) -> 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 generated_files =
-        create_pyproject_file_set(build, target_dir.clone(), ops_id)?.into_names();
+    let kernel_id = KernelIdentifier::new(&kernel_dir, build.general.name.python_name(), unique_id);
+
+    let generated_files = create_pyproject_file_set(build, &kernel_id)?.into_names();
 
     if generated_files.is_empty() {
         eprintln!("No generated artifacts found to clean.");