kernels: switch over to the dependency solver

kernels/src/kernels/cli/__init__.py

@@ -1,18 +1,12 @@
 import argparse
-import dataclasses
-import json
 import sys
 from pathlib import Path
 
+from kernels.cli.download import download_kernels
 from kernels.cli.info import print_kernel_info
+from kernels.cli.lock import lock_kernels
 from kernels.cli.verify_signature import verify_signature
 from kernels.cli.versions import print_kernel_versions
-from kernels.compat import tomllib
-from kernels.install import (
-    install_kernel,
-    install_kernel_all_variants,
-)
-from kernels.locking import KernelLock, get_kernel_locks
 
 
 def main():
@@ -131,43 +125,6 @@ def main():
     args.func(args)
 
 
-def download_kernels(args):
-    lock_path = args.project_dir / "kernels.lock"
-
-    if not lock_path.exists():
-        print(f"No kernels.lock file found in: {args.project_dir}", file=sys.stderr)
-        sys.exit(1)
-
-    with open(args.project_dir / "kernels.lock", "r") as f:
-        lock_json = json.load(f)
-
-    all_successful = True
-
-    for kernel_lock_json in lock_json:
-        kernel_lock = KernelLock.from_json(kernel_lock_json)
-        print(
-            f"Downloading `{kernel_lock.repo_id}` with SHA: {kernel_lock.sha}",
-            file=sys.stderr,
-        )
-        if args.all_variants:
-            install_kernel_all_variants(
-                kernel_lock.repo_id,
-                revision=kernel_lock.sha,
-            )
-        else:
-            try:
-                install_kernel(
-                    kernel_lock.repo_id,
-                    revision=kernel_lock.sha,
-                )
-            except FileNotFoundError as e:
-                print(e, file=sys.stderr)
-                all_successful = False
-
-    if not all_successful:
-        sys.exit(1)
-
-
 def kernel_info(args):
     print_kernel_info(
         args.kernel,
@@ -181,27 +138,6 @@ def kernel_versions(args):
     print_kernel_versions(args.repo_id)
 
 
-def lock_kernels(args):
-    with open(args.project_dir / "pyproject.toml", "rb") as f:
-        data = tomllib.load(f)
-
-    kernel_versions = data.get("tool", {}).get("kernels", {}).get("dependencies", None)
-
-    all_locks = []
-    for kernel, version in kernel_versions.items():
-        all_locks.append(get_kernel_locks(kernel, version))
-
-    with open(args.project_dir / "kernels.lock", "w") as f:
-        json.dump(all_locks, f, cls=_JSONEncoder, indent=2)
-
-
-class _JSONEncoder(json.JSONEncoder):
-    def default(self, o):
-        if dataclasses.is_dataclass(o):
-            return dataclasses.asdict(o)
-        return super().default(o)
-
-
 def _check_moved(_args):
     print(
         "`kernels check` has moved to `kernel-builder check-abi`",

kernels/src/kernels/cli/download.py

@@ -0,0 +1,46 @@
+import sys
+
+from kernels_data import KernelLocks
+
+from kernels.hf_hub import CACHE_DIR, _get_hf_api
+from kernels.resolver import _BYTECODE_IGNORE_PATTERNS, resolve_hub_kernel
+
+
+def download_kernels(args):
+    lock_path = args.project_dir / "kernels.lock"
+
+    if not lock_path.exists():
+        print(f"No kernels.lock file found in: {args.project_dir}", file=sys.stderr)
+        sys.exit(1)
+
+    with open(args.project_dir / "kernels.lock", "r") as f:
+        kernel_locks = KernelLocks.from_json(f.read())
+
+    all_successful = True
+
+    api = _get_hf_api()
+    for dep, lock in kernel_locks.items():
+        print(
+            f"Downloading kernel: {dep.repo_id} (revision: {lock.commit})",
+            file=sys.stderr,
+        )
+
+        try:
+            if args.all_variants:
+                api.snapshot_download(
+                    dep.repo_id,
+                    repo_type="kernel",
+                    allow_patterns="build/*",
+                    ignore_patterns=_BYTECODE_IGNORE_PATTERNS,
+                    cache_dir=CACHE_DIR,
+                    revision=lock.commit,
+                )
+            else:
+                location = resolve_hub_kernel(dep.repo_id, api=api, backend=None, revision=lock.commit)
+                location.install(api=api)
+        except FileNotFoundError as e:
+            print(e, file=sys.stderr)
+            all_successful = False
+
+    if not all_successful:
+        sys.exit(1)

kernels/src/kernels/cli/lock.py

@@ -0,0 +1,24 @@
+from kernels_data import KernelDependency, KernelVersion
+
+from kernels.compat import tomllib
+from kernels.hf_hub import _get_hf_api
+from kernels.locking import extract_dependency_locks
+
+
+def lock_kernels(args):
+    project_dir = args.project_dir
+
+    with open(project_dir / "pyproject.toml", "rb") as f:
+        data = tomllib.load(f)
+
+    kernel_versions = data.get("tool", {}).get("kernels", {}).get("dependencies", None)
+
+    depends = [
+        KernelDependency(repo_id=kernel, version=KernelVersion.Version(version))
+        for kernel, version in kernel_versions.items()
+    ]
+
+    locks = extract_dependency_locks(depends, api=_get_hf_api())
+
+    with open(project_dir / "kernels.lock", "w") as f:
+        f.write(locks.to_json())

kernels/src/kernels/cli/verify_signature.py

@@ -16,7 +16,10 @@
 
 def verify_signature(args: argparse.Namespace) -> None:
     revision = select_revision_or_version(
-        args.repo_id, revision=None, version=args.version, local_files_only=constants.HF_HUB_OFFLINE
+        args.repo_id,
+        revision=None,
+        version=args.version,
+        local_files_only=constants.HF_HUB_OFFLINE,
     )
 
     if args.all_variants:

kernels/src/kernels/importer.py

@@ -116,7 +116,12 @@ def _import_from_path(
     if module is None:
         raise ImportError(f"Cannot load module {module_name} from spec")
     sys.modules[metadata.id] = module
-    spec.loader.exec_module(module)  # type: ignore
+
+    # Avoid an import cycle.
+    from kernels.deps import use_kernel_deps
+
+    with use_kernel_deps(deps):
+        spec.loader.exec_module(module)  # type: ignore
 
     _loaded_kernels[variant_path] = LoadedKernel(
         metadata=metadata,