Kernel loading untangling

kernels/src/kernels/cli/__init__.py

@@ -8,7 +8,7 @@
 from kernels.cli.verify_signature import verify_signature
 from kernels.cli.versions import print_kernel_versions
 from kernels.compat import tomllib
-from kernels.lockfile import KernelLock, get_kernel_locks
+from kernels.locking import KernelLock, get_kernel_locks
 from kernels.utils import (
     install_kernel,
     install_kernel_all_variants,
@@ -153,14 +153,12 @@ def download_kernels(args):
             install_kernel_all_variants(
                 kernel_lock.repo_id,
                 revision=kernel_lock.sha,
-                variant_locks=kernel_lock.variants,
             )
         else:
             try:
                 install_kernel(
                     kernel_lock.repo_id,
                     revision=kernel_lock.sha,
-                    variant_locks=kernel_lock.variants,
                 )
             except FileNotFoundError as e:
                 print(e, file=sys.stderr)

kernels/src/kernels/layer/func.py

@@ -5,9 +5,11 @@
 from typing import TYPE_CHECKING, Protocol, Type
 
 from .._versions import select_revision_or_version
+from ..locking import (
+    get_caller_locked_kernel_revision,
+    get_locked_kernel_revision,
+)
 from ..utils import (
-    _get_caller_locked_kernel,
-    _get_locked_kernel,
     get_kernel,
     get_local_kernel,
 )
@@ -298,10 +300,10 @@ def __init__(
 
     def _resolve_revision(self) -> str:
         if self._lockfile is None:
-            locked_sha = _get_caller_locked_kernel(self._repo_id)
+            locked_sha = get_caller_locked_kernel_revision(self._repo_id)
         else:
             with open(self._lockfile, "r") as f:
-                locked_sha = _get_locked_kernel(self._repo_id, f.read())
+                locked_sha = get_locked_kernel_revision(self._repo_id, f.read())
 
         if locked_sha is None:
             raise ValueError(f"Kernel `{self._repo_id}` is not locked")

kernels/src/kernels/layer/layer.py

@@ -10,9 +10,11 @@
 from typing import TYPE_CHECKING, Callable, Protocol, Type
 
 from .._versions import select_revision_or_version
+from ..locking import (
+    get_caller_locked_kernel_revision,
+    get_locked_kernel_revision,
+)
 from ..utils import (
-    _get_caller_locked_kernel,
-    _get_locked_kernel,
     get_kernel,
     get_local_kernel,
 )
@@ -207,10 +209,10 @@ def __init__(
 
     def _resolve_revision(self) -> str:
         if self._lockfile is None:
-            locked_sha = _get_caller_locked_kernel(self._repo_id)
+            locked_sha = get_caller_locked_kernel_revision(self._repo_id)
         else:
             with open(self._lockfile, "r") as f:
-                locked_sha = _get_locked_kernel(self._repo_id, f.read())
+                locked_sha = get_locked_kernel_revision(self._repo_id, f.read())
 
         if locked_sha is None:
             raise ValueError(f"Kernel `{self._repo_id}` is not locked")

kernels/src/kernels/lockfile.py → kernels/src/kernels/locking.py

@@ -1,6 +1,11 @@
 import hashlib
+import importlib.metadata
+import inspect
+import json
 from dataclasses import dataclass
+from importlib.metadata import Distribution
 from pathlib import Path
+from types import ModuleType
 
 from huggingface_hub.dataclasses import strict
 from huggingface_hub.hf_api import RepoFile
@@ -129,3 +134,47 @@ def write_egg_lockfile(cmd, basename, filename):
         data = open(lock_path, "r").read()
 
     cmd.write_or_delete_file(basename, filename, data)
+
+
+def get_locked_kernel_revision(repo_id: str, lock_json: str) -> str | 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
+    return None
+
+
+def get_caller_locked_kernel_revision(repo_id: str) -> str | None:
+    for dist in _get_caller_distributions():
+        lock_json = dist.read_text("kernels.lock")
+        if lock_json is None:
+            continue
+        locked_sha = get_locked_kernel_revision(repo_id, lock_json)
+        if locked_sha is not None:
+            return locked_sha
+    return None
+
+
+def _get_caller_distributions() -> list[Distribution]:
+    module = _get_caller_module()
+    if module is None:
+        return []
+
+    # Look up all possible distributions that this module could be from.
+    package = module.__name__.split(".")[0]
+    dist_names = importlib.metadata.packages_distributions().get(package)
+    if dist_names is None:
+        return []
+
+    return [importlib.metadata.distribution(dist_name) for dist_name in dist_names]
+
+
+def _get_caller_module() -> ModuleType | None:
+    stack = inspect.stack()
+    # Get first module in the stack that is not the current module.
+    first_module = inspect.getmodule(stack[0][0])
+    for frame in stack[1:]:
+        module = inspect.getmodule(frame[0])
+        if module is not None and module != first_module:
+            return module
+    return first_module