diff --git a/flake.nix b/flake.nix index 1d7430a16..b42ce72e7 100644 --- a/flake.nix +++ b/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}; }; diff --git a/nix-builder/lib/checks.nix b/nix-builder/lib/checks.nix index 1a30a4ae6..2cf6fd137 100644 --- a/nix-builder/lib/checks.nix +++ b/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 + '' diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix index 6ce964f36..cc2f10ff0 100644 --- a/nix-builder/overlay.nix +++ b/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 { }; diff --git a/nix-builder/pkgs/fetch-from-huggingface/default.nix b/nix-builder/pkgs/fetch-from-huggingface/default.nix new file mode 100644 index 000000000..0c5e08d14 --- /dev/null +++ b/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 `-`. + 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" ]; +} diff --git a/nix-builder/pkgs/fetch-from-huggingface/download.py b/nix-builder/pkgs/fetch-from-huggingface/download.py new file mode 100644 index 000000000..d13e9cc1b --- /dev/null +++ b/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)