diff --git a/nix-builder/pkgs/generate-symbols/pyproject.toml b/nix-builder/pkgs/generate-symbols/pyproject.toml new file mode 100644 index 00000000..8ecccdc4 --- /dev/null +++ b/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"] diff --git a/nix-builder/pkgs/generate-symbols/src/generate_symbols.py b/nix-builder/pkgs/generate-symbols/src/generate_symbols.py new file mode 100644 index 00000000..2b220ec1 --- /dev/null +++ b/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() diff --git a/nix-builder/pkgs/generate-symbols/tests/test_generate_symbols.py b/nix-builder/pkgs/generate-symbols/tests/test_generate_symbols.py new file mode 100644 index 00000000..d211e4de --- /dev/null +++ b/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