fetchFromHuggingFace: add

flake.nix

@@ -211,6 +211,7 @@
       rec {
         checks.default = pkgs.callPackage ./nix-builder/lib/checks.nix {
           inherit buildSets self;
+          inherit (buildSet.pkgs) fetchFromHuggingFace;
           inherit (self.lib) genKernelFlakeOutputs;
           build = buildPerSystem.${system};
         };

nix-builder/lib/checks.nix

@@ -8,6 +8,7 @@
 
   build,
   buildSets,
+  fetchFromHuggingFace,
   genKernelFlakeOutputs,
 }:
 
@@ -28,11 +29,36 @@ let
       "Found Torch library registrations that do not use `add_op_namespace_prefix`:"
     ];
   };
+
+  fetchFromHuggingFaceCheck =
+    runCommand "fetch-from-huggingface-check"
+      {
+        relu = fetchFromHuggingFace {
+          repoId = "kernels-community/relu";
+          type = "kernel";
+          revision = "d649efb56fb249ac8f7a57fa1866728ad0c60e52";
+          hash = "sha256-1CM0MGEOCqqnYV989jcUANEbiB727JftjXTdQVUHPwk=";
+        };
+      }
+      ''
+        test -d "$relu/build"
+        # Download metadata must not leak into the output, it would make the
+        # hash unstable.
+        test ! -e "$relu/.cache"
+        touch $out
+      '';
 in
 assert lib.assertMsg (builtins.all (buildSet: buildSet.torch.version == "2.12.0") kernelBuildSets)
   ''
     Torch minver/maxver filtering does not work.
   '';
-runCommand "builder-nix-checks" { buildInputs = [ badRegistrationCheck ]; } ''
-  touch $out
-''
+runCommand "builder-nix-checks"
+  {
+    buildInputs = [
+      badRegistrationCheck
+      fetchFromHuggingFaceCheck
+    ];
+  }
+  ''
+    touch $out
+  ''

nix-builder/overlay.nix

@@ -12,6 +12,8 @@ final: prev:
 
   cmakeNvccThreadsHook = final.callPackage ./pkgs/cmake-nvcc-threads-hook { };
 
+  fetchFromHuggingFace = final.callPackage ./pkgs/fetch-from-huggingface { };
+
   get-kernel-check = final.callPackage ./pkgs/get-kernel-check { };
 
   hash-kernel-hook = final.callPackage ./pkgs/hash-kernel-hook { };

nix-builder/pkgs/fetch-from-huggingface/default.nix

@@ -0,0 +1,68 @@
+{
+  lib,
+  cacert,
+  python3,
+  stdenvNoCC,
+}:
+
+let
+  # `hf download` only accepts the `model`, `dataset` and `space` repository
+  # types, so we use `snapshot_download` (which the CLI wraps) directly.
+  python = python3.withPackages (ps: [ ps.huggingface-hub ]);
+in
+
+{
+  # Repository id, e.g. `kernels-community/relu`.
+  repoId,
+
+  # Repository type, e.g. `model`, `dataset` or `kernel`.
+  type ? "model",
+
+  # Tag or commit SHA.
+  revision,
+
+  # SRI hash of the repository snapshot.
+  hash,
+
+  # Derivation name, defaults to `<repoId>-<revision>`.
+  name ? null,
+}:
+
+stdenvNoCC.mkDerivation {
+  name = lib.strings.sanitizeDerivationName (if name != null then name else "${repoId}-${revision}");
+
+  nativeBuildInputs = [ python ];
+
+  dontUnpack = true;
+  dontConfigure = true;
+  dontBuild = true;
+  dontFixup = true;
+
+  installPhase = ''
+    runHook preInstall
+
+    # Downloading into a local directory bypasses the hub cache, but hf-xet
+    # still needs a writable home for its chunk cache.
+    export HOME="$NIX_BUILD_TOP"
+
+    # We use a custom script for now, since hf-cli does not support
+    # downloading the kernel repo type yet.
+    python ${./download.py} \
+      ${lib.escapeShellArg repoId} \
+      ${lib.escapeShellArg type} \
+      ${lib.escapeShellArg revision} \
+      "$out"
+
+    runHook postInstall
+  '';
+
+  SSL_CERT_FILE = "${cacert}/etc/ssl/certs/ca-bundle.crt";
+  HF_HUB_DISABLE_PROGRESS_BARS = "1";
+  HF_HUB_DISABLE_TELEMETRY = "1";
+
+  outputHash = hash;
+  outputHashAlgo = null;
+  outputHashMode = "recursive";
+
+  impureEnvVars = lib.fetchers.proxyImpureEnvVars ++ [ "HF_TOKEN" ];
+}

nix-builder/pkgs/fetch-from-huggingface/download.py

@@ -0,0 +1,18 @@
+import shutil
+import sys
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+repo_id, repo_type, revision, local_dir = sys.argv[1:]
+
+snapshot_download(
+    repo_id=repo_id,
+    repo_type=repo_type,
+    revision=revision,
+    local_dir=local_dir,
+)
+
+# Downloading leaves per-file metadata (etags, timestamps) behind, which would
+# make the hash of the snapshot unstable.
+shutil.rmtree(Path(local_dir) / ".cache", ignore_errors=True)