kernel-builder: Python script to generate symbols in __all__

nix-builder/pkgs/generate-symbols/pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "generate-symbols"
+version = "0.1.0"
+description = "Generate public API symbols for built kernels"
+requires-python = ">=3.10"
+license = "Apache-2.0"
+dependencies = ["kernels"]
+
+[project.optional-dependencies]
+test = ["pytest>=8"]
+
+[project.scripts]
+generate-symbols = "generate_symbols:main"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+py-modules = ["generate_symbols"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

nix-builder/pkgs/generate-symbols/src/generate_symbols.py

@@ -0,0 +1,139 @@
+"""Inspect a built kernel's public functions and layer classes into symbols.json.
+
+Run this in the build's Python environment, with the kernel's runtime
+dependencies and the kernels package available. The build must include its
+metadata.json. Importing a kernel executes its initialization code.
+The inspection itself does not call exported functions or instantiate classes.
+"""
+
+import argparse
+import inspect
+import json
+from collections.abc import Callable
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+
+def describe_signature(obj: Any) -> dict[str, Any] | None:
+    """Serialize inspect.Signature without evaluating string annotations."""
+    try:
+        signature = inspect.signature(obj, eval_str=False)
+    except (TypeError, ValueError):
+        # Some extension functions and callable objects expose no signature.
+        return None
+
+    return {
+        "parameters": [
+            {
+                "name": parameter.name,
+                "kind": parameter.kind.name,
+                "default": (
+                    None
+                    if parameter.default is inspect.Parameter.empty
+                    else repr(parameter.default)
+                ),
+                "annotation": (
+                    None
+                    if parameter.annotation is inspect.Parameter.empty
+                    else inspect.formatannotation(parameter.annotation)
+                ),
+            }
+            for parameter in signature.parameters.values()
+        ],
+        "return_annotation": (
+            None
+            if signature.return_annotation is inspect.Signature.empty
+            else inspect.formatannotation(signature.return_annotation)
+        ),
+    }
+
+
+def describe_symbol(name: str, obj: Any) -> dict[str, Any]:
+    """Describe a public function or layer class using inspect terminology."""
+    symbol = {
+        "name": name,
+        "qualname": obj.__qualname__,
+        "module": obj.__module__,
+        "kind": "class" if inspect.isclass(obj) else "function",
+        "doc": inspect.getdoc(obj),
+        "signature": describe_signature(obj),
+    }
+    if inspect.isclass(obj):
+        symbol["attributes"] = {
+            attribute: value
+            if isinstance(value := getattr(obj, attribute, None), bool)
+            else None
+            for attribute in ("has_backward", "can_torch_compile")
+        }
+    return symbol
+
+
+def describe_exports(
+    module: ModuleType, predicate: Callable[[Any], bool]
+) -> list[dict[str, Any]]:
+    """Select public exports of the requested kind, preserving __all__ order."""
+    exports = getattr(module, "__all__", ())
+    if not isinstance(exports, (list, tuple)) or not all(
+        isinstance(export, str) for export in exports
+    ):
+        raise ValueError(
+            f"{module.__name__}.__all__ must be a list or tuple of strings"
+        )
+    symbols = []
+    for name in dict.fromkeys(exports):
+        if name.startswith("_") or name == "layers":
+            continue
+        try:
+            obj = getattr(module, name)
+        except AttributeError as error:
+            raise ValueError(
+                f"{module.__name__}.__all__ names a missing symbol: {name!r}"
+            ) from error
+        if not predicate(obj):
+            continue
+        if "_private_for_testing" in obj.__module__.split("."):
+            continue
+        symbols.append(describe_symbol(name, obj))
+    return symbols
+
+
+def generate_symbols(module: ModuleType) -> dict[str, Any]:
+    """Describe top-level functions and optional layers, each scoped by __all__."""
+    functions = describe_exports(module, inspect.isfunction)
+    layers = getattr(module, "layers", None)
+    return {
+        "schema_version": 1,
+        "module": module.__name__,
+        "functions": functions,
+        "layers": describe_exports(layers, inspect.isclass)
+        if inspect.ismodule(layers)
+        else [],
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "kernel_dir",
+        type=Path,
+        help="Local kernel repository or build variant containing metadata.json",
+    )
+    parser.add_argument(
+        "--backend", help="Backend to select when loading a local kernel repository"
+    )
+    parser.add_argument(
+        "--output", type=Path, required=True, help="Destination symbols.json"
+    )
+    args = parser.parse_args()
+
+    from kernels import get_local_kernel
+
+    module = get_local_kernel(args.kernel_dir, backend=args.backend)
+    # Finish inspection and serialization before touching an existing output file.
+    content = json.dumps(generate_symbols(module), indent=2, ensure_ascii=False) + "\n"
+    args.output.write_text(content, encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()

nix-builder/pkgs/generate-symbols/tests/test_generate_symbols.py

@@ -0,0 +1,231 @@
+import importlib
+import json
+from types import ModuleType
+
+import generate_symbols as generator
+import pytest
+
+
+@pytest.fixture
+def kernel():
+    """A small kernel exporting one function and one layer."""
+
+    def relu(x: int) -> int:
+        """Apply ReLU."""
+        raise AssertionError("Inspection must not call the function")
+
+    class ReLU:
+        has_backward = True
+        can_torch_compile = False
+
+        def __init__(self, size: int = 4):
+            raise AssertionError("Inspection must not instantiate the layer")
+
+    layers = ModuleType("example.layers")
+    layers.ReLU = ReLU
+    layers.__all__ = ["ReLU"]
+
+    module = ModuleType("example")
+    module.relu = relu
+    module.layers = layers
+    module.__all__ = ["relu", "layers"]
+    return module
+
+
+def test_public_function_is_included(kernel):
+    result = generator.generate_symbols(kernel)
+
+    assert len(result["functions"]) == 1
+    function = result["functions"][0]
+    assert function["name"] == "relu"
+    assert function["kind"] == "function"
+    assert function["doc"] == "Apply ReLU."
+
+
+def test_public_layer_is_included(kernel):
+    result = generator.generate_symbols(kernel)
+
+    assert len(result["layers"]) == 1
+    layer = result["layers"][0]
+    assert layer["name"] == "ReLU"
+    assert layer["kind"] == "class"
+    assert layer["attributes"] == {
+        "has_backward": True,
+        "can_torch_compile": False,
+    }
+
+
+def test_symbols_not_in_all_are_excluded(kernel):
+    kernel.helper = kernel.relu
+    kernel.layers.OtherLayer = kernel.layers.ReLU
+
+    result = generator.generate_symbols(kernel)
+
+    assert [function["name"] for function in result["functions"]] == ["relu"]
+    assert [layer["name"] for layer in result["layers"]] == ["ReLU"]
+
+
+def test_underscore_names_are_excluded_even_when_exported(kernel):
+    kernel._helper = kernel.relu
+    kernel.__all__.append("_helper")
+    kernel.layers._HiddenLayer = kernel.layers.ReLU
+    kernel.layers.__all__.append("_HiddenLayer")
+
+    result = generator.generate_symbols(kernel)
+
+    assert [function["name"] for function in result["functions"]] == ["relu"]
+    assert [layer["name"] for layer in result["layers"]] == ["ReLU"]
+
+
+def test_function_from_private_testing_module_is_excluded(kernel):
+    # A public export name must not expose a private testing implementation.
+    kernel.relu.__module__ = "example._private_for_testing.helpers"
+
+    result = generator.generate_symbols(kernel)
+
+    assert result["functions"] == []
+    assert [layer["name"] for layer in result["layers"]] == ["ReLU"]
+
+
+def test_layer_from_private_testing_module_is_excluded(kernel):
+    kernel.layers.ReLU.__module__ = "example._private_for_testing"
+
+    result = generator.generate_symbols(kernel)
+
+    assert result["layers"] == []
+    assert [function["name"] for function in result["functions"]] == ["relu"]
+
+
+def test_functions_in_layers_are_not_layers(kernel):
+    kernel.layers.helper = kernel.relu
+    kernel.layers.__all__.append("helper")
+
+    result = generator.generate_symbols(kernel)
+
+    assert [layer["name"] for layer in result["layers"]] == ["ReLU"]
+
+
+def test_layers_are_optional(kernel, monkeypatch):
+    del kernel.layers
+    kernel.__all__ = ["relu"]
+    kernel.__path__ = []  # Model a package without an exposed layers module.
+
+    def unexpected_import(name):
+        pytest.fail(f"Inspection should not import {name}")
+
+    monkeypatch.setattr(importlib, "import_module", unexpected_import)
+
+    result = generator.generate_symbols(kernel)
+
+    assert [function["name"] for function in result["functions"]] == ["relu"]
+    assert result["layers"] == []
+
+
+def test_layers_need_not_appear_in_package_all(kernel):
+    kernel.__all__ = ["relu"]
+
+    result = generator.generate_symbols(kernel)
+
+    assert [layer["name"] for layer in result["layers"]] == ["ReLU"]
+
+
+def test_parameter_kinds_are_preserved():
+    def function(x, /, scale, *args, out=None, **kwargs):
+        pass
+
+    signature = generator.describe_signature(function)
+    parameters = signature["parameters"]
+
+    assert [parameter["name"] for parameter in parameters] == [
+        "x",
+        "scale",
+        "args",
+        "out",
+        "kwargs",
+    ]
+    assert [parameter["kind"] for parameter in parameters] == [
+        "POSITIONAL_ONLY",
+        "POSITIONAL_OR_KEYWORD",
+        "VAR_POSITIONAL",
+        "KEYWORD_ONLY",
+        "VAR_KEYWORD",
+    ]
+
+
+def test_missing_default_differs_from_explicit_none():
+    def function(x, out=None):
+        pass
+
+    signature = generator.describe_signature(function)
+    x, out = signature["parameters"]
+
+    assert x["default"] is None
+    assert out["default"] == "None"
+    assert x["annotation"] is None
+    assert signature["return_annotation"] is None
+
+
+def test_string_annotations_are_not_evaluated():
+    def function(x: "UndefinedTensor") -> "UndefinedTensor":  # noqa: F821
+        pass
+
+    signature = generator.describe_signature(function)
+
+    assert signature["parameters"][0]["annotation"] == "'UndefinedTensor'"
+    assert signature["return_annotation"] == "'UndefinedTensor'"
+
+
+def test_layer_constructor_signature_is_included(kernel):
+    signature = generator.describe_signature(kernel.layers.ReLU)
+
+    assert signature["parameters"] == [
+        {
+            "name": "size",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "4",
+            "annotation": "int",
+        }
+    ]
+
+
+def test_unavailable_signature_is_null(kernel):
+    kernel.layers.ReLU.__signature__ = "invalid signature"
+
+    result = generator.generate_symbols(kernel)
+
+    assert result["layers"][0]["signature"] is None
+
+
+def test_inherited_layer_flags_are_included(kernel):
+    class ChildLayer(kernel.layers.ReLU):
+        pass
+
+    kernel.layers.ChildLayer = ChildLayer
+    kernel.layers.__all__ = ["ChildLayer"]
+
+    result = generator.generate_symbols(kernel)
+
+    assert result["layers"][0]["attributes"] == {
+        "has_backward": True,
+        "can_torch_compile": False,
+    }
+
+
+def test_missing_layer_flags_are_unknown(kernel):
+    del kernel.layers.ReLU.has_backward
+    del kernel.layers.ReLU.can_torch_compile
+
+    result = generator.generate_symbols(kernel)
+
+    assert result["layers"][0]["attributes"] == {
+        "has_backward": None,
+        "can_torch_compile": None,
+    }
+
+
+def test_result_can_be_written_as_json(kernel):
+    result = generator.generate_symbols(kernel)
+
+    assert result["schema_version"] == 1
+    assert result["module"] == "example"
+    assert json.loads(json.dumps(result)) == result