From 19284561a28c536ad3c639b959ec23a3979c5bf2 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 6 Oct 2026 11:47:58 +0530 Subject: [PATCH 01/10] feat: implement kernels describe. --- docs/source/_toctree.yml | 2 + docs/source/cli-describe.md | 61 +++++++ docs/source/cli-info.md | 1 + docs/source/cli.md | 2 +- kernels/src/kernels/cli/__init__.py | 14 ++ kernels/src/kernels/cli/describe.py | 240 ++++++++++++++++++++++++++ kernels/tests/test_cli_describe.py | 255 ++++++++++++++++++++++++++++ 7 files changed, 574 insertions(+), 1 deletion(-) create mode 100644 docs/source/cli-describe.md create mode 100644 kernels/src/kernels/cli/describe.py create mode 100644 kernels/tests/test_cli_describe.py diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 60a88fa6..45094fb4 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -35,6 +35,8 @@ title: Overview - local: cli-info title: kernels info + - local: cli-describe + title: kernels describe - local: cli-variants title: kernels variants - local: cli-versions diff --git a/docs/source/cli-describe.md b/docs/source/cli-describe.md new file mode 100644 index 00000000..80f29c28 --- /dev/null +++ b/docs/source/cli-describe.md @@ -0,0 +1,61 @@ +# kernels describe + +Use `kernels describe` to list a kernel's public functions and layers without +loading the kernel or installing its dependencies. For Hub kernels, the command +downloads only the Python source needed to inspect the API, plus `metadata.json` +when needed to locate an older build's package. It does not download compiled +libraries or a repository snapshot. + +## Usage + +```bash +kernels describe +``` + +Options: + +- `--version `: select an integer kernel version or `latest` (the default). +- `--revision `: select a branch, tag, or commit instead of a version. + +These options are mutually exclusive and apply only to Hub repositories. The +latest numbered version is selected by default, or `main` for an unversioned +repository. + +An existing local directory is treated as a kernel. + +## Examples + +```bash +kernels describe kernels-community/activation +kernels describe kernels-community/activation --version 1 +kernels describe kernels-community/activation --revision main +kernels describe ./my-kernel +``` + +## Export conventions + +Functions must be listed in the top-level module's `__all__`. Layers must be +listed in `layers.py` or `layers/__init__.py`'s own `__all__`. A module without +`__all__` contributes no exports. Private names and symbols from +`_private_for_testing` are purposefully excluded. + +Each layer includes its explicitly declared `has_backward` and +`can_torch_compile` boolean attributes. + +For example, a kernel declaring `gelu` and `Gelu` in the respective `__all__` +lists could produce: + +```text +Repository: example/activation +Revision: v1 + +Functions: + gelu + +Layers: + Gelu has_backward=True can_torch_compile=True +``` + +Use [kernels info](cli-info.md) for metadata such as the license, dependencies, +and supported backends, and [kernels variants](cli-variants.md) for build variants +and compatibility information. diff --git a/docs/source/cli-info.md b/docs/source/cli-info.md index 1de4acd0..14284fe1 100644 --- a/docs/source/cli-info.md +++ b/docs/source/cli-info.md @@ -5,6 +5,7 @@ It prints the kernel metadata (name, version, license, upstream/source repositories, Python dependencies, and supported backends). To list the build variants of a kernel and their compatibility with the current system, use [kernels variants](cli-variants.md). +To list its public functions and layers, use [kernels describe](cli-describe.md). ## Usage diff --git a/docs/source/cli.md b/docs/source/cli.md index 3563e07b..00b41b66 100644 --- a/docs/source/cli.md +++ b/docs/source/cli.md @@ -7,9 +7,9 @@ The `kernels` CLI provides commands for managing compute kernels. | Command | Description | | ------------------------------------------- | ---------------------------------- | | [benchmark](cli-benchmark.md) | Run benchmark results for a kernel | +| [describe](cli-describe.md) | List public functions and layers | | [download](cli-download.md) | Download a kernel | | [info](cli-info.md) | Describe a kernel | | [lock](cli-lock.md) | Lock kernel revisions | | [verify-signature](cli-verify-signature.md) | Verify a kernel signature | | [variants](cli-variants.md) | Show variants and compatibility | - diff --git a/kernels/src/kernels/cli/__init__.py b/kernels/src/kernels/cli/__init__.py index df350769..75dd8cad 100644 --- a/kernels/src/kernels/cli/__init__.py +++ b/kernels/src/kernels/cli/__init__.py @@ -3,6 +3,7 @@ from pathlib import Path from typing import Literal +from kernels.cli.describe import print_kernel_description from kernels.cli.download import download_kernels from kernels.cli.info import print_kernel_info from kernels.cli.lock import lock_kernels @@ -32,6 +33,15 @@ def main(): ) download_parser.set_defaults(func=download_kernels) + describe_parser = subparsers.add_parser("describe", help="Show the public functions and layers of a kernel") + describe_parser.add_argument("kernel", help="The kernel repo ID or a local path to a kernel") + describe_selection = describe_parser.add_mutually_exclusive_group() + describe_selection.add_argument( + "--version", type=_variant_version, help="The kernel version or 'latest' (default: latest version)" + ) + describe_selection.add_argument("--revision", help="The kernel revision (branch, tag, or commit)") + describe_parser.set_defaults(func=kernel_describe) + info_parser = subparsers.add_parser("info", help="Describe a kernel") info_parser.add_argument( "kernel", @@ -162,6 +172,10 @@ def kernel_info(args): ) +def kernel_describe(args): + print_kernel_description(args.kernel, revision=args.revision, version=args.version) + + def kernel_versions(args): print_kernel_versions(args.repo_id) diff --git a/kernels/src/kernels/cli/describe.py b/kernels/src/kernels/cli/describe.py new file mode 100644 index 00000000..765f6f16 --- /dev/null +++ b/kernels/src/kernels/cli/describe.py @@ -0,0 +1,240 @@ +"""Describe kernel exports by reading Python source, without executing it.""" + +import ast +import sys +from pathlib import Path, PurePosixPath +from typing import Literal + +from huggingface_hub import constants +from huggingface_hub.errors import HfHubHTTPError, RemoteEntryNotFoundError + +from kernels._rust import KernelName, Metadata +from kernels._versions import _get_available_versions +from kernels.compat import tomllib +from kernels.hf_hub import _get_cache_dir, _get_hf_api +from kernels.variants import get_variants, get_variants_local + +_Definition = ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef + + +def print_kernel_description( + kernel: str, + *, + version: int | Literal["latest"] | None = None, + revision: str | None = None, +): + """Print the functions and layers explicitly exported by a kernel.""" + try: + if version is not None and revision is not None: + raise ValueError("Only one of revision or version can be specified") + + path = Path(kernel) + if path.is_dir(): + if version is not None or revision is not None: + raise ValueError("revision and version cannot be used with a local path") + source = _local_source(path) + heading = f"Path: {path}" + else: + if revision is None: + versions = _get_available_versions(kernel, local_files_only=constants.HF_HUB_OFFLINE) + if isinstance(version, int): + if version not in versions: + raise ValueError( + f"Version {version} not found, available versions: {', '.join(map(str, sorted(versions)))}" + ) + revision = versions[version].name + else: + revision = versions[max(versions)].name if versions else "main" + source = _hub_source(kernel, revision) + heading = f"Repository: {kernel}\nRevision: {revision}" + + functions = source.exports("__init__.py", (ast.FunctionDef, ast.AsyncFunctionDef)) + layers_file = source.module_file("layers") + layers = source.exports(layers_file, (ast.ClassDef,)) if layers_file else [] + except (ValueError, OSError, SyntaxError, HfHubHTTPError) as error: + print(f"Cannot describe {kernel}: {error}", file=sys.stderr) + sys.exit(1) + + print(heading) + print("\nFunctions:") + for name, _, _ in functions: + print(f" {name}") + if not functions: + print(" No functions declared in __all__.") + + print("\nLayers:") + for name, node, _ in layers: + flags = _assignments(node.body) + capabilities = [] + for flag in ("has_backward", "can_torch_compile"): + value = flags.get(flag) + rendered = ( + str(value.value) if isinstance(value, ast.Constant) and isinstance(value.value, bool) else "unknown" + ) + capabilities.append(f"{flag}={rendered}") + print(f" {name} {' '.join(capabilities)}") + if not layers: + print(" No layers declared in __all__.") + + +def _local_source(path: Path) -> "_Source": + # Accept a package/variant directory as well as a built repository. + if (path / "__init__.py").is_file(): + return _Source(path) + variants = get_variants_local(path / "build") + if variants: + path = path / "build" / min(v.variant_str for v in variants) + elif (path / "build.toml").is_file(): + with (path / "build.toml").open("rb") as config_file: + config = tomllib.load(config_file) + name = config.get("general", {}).get("name") + if not isinstance(name, str): + raise ValueError("Missing general.name in build.toml") + name = KernelName(name).python_name + for directory in ("torch-ext", "tvm-ffi-ext"): + package = path / directory / name + if (package / "__init__.py").is_file(): + return _Source(package) + raise ValueError(f"No Python kernel package found in {path}") + return _source_from_variant(_Source(path)) + + +def _hub_source(repo_id: str, revision: str) -> "_Source": + api = _get_hf_api() + variants = get_variants(api, repo_id=repo_id, revision=revision) + if not variants: + raise ValueError(f"No build variants found (revision: {revision})") + # API source can be inspected on any host; do not resolve compatibility or + # download a snapshot (which would also fetch the compiled kernel). + variant = min(v.variant_str for v in variants) + return _source_from_variant(_Source(Path("build") / variant, repo_id=repo_id, revision=revision)) + + +def _source_from_variant(source: "_Source") -> "_Source": + if source.read("__init__.py") is not None: + return source + metadata_path = source.file("metadata.json") + if metadata_path is not None: + name = Metadata.read_from_file(metadata_path).name.python_name + source = _Source(source.root / name, repo_id=source.repo_id, revision=source.revision) + if source.read("__init__.py") is not None: + return source + raise ValueError("No Python kernel package found") + + +def _assignments(body: list[ast.stmt]) -> dict[str, ast.expr]: + values = {} + for node in body: + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name): + values[target.id] = node.value + elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.value is not None: + values[node.target.id] = node.value + elif isinstance(node, ast.AugAssign) and isinstance(node.target, ast.Name): + if node.target.id in values: + values[node.target.id] = ast.BinOp(values[node.target.id], node.op, node.value) + return values + + +def _export_names(tree: ast.Module, filename: str) -> list[str]: + value = _assignments(tree.body).get("__all__") + if value is None: + return [] + + def strings(node: ast.expr) -> list[str]: + if isinstance(node, (ast.List, ast.Tuple)): + names = [ + item.value for item in node.elts if isinstance(item, ast.Constant) and isinstance(item.value, str) + ] + if len(names) == len(node.elts): + return names + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add): + return strings(node.left) + strings(node.right) + raise ValueError(f"Cannot statically read __all__ in {filename}; expected a literal list or tuple of names") + + return list(dict.fromkeys(strings(value))) + + +class _Source: + """Read only API source files and relative re-exports, never import them.""" + + def __init__(self, root: Path, *, repo_id: str | None = None, revision: str | None = None): + self.root = root + self.repo_id = repo_id + self.revision = revision + self._modules: dict[str, ast.Module | None] = {} + + def file(self, filename: str) -> Path | None: + path = self.root / filename + if self.repo_id is None: + return path if path.is_file() else None + try: + return Path( + _get_hf_api().hf_hub_download( + self.repo_id, + repo_type="kernel", + filename=path.as_posix(), + revision=self.revision, + cache_dir=_get_cache_dir(), + local_files_only=constants.HF_HUB_OFFLINE, + ) + ) + except RemoteEntryNotFoundError: + return None + + def read(self, filename: str) -> ast.Module | None: + if filename not in self._modules: + path = self.file(filename) + self._modules[filename] = ast.parse(path.read_bytes(), filename=filename) if path is not None else None + return self._modules[filename] + + def module_file(self, module: str) -> str | None: + for filename in (f"{module}/__init__.py", f"{module}.py"): + if self.read(filename) is not None: + return filename + return None + + def resolve( + self, filename: str, name: str, seen: frozenset[tuple[str, str]] = frozenset() + ) -> tuple[_Definition, str] | None: + if (filename, name) in seen or "_private_for_testing" in PurePosixPath(filename).parts: + return None + seen = seen | {(filename, name)} + tree = self.read(filename) + if tree is None: + return None + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.name == name: + return node, filename + if isinstance(node, ast.ImportFrom) and node.level and node.module: + for alias in node.names: + if (alias.asname or alias.name) != name: + continue + parts = list(PurePosixPath(filename).parent.parts) + if node.level > len(parts) + 1: + return None + parts = parts[: len(parts) - node.level + 1] + node.module.split(".") + if "_private_for_testing" in parts: + return None + target = self.module_file("/".join(parts)) + return self.resolve(target, alias.name, seen) if target else None + value = _assignments(tree.body).get(name) + if isinstance(value, ast.Name): + return self.resolve(filename, value.id, seen) + return None + + def exports(self, filename: str, kinds: tuple[type, ...]) -> list[tuple[str, _Definition, str]]: + tree = self.read(filename) + if tree is None: + raise ValueError(f"Missing Python source: {filename}") + exports = [] + for name in _export_names(tree, filename): + if name.startswith("_") or name == "layers": + continue + resolved = self.resolve(filename, name) + if resolved is None: + print(f"Cannot statically resolve export {name!r} in {filename}", file=sys.stderr) + elif isinstance(resolved[0], kinds): + exports.append((name, *resolved)) + return exports diff --git a/kernels/tests/test_cli_describe.py b/kernels/tests/test_cli_describe.py new file mode 100644 index 00000000..097be1af --- /dev/null +++ b/kernels/tests/test_cli_describe.py @@ -0,0 +1,255 @@ +import sys +from unittest.mock import Mock, call + +import httpx +import pytest +from huggingface_hub.errors import RemoteEntryNotFoundError +from huggingface_hub.hf_api import GitRefInfo + +from kernels.cli import describe as describe_cli +from kernels.cli import main +from kernels.variants import parse_variant + +REPO_ID = "kernels-community/example" +VARIANT = "torch210-cxx11-cu128-x86_64-linux" +INIT = """ +raise RuntimeError("Describing must not execute kernel code") +from missing_dependency import dependency +from . import layers + +def hidden(): pass +def public(): pass +async def async_public(): pass +class NotAFunction: pass + +__all__ = ["public", "layers", "NotAFunction"] +__all__ += ["async_public"] +""" +LAYERS = """ +raise RuntimeError("Describing must not execute layer code") +class Hidden: pass +class Trainable: + has_backward: bool = True + can_torch_compile = False +class Inference: + has_backward = False + can_torch_compile: bool = True +class Unknown: + can_torch_compile = compute_flag() +__all__ = ("Trainable", "Inference", "Unknown") +""" + + +def write_file(root, filename, content): + path = root / filename + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content) + return path + + +def run_cli(monkeypatch, *args): + monkeypatch.setattr(sys, "argv", ["kernels", "describe", *map(str, args)]) + main() + + +@pytest.fixture +def hub(monkeypatch, tmp_path): + versions = {v: GitRefInfo(name=f"v{v}", ref=f"refs/heads/v{v}", target_commit=str(v) * 40) for v in (1, 2)} + available = Mock(return_value=versions) + variants = Mock(return_value=[parse_variant(VARIANT)]) + write_file(tmp_path, f"build/{VARIANT}/__init__.py", INIT) + write_file(tmp_path, f"build/{VARIANT}/layers.py", LAYERS) + + def download(repo_id, *, filename, **kwargs): + path = tmp_path / filename + if not path.is_file(): + raise RemoteEntryNotFoundError( + filename, response=httpx.Response(404, request=httpx.Request("GET", "https://huggingface.co")) + ) + return str(path) + + # A strict API mock makes any snapshot/binary download fail the test. + api = Mock(spec=["hf_hub_download"]) + api.hf_hub_download.side_effect = download + monkeypatch.setattr(describe_cli, "_get_hf_api", lambda: api) + monkeypatch.setattr(describe_cli, "_get_available_versions", available) + monkeypatch.setattr(describe_cli, "get_variants", variants) + return api, available, variants, tmp_path + + +def assert_api_output(output): + assert "Functions:\n public\n async_public\n" in output + assert "Trainable has_backward=True can_torch_compile=False" in output + assert "Inference has_backward=False can_torch_compile=True" in output + assert "Unknown has_backward=unknown can_torch_compile=unknown" in output + for name in ("hidden", "Hidden", "NotAFunction", "dependency"): + assert name not in output + + +@pytest.mark.parametrize("selection", [[], ["--version", "latest"], ["--version", "2"], ["--version", "1"]]) +def test_hub_versions(monkeypatch, hub, capsys, selection): + api, available, variants, _ = hub + revision = "v1" if selection == ["--version", "1"] else "v2" + run_cli(monkeypatch, REPO_ID, *selection) + output = capsys.readouterr() + assert output.err == "" + assert output.out.startswith(f"Repository: {REPO_ID}\nRevision: {revision}\n") + assert_api_output(output.out) + available.assert_called_once_with(REPO_ID, local_files_only=describe_cli.constants.HF_HUB_OFFLINE) + variants.assert_called_once_with(api, repo_id=REPO_ID, revision=revision) + assert api.hf_hub_download.call_args_list == [ + call( + REPO_ID, + repo_type="kernel", + filename=f"build/{VARIANT}/{filename}", + revision=revision, + cache_dir=describe_cli._get_cache_dir(), + local_files_only=describe_cli.constants.HF_HUB_OFFLINE, + ) + for filename in ("__init__.py", "layers/__init__.py", "layers.py") + ] + + +@pytest.mark.parametrize("revision", ["main", "my-tag", "a" * 40]) +def test_hub_revision(monkeypatch, hub, capsys, revision): + run_cli(monkeypatch, REPO_ID, "--revision", revision) + assert f"Revision: {revision}" in capsys.readouterr().out + hub[1].assert_not_called() + hub[2].assert_called_once_with(hub[0], repo_id=REPO_ID, revision=revision) + + +@pytest.mark.parametrize("selection", [["--version", "3"], []]) +def test_missing_version_or_variants(monkeypatch, hub, capsys, selection): + if not selection: + hub[2].return_value = [] + with pytest.raises(SystemExit) as error: + run_cli(monkeypatch, REPO_ID, *selection) + assert error.value.code == 1 + output = capsys.readouterr() + assert output.out == "" + assert ("Version 3 not found" if selection else "No build variants found") in output.err + hub[0].hf_hub_download.assert_not_called() + + +@pytest.mark.parametrize("selection", [["--version", "1", "--revision", "main"], ["--version", "invalid"]]) +def test_invalid_selection(monkeypatch, hub, capsys, selection): + with pytest.raises(SystemExit) as error: + run_cli(monkeypatch, REPO_ID, *selection) + assert error.value.code == 2 + assert "error:" in capsys.readouterr().err + hub[1].assert_not_called() + hub[2].assert_not_called() + + +@pytest.mark.parametrize("layout", ["flat", "nested", "package", "torch-ext", "tvm-ffi-ext"]) +@pytest.mark.parametrize("layers_file", ["layers.py", "layers/__init__.py"]) +def test_local(monkeypatch, tmp_path, capsys, layout, layers_file): + if layout in ("torch-ext", "tvm-ffi-ext"): + write_file( + tmp_path, + "build.toml", + '[general]\nname = "test-kernel"\nversion = 1\nlicense = "MIT"\nbackends = ["cpu"]\n[torch]\n', + ) + package = tmp_path / layout / "test_kernel" + elif layout == "package": + package = tmp_path + else: + package = tmp_path / "build" / VARIANT + if layout == "nested": + write_file( + package, + "metadata.json", + '{"name": "test-kernel", "id": "test", "version": 1, ' + '"license": "MIT", "python-depends": [], "backend": {"type": "cuda"}}', + ) + package /= "test_kernel" + write_file(package, "__init__.py", INIT) + write_file(package, layers_file, LAYERS) + api = Mock(side_effect=AssertionError("Local inspection must not contact the Hub")) + monkeypatch.setattr(describe_cli, "_get_hf_api", api) + run_cli(monkeypatch, tmp_path) + output = capsys.readouterr() + assert output.err == "" + assert output.out.startswith(f"Path: {tmp_path}\n") + assert_api_output(output.out) + api.assert_not_called() + + +@pytest.mark.parametrize("selection", [["--version", "1"], ["--version", "latest"], ["--revision", "main"]]) +def test_local_rejects_selection(monkeypatch, tmp_path, capsys, selection): + with pytest.raises(SystemExit) as error: + run_cli(monkeypatch, tmp_path, *selection) + assert error.value.code == 1 + assert "cannot be used with a local path" in capsys.readouterr().err + + +def test_no_exports(monkeypatch, tmp_path, capsys): + write_file(tmp_path, "__init__.py", "def hidden(): pass") + write_file(tmp_path, "layers.py", "class Hidden: pass") + run_cli(monkeypatch, tmp_path) + output = capsys.readouterr() + assert output.err == "" + assert "No functions declared in __all__." in output.out + assert "No layers declared in __all__." in output.out + assert "Hidden" not in output.out + + +def test_missing_layers(monkeypatch, tmp_path, capsys): + write_file(tmp_path, "__init__.py", INIT) + run_cli(monkeypatch, tmp_path) + output = capsys.readouterr() + assert " public\n" in output.out + assert "No layers declared in __all__." in output.out + assert output.err == "" + + +def test_relative_reexports_and_private_symbols(monkeypatch, tmp_path, capsys): + write_file( + tmp_path, + "__init__.py", + "from .functions import impl as public\nfrom ._private_for_testing import secret\n" + 'alias = public\n__all__ = ["public", "alias", "secret"]', + ) + write_file(tmp_path, "functions.py", "def impl(): pass\ndef hidden(): pass") + write_file(tmp_path, "_private_for_testing.py", "def secret(): pass") + write_file(tmp_path, "layers/__init__.py", 'from .impl import Layer as PublicLayer\n__all__ = ["PublicLayer"]') + write_file(tmp_path, "layers/impl.py", "from ..implementation import Layer") + write_file(tmp_path, "implementation.py", "class Layer:\n has_backward = False\n can_torch_compile = True") + run_cli(monkeypatch, tmp_path) + output = capsys.readouterr() + assert "Functions:\n public\n alias\n" in output.out + assert "PublicLayer has_backward=False can_torch_compile=True" in output.out + assert "secret" not in output.out + + +def test_download_errors_are_not_hidden(monkeypatch, hub, capsys): + hub[0].hf_hub_download.side_effect = OSError("download failed") + with pytest.raises(SystemExit) as error: + run_cli(monkeypatch, REPO_ID) + assert error.value.code == 1 + assert "download failed" in capsys.readouterr().err + + +def test_hub_fetches_only_reexport_source(monkeypatch, hub, capsys): + api, _, _, root = hub + variant = root / "build" / VARIANT + write_file( + variant, + "__init__.py", + "from .functions import impl as public\nfrom .functions import impl as alias\n" + 'from .unused import hidden\n__all__ = ["public", "alias"]', + ) + write_file(variant, "functions.py", "def impl(): pass") + run_cli(monkeypatch, REPO_ID) + output = capsys.readouterr() + assert "Functions:\n public\n alias\n" in output.out + assert output.err == "" + filenames = [entry.kwargs["filename"] for entry in api.hf_hub_download.call_args_list] + assert filenames.count(f"build/{VARIANT}/functions.py") == 1 + assert all("unused" not in filename for filename in filenames) + + +def test_direct_call_rejects_conflicting_selection(): + with pytest.raises(SystemExit) as error: + describe_cli.print_kernel_description(REPO_ID, version=1, revision="main") + assert error.value.code == 1 From 317ed946733f512130ad22258bff26f00b8a65d8 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 6 Oct 2026 13:05:18 +0530 Subject: [PATCH 02/10] add machine-readability support. --- docs/source/cli-describe.md | 19 +++++++++++- kernels/src/kernels/cli/__init__.py | 3 +- kernels/src/kernels/cli/describe.py | 48 ++++++++++++++++++++--------- kernels/tests/test_cli_describe.py | 25 +++++++++++++++ 4 files changed, 79 insertions(+), 16 deletions(-) diff --git a/docs/source/cli-describe.md b/docs/source/cli-describe.md index 80f29c28..7ae23977 100644 --- a/docs/source/cli-describe.md +++ b/docs/source/cli-describe.md @@ -16,8 +16,9 @@ Options: - `--version `: select an integer kernel version or `latest` (the default). - `--revision `: select a branch, tag, or commit instead of a version. +- `--json`: print the description as JSON for machine-readable output. -These options are mutually exclusive and apply only to Hub repositories. The +`--version` and `--revision` are mutually exclusive and apply only to Hub repositories. The latest numbered version is selected by default, or `main` for an unversioned repository. @@ -30,6 +31,7 @@ kernels describe kernels-community/activation kernels describe kernels-community/activation --version 1 kernels describe kernels-community/activation --revision main kernels describe ./my-kernel +kernels describe kernels-community/activation --json ``` ## Export conventions @@ -56,6 +58,21 @@ Layers: Gelu has_backward=True can_torch_compile=True ``` +With `--json`, the same description is: + +```json +{ + "repo_id": "example/activation", + "revision": "v1", + "functions": ["gelu"], + "layers": [ + {"name": "Gelu", "has_backward": true, "can_torch_compile": true} + ] +} +``` + +Local kernels include `path` instead of `repo_id` and `revision`. + Use [kernels info](cli-info.md) for metadata such as the license, dependencies, and supported backends, and [kernels variants](cli-variants.md) for build variants and compatibility information. diff --git a/kernels/src/kernels/cli/__init__.py b/kernels/src/kernels/cli/__init__.py index 75dd8cad..b22bd43f 100644 --- a/kernels/src/kernels/cli/__init__.py +++ b/kernels/src/kernels/cli/__init__.py @@ -40,6 +40,7 @@ def main(): "--version", type=_variant_version, help="The kernel version or 'latest' (default: latest version)" ) describe_selection.add_argument("--revision", help="The kernel revision (branch, tag, or commit)") + describe_parser.add_argument("--json", action="store_true", help="Print the kernel description as JSON") describe_parser.set_defaults(func=kernel_describe) info_parser = subparsers.add_parser("info", help="Describe a kernel") @@ -173,7 +174,7 @@ def kernel_info(args): def kernel_describe(args): - print_kernel_description(args.kernel, revision=args.revision, version=args.version) + print_kernel_description(args.kernel, revision=args.revision, version=args.version, json_output=args.json) def kernel_versions(args): diff --git a/kernels/src/kernels/cli/describe.py b/kernels/src/kernels/cli/describe.py index 765f6f16..ec5c15c0 100644 --- a/kernels/src/kernels/cli/describe.py +++ b/kernels/src/kernels/cli/describe.py @@ -1,9 +1,10 @@ """Describe kernel exports by reading Python source, without executing it.""" import ast +import json import sys from pathlib import Path, PurePosixPath -from typing import Literal +from typing import Any, Literal from huggingface_hub import constants from huggingface_hub.errors import HfHubHTTPError, RemoteEntryNotFoundError @@ -22,6 +23,7 @@ def print_kernel_description( *, version: int | Literal["latest"] | None = None, revision: str | None = None, + json_output: bool = False, ): """Print the functions and layers explicitly exported by a kernel.""" try: @@ -33,7 +35,7 @@ def print_kernel_description( if version is not None or revision is not None: raise ValueError("revision and version cannot be used with a local path") source = _local_source(path) - heading = f"Path: {path}" + info: dict[str, Any] = {"path": str(path)} else: if revision is None: versions = _get_available_versions(kernel, local_files_only=constants.HF_HUB_OFFLINE) @@ -46,7 +48,7 @@ def print_kernel_description( else: revision = versions[max(versions)].name if versions else "main" source = _hub_source(kernel, revision) - heading = f"Repository: {kernel}\nRevision: {revision}" + info = {"repo_id": kernel, "revision": revision} functions = source.exports("__init__.py", (ast.FunctionDef, ast.AsyncFunctionDef)) layers_file = source.module_file("layers") @@ -55,25 +57,43 @@ def print_kernel_description( print(f"Cannot describe {kernel}: {error}", file=sys.stderr) sys.exit(1) - print(heading) + info["functions"] = [name for name, _, _ in functions] + info["layers"] = [] + for name, node, _ in layers: + flags = _assignments(node.body) + layer: dict[str, Any] = {"name": name} + for flag in ("has_backward", "can_torch_compile"): + value = flags.get(flag) + layer[flag] = value.value if isinstance(value, ast.Constant) and isinstance(value.value, bool) else None + info["layers"].append(layer) + + if json_output: + print(json.dumps(info, indent=2)) + else: + _print_human(info) + + +def _print_human(info: dict): + if "repo_id" in info: + print(f"Repository: {info['repo_id']}") + print(f"Revision: {info['revision']}") + else: + print(f"Path: {info['path']}") + print("\nFunctions:") - for name, _, _ in functions: + for name in info["functions"]: print(f" {name}") - if not functions: + if not info["functions"]: print(" No functions declared in __all__.") print("\nLayers:") - for name, node, _ in layers: - flags = _assignments(node.body) + for layer in info["layers"]: capabilities = [] for flag in ("has_backward", "can_torch_compile"): - value = flags.get(flag) - rendered = ( - str(value.value) if isinstance(value, ast.Constant) and isinstance(value.value, bool) else "unknown" - ) + rendered = "unknown" if layer[flag] is None else str(layer[flag]) capabilities.append(f"{flag}={rendered}") - print(f" {name} {' '.join(capabilities)}") - if not layers: + print(f" {layer['name']} {' '.join(capabilities)}") + if not info["layers"]: print(" No layers declared in __all__.") diff --git a/kernels/tests/test_cli_describe.py b/kernels/tests/test_cli_describe.py index 097be1af..d3010d4d 100644 --- a/kernels/tests/test_cli_describe.py +++ b/kernels/tests/test_cli_describe.py @@ -1,3 +1,4 @@ +import json import sys from unittest.mock import Mock, call @@ -118,6 +119,30 @@ def test_hub_revision(monkeypatch, hub, capsys, revision): hub[2].assert_called_once_with(hub[0], repo_id=REPO_ID, revision=revision) +@pytest.mark.parametrize("local", [False, True]) +def test_json(monkeypatch, hub, capsys, local): + path = hub[3] / "build" / VARIANT + run_cli(monkeypatch, path if local else REPO_ID, "--json") + output = capsys.readouterr() + assert output.err == "" + expected = {"path": str(path)} if local else {"repo_id": REPO_ID, "revision": "v2"} + assert json.loads(output.out) == { + **expected, + "functions": ["public", "async_public"], + "layers": [ + {"name": "Trainable", "has_backward": True, "can_torch_compile": False}, + {"name": "Inference", "has_backward": False, "can_torch_compile": True}, + {"name": "Unknown", "has_backward": None, "can_torch_compile": None}, + ], + } + + +def test_json_no_exports(monkeypatch, tmp_path, capsys): + write_file(tmp_path, "__init__.py", "__all__ = []") + run_cli(monkeypatch, tmp_path, "--json") + assert json.loads(capsys.readouterr().out) == {"path": str(tmp_path), "functions": [], "layers": []} + + @pytest.mark.parametrize("selection", [["--version", "3"], []]) def test_missing_version_or_variants(monkeypatch, hub, capsys, selection): if not selection: From 0e9cdd5bd1285d91612ae47985db21e09665e412 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 7 Oct 2026 14:53:55 +0530 Subject: [PATCH 03/10] rename api to describe-api. --- docs/source/_toctree.yml | 4 ++-- .../{cli-describe.md => cli-describe-api.md} | 16 +++++++-------- docs/source/cli-info.md | 2 +- docs/source/cli.md | 2 +- kernels/src/kernels/cli/__init__.py | 20 ++++++++++--------- .../cli/{describe.py => describe_api.py} | 0 ...i_describe.py => test_cli_describe_api.py} | 20 +++++++++---------- 7 files changed, 33 insertions(+), 31 deletions(-) rename docs/source/{cli-describe.md => cli-describe-api.md} (82%) rename kernels/src/kernels/cli/{describe.py => describe_api.py} (100%) rename kernels/tests/{test_cli_describe.py => test_cli_describe_api.py} (93%) diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 45094fb4..50785888 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -35,8 +35,8 @@ title: Overview - local: cli-info title: kernels info - - local: cli-describe - title: kernels describe + - local: cli-describe-api + title: kernels describe-api - local: cli-variants title: kernels variants - local: cli-versions diff --git a/docs/source/cli-describe.md b/docs/source/cli-describe-api.md similarity index 82% rename from docs/source/cli-describe.md rename to docs/source/cli-describe-api.md index 7ae23977..68cd894d 100644 --- a/docs/source/cli-describe.md +++ b/docs/source/cli-describe-api.md @@ -1,6 +1,6 @@ -# kernels describe +# kernels describe-api -Use `kernels describe` to list a kernel's public functions and layers without +Use `kernels describe-api` to list a kernel's public functions and layers without loading the kernel or installing its dependencies. For Hub kernels, the command downloads only the Python source needed to inspect the API, plus `metadata.json` when needed to locate an older build's package. It does not download compiled @@ -9,7 +9,7 @@ libraries or a repository snapshot. ## Usage ```bash -kernels describe +kernels describe-api ``` Options: @@ -27,11 +27,11 @@ An existing local directory is treated as a kernel. ## Examples ```bash -kernels describe kernels-community/activation -kernels describe kernels-community/activation --version 1 -kernels describe kernels-community/activation --revision main -kernels describe ./my-kernel -kernels describe kernels-community/activation --json +kernels describe-api kernels-community/activation +kernels describe-api kernels-community/activation --version 1 +kernels describe-api kernels-community/activation --revision main +kernels describe-api ./my-kernel +kernels describe-api kernels-community/activation --json ``` ## Export conventions diff --git a/docs/source/cli-info.md b/docs/source/cli-info.md index 14284fe1..f359eaf2 100644 --- a/docs/source/cli-info.md +++ b/docs/source/cli-info.md @@ -5,7 +5,7 @@ It prints the kernel metadata (name, version, license, upstream/source repositories, Python dependencies, and supported backends). To list the build variants of a kernel and their compatibility with the current system, use [kernels variants](cli-variants.md). -To list its public functions and layers, use [kernels describe](cli-describe.md). +To list its public functions and layers, use [kernels describe-api](cli-describe-api.md). ## Usage diff --git a/docs/source/cli.md b/docs/source/cli.md index 00b41b66..e4b3d3b8 100644 --- a/docs/source/cli.md +++ b/docs/source/cli.md @@ -7,7 +7,7 @@ The `kernels` CLI provides commands for managing compute kernels. | Command | Description | | ------------------------------------------- | ---------------------------------- | | [benchmark](cli-benchmark.md) | Run benchmark results for a kernel | -| [describe](cli-describe.md) | List public functions and layers | +| [describe-api](cli-describe-api.md) | List public functions and layers | | [download](cli-download.md) | Download a kernel | | [info](cli-info.md) | Describe a kernel | | [lock](cli-lock.md) | Lock kernel revisions | diff --git a/kernels/src/kernels/cli/__init__.py b/kernels/src/kernels/cli/__init__.py index 8b75922b..524627ad 100644 --- a/kernels/src/kernels/cli/__init__.py +++ b/kernels/src/kernels/cli/__init__.py @@ -3,7 +3,7 @@ from pathlib import Path from typing import Literal -from kernels.cli.describe import print_kernel_description +from kernels.cli.describe_api import print_kernel_description from kernels.cli.download import download_kernels from kernels.cli.info import print_kernel_info from kernels.cli.lock import lock_kernels @@ -33,15 +33,17 @@ def main(): ) download_parser.set_defaults(func=download_kernels) - describe_parser = subparsers.add_parser("describe", help="Show the public functions and layers of a kernel") - describe_parser.add_argument("kernel", help="The kernel repo ID or a local path to a kernel") - describe_selection = describe_parser.add_mutually_exclusive_group() - describe_selection.add_argument( + describe_api_parser = subparsers.add_parser( + "describe-api", help="Show the public functions and layers of a kernel" + ) + describe_api_parser.add_argument("kernel", help="The kernel repo ID or a local path to a kernel") + describe_api_selection = describe_api_parser.add_mutually_exclusive_group() + describe_api_selection.add_argument( "--version", type=_variant_version, help="The kernel version or 'latest' (default: latest version)" ) - describe_selection.add_argument("--revision", help="The kernel revision (branch, tag, or commit)") - describe_parser.add_argument("--json", action="store_true", help="Print the kernel description as JSON") - describe_parser.set_defaults(func=kernel_describe) + describe_api_selection.add_argument("--revision", help="The kernel revision (branch, tag, or commit)") + describe_api_parser.add_argument("--json", action="store_true", help="Print the kernel description as JSON") + describe_api_parser.set_defaults(func=kernel_describe_api) info_parser = subparsers.add_parser("info", help="Describe a kernel") info_parser.add_argument( @@ -176,7 +178,7 @@ def kernel_info(args): ) -def kernel_describe(args): +def kernel_describe_api(args): print_kernel_description(args.kernel, revision=args.revision, version=args.version, json_output=args.json) diff --git a/kernels/src/kernels/cli/describe.py b/kernels/src/kernels/cli/describe_api.py similarity index 100% rename from kernels/src/kernels/cli/describe.py rename to kernels/src/kernels/cli/describe_api.py diff --git a/kernels/tests/test_cli_describe.py b/kernels/tests/test_cli_describe_api.py similarity index 93% rename from kernels/tests/test_cli_describe.py rename to kernels/tests/test_cli_describe_api.py index d3010d4d..1d9c1ca5 100644 --- a/kernels/tests/test_cli_describe.py +++ b/kernels/tests/test_cli_describe_api.py @@ -7,7 +7,7 @@ from huggingface_hub.errors import RemoteEntryNotFoundError from huggingface_hub.hf_api import GitRefInfo -from kernels.cli import describe as describe_cli +from kernels.cli import describe_api as describe_api_cli from kernels.cli import main from kernels.variants import parse_variant @@ -49,7 +49,7 @@ def write_file(root, filename, content): def run_cli(monkeypatch, *args): - monkeypatch.setattr(sys, "argv", ["kernels", "describe", *map(str, args)]) + monkeypatch.setattr(sys, "argv", ["kernels", "describe-api", *map(str, args)]) main() @@ -72,9 +72,9 @@ def download(repo_id, *, filename, **kwargs): # A strict API mock makes any snapshot/binary download fail the test. api = Mock(spec=["hf_hub_download"]) api.hf_hub_download.side_effect = download - monkeypatch.setattr(describe_cli, "_get_hf_api", lambda: api) - monkeypatch.setattr(describe_cli, "_get_available_versions", available) - monkeypatch.setattr(describe_cli, "get_variants", variants) + monkeypatch.setattr(describe_api_cli, "_get_hf_api", lambda: api) + monkeypatch.setattr(describe_api_cli, "_get_available_versions", available) + monkeypatch.setattr(describe_api_cli, "get_variants", variants) return api, available, variants, tmp_path @@ -96,7 +96,7 @@ def test_hub_versions(monkeypatch, hub, capsys, selection): assert output.err == "" assert output.out.startswith(f"Repository: {REPO_ID}\nRevision: {revision}\n") assert_api_output(output.out) - available.assert_called_once_with(REPO_ID, local_files_only=describe_cli.constants.HF_HUB_OFFLINE) + available.assert_called_once_with(REPO_ID, local_files_only=describe_api_cli.constants.HF_HUB_OFFLINE) variants.assert_called_once_with(api, repo_id=REPO_ID, revision=revision) assert api.hf_hub_download.call_args_list == [ call( @@ -104,8 +104,8 @@ def test_hub_versions(monkeypatch, hub, capsys, selection): repo_type="kernel", filename=f"build/{VARIANT}/{filename}", revision=revision, - cache_dir=describe_cli._get_cache_dir(), - local_files_only=describe_cli.constants.HF_HUB_OFFLINE, + cache_dir=describe_api_cli._get_cache_dir(), + local_files_only=describe_api_cli.constants.HF_HUB_OFFLINE, ) for filename in ("__init__.py", "layers/__init__.py", "layers.py") ] @@ -191,7 +191,7 @@ def test_local(monkeypatch, tmp_path, capsys, layout, layers_file): write_file(package, "__init__.py", INIT) write_file(package, layers_file, LAYERS) api = Mock(side_effect=AssertionError("Local inspection must not contact the Hub")) - monkeypatch.setattr(describe_cli, "_get_hf_api", api) + monkeypatch.setattr(describe_api_cli, "_get_hf_api", api) run_cli(monkeypatch, tmp_path) output = capsys.readouterr() assert output.err == "" @@ -276,5 +276,5 @@ def test_hub_fetches_only_reexport_source(monkeypatch, hub, capsys): def test_direct_call_rejects_conflicting_selection(): with pytest.raises(SystemExit) as error: - describe_cli.print_kernel_description(REPO_ID, version=1, revision="main") + describe_api_cli.print_kernel_description(REPO_ID, version=1, revision="main") assert error.value.code == 1 From 843e4e94aa30e70cbf22fb4148e9c26395fb8ee8 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 8 Oct 2026 07:11:08 +0000 Subject: [PATCH 04/10] kernel-builder: first stab at symbol generation python script. --- docs/source/_toctree.yml | 2 - docs/source/cli-describe-api.md | 78 ---- docs/source/cli-info.md | 1 - docs/source/cli.md | 2 +- kernels/src/kernels/cli/__init__.py | 17 - kernels/src/kernels/cli/describe_api.py | 260 ------------- kernels/tests/test_cli_describe_api.py | 280 -------------- nix-builder/scripts/generate_symbols.py | 181 +++++++++ nix-builder/tests/test_generate_symbols.py | 421 +++++++++++++++++++++ 9 files changed, 603 insertions(+), 639 deletions(-) delete mode 100644 docs/source/cli-describe-api.md delete mode 100644 kernels/src/kernels/cli/describe_api.py delete mode 100644 kernels/tests/test_cli_describe_api.py create mode 100644 nix-builder/scripts/generate_symbols.py create mode 100644 nix-builder/tests/test_generate_symbols.py diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 50785888..60a88fa6 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -35,8 +35,6 @@ title: Overview - local: cli-info title: kernels info - - local: cli-describe-api - title: kernels describe-api - local: cli-variants title: kernels variants - local: cli-versions diff --git a/docs/source/cli-describe-api.md b/docs/source/cli-describe-api.md deleted file mode 100644 index 68cd894d..00000000 --- a/docs/source/cli-describe-api.md +++ /dev/null @@ -1,78 +0,0 @@ -# kernels describe-api - -Use `kernels describe-api` to list a kernel's public functions and layers without -loading the kernel or installing its dependencies. For Hub kernels, the command -downloads only the Python source needed to inspect the API, plus `metadata.json` -when needed to locate an older build's package. It does not download compiled -libraries or a repository snapshot. - -## Usage - -```bash -kernels describe-api -``` - -Options: - -- `--version `: select an integer kernel version or `latest` (the default). -- `--revision `: select a branch, tag, or commit instead of a version. -- `--json`: print the description as JSON for machine-readable output. - -`--version` and `--revision` are mutually exclusive and apply only to Hub repositories. The -latest numbered version is selected by default, or `main` for an unversioned -repository. - -An existing local directory is treated as a kernel. - -## Examples - -```bash -kernels describe-api kernels-community/activation -kernels describe-api kernels-community/activation --version 1 -kernels describe-api kernels-community/activation --revision main -kernels describe-api ./my-kernel -kernels describe-api kernels-community/activation --json -``` - -## Export conventions - -Functions must be listed in the top-level module's `__all__`. Layers must be -listed in `layers.py` or `layers/__init__.py`'s own `__all__`. A module without -`__all__` contributes no exports. Private names and symbols from -`_private_for_testing` are purposefully excluded. - -Each layer includes its explicitly declared `has_backward` and -`can_torch_compile` boolean attributes. - -For example, a kernel declaring `gelu` and `Gelu` in the respective `__all__` -lists could produce: - -```text -Repository: example/activation -Revision: v1 - -Functions: - gelu - -Layers: - Gelu has_backward=True can_torch_compile=True -``` - -With `--json`, the same description is: - -```json -{ - "repo_id": "example/activation", - "revision": "v1", - "functions": ["gelu"], - "layers": [ - {"name": "Gelu", "has_backward": true, "can_torch_compile": true} - ] -} -``` - -Local kernels include `path` instead of `repo_id` and `revision`. - -Use [kernels info](cli-info.md) for metadata such as the license, dependencies, -and supported backends, and [kernels variants](cli-variants.md) for build variants -and compatibility information. diff --git a/docs/source/cli-info.md b/docs/source/cli-info.md index f359eaf2..1de4acd0 100644 --- a/docs/source/cli-info.md +++ b/docs/source/cli-info.md @@ -5,7 +5,6 @@ It prints the kernel metadata (name, version, license, upstream/source repositories, Python dependencies, and supported backends). To list the build variants of a kernel and their compatibility with the current system, use [kernels variants](cli-variants.md). -To list its public functions and layers, use [kernels describe-api](cli-describe-api.md). ## Usage diff --git a/docs/source/cli.md b/docs/source/cli.md index e4b3d3b8..3563e07b 100644 --- a/docs/source/cli.md +++ b/docs/source/cli.md @@ -7,9 +7,9 @@ The `kernels` CLI provides commands for managing compute kernels. | Command | Description | | ------------------------------------------- | ---------------------------------- | | [benchmark](cli-benchmark.md) | Run benchmark results for a kernel | -| [describe-api](cli-describe-api.md) | List public functions and layers | | [download](cli-download.md) | Download a kernel | | [info](cli-info.md) | Describe a kernel | | [lock](cli-lock.md) | Lock kernel revisions | | [verify-signature](cli-verify-signature.md) | Verify a kernel signature | | [variants](cli-variants.md) | Show variants and compatibility | + diff --git a/kernels/src/kernels/cli/__init__.py b/kernels/src/kernels/cli/__init__.py index 524627ad..aafcc088 100644 --- a/kernels/src/kernels/cli/__init__.py +++ b/kernels/src/kernels/cli/__init__.py @@ -3,7 +3,6 @@ from pathlib import Path from typing import Literal -from kernels.cli.describe_api import print_kernel_description from kernels.cli.download import download_kernels from kernels.cli.info import print_kernel_info from kernels.cli.lock import lock_kernels @@ -33,18 +32,6 @@ def main(): ) download_parser.set_defaults(func=download_kernels) - describe_api_parser = subparsers.add_parser( - "describe-api", help="Show the public functions and layers of a kernel" - ) - describe_api_parser.add_argument("kernel", help="The kernel repo ID or a local path to a kernel") - describe_api_selection = describe_api_parser.add_mutually_exclusive_group() - describe_api_selection.add_argument( - "--version", type=_variant_version, help="The kernel version or 'latest' (default: latest version)" - ) - describe_api_selection.add_argument("--revision", help="The kernel revision (branch, tag, or commit)") - describe_api_parser.add_argument("--json", action="store_true", help="Print the kernel description as JSON") - describe_api_parser.set_defaults(func=kernel_describe_api) - info_parser = subparsers.add_parser("info", help="Describe a kernel") info_parser.add_argument( "kernel", @@ -178,10 +165,6 @@ def kernel_info(args): ) -def kernel_describe_api(args): - print_kernel_description(args.kernel, revision=args.revision, version=args.version, json_output=args.json) - - def kernel_versions(args): print_kernel_versions(args.repo_id) diff --git a/kernels/src/kernels/cli/describe_api.py b/kernels/src/kernels/cli/describe_api.py deleted file mode 100644 index ec5c15c0..00000000 --- a/kernels/src/kernels/cli/describe_api.py +++ /dev/null @@ -1,260 +0,0 @@ -"""Describe kernel exports by reading Python source, without executing it.""" - -import ast -import json -import sys -from pathlib import Path, PurePosixPath -from typing import Any, Literal - -from huggingface_hub import constants -from huggingface_hub.errors import HfHubHTTPError, RemoteEntryNotFoundError - -from kernels._rust import KernelName, Metadata -from kernels._versions import _get_available_versions -from kernels.compat import tomllib -from kernels.hf_hub import _get_cache_dir, _get_hf_api -from kernels.variants import get_variants, get_variants_local - -_Definition = ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef - - -def print_kernel_description( - kernel: str, - *, - version: int | Literal["latest"] | None = None, - revision: str | None = None, - json_output: bool = False, -): - """Print the functions and layers explicitly exported by a kernel.""" - try: - if version is not None and revision is not None: - raise ValueError("Only one of revision or version can be specified") - - path = Path(kernel) - if path.is_dir(): - if version is not None or revision is not None: - raise ValueError("revision and version cannot be used with a local path") - source = _local_source(path) - info: dict[str, Any] = {"path": str(path)} - else: - if revision is None: - versions = _get_available_versions(kernel, local_files_only=constants.HF_HUB_OFFLINE) - if isinstance(version, int): - if version not in versions: - raise ValueError( - f"Version {version} not found, available versions: {', '.join(map(str, sorted(versions)))}" - ) - revision = versions[version].name - else: - revision = versions[max(versions)].name if versions else "main" - source = _hub_source(kernel, revision) - info = {"repo_id": kernel, "revision": revision} - - functions = source.exports("__init__.py", (ast.FunctionDef, ast.AsyncFunctionDef)) - layers_file = source.module_file("layers") - layers = source.exports(layers_file, (ast.ClassDef,)) if layers_file else [] - except (ValueError, OSError, SyntaxError, HfHubHTTPError) as error: - print(f"Cannot describe {kernel}: {error}", file=sys.stderr) - sys.exit(1) - - info["functions"] = [name for name, _, _ in functions] - info["layers"] = [] - for name, node, _ in layers: - flags = _assignments(node.body) - layer: dict[str, Any] = {"name": name} - for flag in ("has_backward", "can_torch_compile"): - value = flags.get(flag) - layer[flag] = value.value if isinstance(value, ast.Constant) and isinstance(value.value, bool) else None - info["layers"].append(layer) - - if json_output: - print(json.dumps(info, indent=2)) - else: - _print_human(info) - - -def _print_human(info: dict): - if "repo_id" in info: - print(f"Repository: {info['repo_id']}") - print(f"Revision: {info['revision']}") - else: - print(f"Path: {info['path']}") - - print("\nFunctions:") - for name in info["functions"]: - print(f" {name}") - if not info["functions"]: - print(" No functions declared in __all__.") - - print("\nLayers:") - for layer in info["layers"]: - capabilities = [] - for flag in ("has_backward", "can_torch_compile"): - rendered = "unknown" if layer[flag] is None else str(layer[flag]) - capabilities.append(f"{flag}={rendered}") - print(f" {layer['name']} {' '.join(capabilities)}") - if not info["layers"]: - print(" No layers declared in __all__.") - - -def _local_source(path: Path) -> "_Source": - # Accept a package/variant directory as well as a built repository. - if (path / "__init__.py").is_file(): - return _Source(path) - variants = get_variants_local(path / "build") - if variants: - path = path / "build" / min(v.variant_str for v in variants) - elif (path / "build.toml").is_file(): - with (path / "build.toml").open("rb") as config_file: - config = tomllib.load(config_file) - name = config.get("general", {}).get("name") - if not isinstance(name, str): - raise ValueError("Missing general.name in build.toml") - name = KernelName(name).python_name - for directory in ("torch-ext", "tvm-ffi-ext"): - package = path / directory / name - if (package / "__init__.py").is_file(): - return _Source(package) - raise ValueError(f"No Python kernel package found in {path}") - return _source_from_variant(_Source(path)) - - -def _hub_source(repo_id: str, revision: str) -> "_Source": - api = _get_hf_api() - variants = get_variants(api, repo_id=repo_id, revision=revision) - if not variants: - raise ValueError(f"No build variants found (revision: {revision})") - # API source can be inspected on any host; do not resolve compatibility or - # download a snapshot (which would also fetch the compiled kernel). - variant = min(v.variant_str for v in variants) - return _source_from_variant(_Source(Path("build") / variant, repo_id=repo_id, revision=revision)) - - -def _source_from_variant(source: "_Source") -> "_Source": - if source.read("__init__.py") is not None: - return source - metadata_path = source.file("metadata.json") - if metadata_path is not None: - name = Metadata.read_from_file(metadata_path).name.python_name - source = _Source(source.root / name, repo_id=source.repo_id, revision=source.revision) - if source.read("__init__.py") is not None: - return source - raise ValueError("No Python kernel package found") - - -def _assignments(body: list[ast.stmt]) -> dict[str, ast.expr]: - values = {} - for node in body: - if isinstance(node, ast.Assign): - for target in node.targets: - if isinstance(target, ast.Name): - values[target.id] = node.value - elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.value is not None: - values[node.target.id] = node.value - elif isinstance(node, ast.AugAssign) and isinstance(node.target, ast.Name): - if node.target.id in values: - values[node.target.id] = ast.BinOp(values[node.target.id], node.op, node.value) - return values - - -def _export_names(tree: ast.Module, filename: str) -> list[str]: - value = _assignments(tree.body).get("__all__") - if value is None: - return [] - - def strings(node: ast.expr) -> list[str]: - if isinstance(node, (ast.List, ast.Tuple)): - names = [ - item.value for item in node.elts if isinstance(item, ast.Constant) and isinstance(item.value, str) - ] - if len(names) == len(node.elts): - return names - if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add): - return strings(node.left) + strings(node.right) - raise ValueError(f"Cannot statically read __all__ in {filename}; expected a literal list or tuple of names") - - return list(dict.fromkeys(strings(value))) - - -class _Source: - """Read only API source files and relative re-exports, never import them.""" - - def __init__(self, root: Path, *, repo_id: str | None = None, revision: str | None = None): - self.root = root - self.repo_id = repo_id - self.revision = revision - self._modules: dict[str, ast.Module | None] = {} - - def file(self, filename: str) -> Path | None: - path = self.root / filename - if self.repo_id is None: - return path if path.is_file() else None - try: - return Path( - _get_hf_api().hf_hub_download( - self.repo_id, - repo_type="kernel", - filename=path.as_posix(), - revision=self.revision, - cache_dir=_get_cache_dir(), - local_files_only=constants.HF_HUB_OFFLINE, - ) - ) - except RemoteEntryNotFoundError: - return None - - def read(self, filename: str) -> ast.Module | None: - if filename not in self._modules: - path = self.file(filename) - self._modules[filename] = ast.parse(path.read_bytes(), filename=filename) if path is not None else None - return self._modules[filename] - - def module_file(self, module: str) -> str | None: - for filename in (f"{module}/__init__.py", f"{module}.py"): - if self.read(filename) is not None: - return filename - return None - - def resolve( - self, filename: str, name: str, seen: frozenset[tuple[str, str]] = frozenset() - ) -> tuple[_Definition, str] | None: - if (filename, name) in seen or "_private_for_testing" in PurePosixPath(filename).parts: - return None - seen = seen | {(filename, name)} - tree = self.read(filename) - if tree is None: - return None - for node in tree.body: - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.name == name: - return node, filename - if isinstance(node, ast.ImportFrom) and node.level and node.module: - for alias in node.names: - if (alias.asname or alias.name) != name: - continue - parts = list(PurePosixPath(filename).parent.parts) - if node.level > len(parts) + 1: - return None - parts = parts[: len(parts) - node.level + 1] + node.module.split(".") - if "_private_for_testing" in parts: - return None - target = self.module_file("/".join(parts)) - return self.resolve(target, alias.name, seen) if target else None - value = _assignments(tree.body).get(name) - if isinstance(value, ast.Name): - return self.resolve(filename, value.id, seen) - return None - - def exports(self, filename: str, kinds: tuple[type, ...]) -> list[tuple[str, _Definition, str]]: - tree = self.read(filename) - if tree is None: - raise ValueError(f"Missing Python source: {filename}") - exports = [] - for name in _export_names(tree, filename): - if name.startswith("_") or name == "layers": - continue - resolved = self.resolve(filename, name) - if resolved is None: - print(f"Cannot statically resolve export {name!r} in {filename}", file=sys.stderr) - elif isinstance(resolved[0], kinds): - exports.append((name, *resolved)) - return exports diff --git a/kernels/tests/test_cli_describe_api.py b/kernels/tests/test_cli_describe_api.py deleted file mode 100644 index 1d9c1ca5..00000000 --- a/kernels/tests/test_cli_describe_api.py +++ /dev/null @@ -1,280 +0,0 @@ -import json -import sys -from unittest.mock import Mock, call - -import httpx -import pytest -from huggingface_hub.errors import RemoteEntryNotFoundError -from huggingface_hub.hf_api import GitRefInfo - -from kernels.cli import describe_api as describe_api_cli -from kernels.cli import main -from kernels.variants import parse_variant - -REPO_ID = "kernels-community/example" -VARIANT = "torch210-cxx11-cu128-x86_64-linux" -INIT = """ -raise RuntimeError("Describing must not execute kernel code") -from missing_dependency import dependency -from . import layers - -def hidden(): pass -def public(): pass -async def async_public(): pass -class NotAFunction: pass - -__all__ = ["public", "layers", "NotAFunction"] -__all__ += ["async_public"] -""" -LAYERS = """ -raise RuntimeError("Describing must not execute layer code") -class Hidden: pass -class Trainable: - has_backward: bool = True - can_torch_compile = False -class Inference: - has_backward = False - can_torch_compile: bool = True -class Unknown: - can_torch_compile = compute_flag() -__all__ = ("Trainable", "Inference", "Unknown") -""" - - -def write_file(root, filename, content): - path = root / filename - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(content) - return path - - -def run_cli(monkeypatch, *args): - monkeypatch.setattr(sys, "argv", ["kernels", "describe-api", *map(str, args)]) - main() - - -@pytest.fixture -def hub(monkeypatch, tmp_path): - versions = {v: GitRefInfo(name=f"v{v}", ref=f"refs/heads/v{v}", target_commit=str(v) * 40) for v in (1, 2)} - available = Mock(return_value=versions) - variants = Mock(return_value=[parse_variant(VARIANT)]) - write_file(tmp_path, f"build/{VARIANT}/__init__.py", INIT) - write_file(tmp_path, f"build/{VARIANT}/layers.py", LAYERS) - - def download(repo_id, *, filename, **kwargs): - path = tmp_path / filename - if not path.is_file(): - raise RemoteEntryNotFoundError( - filename, response=httpx.Response(404, request=httpx.Request("GET", "https://huggingface.co")) - ) - return str(path) - - # A strict API mock makes any snapshot/binary download fail the test. - api = Mock(spec=["hf_hub_download"]) - api.hf_hub_download.side_effect = download - monkeypatch.setattr(describe_api_cli, "_get_hf_api", lambda: api) - monkeypatch.setattr(describe_api_cli, "_get_available_versions", available) - monkeypatch.setattr(describe_api_cli, "get_variants", variants) - return api, available, variants, tmp_path - - -def assert_api_output(output): - assert "Functions:\n public\n async_public\n" in output - assert "Trainable has_backward=True can_torch_compile=False" in output - assert "Inference has_backward=False can_torch_compile=True" in output - assert "Unknown has_backward=unknown can_torch_compile=unknown" in output - for name in ("hidden", "Hidden", "NotAFunction", "dependency"): - assert name not in output - - -@pytest.mark.parametrize("selection", [[], ["--version", "latest"], ["--version", "2"], ["--version", "1"]]) -def test_hub_versions(monkeypatch, hub, capsys, selection): - api, available, variants, _ = hub - revision = "v1" if selection == ["--version", "1"] else "v2" - run_cli(monkeypatch, REPO_ID, *selection) - output = capsys.readouterr() - assert output.err == "" - assert output.out.startswith(f"Repository: {REPO_ID}\nRevision: {revision}\n") - assert_api_output(output.out) - available.assert_called_once_with(REPO_ID, local_files_only=describe_api_cli.constants.HF_HUB_OFFLINE) - variants.assert_called_once_with(api, repo_id=REPO_ID, revision=revision) - assert api.hf_hub_download.call_args_list == [ - call( - REPO_ID, - repo_type="kernel", - filename=f"build/{VARIANT}/{filename}", - revision=revision, - cache_dir=describe_api_cli._get_cache_dir(), - local_files_only=describe_api_cli.constants.HF_HUB_OFFLINE, - ) - for filename in ("__init__.py", "layers/__init__.py", "layers.py") - ] - - -@pytest.mark.parametrize("revision", ["main", "my-tag", "a" * 40]) -def test_hub_revision(monkeypatch, hub, capsys, revision): - run_cli(monkeypatch, REPO_ID, "--revision", revision) - assert f"Revision: {revision}" in capsys.readouterr().out - hub[1].assert_not_called() - hub[2].assert_called_once_with(hub[0], repo_id=REPO_ID, revision=revision) - - -@pytest.mark.parametrize("local", [False, True]) -def test_json(monkeypatch, hub, capsys, local): - path = hub[3] / "build" / VARIANT - run_cli(monkeypatch, path if local else REPO_ID, "--json") - output = capsys.readouterr() - assert output.err == "" - expected = {"path": str(path)} if local else {"repo_id": REPO_ID, "revision": "v2"} - assert json.loads(output.out) == { - **expected, - "functions": ["public", "async_public"], - "layers": [ - {"name": "Trainable", "has_backward": True, "can_torch_compile": False}, - {"name": "Inference", "has_backward": False, "can_torch_compile": True}, - {"name": "Unknown", "has_backward": None, "can_torch_compile": None}, - ], - } - - -def test_json_no_exports(monkeypatch, tmp_path, capsys): - write_file(tmp_path, "__init__.py", "__all__ = []") - run_cli(monkeypatch, tmp_path, "--json") - assert json.loads(capsys.readouterr().out) == {"path": str(tmp_path), "functions": [], "layers": []} - - -@pytest.mark.parametrize("selection", [["--version", "3"], []]) -def test_missing_version_or_variants(monkeypatch, hub, capsys, selection): - if not selection: - hub[2].return_value = [] - with pytest.raises(SystemExit) as error: - run_cli(monkeypatch, REPO_ID, *selection) - assert error.value.code == 1 - output = capsys.readouterr() - assert output.out == "" - assert ("Version 3 not found" if selection else "No build variants found") in output.err - hub[0].hf_hub_download.assert_not_called() - - -@pytest.mark.parametrize("selection", [["--version", "1", "--revision", "main"], ["--version", "invalid"]]) -def test_invalid_selection(monkeypatch, hub, capsys, selection): - with pytest.raises(SystemExit) as error: - run_cli(monkeypatch, REPO_ID, *selection) - assert error.value.code == 2 - assert "error:" in capsys.readouterr().err - hub[1].assert_not_called() - hub[2].assert_not_called() - - -@pytest.mark.parametrize("layout", ["flat", "nested", "package", "torch-ext", "tvm-ffi-ext"]) -@pytest.mark.parametrize("layers_file", ["layers.py", "layers/__init__.py"]) -def test_local(monkeypatch, tmp_path, capsys, layout, layers_file): - if layout in ("torch-ext", "tvm-ffi-ext"): - write_file( - tmp_path, - "build.toml", - '[general]\nname = "test-kernel"\nversion = 1\nlicense = "MIT"\nbackends = ["cpu"]\n[torch]\n', - ) - package = tmp_path / layout / "test_kernel" - elif layout == "package": - package = tmp_path - else: - package = tmp_path / "build" / VARIANT - if layout == "nested": - write_file( - package, - "metadata.json", - '{"name": "test-kernel", "id": "test", "version": 1, ' - '"license": "MIT", "python-depends": [], "backend": {"type": "cuda"}}', - ) - package /= "test_kernel" - write_file(package, "__init__.py", INIT) - write_file(package, layers_file, LAYERS) - api = Mock(side_effect=AssertionError("Local inspection must not contact the Hub")) - monkeypatch.setattr(describe_api_cli, "_get_hf_api", api) - run_cli(monkeypatch, tmp_path) - output = capsys.readouterr() - assert output.err == "" - assert output.out.startswith(f"Path: {tmp_path}\n") - assert_api_output(output.out) - api.assert_not_called() - - -@pytest.mark.parametrize("selection", [["--version", "1"], ["--version", "latest"], ["--revision", "main"]]) -def test_local_rejects_selection(monkeypatch, tmp_path, capsys, selection): - with pytest.raises(SystemExit) as error: - run_cli(monkeypatch, tmp_path, *selection) - assert error.value.code == 1 - assert "cannot be used with a local path" in capsys.readouterr().err - - -def test_no_exports(monkeypatch, tmp_path, capsys): - write_file(tmp_path, "__init__.py", "def hidden(): pass") - write_file(tmp_path, "layers.py", "class Hidden: pass") - run_cli(monkeypatch, tmp_path) - output = capsys.readouterr() - assert output.err == "" - assert "No functions declared in __all__." in output.out - assert "No layers declared in __all__." in output.out - assert "Hidden" not in output.out - - -def test_missing_layers(monkeypatch, tmp_path, capsys): - write_file(tmp_path, "__init__.py", INIT) - run_cli(monkeypatch, tmp_path) - output = capsys.readouterr() - assert " public\n" in output.out - assert "No layers declared in __all__." in output.out - assert output.err == "" - - -def test_relative_reexports_and_private_symbols(monkeypatch, tmp_path, capsys): - write_file( - tmp_path, - "__init__.py", - "from .functions import impl as public\nfrom ._private_for_testing import secret\n" - 'alias = public\n__all__ = ["public", "alias", "secret"]', - ) - write_file(tmp_path, "functions.py", "def impl(): pass\ndef hidden(): pass") - write_file(tmp_path, "_private_for_testing.py", "def secret(): pass") - write_file(tmp_path, "layers/__init__.py", 'from .impl import Layer as PublicLayer\n__all__ = ["PublicLayer"]') - write_file(tmp_path, "layers/impl.py", "from ..implementation import Layer") - write_file(tmp_path, "implementation.py", "class Layer:\n has_backward = False\n can_torch_compile = True") - run_cli(monkeypatch, tmp_path) - output = capsys.readouterr() - assert "Functions:\n public\n alias\n" in output.out - assert "PublicLayer has_backward=False can_torch_compile=True" in output.out - assert "secret" not in output.out - - -def test_download_errors_are_not_hidden(monkeypatch, hub, capsys): - hub[0].hf_hub_download.side_effect = OSError("download failed") - with pytest.raises(SystemExit) as error: - run_cli(monkeypatch, REPO_ID) - assert error.value.code == 1 - assert "download failed" in capsys.readouterr().err - - -def test_hub_fetches_only_reexport_source(monkeypatch, hub, capsys): - api, _, _, root = hub - variant = root / "build" / VARIANT - write_file( - variant, - "__init__.py", - "from .functions import impl as public\nfrom .functions import impl as alias\n" - 'from .unused import hidden\n__all__ = ["public", "alias"]', - ) - write_file(variant, "functions.py", "def impl(): pass") - run_cli(monkeypatch, REPO_ID) - output = capsys.readouterr() - assert "Functions:\n public\n alias\n" in output.out - assert output.err == "" - filenames = [entry.kwargs["filename"] for entry in api.hf_hub_download.call_args_list] - assert filenames.count(f"build/{VARIANT}/functions.py") == 1 - assert all("unused" not in filename for filename in filenames) - - -def test_direct_call_rejects_conflicting_selection(): - with pytest.raises(SystemExit) as error: - describe_api_cli.print_kernel_description(REPO_ID, version=1, revision="main") - assert error.value.code == 1 diff --git a/nix-builder/scripts/generate_symbols.py b/nix-builder/scripts/generate_symbols.py new file mode 100644 index 00000000..3d779721 --- /dev/null +++ b/nix-builder/scripts/generate_symbols.py @@ -0,0 +1,181 @@ +"""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 available. Importing a kernel executes its initialization code. +The inspection itself does not call exported functions or instantiate classes. +""" + +import argparse +import importlib +import importlib.util +import inspect +import json +import sys +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__. + + As in describe-api, layers need not appear in the package's __all__. + Only Python functions (including async functions) and layer classes are + included; other exported modules, data, and callable objects are ignored. + """ + functions = describe_exports(module, inspect.isfunction) + layers = getattr(module, "layers", None) + if layers is None and hasattr(module, "__path__"): + # The package need not import its optional layers submodule itself. + layers_name = f"{module.__name__}.layers" + try: + layers = importlib.import_module(layers_name) + except ModuleNotFoundError as error: + if error.name != layers_name: + raise + return { + "schema_version": 1, + "module": module.__name__, + "functions": functions, + "layers": describe_exports(layers, inspect.isclass) + if inspect.ismodule(layers) + else [], + } + + +def load_package(module_name: str, package_dir: Path) -> ModuleType: + """Import a package from a build variant, preserving relative imports.""" + init_path = package_dir.resolve() / "__init__.py" + if not init_path.is_file(): + raise FileNotFoundError(f"No kernel package found at {init_path}") + if module_name in sys.modules: + raise ValueError(f"Module {module_name!r} is already loaded") + spec = importlib.util.spec_from_file_location(module_name, init_path) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load {module_name!r} from {init_path}") + module = importlib.util.module_from_spec(spec) + # Registration is needed by relative imports and decorators such as dataclass. + sys.modules[module_name] = module + try: + spec.loader.exec_module(module) + except BaseException: + # A failed import can leave submodules behind as well as the package. + for name in list(sys.modules): + if name == module_name or name.startswith(f"{module_name}."): + del sys.modules[name] + raise + return module + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("module", help="Python module name to import") + parser.add_argument( + "--package-dir", + type=Path, + help=( + "Built package directory containing __init__.py " + "(otherwise use Python's import path)" + ), + ) + parser.add_argument( + "--output", type=Path, required=True, help="Destination symbols.json" + ) + args = parser.parse_args() + + module = ( + load_package(args.module, args.package_dir) + if args.package_dir is not None + else importlib.import_module(args.module) + ) + # 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/tests/test_generate_symbols.py b/nix-builder/tests/test_generate_symbols.py new file mode 100644 index 00000000..aa7db365 --- /dev/null +++ b/nix-builder/tests/test_generate_symbols.py @@ -0,0 +1,421 @@ +import functools +import importlib.util +import json +import os +import subprocess +import sys +import textwrap +from pathlib import Path +from types import ModuleType + +import pytest + +SCRIPT = Path(__file__).parents[1] / "scripts" / "generate_symbols.py" +SPEC = importlib.util.spec_from_file_location("generate_symbols", SCRIPT) +generator = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(generator) + + +def make_module(name="example", **exports): + module = ModuleType(name) + module.__dict__.update(exports) + module.__all__ = list(exports) + return module + + +def test_runtime_exports_aliases_and_layers(): + def implementation(x: int, /, *, scale: float = 1.0) -> int: + """Scale the input.""" + raise AssertionError("Must not run the kernel") + + class Base: + has_backward = True + + class Layer(Base): + can_torch_compile = False + + def __init__(self, size: int = 4): + raise AssertionError("Must not instantiate layers") + + layers = make_module("example.layers", Layer=Layer) + module = make_module(layers=layers, alias=implementation) + module.unexported = implementation + module._private_for_testing = make_module(hidden=implementation) + # __all__ need not be statically resolvable, and duplicates are harmless. + module.__all__ = ["alias"] + [name for name in ("layers", "alias")] + + result = generator.generate_symbols(module) + + assert result["schema_version"] == 1 + (alias,) = result["functions"] + assert alias["name"] == "alias" + assert alias["qualname"].endswith("implementation") + assert alias["kind"] == "function" + assert alias["doc"] == "Scale the input." + layer = result["layers"][0] + assert layer["name"] == "Layer" + assert layer["kind"] == "class" + assert layer["attributes"] == {"has_backward": True, "can_torch_compile": False} + assert layer["signature"]["parameters"] == [ + { + "name": "size", + "kind": "POSITIONAL_OR_KEYWORD", + "default": "4", + "annotation": "int", + } + ] + assert json.loads(json.dumps(result)) == result + + +def test_signature_preserves_parameter_kinds_and_empty_values(): + def function(a, /, b: int = None, *args: float, c="é", **kwargs) -> None: + pass + + assert generator.describe_signature(function) == { + "parameters": [ + { + "name": "a", + "kind": "POSITIONAL_ONLY", + "default": None, + "annotation": None, + }, + { + "name": "b", + "kind": "POSITIONAL_OR_KEYWORD", + "default": "None", + "annotation": "int", + }, + { + "name": "args", + "kind": "VAR_POSITIONAL", + "default": None, + "annotation": "float", + }, + {"name": "c", "kind": "KEYWORD_ONLY", "default": "'é'", "annotation": None}, + { + "name": "kwargs", + "kind": "VAR_KEYWORD", + "default": None, + "annotation": None, + }, + ], + "return_annotation": "None", + } + + +def test_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_only_functions_and_layer_classes_are_included(): + def original(x: int, y: int = 2): + pass + + @functools.wraps(original) + def wrapper(*args, **kwargs): + raise AssertionError("Must not call wrappers") + + async def async_function(): + raise AssertionError("Must not call async functions") + + class Layer: + def __call__(self): + raise AssertionError("Must not call exported objects") + + def method(self): + pass + + layers = make_module("example.layers", Layer=Layer, function=original) + layers.unexported = type("HiddenLayer", (), {}) + other_module = make_module("example.other", function=original, Layer=Layer) + module = make_module( + wrapped=wrapper, + async_function=async_function, + partial=functools.partial(original, y=3), + callable=Layer(), + method=Layer().method, + builtin=len, + data=42, + TopLevelClass=Layer, + other=other_module, + layers=layers, + ) + result = generator.generate_symbols(module) + assert [item["name"] for item in result["functions"]] == [ + "wrapped", + "async_function", + ] + assert result["functions"][0]["signature"] == generator.describe_signature(original) + assert [item["name"] for item in result["layers"]] == ["Layer"] + + +def test_unavailable_signature_is_null(): + class Layer: + __signature__ = "not an inspect.Signature" + + result = generator.generate_symbols(make_module(layers=make_module(Layer=Layer))) + assert result["layers"][0]["signature"] is None + + +def test_non_boolean_layer_attributes_are_unknown(): + class Layer: + has_backward = "True" + can_torch_compile = 1 + + symbol = generator.describe_symbol("Layer", Layer) + assert symbol["attributes"] == {"has_backward": None, "can_torch_compile": None} + + +@pytest.mark.parametrize("exports", [None, "function", [1], {"function"}]) +def test_invalid_all_fails(exports): + module = make_module() + module.__all__ = exports + with pytest.raises(ValueError, match=r"example\.__all__ must be"): + generator.generate_symbols(module) + + +def test_missing_export_fails_with_context(): + module = make_module() + module.__all__ = ["missing"] + with pytest.raises( + ValueError, match=r"example\.__all__ names a missing symbol: 'missing'" + ): + generator.generate_symbols(module) + + +def test_no_all_exports_nothing(): + module = ModuleType("example") + module.public = lambda: None + module.layers = ModuleType("example.layers") + module.layers.Layer = type("Layer", (), {}) + result = generator.generate_symbols(module) + assert result["functions"] == [] + assert result["layers"] == [] + + +def test_private_exports_and_testing_reexports_are_excluded(): + def public(): + pass + + def secret(): + pass + + class PrivateLayer: + pass + + secret.__module__ = "example._private_for_testing.helpers" + PrivateLayer.__module__ = "example._private_for_testing" + module = make_module( + _hidden=public, + secret=secret, + public=public, + layers=make_module(_Hidden=type("Hidden", (), {}), Alias=PrivateLayer), + ) + result = generator.generate_symbols(module) + assert [item["name"] for item in result["functions"]] == ["public"] + assert result["layers"] == [] + + +def test_lazy_exports_not_in_dir(): + def public(): + pass + + module = make_module() + module.__all__ = ["lazy"] + module.__getattr__ = lambda name: public if name == "lazy" else None + assert generator.generate_symbols(module)["functions"][0]["name"] == "lazy" + + +def test_layers_need_not_be_exported_by_package(): + module = make_module() + module.layers = make_module(Layer=type("Layer", (), {})) + result = generator.generate_symbols(module) + assert result["functions"] == [] + assert [item["name"] for item in result["layers"]] == ["Layer"] + + +def test_missing_layers_is_supported(): + module = make_module(public=lambda: None) + result = generator.generate_symbols(module) + assert [item["name"] for item in result["functions"]] == ["public"] + assert result["layers"] == [] + + +def test_export_order_and_aliases_are_preserved(): + def implementation(): + pass + + module = make_module(z=implementation, a=implementation) + module.__all__ = ("z", "a", "z") + assert [ + item["name"] for item in generator.generate_symbols(module)["functions"] + ] == ["z", "a"] + + +@pytest.fixture +def built_package(tmp_path): + package = tmp_path / "variant" + package.mkdir() + (package / "__init__.py").write_text( + textwrap.dedent("""\ + from . import layers + from .implementation import function as alias + print("Initialization output") + __all__ = ["alias", "layers"] + """), + encoding="utf-8", + ) + (package / "implementation.py").write_text( + 'def function(x: "UndefinedTensor", *, scale=2):\n """Scale é."""\n', + encoding="utf-8", + ) + (package / "layers.py").write_text( + textwrap.dedent("""\ + from dataclasses import dataclass + @dataclass + class Layer: + size: int = 4 + has_backward = True + can_torch_compile = True + __all__ = ["Layer"] + """), + encoding="utf-8", + ) + return package + + +def test_cli_loads_built_package_and_writes_json(built_package): + output = built_package / "symbols.json" + command = [ + sys.executable, + str(SCRIPT), + "example_build", + "--package-dir", + str(built_package), + "--output", + str(output), + ] + result = subprocess.run(command, capture_output=True, text=True, check=True) + content = output.read_text(encoding="utf-8") + document = json.loads(content) + assert "Initialization output" in result.stdout + assert document["module"] == "example_build" + (alias,) = document["functions"] + assert alias["module"] == "example_build.implementation" + assert alias["doc"] == "Scale é." + assert document["layers"][0]["attributes"]["has_backward"] is True + subprocess.run(command, capture_output=True, text=True, check=True) + assert output.read_text(encoding="utf-8") == content + + +def test_cli_imports_from_python_path(built_package, tmp_path): + output = tmp_path / "symbols.json" + subprocess.run( + [sys.executable, str(SCRIPT), "variant", "--output", str(output)], + env={**os.environ, "PYTHONPATH": str(tmp_path)}, + capture_output=True, + text=True, + check=True, + ) + assert json.loads(output.read_text())["module"] == "variant" + + +@pytest.mark.parametrize("has_layers", [False, True]) +def test_cli_discovers_optional_layers(built_package, has_layers): + # Match the PR: layers need not be imported or exported by __init__.py. + (built_package / "__init__.py").write_text( + "from .implementation import function\n__all__ = ['function']\n", + encoding="utf-8", + ) + if not has_layers: + (built_package / "layers.py").unlink() + output = built_package / "symbols.json" + subprocess.run( + [ + sys.executable, + str(SCRIPT), + "example_build", + "--package-dir", + str(built_package), + "--output", + str(output), + ], + capture_output=True, + text=True, + check=True, + ) + result = json.loads(output.read_text()) + assert [item["name"] for item in result["functions"]] == ["function"] + assert [item["name"] for item in result["layers"]] == ( + ["Layer"] if has_layers else [] + ) + + +def test_optional_layers_import_errors_are_not_hidden(built_package): + (built_package / "__init__.py").write_text("", encoding="utf-8") + (built_package / "layers.py").write_text( + "import missing_layer_dependency_for_test\n", encoding="utf-8" + ) + output = built_package / "symbols.json" + result = subprocess.run( + [ + sys.executable, + str(SCRIPT), + "example_build", + "--package-dir", + str(built_package), + "--output", + str(output), + ], + capture_output=True, + text=True, + ) + assert result.returncode != 0 + assert "missing_layer_dependency_for_test" in result.stderr + assert not output.exists() + + +@pytest.mark.parametrize( + "source", ["raise RuntimeError('broken import')", "__all__ = ['missing']"] +) +def test_cli_failure_preserves_existing_output(built_package, source): + output = built_package / "symbols.json" + output.write_text("previous output", encoding="utf-8") + (built_package / "__init__.py").write_text(source, encoding="utf-8") + result = subprocess.run( + [ + sys.executable, + str(SCRIPT), + "example_build", + "--package-dir", + str(built_package), + "--output", + str(output), + ], + capture_output=True, + text=True, + ) + assert result.returncode != 0 + assert output.read_text(encoding="utf-8") == "previous output" + + +def test_missing_package_is_an_error(tmp_path): + with pytest.raises(FileNotFoundError, match="No kernel package found"): + generator.load_package("missing_package", tmp_path) + + +def test_failed_import_cleans_up_submodules(built_package): + (built_package / "__init__.py").write_text( + "from . import layers\nraise RuntimeError('broken import')\n", encoding="utf-8" + ) + with pytest.raises(RuntimeError, match="broken import"): + generator.load_package("failed_build", built_package) + assert not any( + name == "failed_build" or name.startswith("failed_build.") + for name in sys.modules + ) From 55c573de7799f1dfef9b725e365315277e2c7095 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 8 Oct 2026 08:41:04 +0000 Subject: [PATCH 05/10] kernel-builder: remove unneeded tests. --- nix-builder/tests/test_generate_symbols.py | 157 +-------------------- 1 file changed, 4 insertions(+), 153 deletions(-) diff --git a/nix-builder/tests/test_generate_symbols.py b/nix-builder/tests/test_generate_symbols.py index aa7db365..95037a74 100644 --- a/nix-builder/tests/test_generate_symbols.py +++ b/nix-builder/tests/test_generate_symbols.py @@ -1,10 +1,7 @@ import functools import importlib.util import json -import os -import subprocess import sys -import textwrap from pathlib import Path from types import ModuleType @@ -257,164 +254,18 @@ def implementation(): ] == ["z", "a"] -@pytest.fixture -def built_package(tmp_path): - package = tmp_path / "variant" - package.mkdir() - (package / "__init__.py").write_text( - textwrap.dedent("""\ - from . import layers - from .implementation import function as alias - print("Initialization output") - __all__ = ["alias", "layers"] - """), - encoding="utf-8", - ) - (package / "implementation.py").write_text( - 'def function(x: "UndefinedTensor", *, scale=2):\n """Scale é."""\n', - encoding="utf-8", - ) - (package / "layers.py").write_text( - textwrap.dedent("""\ - from dataclasses import dataclass - @dataclass - class Layer: - size: int = 4 - has_backward = True - can_torch_compile = True - __all__ = ["Layer"] - """), - encoding="utf-8", - ) - return package - - -def test_cli_loads_built_package_and_writes_json(built_package): - output = built_package / "symbols.json" - command = [ - sys.executable, - str(SCRIPT), - "example_build", - "--package-dir", - str(built_package), - "--output", - str(output), - ] - result = subprocess.run(command, capture_output=True, text=True, check=True) - content = output.read_text(encoding="utf-8") - document = json.loads(content) - assert "Initialization output" in result.stdout - assert document["module"] == "example_build" - (alias,) = document["functions"] - assert alias["module"] == "example_build.implementation" - assert alias["doc"] == "Scale é." - assert document["layers"][0]["attributes"]["has_backward"] is True - subprocess.run(command, capture_output=True, text=True, check=True) - assert output.read_text(encoding="utf-8") == content - - -def test_cli_imports_from_python_path(built_package, tmp_path): - output = tmp_path / "symbols.json" - subprocess.run( - [sys.executable, str(SCRIPT), "variant", "--output", str(output)], - env={**os.environ, "PYTHONPATH": str(tmp_path)}, - capture_output=True, - text=True, - check=True, - ) - assert json.loads(output.read_text())["module"] == "variant" - - -@pytest.mark.parametrize("has_layers", [False, True]) -def test_cli_discovers_optional_layers(built_package, has_layers): - # Match the PR: layers need not be imported or exported by __init__.py. - (built_package / "__init__.py").write_text( - "from .implementation import function\n__all__ = ['function']\n", - encoding="utf-8", - ) - if not has_layers: - (built_package / "layers.py").unlink() - output = built_package / "symbols.json" - subprocess.run( - [ - sys.executable, - str(SCRIPT), - "example_build", - "--package-dir", - str(built_package), - "--output", - str(output), - ], - capture_output=True, - text=True, - check=True, - ) - result = json.loads(output.read_text()) - assert [item["name"] for item in result["functions"]] == ["function"] - assert [item["name"] for item in result["layers"]] == ( - ["Layer"] if has_layers else [] - ) - - -def test_optional_layers_import_errors_are_not_hidden(built_package): - (built_package / "__init__.py").write_text("", encoding="utf-8") - (built_package / "layers.py").write_text( - "import missing_layer_dependency_for_test\n", encoding="utf-8" - ) - output = built_package / "symbols.json" - result = subprocess.run( - [ - sys.executable, - str(SCRIPT), - "example_build", - "--package-dir", - str(built_package), - "--output", - str(output), - ], - capture_output=True, - text=True, - ) - assert result.returncode != 0 - assert "missing_layer_dependency_for_test" in result.stderr - assert not output.exists() - - -@pytest.mark.parametrize( - "source", ["raise RuntimeError('broken import')", "__all__ = ['missing']"] -) -def test_cli_failure_preserves_existing_output(built_package, source): - output = built_package / "symbols.json" - output.write_text("previous output", encoding="utf-8") - (built_package / "__init__.py").write_text(source, encoding="utf-8") - result = subprocess.run( - [ - sys.executable, - str(SCRIPT), - "example_build", - "--package-dir", - str(built_package), - "--output", - str(output), - ], - capture_output=True, - text=True, - ) - assert result.returncode != 0 - assert output.read_text(encoding="utf-8") == "previous output" - - def test_missing_package_is_an_error(tmp_path): with pytest.raises(FileNotFoundError, match="No kernel package found"): generator.load_package("missing_package", tmp_path) -def test_failed_import_cleans_up_submodules(built_package): - (built_package / "__init__.py").write_text( +def test_failed_import_cleans_up_submodules(tmp_path): + (tmp_path / "layers.py").write_text("", encoding="utf-8") + (tmp_path / "__init__.py").write_text( "from . import layers\nraise RuntimeError('broken import')\n", encoding="utf-8" ) with pytest.raises(RuntimeError, match="broken import"): - generator.load_package("failed_build", built_package) + generator.load_package("failed_build", tmp_path) assert not any( name == "failed_build" or name.startswith("failed_build.") for name in sys.modules From ed17a222c69201402a8becb2596996ed287e803f Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 8 Oct 2026 10:18:58 +0000 Subject: [PATCH 06/10] simplify tests. --- nix-builder/scripts/generate_symbols.py | 48 +-- nix-builder/tests/test_generate_symbols.py | 373 +++++++++------------ 2 files changed, 176 insertions(+), 245 deletions(-) diff --git a/nix-builder/scripts/generate_symbols.py b/nix-builder/scripts/generate_symbols.py index 3d779721..edbd52a2 100644 --- a/nix-builder/scripts/generate_symbols.py +++ b/nix-builder/scripts/generate_symbols.py @@ -1,16 +1,15 @@ """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 available. Importing a kernel executes its initialization code. +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 importlib -import importlib.util import inspect import json -import sys from collections.abc import Callable from pathlib import Path from types import ModuleType @@ -127,51 +126,24 @@ def generate_symbols(module: ModuleType) -> dict[str, Any]: } -def load_package(module_name: str, package_dir: Path) -> ModuleType: - """Import a package from a build variant, preserving relative imports.""" - init_path = package_dir.resolve() / "__init__.py" - if not init_path.is_file(): - raise FileNotFoundError(f"No kernel package found at {init_path}") - if module_name in sys.modules: - raise ValueError(f"Module {module_name!r} is already loaded") - spec = importlib.util.spec_from_file_location(module_name, init_path) - if spec is None or spec.loader is None: - raise ImportError(f"Cannot load {module_name!r} from {init_path}") - module = importlib.util.module_from_spec(spec) - # Registration is needed by relative imports and decorators such as dataclass. - sys.modules[module_name] = module - try: - spec.loader.exec_module(module) - except BaseException: - # A failed import can leave submodules behind as well as the package. - for name in list(sys.modules): - if name == module_name or name.startswith(f"{module_name}."): - del sys.modules[name] - raise - return module - - def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("module", help="Python module name to import") parser.add_argument( - "--package-dir", + "kernel_dir", type=Path, - help=( - "Built package directory containing __init__.py " - "(otherwise use Python's import 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() - module = ( - load_package(args.module, args.package_dir) - if args.package_dir is not None - else importlib.import_module(args.module) - ) + 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") diff --git a/nix-builder/tests/test_generate_symbols.py b/nix-builder/tests/test_generate_symbols.py index 95037a74..4b369dd3 100644 --- a/nix-builder/tests/test_generate_symbols.py +++ b/nix-builder/tests/test_generate_symbols.py @@ -1,272 +1,231 @@ -import functools import importlib.util import json -import sys from pathlib import Path from types import ModuleType import pytest +# The generator is a standalone script, so load it by its file path. SCRIPT = Path(__file__).parents[1] / "scripts" / "generate_symbols.py" SPEC = importlib.util.spec_from_file_location("generate_symbols", SCRIPT) generator = importlib.util.module_from_spec(SPEC) SPEC.loader.exec_module(generator) -def make_module(name="example", **exports): - module = ModuleType(name) - module.__dict__.update(exports) - module.__all__ = list(exports) - return module - +@pytest.fixture +def kernel(): + """A small kernel exporting one function and one layer.""" -def test_runtime_exports_aliases_and_layers(): - def implementation(x: int, /, *, scale: float = 1.0) -> int: - """Scale the input.""" - raise AssertionError("Must not run the kernel") + def relu(x: int) -> int: + """Apply ReLU.""" + raise AssertionError("Inspection must not call the function") - class Base: + class ReLU: has_backward = True - - class Layer(Base): can_torch_compile = False def __init__(self, size: int = 4): - raise AssertionError("Must not instantiate layers") + 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 - layers = make_module("example.layers", Layer=Layer) - module = make_module(layers=layers, alias=implementation) - module.unexported = implementation - module._private_for_testing = make_module(hidden=implementation) - # __all__ need not be statically resolvable, and duplicates are harmless. - module.__all__ = ["alias"] + [name for name in ("layers", "alias")] - result = generator.generate_symbols(module) +def test_public_function_is_included(kernel): + result = generator.generate_symbols(kernel) - assert result["schema_version"] == 1 - (alias,) = result["functions"] - assert alias["name"] == "alias" - assert alias["qualname"].endswith("implementation") - assert alias["kind"] == "function" - assert alias["doc"] == "Scale the input." + 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"] == "Layer" + assert layer["name"] == "ReLU" assert layer["kind"] == "class" - assert layer["attributes"] == {"has_backward": True, "can_torch_compile": False} - assert layer["signature"]["parameters"] == [ - { - "name": "size", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "4", - "annotation": "int", - } - ] - assert json.loads(json.dumps(result)) == result + assert layer["attributes"] == { + "has_backward": True, + "can_torch_compile": False, + } -def test_signature_preserves_parameter_kinds_and_empty_values(): - def function(a, /, b: int = None, *args: float, c="é", **kwargs) -> None: - pass +def test_symbols_not_in_all_are_excluded(kernel): + kernel.helper = kernel.relu + kernel.layers.OtherLayer = kernel.layers.ReLU - assert generator.describe_signature(function) == { - "parameters": [ - { - "name": "a", - "kind": "POSITIONAL_ONLY", - "default": None, - "annotation": None, - }, - { - "name": "b", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "int", - }, - { - "name": "args", - "kind": "VAR_POSITIONAL", - "default": None, - "annotation": "float", - }, - {"name": "c", "kind": "KEYWORD_ONLY", "default": "'é'", "annotation": None}, - { - "name": "kwargs", - "kind": "VAR_KEYWORD", - "default": None, - "annotation": None, - }, - ], - "return_annotation": "None", - } + 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_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_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) -def test_only_functions_and_layer_classes_are_included(): - def original(x: int, y: int = 2): - pass + assert [function["name"] for function in result["functions"]] == ["relu"] + assert [layer["name"] for layer in result["layers"]] == ["ReLU"] - @functools.wraps(original) - def wrapper(*args, **kwargs): - raise AssertionError("Must not call wrappers") - - async def async_function(): - raise AssertionError("Must not call async functions") - - class Layer: - def __call__(self): - raise AssertionError("Must not call exported objects") - - def method(self): - pass - - layers = make_module("example.layers", Layer=Layer, function=original) - layers.unexported = type("HiddenLayer", (), {}) - other_module = make_module("example.other", function=original, Layer=Layer) - module = make_module( - wrapped=wrapper, - async_function=async_function, - partial=functools.partial(original, y=3), - callable=Layer(), - method=Layer().method, - builtin=len, - data=42, - TopLevelClass=Layer, - other=other_module, - layers=layers, - ) - result = generator.generate_symbols(module) - assert [item["name"] for item in result["functions"]] == [ - "wrapped", - "async_function", - ] - assert result["functions"][0]["signature"] == generator.describe_signature(original) - assert [item["name"] for item in result["layers"]] == ["Layer"] +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" -def test_unavailable_signature_is_null(): - class Layer: - __signature__ = "not an inspect.Signature" + result = generator.generate_symbols(kernel) - result = generator.generate_symbols(make_module(layers=make_module(Layer=Layer))) - assert result["layers"][0]["signature"] is None + assert result["functions"] == [] + assert [layer["name"] for layer in result["layers"]] == ["ReLU"] -def test_non_boolean_layer_attributes_are_unknown(): - class Layer: - has_backward = "True" - can_torch_compile = 1 +def test_layer_from_private_testing_module_is_excluded(kernel): + kernel.layers.ReLU.__module__ = "example._private_for_testing" - symbol = generator.describe_symbol("Layer", Layer) - assert symbol["attributes"] == {"has_backward": None, "can_torch_compile": None} + result = generator.generate_symbols(kernel) + assert result["layers"] == [] + assert [function["name"] for function in result["functions"]] == ["relu"] -@pytest.mark.parametrize("exports", [None, "function", [1], {"function"}]) -def test_invalid_all_fails(exports): - module = make_module() - module.__all__ = exports - with pytest.raises(ValueError, match=r"example\.__all__ must be"): - generator.generate_symbols(module) +def test_functions_in_layers_are_not_layers(kernel): + kernel.layers.helper = kernel.relu + kernel.layers.__all__.append("helper") -def test_missing_export_fails_with_context(): - module = make_module() - module.__all__ = ["missing"] - with pytest.raises( - ValueError, match=r"example\.__all__ names a missing symbol: 'missing'" - ): - generator.generate_symbols(module) + result = generator.generate_symbols(kernel) + assert [layer["name"] for layer in result["layers"]] == ["ReLU"] -def test_no_all_exports_nothing(): - module = ModuleType("example") - module.public = lambda: None - module.layers = ModuleType("example.layers") - module.layers.Layer = type("Layer", (), {}) - result = generator.generate_symbols(module) - assert result["functions"] == [] + +def test_layers_are_optional(kernel): + del kernel.layers + kernel.__all__ = ["relu"] + + result = generator.generate_symbols(kernel) + + assert [function["name"] for function in result["functions"]] == ["relu"] assert result["layers"] == [] -def test_private_exports_and_testing_reexports_are_excluded(): - def public(): - pass +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 secret(): + +def test_parameter_kinds_are_preserved(): + def function(x, /, scale, *args, out=None, **kwargs): pass - class PrivateLayer: + 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 - secret.__module__ = "example._private_for_testing.helpers" - PrivateLayer.__module__ = "example._private_for_testing" - module = make_module( - _hidden=public, - secret=secret, - public=public, - layers=make_module(_Hidden=type("Hidden", (), {}), Alias=PrivateLayer), - ) - result = generator.generate_symbols(module) - assert [item["name"] for item in result["functions"]] == ["public"] - assert result["layers"] == [] + 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_lazy_exports_not_in_dir(): - def public(): + +def test_string_annotations_are_not_evaluated(): + def function(x: "UndefinedTensor") -> "UndefinedTensor": # noqa: F821 pass - module = make_module() - module.__all__ = ["lazy"] - module.__getattr__ = lambda name: public if name == "lazy" else None - assert generator.generate_symbols(module)["functions"][0]["name"] == "lazy" + signature = generator.describe_signature(function) + assert signature["parameters"][0]["annotation"] == "'UndefinedTensor'" + assert signature["return_annotation"] == "'UndefinedTensor'" -def test_layers_need_not_be_exported_by_package(): - module = make_module() - module.layers = make_module(Layer=type("Layer", (), {})) - result = generator.generate_symbols(module) - assert result["functions"] == [] - assert [item["name"] for item in result["layers"]] == ["Layer"] +def test_layer_constructor_signature_is_included(kernel): + signature = generator.describe_signature(kernel.layers.ReLU) -def test_missing_layers_is_supported(): - module = make_module(public=lambda: None) - result = generator.generate_symbols(module) - assert [item["name"] for item in result["functions"]] == ["public"] - assert result["layers"] == [] + 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_export_order_and_aliases_are_preserved(): - def implementation(): +def test_inherited_layer_flags_are_included(kernel): + class ChildLayer(kernel.layers.ReLU): pass - module = make_module(z=implementation, a=implementation) - module.__all__ = ("z", "a", "z") - assert [ - item["name"] for item in generator.generate_symbols(module)["functions"] - ] == ["z", "a"] - - -def test_missing_package_is_an_error(tmp_path): - with pytest.raises(FileNotFoundError, match="No kernel package found"): - generator.load_package("missing_package", tmp_path) - - -def test_failed_import_cleans_up_submodules(tmp_path): - (tmp_path / "layers.py").write_text("", encoding="utf-8") - (tmp_path / "__init__.py").write_text( - "from . import layers\nraise RuntimeError('broken import')\n", encoding="utf-8" - ) - with pytest.raises(RuntimeError, match="broken import"): - generator.load_package("failed_build", tmp_path) - assert not any( - name == "failed_build" or name.startswith("failed_build.") - for name in sys.modules - ) + 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 From 904b52b77be154ad441b485677df46ebc90343b0 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 8 Oct 2026 12:12:51 +0000 Subject: [PATCH 07/10] kernel-builder: move script and adjust paths. --- .../{scripts => pkgs/generate-symbols}/generate_symbols.py | 0 .../{tests => pkgs/generate-symbols}/test_generate_symbols.py | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename nix-builder/{scripts => pkgs/generate-symbols}/generate_symbols.py (100%) rename nix-builder/{tests => pkgs/generate-symbols}/test_generate_symbols.py (98%) diff --git a/nix-builder/scripts/generate_symbols.py b/nix-builder/pkgs/generate-symbols/generate_symbols.py similarity index 100% rename from nix-builder/scripts/generate_symbols.py rename to nix-builder/pkgs/generate-symbols/generate_symbols.py diff --git a/nix-builder/tests/test_generate_symbols.py b/nix-builder/pkgs/generate-symbols/test_generate_symbols.py similarity index 98% rename from nix-builder/tests/test_generate_symbols.py rename to nix-builder/pkgs/generate-symbols/test_generate_symbols.py index 4b369dd3..da3534e0 100644 --- a/nix-builder/tests/test_generate_symbols.py +++ b/nix-builder/pkgs/generate-symbols/test_generate_symbols.py @@ -6,7 +6,7 @@ import pytest # The generator is a standalone script, so load it by its file path. -SCRIPT = Path(__file__).parents[1] / "scripts" / "generate_symbols.py" +SCRIPT = Path(__file__).with_name("generate_symbols.py") SPEC = importlib.util.spec_from_file_location("generate_symbols", SCRIPT) generator = importlib.util.module_from_spec(SPEC) SPEC.loader.exec_module(generator) From 21f99a7e963b82f674153931ce9bf4a9e5872cd6 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 8 Oct 2026 12:24:20 +0000 Subject: [PATCH 08/10] remove importlib fallback. --- .../pkgs/generate-symbols/generate_symbols.py | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/nix-builder/pkgs/generate-symbols/generate_symbols.py b/nix-builder/pkgs/generate-symbols/generate_symbols.py index edbd52a2..2b220ec1 100644 --- a/nix-builder/pkgs/generate-symbols/generate_symbols.py +++ b/nix-builder/pkgs/generate-symbols/generate_symbols.py @@ -7,7 +7,6 @@ """ import argparse -import importlib import inspect import json from collections.abc import Callable @@ -100,22 +99,9 @@ def describe_exports( def generate_symbols(module: ModuleType) -> dict[str, Any]: - """Describe top-level functions and optional layers, each scoped by __all__. - - As in describe-api, layers need not appear in the package's __all__. - Only Python functions (including async functions) and layer classes are - included; other exported modules, data, and callable objects are ignored. - """ + """Describe top-level functions and optional layers, each scoped by __all__.""" functions = describe_exports(module, inspect.isfunction) layers = getattr(module, "layers", None) - if layers is None and hasattr(module, "__path__"): - # The package need not import its optional layers submodule itself. - layers_name = f"{module.__name__}.layers" - try: - layers = importlib.import_module(layers_name) - except ModuleNotFoundError as error: - if error.name != layers_name: - raise return { "schema_version": 1, "module": module.__name__, From d313f02f061688a43dcf96c7bd22322e65be33b8 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 8 Oct 2026 12:32:05 +0000 Subject: [PATCH 09/10] kernel-builder: test for no export. --- .../pkgs/generate-symbols/test_generate_symbols.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/nix-builder/pkgs/generate-symbols/test_generate_symbols.py b/nix-builder/pkgs/generate-symbols/test_generate_symbols.py index da3534e0..2ecff8e2 100644 --- a/nix-builder/pkgs/generate-symbols/test_generate_symbols.py +++ b/nix-builder/pkgs/generate-symbols/test_generate_symbols.py @@ -111,9 +111,15 @@ def test_functions_in_layers_are_not_layers(kernel): assert [layer["name"] for layer in result["layers"]] == ["ReLU"] -def test_layers_are_optional(kernel): +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) From 70a1f8c3d7de4571a862582a34e8ae876e205747 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 8 Oct 2026 12:36:13 +0000 Subject: [PATCH 10/10] kernel-builder: turn symbol generation into a proper python project --- .../pkgs/generate-symbols/pyproject.toml | 24 +++++++++++++++++++ .../{ => src}/generate_symbols.py | 0 .../{ => tests}/test_generate_symbols.py | 10 ++------ 3 files changed, 26 insertions(+), 8 deletions(-) create mode 100644 nix-builder/pkgs/generate-symbols/pyproject.toml rename nix-builder/pkgs/generate-symbols/{ => src}/generate_symbols.py (100%) rename nix-builder/pkgs/generate-symbols/{ => tests}/test_generate_symbols.py (95%) 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/generate_symbols.py b/nix-builder/pkgs/generate-symbols/src/generate_symbols.py similarity index 100% rename from nix-builder/pkgs/generate-symbols/generate_symbols.py rename to nix-builder/pkgs/generate-symbols/src/generate_symbols.py diff --git a/nix-builder/pkgs/generate-symbols/test_generate_symbols.py b/nix-builder/pkgs/generate-symbols/tests/test_generate_symbols.py similarity index 95% rename from nix-builder/pkgs/generate-symbols/test_generate_symbols.py rename to nix-builder/pkgs/generate-symbols/tests/test_generate_symbols.py index 2ecff8e2..d211e4de 100644 --- a/nix-builder/pkgs/generate-symbols/test_generate_symbols.py +++ b/nix-builder/pkgs/generate-symbols/tests/test_generate_symbols.py @@ -1,16 +1,10 @@ -import importlib.util +import importlib import json -from pathlib import Path from types import ModuleType +import generate_symbols as generator import pytest -# The generator is a standalone script, so load it by its file path. -SCRIPT = Path(__file__).with_name("generate_symbols.py") -SPEC = importlib.util.spec_from_file_location("generate_symbols", SCRIPT) -generator = importlib.util.module_from_spec(SPEC) -SPEC.loader.exec_module(generator) - @pytest.fixture def kernel():