kernels 0.2.0__tar.gz → 0.2.1__tar.gz

{kernels-0.2.0 → kernels-0.2.1}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.2
 Name: kernels
-Version: 0.2.0
-Summary: Download cuda kernels
+Version: 0.2.1
+Summary: Download compute kernels
 Author-email: OlivierDehaene <olivier@huggingface.co>, Daniel de Kok <daniel@huggingface.co>, David Holtz <david@huggingface.co>, Nicolas Patry <nicolas@huggingface.co>
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown

{kernels-0.2.0 → kernels-0.2.1}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "kernels"
-version = "0.2.0"
-description = "Download cuda kernels"
+version = "0.2.1"
+description = "Download compute kernels"
 authors = [
   { name = "OlivierDehaene", email = "olivier@huggingface.co" },
   { name = "Daniel de Kok", email = "daniel@huggingface.co" },

{kernels-0.2.0 → kernels-0.2.1}/src/kernels/utils.py

@@ -144,9 +144,18 @@ def get_kernel(repo_id: str, revision: str = "main") -> ModuleType:
     return import_from_path(package_name, package_path / package_name / "__init__.py")
 
 
-def load_kernel(repo_id: str) -> ModuleType:
-    """Get a pre-downloaded, locked kernel."""
-    locked_sha = _get_caller_locked_kernel(repo_id)
+def load_kernel(repo_id: str, *, lockfile: Optional[Path] = None) -> ModuleType:
+    """
+    Get a pre-downloaded, locked kernel.
+
+    If `lockfile` is not specified, the lockfile will be loaded from the
+    caller's package metadata.
+    """
+    if lockfile is None:
+        locked_sha = _get_caller_locked_kernel(repo_id)
+    else:
+        with open(lockfile, "r") as f:
+            locked_sha = _get_locked_kernel(repo_id, f.read())
 
     if locked_sha is None:
         raise ValueError(
@@ -163,6 +172,7 @@ def load_kernel(repo_id: str) -> ModuleType:
             repo_id,
             allow_patterns=[f"build/{variant}/*", f"build/{universal_variant}/*"],
             cache_dir=CACHE_DIR,
+            revision=locked_sha,
             local_files_only=True,
         )
     )
@@ -200,11 +210,19 @@ def get_locked_kernel(repo_id: str, local_files_only: bool = False) -> ModuleTyp
 def _get_caller_locked_kernel(repo_id: str) -> Optional[str]:
     for dist in _get_caller_distributions():
         lock_json = dist.read_text("kernels.lock")
-        if lock_json is not None:
-            for kernel_lock_json in json.loads(lock_json):
-                kernel_lock = KernelLock.from_json(kernel_lock_json)
-                if kernel_lock.repo_id == repo_id:
-                    return kernel_lock.sha
+        if lock_json is None:
+            continue
+        locked_sha = _get_locked_kernel(repo_id, lock_json)
+        if locked_sha is not None:
+            return locked_sha
+    return None
+
+
+def _get_locked_kernel(repo_id: str, lock_json: str) -> Optional[str]:
+    for kernel_lock_json in json.loads(lock_json):
+        kernel_lock = KernelLock.from_json(kernel_lock_json)
+        if kernel_lock.repo_id == repo_id:
+            return kernel_lock.sha
     return None
 
 

{kernels-0.2.0 → kernels-0.2.1}/src/kernels.egg-info/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.2
 Name: kernels
-Version: 0.2.0
-Summary: Download cuda kernels
+Version: 0.2.1
+Summary: Download compute kernels
 Author-email: OlivierDehaene <olivier@huggingface.co>, Daniel de Kok <daniel@huggingface.co>, David Holtz <david@huggingface.co>, Nicolas Patry <nicolas@huggingface.co>
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown

{kernels-0.2.0 → kernels-0.2.1}/src/kernels.egg-info/SOURCES.txt

@@ -13,4 +13,4 @@ src/kernels.egg-info/requires.txt
 src/kernels.egg-info/top_level.txt
 tests/test_basic.py
 tests/test_benchmarks.py
-tests/test_hash_validation.py
+tests/test_kernel_locking.py

kernels-0.2.0/tests/test_hash_validation.py → kernels-0.2.1/tests/test_kernel_locking.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+from kernels import load_kernel
 from kernels.cli import download_kernels
 
 
@@ -11,11 +12,13 @@ class DownloadArgs:
     project_dir: Path
 
 
-def test_download_hash_validation():
-    project_dir = Path(__file__).parent / "hash_validation"
-    download_kernels(DownloadArgs(all_variants=False, project_dir=project_dir))
-
-
 def test_download_all_hash_validation():
-    project_dir = Path(__file__).parent / "hash_validation"
+    project_dir = Path(__file__).parent / "kernel_locking"
     download_kernels(DownloadArgs(all_variants=True, project_dir=project_dir))
+
+
+def test_load_locked():
+    project_dir = Path(__file__).parent / "kernel_locking"
+    # Also validates that hashing works correctly.
+    download_kernels(DownloadArgs(all_variants=False, project_dir=project_dir))
+    load_kernel("kernels-community/activation", lockfile=project_dir / "kernels.lock")