diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index 30679defb..eb36c60f8 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -2,7 +2,7 @@ > [!TIP] > Refer to the [kernel requirements' page](../kernel-requirements.md) to -> get an idea of what is expected from the kernel structure and content. +> get an idea of what is expected from the kernel structure and content. ## Introduction @@ -320,7 +320,7 @@ The following options can be set for a kernel: supported backends are `cpu`, `cuda`, `metal`, `rocm`, and `xpu`. **The `cpu` backend is currently experimental and might still change.** - `depends` (required): a list of dependencies. The supported dependencies - are listed in [`deps.nix`](https://github.com/huggingface/kernels/blob/main/builder/lib/deps.nix). + are listed in [`cpp-deps.nix`](https://github.com/huggingface/kernels/blob/main/builder/lib/cpp-deps.nix). - `src` (required): a list of source files and headers. - `include` (optional): include directories relative to the project root. Default: `[]`. diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix index f6fd9f9b5..6a1cf836b 100644 --- a/examples/kernels/flake.nix +++ b/examples/kernels/flake.nix @@ -207,6 +207,11 @@ path = ./gemm-triton-autotune; drv = sys: out: out.packages.${sys}.redistributable.torch-cuda; } + { + name = "kernel-deps-kernel"; + path = ./kernel-deps; + drv = sys: out: out.packages.${sys}.redistributable.torch-cuda; + } ]; # ROCm kernels to build in CI. diff --git a/examples/kernels/kernel-deps/CARD.md b/examples/kernels/kernel-deps/CARD.md new file mode 100644 index 000000000..927ad93fd --- /dev/null +++ b/examples/kernels/kernel-deps/CARD.md @@ -0,0 +1,66 @@ +--- +library_name: kernels +{% if license %}license: {{ license }} +{% endif %}--- + +This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated. + +## How to use +{% if functions %} + +```python +# make sure `kernels` is installed: `pip install -U kernels` +from kernels import get_kernel + +# If the org / user isn't a trusted publisher, pass `trust_remote_code=True` to the +# `get_kernel` call. You can find whether this kernel is from a trusted publisher +# by going to the kernel's Hub page and finding the "Trusted publisher" status at +# the top of the page. +kernel_module = get_kernel("{{ repo_id }}", version={{ version }}) +{{ functions[0] }} = kernel_module.{{ functions[0] }} + +{{ functions[0] }}(...) +``` +{% else %} + +Usage example not available. +{% endif %} + +## Available functions +{% if functions %} +{% for func in functions %} +- `{{ func }}` +{% endfor %} +{% else %} + +Function list not available. +{% endif %} +{% if layers %} + +## Available layers +{% for layer in layers %} +- `{{ layer }}` +{% endfor %} +{% endif %} + +## Benchmarks +{% if has_benchmark %} + +Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`. +{% else %} + +No benchmark available yet. +{% endif %} +{% if upstream %} + +## Upstream + +The original source code for this kernel comes from {{ upstream }}. +{% endif %} +{% if source %} + +## Source + +The kernel-builder formatted source for this kernel is available at {{ source }}. +{% endif %} + diff --git a/examples/kernels/kernel-deps/build.toml b/examples/kernels/kernel-deps/build.toml new file mode 100644 index 000000000..3f072bbab --- /dev/null +++ b/examples/kernels/kernel-deps/build.toml @@ -0,0 +1,24 @@ +[general] +name = "kernel-deps" +version = 1 +edition = 5 +license = "apache-2.0" +backends = [ + "cpu", + "cuda", + "metal", + "rocm", + "xpu", +] + +kernel-depends = [ + { repo-id = "kernels-community/activation", version = 1 }, + { repo-id = "kernels-community/einops", version = 1 }, +] + +[general.hub] +repo-id = "kernels-test/kernel-deps" + +[torch-noarch] + +[kernel] diff --git a/examples/kernels/kernel-deps/flake.nix b/examples/kernels/kernel-deps/flake.nix new file mode 100644 index 000000000..29c0cc8f8 --- /dev/null +++ b/examples/kernels/kernel-deps/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for kernels tests"; + + inputs = { + kernel-builder.url = "path:../../.."; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/examples/kernels/kernel-deps/kernels.lock b/examples/kernels/kernel-deps/kernels.lock new file mode 100644 index 000000000..0d80241c3 --- /dev/null +++ b/examples/kernels/kernel-deps/kernels.lock @@ -0,0 +1,22 @@ +[ + { + "dependency": { + "repo-id": "kernels-community/activation", + "version": 1 + }, + "lock": { + "commit": "36f4b285ab9125127d1b7364a786089ef10bcf8f", + "hash": "sha256-S2bSgxoK5JbuyBnzZSXNcV7NmMGtccuqfsjPgbwki7o=" + } + }, + { + "dependency": { + "repo-id": "kernels-community/einops", + "version": 1 + }, + "lock": { + "commit": "94986c119039ea0cb65592c210c57255de1089eb", + "hash": "sha256-qeQZSajDKBw4gfjZEE6wreqvdSKGfZmkwi/KnkjPj5g=" + } + } +] diff --git a/examples/kernels/kernel-deps/tests/test_kernel_deps.py b/examples/kernels/kernel-deps/tests/test_kernel_deps.py new file mode 100644 index 000000000..fd7a6ab38 --- /dev/null +++ b/examples/kernels/kernel-deps/tests/test_kernel_deps.py @@ -0,0 +1,18 @@ +import torch +import kernels + +kernel_deps = kernels.get_kernel("kernels-test/kernel-deps", version=1) + + +def test_silu(): + torch.manual_seed(42) + x = torch.randn(4, 64, device="cuda") + torch.testing.assert_close(kernel_deps.silu(x), torch.nn.functional.silu(x)) + + +def test_swap_last_dimensions(): + torch.manual_seed(42) + x = torch.randn(3, 4, 5, 6) + y = kernel_deps.swap_last_dimensions(x) + assert y.shape == (3, 4, 6, 5) + torch.testing.assert_close(y, x.transpose(-1, -2)) diff --git a/examples/kernels/kernel-deps/torch-ext/kernel_deps/__init__.py b/examples/kernels/kernel-deps/torch-ext/kernel_deps/__init__.py new file mode 100644 index 000000000..6e630cfa1 --- /dev/null +++ b/examples/kernels/kernel-deps/torch-ext/kernel_deps/__init__.py @@ -0,0 +1,33 @@ +import kernels +import torch + + +activation = kernels.get_kernel_dep("kernels-community/activation") +einops = kernels.get_kernel_dep("kernels-community/einops") + + +def silu(x: torch.Tensor) -> torch.Tensor: + out = torch.empty_like(x) + activation.silu(out, x) + return out + + +def swap_last_dimensions(x: torch.Tensor) -> torch.Tensor: + """ + Swap the last two dimensions of a tensor. + + Shapes: + x: (..., M, N) with at least two dimensions + return: (..., N, M) + + Raises: + ValueError: if the tensor has fewer than two dimensions. + """ + if x.ndim < 2: + raise ValueError( + f"Expected a tensor with at least two dimensions, got {x.ndim}" + ) + return einops.rearrange(x, "... a b -> ... b a") + + +__all__ = ["swap_last_dimensions"] diff --git a/nix-builder/lib/build.nix b/nix-builder/lib/build.nix index 917a7cec4..4d2e450ee 100644 --- a/nix-builder/lib/build.nix +++ b/nix-builder/lib/build.nix @@ -136,6 +136,7 @@ rec { _: kernel: builtins.length (kernel.cuda-capabilities or supportedCudaCapabilities) ) kernelConfig.toml.kernel ); + kernelDeps = pkgs.fetchKernelDeps src; pythonDeps = (kernelConfig.toml.general.python-depends or [ ]); backendPythonDeps = lib.attrByPath [ buildConfig.backend "python-depends" ] [ ] @@ -150,6 +151,7 @@ rec { src rev doGetKernelCheck + kernelDeps pythonDeps backendPythonDeps kernelProvenance @@ -166,6 +168,7 @@ rec { src stripRPath rev + kernelDeps pythonDeps backendPythonDeps kernelProvenance @@ -184,6 +187,7 @@ rec { src stripRPath rev + kernelDeps pythonDeps backendPythonDeps kernelProvenance diff --git a/nix-builder/lib/cpp-deps.nix b/nix-builder/lib/cpp-deps.nix new file mode 100644 index 000000000..f5c38ec0e --- /dev/null +++ b/nix-builder/lib/cpp-deps.nix @@ -0,0 +1,47 @@ +{ + lib, + pkgs, + stdenv, + torch, +}: + +let + deps = { + "cutlass_2_10" = [ + pkgs.cutlass_2_10 + ]; + "cutlass_3_5" = [ + pkgs.cutlass_3_5 + ]; + "cutlass_3_6" = [ + pkgs.cutlass_3_6 + ]; + "cutlass_3_8" = [ + pkgs.cutlass_3_8 + ]; + "cutlass_3_9" = [ + pkgs.cutlass_3_9 + ]; + "cutlass_4_0" = [ + pkgs.cutlass_4_0 + ]; + "cutlass_4_4" = [ + pkgs.cutlass_4_4 + ]; + "cutlass_4_5" = [ + pkgs.cutlass_4_5 + ]; + "torch" = [ + torch + ]; + "sycl_tla" = [ + (torch.xpuPackages.sycl-tla.override { inherit stdenv; }) + ]; + "metal-cpp" = [ + (pkgs.metal-cpp.override { inherit stdenv; }).dev + ]; + }; + + getCppDep = dep: deps.${dep} or (throw "Unknown dependency: ${dep}"); +in +deps: lib.flatten (map getCppDep deps) diff --git a/nix-builder/lib/extension/default.nix b/nix-builder/lib/extension/default.nix index 286f6e023..d5be56c30 100644 --- a/nix-builder/lib/extension/default.nix +++ b/nix-builder/lib/extension/default.nix @@ -91,13 +91,12 @@ in mkTorchNoArchExtension = callPackage ./torch/no-arch.nix { inherit torch; }; - inherit - (import ../deps.nix { + resolveCppDeps = ( + import ../cpp-deps.nix { inherit lib pkgs torch; stdenv = effectiveStdenv; - }) - resolveCppDeps - ; + } + ); stdenv = effectiveStdenv; } diff --git a/nix-builder/lib/extension/torch/arch.nix b/nix-builder/lib/extension/torch/arch.nix index 004cc55a5..e28eae3dc 100644 --- a/nix-builder/lib/extension/torch/arch.nix +++ b/nix-builder/lib/extension/torch/arch.nix @@ -55,6 +55,10 @@ nvccThreads, + # Dependencies on other kernels. Path to a JSON file that maps + # kernel dependencies to output paths. + kernelDeps, + # A stringly-typed list of Python dependencies. Ideally we'd take a # list of derivations, but we also need to write the dependencies to # the output. @@ -86,7 +90,7 @@ assert (buildConfig.metal or false) -> stdenv.hostPlatform.isDarwin; let inherit - (import ../../deps.nix { + (import ../../python-deps.nix { inherit lib pkgs @@ -125,6 +129,7 @@ stdenv.mkDerivation (prevAttrs: { inherit doAbiCheck doKernelBuildCheck + kernelDeps moduleName nvccThreads ; diff --git a/nix-builder/lib/extension/torch/no-arch.nix b/nix-builder/lib/extension/torch/no-arch.nix index e0f2b132d..a1b1d36d7 100644 --- a/nix-builder/lib/extension/torch/no-arch.nix +++ b/nix-builder/lib/extension/torch/no-arch.nix @@ -35,6 +35,10 @@ src, + # Dependencies on other kernels. Path to a JSON file that maps + # kernel dependencies to output paths. + kernelDeps, + # A stringly-typed list of Python dependencies. Ideally we'd take a # list of derivations, but we also need to write the dependencies to # the output. @@ -55,7 +59,7 @@ assert (buildConfig.metal or false) -> stdenv.hostPlatform.isDarwin; let inherit - (import ../../deps.nix { + (import ../../python-deps.nix { inherit lib pkgs @@ -78,7 +82,7 @@ in stdenv.mkDerivation (prevAttrs: { name = "${kernelName}-torch-ext"; - inherit doKernelBuildCheck moduleName; + inherit doKernelBuildCheck kernelDeps moduleName; src = pkgs.runCommand "source" { } '' mkdir -p $out diff --git a/nix-builder/lib/extension/tvm-ffi/arch.nix b/nix-builder/lib/extension/tvm-ffi/arch.nix index 879076e88..bfdd06041 100644 --- a/nix-builder/lib/extension/tvm-ffi/arch.nix +++ b/nix-builder/lib/extension/tvm-ffi/arch.nix @@ -55,6 +55,10 @@ nvccThreads, + # Dependencies on other kernels. Path to a JSON file that maps + # kernel dependencies to output paths. + kernelDeps, + # A stringly-typed list of Python dependencies. Ideally we'd take a # list of derivations, but we also need to write the dependencies to # the output. @@ -83,7 +87,7 @@ assert (buildConfig.metal or false) -> stdenv.hostPlatform.isDarwin; let inherit - (import ../../deps.nix { + (import ../../python-deps.nix { inherit lib pkgs @@ -125,6 +129,7 @@ stdenv.mkDerivation (prevAttrs: { inherit doAbiCheck doKernelBuildCheck + kernelDeps moduleName nvccThreads ; diff --git a/nix-builder/lib/gen-flake-outputs.nix b/nix-builder/lib/gen-flake-outputs.nix index 335dd8250..53b05f7f9 100644 --- a/nix-builder/lib/gen-flake-outputs.nix +++ b/nix-builder/lib/gen-flake-outputs.nix @@ -241,6 +241,15 @@ in ${uploadStr} ''; + lock-kernel-deps = writeScriptBin "lock-kernel-deps" '' + #!/usr/bin/env bash + set -euo pipefail + + kernel_dir="''${1:-.}" + + ${pkgs.lock-kernel-deps}/bin/lock-kernel-deps "$kernel_dir" -o "$kernel_dir/kernels.lock" + ''; + ci = let setsWithFramework = diff --git a/nix-builder/lib/deps.nix b/nix-builder/lib/python-deps.nix similarity index 60% rename from nix-builder/lib/deps.nix rename to nix-builder/lib/python-deps.nix index af2be5db1..d04fb786b 100644 --- a/nix-builder/lib/deps.nix +++ b/nix-builder/lib/python-deps.nix @@ -6,42 +6,6 @@ }: let - cppDeps = { - "cutlass_2_10" = [ - pkgs.cutlass_2_10 - ]; - "cutlass_3_5" = [ - pkgs.cutlass_3_5 - ]; - "cutlass_3_6" = [ - pkgs.cutlass_3_6 - ]; - "cutlass_3_8" = [ - pkgs.cutlass_3_8 - ]; - "cutlass_3_9" = [ - pkgs.cutlass_3_9 - ]; - "cutlass_4_0" = [ - pkgs.cutlass_4_0 - ]; - "cutlass_4_4" = [ - pkgs.cutlass_4_4 - ]; - "cutlass_4_5" = [ - pkgs.cutlass_4_5 - ]; - "torch" = [ - torch - ]; - "sycl_tla" = [ - (torch.xpuPackages.sycl-tla.override { inherit stdenv; }) - ]; - "metal-cpp" = [ - (pkgs.metal-cpp.override { inherit stdenv; }).dev - ]; - }; - pythonDeps = let depsJson = builtins.fromJSON (builtins.readFile ../../kernels-data/src/python_dependencies.json); @@ -55,7 +19,6 @@ let backends = lib.mapAttrs updateBackend depsJson.backends; }; - getCppDep = dep: cppDeps.${dep} or (throw "Unknown dependency: ${dep}"); getPythonDep = dep: lib.attrByPath [ "general" dep "nix" ] (throw "Unknown Python dependency: ${dep}") pythonDeps; getBackendPythonDep = @@ -72,7 +35,6 @@ let ] (throw "Unknown Python dependency for backend `${backend}`: ${dep}") backendDeps; in { - resolveCppDeps = deps: lib.flatten (map getCppDep deps); resolvePythonDeps = deps: lib.flatten (map getPythonDep deps); resolveBackendPythonDeps = backend: deps: lib.flatten (map (getBackendPythonDep backend) deps); } diff --git a/nix-builder/lib/source-set.nix b/nix-builder/lib/source-set.nix index 0fda3ad8d..f9683d307 100644 --- a/nix-builder/lib/source-set.nix +++ b/nix-builder/lib/source-set.nix @@ -15,9 +15,12 @@ let "pyi" ]; pyFilter = file: builtins.any (ext: file.hasExt ext) pyExt; - extSrc = extConfig.src or [ ] ++ [ "build.toml" ]; + extSrc = extConfig.src or [ ] ++ [ + "build.toml" + ]; torchExtPath = path + "/torch-ext"; tvmFfiExtPath = path + "/tvm-ffi-ext"; + lockSet = fileset.maybeMissing (path + "/kernels.lock"); pySrcSet = let path = @@ -46,6 +49,7 @@ fileset.toSource { root = path; fileset = fileset.unions [ kernelsSrc + lockSet srcSet pySrcSet pyTestsSet diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix index 513003e16..3e4f5f689 100644 --- a/nix-builder/overlay.nix +++ b/nix-builder/overlay.nix @@ -14,12 +14,16 @@ final: prev: fetchFromHuggingFace = final.callPackage ./pkgs/fetch-from-huggingface { }; + fetchKernelDeps = final.callPackage ./pkgs/fetch-kernel-deps { }; + get-kernel-check = final.callPackage ./pkgs/get-kernel-check { }; hash-kernel-hook = final.callPackage ./pkgs/hash-kernel-hook { }; kernel-layout-check = final.callPackage ./pkgs/kernel-layout-check { }; + lock-kernel-deps = final.callPackage ./pkgs/lock-kernel-deps { }; + nvtx = final.callPackage ./pkgs/nvtx { }; metal-cpp = final.callPackage ./pkgs/metal-cpp { }; diff --git a/nix-builder/pkgs/fetch-kernel-deps/default.nix b/nix-builder/pkgs/fetch-kernel-deps/default.nix new file mode 100644 index 000000000..2ec5e72c8 --- /dev/null +++ b/nix-builder/pkgs/fetch-kernel-deps/default.nix @@ -0,0 +1,27 @@ +{ + fetchFromHuggingFace, + writeText, +}: + +# Path to the kernel (flake). +path: + +let + locks = + if builtins.pathExists "${path}/kernels.lock" then + builtins.fromJSON (builtins.readFile "${path}/kernels.lock") + else + [ ]; + + # Map kernel locks to Nix store paths by fetching the kernel. + paths = map (lock: { + dependency = lock.dependency; + path = fetchFromHuggingFace { + repoId = lock.dependency.repo-id; + type = "kernel"; + revision = lock.lock.commit; + inherit (lock.lock) hash; + }; + }) locks; +in +writeText "kernel-deps" (builtins.toJSON paths) diff --git a/nix-builder/pkgs/get-kernel-check/default.nix b/nix-builder/pkgs/get-kernel-check/default.nix index dc1d8eb98..6a5c05247 100644 --- a/nix-builder/pkgs/get-kernel-check/default.nix +++ b/nix-builder/pkgs/get-kernel-check/default.nix @@ -17,6 +17,7 @@ makeSetupHook { python3 = "${python3}/bin/python"; kernels = "${with python3.pkgs; makePythonPath [ kernels ]}"; proot = lib.optionalString useFakeSys "${proot}/bin/proot"; + pyhook = ./get-kernel-check-hook.py; useFakeSys = lib.optionalString useFakeSys "1"; }; } ./get-kernel-check-hook.sh diff --git a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py new file mode 100755 index 000000000..f93db419b --- /dev/null +++ b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.py @@ -0,0 +1,32 @@ +import os +from pathlib import Path + +from kernels.load import get_kernel_with_resolver +from kernels.hf_hub import _get_hf_api +from kernels.resolver import KernelPathsResolver, RepoPathsResolver, SequentialResolver +from kernels_data import KernelDependency, KernelPaths, KernelVersion + +out = os.getenv("out") +if not out: + raise ValueError("`out` environment variable is not set by Nix derivation") + +kernelDeps = os.getenv("kernelDeps") +if not kernelDeps: + raise ValueError("`kernelDeps` environment variable is not set by Nix derivation") + + +with open(kernelDeps) as f: + kernel_paths = KernelPaths.from_json(f.read()) + +kernel = KernelDependency(repo_id=out, version=KernelVersion.Version(0)) +resolvers = [ + RepoPathsResolver(local_kernels={out: Path(out)}), + KernelPathsResolver(kernel_paths=kernel_paths), +] + +get_kernel_with_resolver( + api=_get_hf_api(), + backend=None, + kernel=kernel, + resolver=SequentialResolver(resolvers=resolvers), +) diff --git a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh index 5d229f8ef..6b7659878 100755 --- a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh +++ b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh @@ -10,6 +10,11 @@ _getKernelCheckHook() { exit 1 fi + if [ -z ${kernelDeps+x} ]; then + echo "kernelDeps must be set in derivation" + exit 1 + fi + echo "Check whether the kernel can be loaded with get-kernel: ${moduleName}" # We strip the full library paths from the extension. Unfortunately, @@ -42,7 +47,7 @@ _getKernelCheckHook() { PYTHONPATH="@kernels@" \ ${prootCmd} \ - @python3@ -c "from pathlib import Path; import kernels; kernels.get_local_kernel(Path('${out}'))" + @python3@ @pyhook@ } postInstallCheckHooks+=(_getKernelCheckHook) diff --git a/nix-builder/pkgs/lock-kernel-deps/default.nix b/nix-builder/pkgs/lock-kernel-deps/default.nix new file mode 100644 index 000000000..6516acf55 --- /dev/null +++ b/nix-builder/pkgs/lock-kernel-deps/default.nix @@ -0,0 +1,15 @@ +{ + replaceVars, + runCommand, +}: + +let + lockKernelDeps = replaceVars ./lock_kernel_deps.py { + download = ../fetch-from-huggingface/download.py; + }; +in +runCommand "lock-kernel-deps" { } '' + mkdir -p $out/bin + install -Dm0755 ${lockKernelDeps} $out/bin/lock-kernel-deps + patchShebangs $out/bin/lock-kernel-deps +'' diff --git a/nix-builder/pkgs/lock-kernel-deps/lock_kernel_deps.py b/nix-builder/pkgs/lock-kernel-deps/lock_kernel_deps.py new file mode 100755 index 000000000..79d1068e0 --- /dev/null +++ b/nix-builder/pkgs/lock-kernel-deps/lock_kernel_deps.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Lock a kernel's dependencies and compute the Nix hashes. + +Reads the kernel's `build.toml`, resolves its dependencies to commits, and +outputs a lock file with a `hash` field added to every lock. The hash is the +SRI hash of a snapshot of the kernel repository, as used by the +`fetchFromHuggingFace` fixed-output derivation. +""" + +import argparse +import subprocess +import sys +import tempfile +from pathlib import Path + +from huggingface_hub.hf_api import HfApi +from kernels.locking import extract_dependency_locks +from kernels_data import Build, KernelLocks, NixKernelLock, NixKernelLocks + + +def nix_hash(repo_id: str, revision: str) -> str: + """Download a kernel repository and return the SRI hash of the snapshot.""" + with tempfile.TemporaryDirectory() as tmp_dir: + snapshot = Path(tmp_dir) / "snapshot" + subprocess.run( + [ + sys.executable, + str("@download@"), + repo_id, + "kernel", + revision, + str(snapshot), + ], + check=True, + ) + # `--type sha256 --sri` matches `outputHashMode = "recursive"`, which + # hashes the NAR serialization of the snapshot. + hash = subprocess.run( + ["nix", "hash", "path", "--type", "sha256", "--sri", str(snapshot)], + check=True, + capture_output=True, + text=True, + ) + + return hash.stdout.strip() + + +def lock_kernel_deps(kernel_dir: Path) -> KernelLocks: + """Resolve the dependencies of the kernel in `kernel_dir` to commits.""" + build_toml = kernel_dir / "build.toml" + if not build_toml.exists(): + raise FileNotFoundError(f"build.toml not found in {kernel_dir}") + + build = Build.open(kernel_dir) + api = HfApi() + + kernel_locks = {} + + for backend in build.general.backends: + for dep, lock in extract_dependency_locks( + build.all_kernel_depends(backend), + api=api, + ).items(): + kernel_locks[dep] = lock + + return KernelLocks(locks=kernel_locks) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "kernel_dir", + type=Path, + nargs="?", + default=Path("."), + help="kernel directory to lock dependencies for, defaults to the current directory", + ) + parser.add_argument( + "-o", + "--output", + type=Path, + help="file to write to, defaults to standard output", + ) + args = parser.parse_args() + + locks = lock_kernel_deps(args.kernel_dir) + + nix_locks = NixKernelLocks( + { + dep: NixKernelLock(lock.commit, nix_hash(dep.repo_id, lock.commit)) + for dep, lock in _report(locks.items()) + } + ) + + if args.output is None: + print(nix_locks.to_json()) + else: + args.output.write_text(nix_locks.to_json() + "\n") + + +def _report(items): + """Report progress on standard error, so that standard output stays clean.""" + for n, (dep, lock) in enumerate(items, start=1): + print( + f"[{n}/{len(items)}] Hashing {dep.repo_id} at {lock.commit}", + file=sys.stderr, + ) + yield dep, lock + + +if __name__ == "__main__": + main()