Project rename + build primitive.

kernels/__init__.py

@@ -1,3 +0,0 @@
-from kernels.utils import get_kernel
-
-__all__ = ["get_kernel"]

pyproject.toml

@@ -1,17 +1,26 @@
-[tool.poetry]
-name = "kernels"
+[project]
+name = "hf-kernels"
 version = "0.1.0"
-description = ""
-authors = ["OlivierDehaene <olivier@huggingface.co>"]
+description = "Download cuda kernels"
+authors = [
+    {name = "OlivierDehaene", email = "olivier@huggingface.co"},
+    {name = "Daniel de Kok", email = "daniel@huggingface.co"},
+    {name = "David Holtz", email = "david@huggingface.co"},
+    {name = "Nicolas Patry", email = "nicolas@huggingface.co"}
+]
 readme = "README.md"
 
-[tool.poetry.dependencies]
+[dependencies]
 python = "^3.9"
 huggingface-hub = "^0.26.3"
 packaging = "^24.2"
 tomli = { version = "^2.0.1", python = "<3.11" }
 
+[tool.kernels.dependencies]
+"kernels-community/activation" = "^0.0.1"
+
 
 [build-system]
-requires = ["poetry-core"]
-build-backend = "poetry.core.masonry.api"
+requires = ["torch", "huggingface_hub", "numpy"]
+build-backend = "hf_kernels.build"
+backend-path = ["src"]

src/hf_kernels/__init__.py

@@ -0,0 +1,3 @@
+from hf_kernels.utils import get_kernel
+
+__all__ = ["get_kernel", "load_kernel", "install_kernel"]

src/hf_kernels/build.py

@@ -0,0 +1,146 @@
+"""
+Python shims for the PEP 517 and PEP 660 build backend.
+
+Major imports in this module are required to be lazy:
+```
+$ hyperfine \
+     "/usr/bin/python3 -c \"print('hi')\"" \
+     "/usr/bin/python3 -c \"from subprocess import check_call; print('hi')\""
+Base: Time (mean ± σ):      11.0 ms ±   1.7 ms    [User: 8.5 ms, System: 2.5 ms]
+With import: Time (mean ± σ):      15.2 ms ±   2.0 ms    [User: 12.3 ms, System: 2.9 ms]
+Base 1.38 ± 0.28 times faster than with import
+```
+
+The same thing goes for the typing module, so we use Python 3.10 type annotations that
+don't require importing typing but then quote them so earlier Python version ignore
+them while IDEs and type checker can see through the quotes.
+"""
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence  # noqa:I001
+    from typing import Any  # noqa:I001
+
+
+def warn_config_settings(config_settings: "Mapping[Any, Any] | None" = None) -> None:
+    import sys
+
+    if config_settings:
+        print("Warning: Config settings are not supported", file=sys.stderr)
+
+
+def call(
+    args: "Sequence[str]", config_settings: "Mapping[Any, Any] | None" = None
+) -> str:
+    """Invoke a uv subprocess and return the filename from stdout."""
+    import shutil
+    import subprocess
+    import sys
+
+    print(f" === Args {args}", file=sys.stderr)
+
+    warn_config_settings(config_settings)
+    # Unlike `find_uv_bin`, this mechanism must work according to PEP 517
+    uv_bin = shutil.which("uv")
+    if uv_bin is None:
+        raise RuntimeError("uv was not properly installed")
+    # Forward stderr, capture stdout for the filename
+    result = subprocess.run([uv_bin, *args], stdout=subprocess.PIPE)
+    if result.returncode != 0:
+        sys.exit(result.returncode)
+    # If there was extra stdout, forward it (there should not be extra stdout)
+    stdout = result.stdout.decode("utf-8").strip().splitlines(keepends=True)
+    sys.stdout.writelines(stdout[:-1])
+    # Fail explicitly instead of an irrelevant stacktrace
+    if not stdout:
+        print("uv subprocess did not return a filename on stdout", file=sys.stderr)
+        sys.exit(1)
+    return stdout[-1].strip()
+
+
+def build_sdist(
+    sdist_directory: str, config_settings: "Mapping[Any, Any] | None" = None
+) -> str:
+    """PEP 517 hook `build_sdist`."""
+    args = ["build-backend", "build-sdist", sdist_directory]
+    return call(args, config_settings)
+
+
+def build_wheel(
+    wheel_directory: str,
+    config_settings: "Mapping[Any, Any] | None" = None,
+    metadata_directory: "str | None" = None,
+) -> str:
+    """PEP 517 hook `build_wheel`."""
+    args = ["build-backend", "build-wheel", wheel_directory]
+    if metadata_directory:
+        args.extend(["--metadata-directory", metadata_directory])
+    return call(args, config_settings)
+
+
+def get_requires_for_build_sdist(
+    config_settings: "Mapping[Any, Any] | None" = None,
+) -> "Sequence[str]":
+    """PEP 517 hook `get_requires_for_build_sdist`."""
+    warn_config_settings(config_settings)
+    return []
+
+
+def get_requires_for_build_wheel(
+    config_settings: "Mapping[Any, Any] | None" = None,
+) -> "Sequence[str]":
+    """PEP 517 hook `get_requires_for_build_wheel`."""
+    warn_config_settings(config_settings)
+    return []
+
+
+def prepare_metadata_for_build_wheel(
+    metadata_directory: str, config_settings: "Mapping[Any, Any] | None" = None
+) -> str:
+    """PEP 517 hook `prepare_metadata_for_build_wheel`."""
+    args = ["build-backend", "prepare-metadata-for-build-wheel", metadata_directory]
+    return call(args, config_settings)
+
+
+def build_editable(
+    wheel_directory: str,
+    config_settings: "Mapping[Any, Any] | None" = None,
+    metadata_directory: "str | None" = None,
+) -> str:
+    """PEP 660 hook `build_editable`."""
+    args = ["build-backend", "build-editable", wheel_directory]
+    import os
+    import sys
+    import tomllib
+
+    cwd = os.getcwd()
+    filename = os.path.join(cwd, "pyproject.toml")
+    with open(filename, "rb") as f:
+        data = tomllib.load(f)
+
+    for kernel, _ in (
+        data.get("tool", {}).get("kernels", {}).get("dependencies", {}).items()
+    ):
+        from hf_kernels.utils import install_kernel
+
+        install_kernel(kernel, revision="main")
+
+    if metadata_directory:
+        args.extend(["--metadata-directory", metadata_directory])
+    return call(args, config_settings)
+
+
+def get_requires_for_build_editable(
+    config_settings: "Mapping[Any, Any] | None" = None,
+) -> "Sequence[str]":
+    """PEP 660 hook `get_requires_for_build_editable`."""
+    warn_config_settings(config_settings)
+    return []
+
+
+def prepare_metadata_for_build_editable(
+    metadata_directory: str, config_settings: "Mapping[Any, Any] | None" = None
+) -> str:
+    """PEP 660 hook `prepare_metadata_for_build_editable`."""
+    args = ["build-backend", "prepare-metadata-for-build-editable", metadata_directory]
+    return call(args, config_settings)

src/hf_kernels/utils.py

@@ -1,6 +1,7 @@
 import importlib
 import platform
 import sys
+import os
 
 import torch
 from huggingface_hub import hf_hub_download, snapshot_download
@@ -30,8 +31,10 @@ def import_from_path(module_name: str, file_path):
     return module
 
 
-def get_package_path(repo_id: str):
-    repo_path = snapshot_download(repo_id, allow_patterns=f"build/{build_variant()}/*")
+def install_kernel(repo_id: str, revision: str):
+    repo_path = snapshot_download(
+        repo_id, allow_patterns=f"build/{build_variant()}/*", revision=revision
+    )
     return f"{repo_path}/build/{build_variant()}"
 
 
@@ -40,7 +43,19 @@ def get_metadata(repo_id: str):
         return tomllib.load(f)
 
 
-def get_kernel(repo_id: str):
+def get_kernel(repo_id: str, revision: str = "main"):
     package_name = get_metadata(repo_id)["torch"]["name"]
-    package_path = get_package_path(repo_id)
+    package_path = install_kernel(repo_id, revision=revision)
+    return import_from_path(package_name, f"{package_path}/{package_name}/__init__.py")
+
+
+def load_kernel(repo_id: str, revision: str = "main"):
+    filename = hf_hub_download(
+        repo_id, "build.toml", local_files_only=True, revision=revision
+    )
+    with open(filename, "rb") as f:
+        metadata = tomllib.load(f)
+    package_name = metadata["torch"]["name"]
+    repo_path = os.path.dirname(filename)
+    package_path = f"{repo_path}/build/{build_variant()}"
     return import_from_path(package_name, f"{package_path}/{package_name}/__init__.py")