diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index c5a32d5..0144f7a 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -320,3 +320,14 @@ $ nix run .#ciTests.torch210-cxx11-cpu-x86_64-linux When running the tests on a non-NixOS systems, make sure that [the CUDA driver library can be found](https://danieldk.eu/Software/Nix/Nix-CUDA-on-non-NixOS-systems#solutions). + +## Kernel docs + +We provide a utility to generate a system card for a given kernel, utilizing +information from the `build.toml` and metadata. This system card provides a +reasonable starting point and is meant to be edited afterward by the kernel +developer. + +To generate the card, use the `kernels create-and-upload-card ...` command. +Alternatively, specify `--create_card` while uploading the kernel builds +through `kernels upload`. Refer to the [documentation](../cli.md) to know more. \ No newline at end of file diff --git a/kernels/src/kernels/cli/__init__.py b/kernels/src/kernels/cli/__init__.py index f172678..e40445d 100644 --- a/kernels/src/kernels/cli/__init__.py +++ b/kernels/src/kernels/cli/__init__.py @@ -25,6 +25,8 @@ _update_kernel_card_usage, ) +SYSTEM_CARD_PATH = "CARD.md" + def main(): parser = argparse.ArgumentParser( @@ -247,36 +249,45 @@ def main(): ) init_parser.set_defaults(func=run_init) - repocard_parser = subparsers.add_parser( - "create-and-upload-card", - help="Create and optionally upload a kernel card.", + init_card_parser = subparsers.add_parser( + "init-card", + help="Initialize a kernel system card template inside the build directory of the kernel.", ) - repocard_parser.add_argument( + init_card_parser.add_argument( "kernel_dir", type=str, help="Path to the kernels source.", ) - repocard_parser.add_argument( - "--card-path", type=str, required=True, help="Path to save the card to." + init_card_parser.add_argument( + "--repo_id", + type=str, + default=None, + help="When specified, existing card content is reused. Specific parts are updated based on the build information.", + ) + init_card_parser.set_defaults(func=initialize_card) + + fill_card_parser = subparsers.add_parser( + "fill-card", + help="Fill a system card template based on the `build` information and save it.", + ) + fill_card_parser.add_argument( + "kernel_dir", + type=str, + help="Path to the kernels source.", ) - repocard_parser.add_argument( + fill_card_parser.add_argument( "--description", type=str, default=None, help="Description to introduce the kernel.", ) - repocard_parser.add_argument( + fill_card_parser.add_argument( "--repo-id", type=str, default=None, help="If specified it will be pushed to a repository on the Hub.", ) - repocard_parser.add_argument( - "--create-pr", - action="store_true", - help="If specified it will create a PR on the `repo_id`.", - ) - repocard_parser.set_defaults(func=create_and_upload_card) + fill_card_parser.set_defaults(func=fill_kernel_card) args = parser.parse_args() args.func(args) @@ -346,15 +357,24 @@ def upload_kernels(args): ) -def create_and_upload_card(args): - if not args.repo_id and args.create_pr: - raise ValueError("`create_pr` cannot be True when `repo_id` is None.") - +def initialize_card(args): kernel_dir = Path(args.kernel_dir).resolve() kernel_card = _load_or_create_kernel_card( - kernel_description=args.description, license="apache-2.0" + repo_id_or_path=args.repo_id, + kernel_description=args.description, + license="apache-2.0", ) + kernel_card.save(kernel_dir / "build" / SYSTEM_CARD_PATH) + +def fill_kernel_card(args): + kernel_dir = Path(args.kernel_dir).resolve() + card_path_from_kernel_build = kernel_dir / "build" / SYSTEM_CARD_PATH + kernel_card = _load_or_create_kernel_card( + repo_id_or_path=card_path_from_kernel_build, + kernel_description=args.description, + license="apache-2.0", + ) updated_card = _update_kernel_card_usage( kernel_card=kernel_card, local_path=kernel_dir, repo_id=args.repo_id ) @@ -368,12 +388,7 @@ def create_and_upload_card(args): updated_card = _update_kernel_card_license( kernel_card=kernel_card, local_path=kernel_dir ) - - card_path = args.card_path - updated_card.save(card_path) - - if args.repo_id: - updated_card.push_to_hub(repo_id=args.repo_id, create_pr=args.create_pr) + updated_card.save(card_path_from_kernel_build) class _JSONEncoder(json.JSONEncoder): diff --git a/kernels/src/kernels/cli/upload.py b/kernels/src/kernels/cli/upload.py index 52f37ff..b91df11 100644 --- a/kernels/src/kernels/cli/upload.py +++ b/kernels/src/kernels/cli/upload.py @@ -70,6 +70,15 @@ def upload_kernels_dir( allow_patterns=["benchmark*.py"], ) + card_path = build_dir / "CARD.md" + if (card_path).exists(): + api.upload_file( + path_or_fileobj=card_path, + path_in_repo="README.md", + revision=branch, + commit_message="File uploaded using `kernels`.", + ) + api.upload_folder( repo_id=repo_id, folder_path=build_dir, diff --git a/kernels/tests/test_kernel_card.py b/kernels/tests/test_kernel_card.py index d8bf5e3..2899455 100644 --- a/kernels/tests/test_kernel_card.py +++ b/kernels/tests/test_kernel_card.py @@ -1,20 +1,23 @@ import tempfile from pathlib import Path from dataclasses import dataclass -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest +from huggingface_hub import ModelCard, ModelCardData +from huggingface_hub.errors import RepositoryNotFoundError -from kernels.cli import create_and_upload_card +from kernels.cli import initialize_card, fill_kernel_card +from kernels.kernel_card_utils import KERNEL_CARD_TEMPLATE_PATH + +SYSTEM_CARD_PATH = "CARD.md" @dataclass class CardArgs: kernel_dir: str - card_path: str - description: str | None = None repo_id: str | None = None - create_pr: bool = False + description: str | None = None @pytest.fixture @@ -22,17 +25,13 @@ def mock_kernel_dir(): with tempfile.TemporaryDirectory() as tmpdir: kernel_dir = Path(tmpdir) - build_toml = kernel_dir / "build.toml" - build_toml.write_text( + (kernel_dir / "build.toml").write_text( """[general] name = "test_kernel" backends = ["cuda", "metal"] license = "apache-2.0" version = 1 -[general.hub] -repo-id = "test-org/test-kernel" - [kernel._test] backend = "cuda" cuda-capabilities = ["8.0", "8.9"] @@ -41,256 +40,141 @@ def mock_kernel_dir(): torch_ext_dir = kernel_dir / "torch-ext" / "test_kernel" torch_ext_dir.mkdir(parents=True) - - init_file = torch_ext_dir / "__init__.py" - init_file.write_text( - """from .core import func1, func2 - -__all__ = ["func1", "func2", "func3"] -""" + (torch_ext_dir / "__init__.py").write_text( + 'from .core import func1, func2\n\n__all__ = ["func1", "func2"]\n' ) - - core_file = torch_ext_dir / "core.py" - core_file.write_text( - """def func1(): - pass - -def func2(): - pass - -def func3(): - pass -""" + (torch_ext_dir / "core.py").write_text( + "def func1():\n pass\n\ndef func2():\n pass\n" ) + (kernel_dir / "build").mkdir() yield kernel_dir @pytest.fixture -def mock_kernel_dir_with_benchmark(mock_kernel_dir): - benchmarks_dir = mock_kernel_dir / "benchmarks" - benchmarks_dir.mkdir() - - benchmark_file = benchmarks_dir / "benchmark.py" - benchmark_file.write_text( - """import time - -def benchmark(): - # Simple benchmark - start = time.time() - # ... benchmark code ... - end = time.time() - return end - start -""" +def initialized_kernel_dir(mock_kernel_dir): + card = ModelCard.from_template( + card_data=ModelCardData(license="apache-2.0", library_name="kernels"), + template_path=str(KERNEL_CARD_TEMPLATE_PATH), + model_description="Test kernel.", ) - + card.save(mock_kernel_dir / "build" / SYSTEM_CARD_PATH) return mock_kernel_dir -@pytest.fixture -def mock_kernel_dir_minimal(): - with tempfile.TemporaryDirectory() as tmpdir: - kernel_dir = Path(tmpdir) +def test_initialize_card_creates_file(mock_kernel_dir): + args = CardArgs(kernel_dir=str(mock_kernel_dir)) + with patch( + "huggingface_hub.ModelCard.load", + side_effect=RepositoryNotFoundError("test", response=MagicMock()), + ): + initialize_card(args) + assert (mock_kernel_dir / "build" / SYSTEM_CARD_PATH).exists() - build_toml = kernel_dir / "build.toml" - build_toml.write_text( - """[general] -name = "minimal_kernel" -backends = ["cuda"] -""" - ) - yield kernel_dir +def test_initialize_card_with_description(mock_kernel_dir): + description = "A test kernel." + args = CardArgs(kernel_dir=str(mock_kernel_dir), description=description) + with patch( + "huggingface_hub.ModelCard.load", + side_effect=RepositoryNotFoundError("test", response=MagicMock()), + ): + initialize_card(args) + content = (mock_kernel_dir / "build" / SYSTEM_CARD_PATH).read_text() + assert "---" in content + assert description in content -def test_create_and_upload_card_basic(mock_kernel_dir): - with tempfile.TemporaryDirectory() as tmpdir: - card_path = Path(tmpdir) / "README.md" - - args = CardArgs( - kernel_dir=str(mock_kernel_dir), - card_path=str(card_path), - description="This is a test kernel for testing purposes.", - ) - - create_and_upload_card(args) - - assert card_path.exists() - - card_content = card_path.read_text() - - assert "---" in card_content - assert "This is a test kernel for testing purposes." in card_content +def test_fill_kernel_card_backends(initialized_kernel_dir): + args = CardArgs(kernel_dir=str(initialized_kernel_dir)) + fill_kernel_card(args) + content = (initialized_kernel_dir / "build" / SYSTEM_CARD_PATH).read_text() + assert "- cuda" in content + assert "- metal" in content -def test_create_and_upload_card_updates_usage(mock_kernel_dir): - """Test that usage code snippet is properly generated.""" - with tempfile.TemporaryDirectory() as tmpdir: - card_path = Path(tmpdir) / "README.md" - - args = CardArgs( - kernel_dir=str(mock_kernel_dir), - card_path=str(card_path), - ) +def test_fill_kernel_card_cuda_capabilities(initialized_kernel_dir): + args = CardArgs(kernel_dir=str(initialized_kernel_dir)) + fill_kernel_card(args) + content = (initialized_kernel_dir / "build" / SYSTEM_CARD_PATH).read_text() + assert "## CUDA Capabilities" in content + assert "- 8.0" in content or "- 8.9" in content - create_and_upload_card(args) - card_content = card_path.read_text() +def test_fill_kernel_card_available_funcs(initialized_kernel_dir): + args = CardArgs(kernel_dir=str(initialized_kernel_dir)) + fill_kernel_card(args) + content = (initialized_kernel_dir / "build" / SYSTEM_CARD_PATH).read_text() + assert "- `func1`" in content + assert "- `func2`" in content - assert "## How to use" in card_content - assert "from kernels import get_kernel" in card_content - assert "func1" in card_content - assert "TODO: add an example code snippet" not in card_content - - -def test_create_and_upload_card_updates_available_functions(mock_kernel_dir): - with tempfile.TemporaryDirectory() as tmpdir: - card_path = Path(tmpdir) / "README.md" - - args = CardArgs( - kernel_dir=str(mock_kernel_dir), - card_path=str(card_path), - ) - - create_and_upload_card(args) - - card_content = card_path.read_text() - - assert "## Available functions" in card_content - assert "- `func1`" in card_content - assert "- `func2`" in card_content - assert "- `func3`" in card_content - assert ( - "[TODO: add the functions available through this kernel]" - not in card_content - ) - - -def test_create_and_upload_card_updates_backends(mock_kernel_dir): - with tempfile.TemporaryDirectory() as tmpdir: - card_path = Path(tmpdir) / "README.md" - - args = CardArgs( - kernel_dir=str(mock_kernel_dir), - card_path=str(card_path), - ) - create_and_upload_card(args) +def test_fill_kernel_card_usage_with_repo_id(initialized_kernel_dir): + repo_id = "test-org/test-kernel" + args = CardArgs(kernel_dir=str(initialized_kernel_dir), repo_id=repo_id) + fill_kernel_card(args) + content = (initialized_kernel_dir / "build" / SYSTEM_CARD_PATH).read_text() + assert f'get_kernel("{repo_id}")' in content - card_content = card_path.read_text() - assert "## Supported backends" in card_content - assert "- cuda" in card_content - assert "- metal" in card_content - assert "[TODO: add the backends this kernel supports]" not in card_content +def test_fill_kernel_card_license(initialized_kernel_dir): + args = CardArgs(kernel_dir=str(initialized_kernel_dir)) + fill_kernel_card(args) + content = (initialized_kernel_dir / "build" / SYSTEM_CARD_PATH).read_text() + assert "license: apache-2.0" in content -def test_create_and_upload_card_updates_cuda_capabilities(mock_kernel_dir): - with tempfile.TemporaryDirectory() as tmpdir: - card_path = Path(tmpdir) / "README.md" - - args = CardArgs( - kernel_dir=str(mock_kernel_dir), - card_path=str(card_path), - ) - - create_and_upload_card(args) - - card_content = card_path.read_text() - - assert "## CUDA Capabilities" in card_content - assert "- 8.0" in card_content or "- 8.9" in card_content - - -def test_create_and_upload_card_updates_license(mock_kernel_dir): - with tempfile.TemporaryDirectory() as tmpdir: - card_path = Path(tmpdir) / "README.md" - - args = CardArgs( - kernel_dir=str(mock_kernel_dir), - card_path=str(card_path), - ) - - create_and_upload_card(args) - - card_content = card_path.read_text() - - assert "license: apache-2.0" in card_content - - -def test_create_and_upload_card_with_benchmark(mock_kernel_dir_with_benchmark): - with tempfile.TemporaryDirectory() as tmpdir: - card_path = Path(tmpdir) / "README.md" - - args = CardArgs( - kernel_dir=str(mock_kernel_dir_with_benchmark), - card_path=str(card_path), - ) - - create_and_upload_card(args) - - card_content = card_path.read_text() - - assert "## Benchmarks" in card_content - assert "Benchmarking script is available for this kernel" in card_content - assert "kernels benchmark" in card_content - - -def test_create_and_upload_card_minimal_structure(mock_kernel_dir_minimal): - with tempfile.TemporaryDirectory() as tmpdir: - card_path = Path(tmpdir) / "README.md" - - args = CardArgs( - kernel_dir=str(mock_kernel_dir_minimal), - card_path=str(card_path), - ) - - create_and_upload_card(args) - - assert card_path.exists() - - card_content = card_path.read_text() - - assert "---" in card_content - assert "## How to use" in card_content - assert "## Available functions" in card_content - assert "## Supported backends" in card_content - - -def test_create_and_upload_card_custom_description(mock_kernel_dir): - with tempfile.TemporaryDirectory() as tmpdir: - card_path = Path(tmpdir) / "README.md" - - custom_desc = "My custom kernel description with special features." - - args = CardArgs( - kernel_dir=str(mock_kernel_dir), - card_path=str(card_path), - description=custom_desc, - ) - - create_and_upload_card(args) - - card_content = card_path.read_text() - - assert custom_desc in card_content - - -def test_create_and_upload_card_usage_with_repo_id(mock_kernel_dir): - """Test that usage code snippet includes the provided repo_id.""" - with tempfile.TemporaryDirectory() as tmpdir: - card_path = Path(tmpdir) / "README.md" +def test_fill_kernel_card_benchmark(initialized_kernel_dir): + benchmarks_dir = initialized_kernel_dir / "benchmarks" + benchmarks_dir.mkdir() + (benchmarks_dir / "benchmark.py").write_text("def benchmark(): pass\n") + args = CardArgs(kernel_dir=str(initialized_kernel_dir)) + fill_kernel_card(args) + content = (initialized_kernel_dir / "build" / SYSTEM_CARD_PATH).read_text() + assert "## Benchmarks" in content + assert "Benchmarking script is available" in content + + +def test_fill_kernel_card_preserves_existing_content(mock_kernel_dir): + existing_description = "A hand-written description of this kernel." + existing_notes = "Custom notes that should not be overwritten." + existing_source = "https://github.com/example/kernel-source" + + card_path = mock_kernel_dir / "build" / SYSTEM_CARD_PATH + card_path.write_text( + "---\n" + "license: mit\n" + "library_name: kernels\n" + "---\n\n" + f"{existing_description}\n\n" + "## How to use\n\n" + "```python\n" + "# TODO: add an example code snippet for running this kernel\n" + "```\n\n" + "## Available functions\n\n" + "[TODO: add the functions available through this kernel]\n\n" + "## Supported backends\n\n" + "[TODO: add the backends this kernel supports]\n\n" + "## Benchmarks\n\n" + "[TODO: provide benchmarks if available]\n\n" + f"## Source code\n\n{existing_source}\n\n" + f"## Notes\n\n{existing_notes}\n" + ) - args = CardArgs( - kernel_dir=str(mock_kernel_dir), - card_path=str(card_path), - repo_id="my-org/my-kernel", - ) + repo_id = "test-org/test-kernel" + args = CardArgs(kernel_dir=str(mock_kernel_dir), repo_id=repo_id) + fill_kernel_card(args) - with patch("huggingface_hub.ModelCard.push_to_hub"): - create_and_upload_card(args) + content = card_path.read_text() - card_content = card_path.read_text() + assert "- `func1`" in content + assert "- `func2`" in content + assert "- cuda" in content + assert "- metal" in content + assert f'get_kernel("{repo_id}")' in content + assert "[TODO: add the functions available through this kernel]" not in content + assert "[TODO: add the backends this kernel supports]" not in content - assert "## How to use" in card_content - assert 'get_kernel("my-org/my-kernel")' in card_content + assert existing_description in content + assert existing_notes in content + assert existing_source in content diff --git a/kernels/tests/test_kernel_upload.py b/kernels/tests/test_kernel_upload.py index 11928bf..866bf4b 100644 --- a/kernels/tests/test_kernel_upload.py +++ b/kernels/tests/test_kernel_upload.py @@ -4,10 +4,12 @@ import tempfile from dataclasses import dataclass from pathlib import Path +from unittest.mock import MagicMock, patch import pytest from kernels.cli import upload_kernels +from kernels.cli.upload import upload_kernels_dir from kernels.utils import _get_hf_api REPO_ID = "valid_org/kernels-upload-test" @@ -30,6 +32,10 @@ class UploadArgs: repo_id: None private: False branch: None + create_card: bool = False + card_path: str = None + description: str = None + create_pr: bool = False def next_filename(path: Path) -> Path: @@ -120,3 +126,30 @@ def test_kernel_upload_deletes_as_expected(): str(filename_to_change) in k for k in repo_filenames ), f"{repo_filenames=}" _get_hf_api().delete_repo(repo_id=REPO_ID) + + +def test_upload_includes_card_as_readme(): + with tempfile.TemporaryDirectory() as tmpdir: + kernel_dir = Path(tmpdir).resolve() + + variant_dir = kernel_dir / "build" / "torch-cuda" + variant_dir.mkdir(parents=True) + (variant_dir / "metadata.json").write_text( + '{"version": 1, "python-depends": []}' + ) + + card_path = kernel_dir / "build" / "CARD.md" + card_path.write_text("# Test Kernel\n") + + mock_api = MagicMock() + mock_api.create_repo.return_value.repo_id = REPO_ID + + with patch("kernels.cli.upload._get_hf_api", return_value=mock_api): + upload_kernels_dir(kernel_dir, repo_id=REPO_ID, branch=None, private=False) + + mock_api.upload_file.assert_called_once_with( + path_or_fileobj=card_path, + path_in_repo="README.md", + revision="v1", + commit_message="File uploaded using `kernels`.", + )