Set version to 0.1.6 (#23)

So far we have been using the extension name in `build.toml` as
the module name. However, this can cause naming conflicts. For
instance, if a kernel named `moe` is loaded through `hf_kernels`,
it would be registered as the `moe` module. This would cause
subsequent imports of a module named `moe` from the Python path
to be resolved as the module loaded through `hf_kernels`.

Solve this issue by adding some unique material to the module
name (hex-encoded hash of the kernel path).

.github/workflows/test.yml

@@ -23,8 +23,8 @@ jobs:
     strategy:
       max-parallel: 4
       matrix:
-        python: ["3.10", "3.12"]
-        torch: ["2.4.0", "2.5.0"]
+        python-version: ["3.10", "3.12"]
+        torch-version: ["2.5.1", "2.6.0"]
 
     steps:
       - name: Checkout code
@@ -35,6 +35,9 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Lock Torch version
+        run: uv lock --upgrade-package "torch==${{ matrix.torch-version }}"
+
       - name: Install the project
         run: uv sync --all-extras --dev
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hf-kernels"
-version = "0.1.1"
+version = "0.1.6"
 description = "Download cuda kernels"
 authors = [
   { name = "OlivierDehaene", email = "olivier@huggingface.co" },
@@ -9,7 +9,7 @@ authors = [
   { name = "Nicolas Patry", email = "nicolas@huggingface.co" },
 ]
 readme = "README.md"
-requires-python = ">= 3.8"
+requires-python = ">= 3.9"
 dependencies = [
   "huggingface-hub>=0.26.3",
   "packaging>=24.2",

src/hf_kernels/cli.py

@@ -6,14 +6,14 @@ from pathlib import Path
 
 from hf_kernels.compat import tomllib
 from hf_kernels.lockfile import KernelLock, get_kernel_locks
-from hf_kernels.utils import install_kernel
+from hf_kernels.utils import install_kernel, install_kernel_all_variants
 
 
 def main():
     parser = argparse.ArgumentParser(
         prog="hf-kernel", description="Manage compute kernels"
     )
-    subparsers = parser.add_subparsers()
+    subparsers = parser.add_subparsers(required=True)
 
     download_parser = subparsers.add_parser("download", help="Download locked kernels")
     download_parser.add_argument(
@@ -21,6 +21,11 @@ def main():
         type=Path,
         help="The project directory",
     )
+    download_parser.add_argument(
+        "--all-variants",
+        action="store_true",
+        help="Download all build variants of the kernel",
+    )
     download_parser.set_defaults(func=download_kernels)
 
     lock_parser = subparsers.add_parser("lock", help="Lock kernel revisions")
@@ -51,7 +56,10 @@ def download_kernels(args):
             f"Downloading `{kernel_lock.repo_id}` at with SHA: {kernel_lock.sha}",
             file=sys.stderr,
         )
-        install_kernel(kernel_lock.repo_id, kernel_lock.sha)
+        if args.all_variants:
+            install_kernel_all_variants(kernel_lock.repo_id, kernel_lock.sha)
+        else:
+            install_kernel(kernel_lock.repo_id, kernel_lock.sha)
 
 
 def lock_kernels(args):

src/hf_kernels/lockfile.py

@@ -89,7 +89,12 @@ def write_egg_lockfile(cmd, basename, filename):
     import logging
 
     cwd = Path.cwd()
-    with open(cwd / "pyproject.toml", "rb") as f:
+    pyproject_path = cwd / "pyproject.toml"
+    if not pyproject_path.exists():
+        # Nothing to do if the project doesn't have pyproject.toml.
+        return
+
+    with open(pyproject_path, "rb") as f:
         data = tomllib.load(f)
 
     kernel_versions = data.get("tool", {}).get("kernels", {}).get("dependencies", None)

src/hf_kernels/utils.py

@@ -1,3 +1,4 @@
+import ctypes
 import importlib
 import importlib.metadata
 import inspect
@@ -9,15 +10,18 @@ from importlib.metadata import Distribution
 from types import ModuleType
 from typing import List, Optional
 
-import torch
 from huggingface_hub import hf_hub_download, snapshot_download
 from packaging.version import parse
 
 from hf_kernels.compat import tomllib
 from hf_kernels.lockfile import KernelLock
 
+CACHE_DIR: Optional[str] = os.environ.get("HF_KERNELS_CACHE", None)
+
 
 def build_variant():
+    import torch
+
     torch_version = parse(torch.__version__)
     cuda_version = parse(torch.version.cuda)
     cxxabi = "cxx11" if torch.compiled_with_cxx11_abi() else "cxx98"
@@ -28,6 +32,12 @@ def build_variant():
 
 
 def import_from_path(module_name: str, file_path):
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path)).value)
+    module_name = f"{module_name}_{path_hash}"
     spec = importlib.util.spec_from_file_location(module_name, file_path)
     module = importlib.util.module_from_spec(spec)
     sys.modules[module_name] = module
@@ -42,16 +52,33 @@ def install_kernel(repo_id: str, revision: str, local_files_only: bool = False):
     repo_path = snapshot_download(
         repo_id,
         allow_patterns=f"build/{build_variant()}/*",
+        cache_dir=CACHE_DIR,
         revision=revision,
         local_files_only=local_files_only,
     )
     return package_name, f"{repo_path}/build/{build_variant()}"
 
 
+def install_kernel_all_variants(
+    repo_id: str, revision: str, local_files_only: bool = False
+):
+    snapshot_download(
+        repo_id,
+        allow_patterns="build/*",
+        cache_dir=CACHE_DIR,
+        revision=revision,
+        local_files_only=local_files_only,
+    )
+
+
 def get_metadata(repo_id: str, revision: str, local_files_only: bool = False):
     with open(
         hf_hub_download(
-            repo_id, "build.toml", revision=revision, local_files_only=local_files_only
+            repo_id,
+            "build.toml",
+            cache_dir=CACHE_DIR,
+            revision=revision,
+            local_files_only=local_files_only,
         ),
         "rb",
     ) as f:
@@ -71,7 +98,11 @@ def load_kernel(repo_id: str):
         raise ValueError(f"Kernel `{repo_id}` is not locked")
 
     filename = hf_hub_download(
-        repo_id, "build.toml", local_files_only=True, revision=locked_sha
+        repo_id,
+        "build.toml",
+        cache_dir=CACHE_DIR,
+        local_files_only=True,
+        revision=locked_sha,
     )
     with open(filename, "rb") as f:
         metadata = tomllib.load(f)