kernels-data: add and expose supported backend archs (#480)

* kernels-data: add and expose supported backend archs

* Fix stub

Author: danieldk

kernel-builder/src/pyproject/common.rs

@@ -38,6 +38,7 @@ pub fn write_metadata(general: &General, file_set: &mut FileSet) -> Result<()> {
             upstream: general.upstream.clone(),
             python_depends,
             backend: BackendInfo {
+                archs: None,
                 backend_type: *backend,
             },
         };

kernels-data/bindings/python/kernels_data.pyi

@@ -1,10 +1,10 @@
 """Type stubs for kernels_data module."""
 
+import os
 from enum import Enum
 from typing import Optional, final
-import os
 
-__all__ = ["Backend", "KernelName", "Metadata", "Version", "__version__"]
+__all__ = ["Backend", "BackendInfo", "KernelName", "Metadata", "Version", "__version__"]
 
 __version__: str
 
@@ -36,6 +36,22 @@ class Backend(Enum):
     def __str__(self) -> str: ...
     def __repr__(self) -> str: ...
 
+@final
+class BackendInfo:
+    """Backend information."""
+
+    @property
+    def backend_type(self) -> Backend:
+        """Return the backend type."""
+        ...
+
+    @property
+    def archs(self) -> Optional[list[str]]:
+        """Optional list of target architectures."""
+        ...
+
+    def __repr__(self) -> str: ...
+
 @final
 class Version:
     """A dotted numeric version (e.g. ``12.8.0``).
@@ -105,5 +121,5 @@ class Metadata:
     @property
     def python_depends(self) -> list[str]: ...
     @property
-    def backend(self) -> Backend: ...
+    def backend(self) -> BackendInfo: ...
     def __repr__(self) -> str: ...

kernels-data/bindings/python/src/lib.rs

@@ -2,7 +2,7 @@ use std::path::PathBuf;
 use std::str::FromStr;
 
 use kernels_data::config::{Backend, KernelName};
-use kernels_data::metadata::{Metadata, parse_metadata};
+use kernels_data::metadata::{BackendInfo, Metadata, parse_metadata};
 use kernels_data::version::Version;
 use pyo3::Bound as PyBound;
 use pyo3::exceptions::PyValueError;
@@ -146,6 +146,44 @@ impl PyBackend {
     }
 }
 
+/// Backend information
+#[pyclass(name = "BackendInfo", frozen)]
+#[derive(Clone, Debug)]
+struct PyBackendInfo {
+    backend_type: PyBackend,
+    archs: Option<Vec<String>>,
+}
+
+impl From<BackendInfo> for PyBackendInfo {
+    fn from(backend_info: BackendInfo) -> Self {
+        Self {
+            backend_type: backend_info.backend_type.into(),
+            archs: backend_info.archs,
+        }
+    }
+}
+
+#[pymethods]
+impl PyBackendInfo {
+    fn __repr__(&self) -> String {
+        format!(
+            "BackendInfo(backend_type={}, archs={:?})",
+            self.backend_type.__repr__(),
+            self.archs
+        )
+    }
+
+    #[getter]
+    fn backend_type(&self) -> PyBackend {
+        self.backend_type
+    }
+
+    #[getter]
+    fn archs(&self) -> Option<&[String]> {
+        self.archs.as_deref()
+    }
+}
+
 /// Parsed `metadata.json` for a kernel build variant.
 #[pyclass(name = "Metadata", frozen)]
 #[derive(Clone, Debug)]
@@ -154,7 +192,7 @@ struct PyMetadata {
     license: Option<String>,
     upstream: Option<String>,
     python_depends: Vec<String>,
-    backend: PyBackend,
+    backend: PyBackendInfo,
 }
 
 impl From<Metadata> for PyMetadata {
@@ -164,7 +202,7 @@ impl From<Metadata> for PyMetadata {
             license: m.license,
             upstream: m.upstream.map(|u| u.to_string()),
             python_depends: m.python_depends,
-            backend: m.backend.backend_type.into(),
+            backend: m.backend.into(),
         }
     }
 }
@@ -202,8 +240,8 @@ impl PyMetadata {
     }
 
     #[getter]
-    fn backend(&self) -> PyBackend {
-        self.backend
+    fn backend(&self) -> PyBackendInfo {
+        self.backend.clone()
     }
 
     fn __repr__(&self) -> String {
@@ -221,6 +259,7 @@ impl PyMetadata {
 #[pyo3::pymodule(name = "kernels_data")]
 fn kernels_data_py(m: &PyBound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyBackend>()?;
+    m.add_class::<PyBackendInfo>()?;
     m.add_class::<PyKernelName>()?;
     m.add_class::<PyMetadata>()?;
     m.add_class::<PyVersion>()?;

kernels-data/bindings/python/tests/test_kernels_data.py

@@ -94,7 +94,7 @@ def test_metadata_load_full(tmp_path):
                 "license": "Apache-2.0",
                 "upstream": "https://github.com/example/kernel",
                 "python-depends": ["torch"],
-                "backend": {"type": "cuda"},
+                "backend": {"type": "cuda", "archs": ["9.0", "10.0"]},
             }
         )
     )
@@ -103,28 +103,25 @@ def test_metadata_load_full(tmp_path):
     assert m.license == "Apache-2.0"
     assert m.upstream == "https://github.com/example/kernel"
     assert m.python_depends == ["torch"]
-    assert m.backend == Backend.CUDA
+    assert m.backend.backend_type == Backend.CUDA
+    assert m.backend.archs == ["9.0", "10.0"]
 
 
 def test_metadata_load_minimal(tmp_path):
     path = tmp_path / "metadata.json"
-    path.write_text(
-        json.dumps({"python-depends": [], "backend": {"type": "cpu"}})
-    )
+    path.write_text(json.dumps({"python-depends": [], "backend": {"type": "cpu"}}))
     m = Metadata.load(path)
     assert m.version is None
     assert m.license is None
     assert m.upstream is None
     assert m.python_depends == []
-    assert m.backend == Backend.CPU
+    assert m.backend.backend_type == Backend.CPU
 
 
 def test_metadata_load_cann(tmp_path):
     path = tmp_path / "metadata.json"
-    path.write_text(
-        json.dumps({"python-depends": [], "backend": {"type": "cann"}})
-    )
-    assert Metadata.load(path).backend == Backend.CANN
+    path.write_text(json.dumps({"python-depends": [], "backend": {"type": "cann"}}))
+    assert Metadata.load(path).backend.backend_type == Backend.CANN
 
 
 def test_metadata_load_unknown_field_rejected(tmp_path):
@@ -155,7 +152,7 @@ def test_metadata_load(tmp_path):
         **{"python-depends": ["torch"], "backend": {"type": "cuda"}},
     )
     m = Metadata.load(path)
-    assert m.backend == Backend.CUDA
+    assert m.backend.backend_type == Backend.CUDA
 
 
 def test_metadata_load_missing_file(tmp_path):

kernels-data/src/metadata.rs

@@ -10,6 +10,7 @@ use crate::config::Backend;
 pub struct BackendInfo {
     #[serde(rename = "type")]
     pub backend_type: Backend,
+    pub archs: Option<Vec<String>>,
 }
 
 #[derive(Debug, Deserialize, Serialize)]