flake.nix: migrate to pname+version

This commit is contained in:
Someone Serge 2023-12-25 17:02:36 +00:00
parent 286bb392a5
commit cc37631f7c
No known key found for this signature in database
GPG key ID: 7B0E3B1390D61DA4
3 changed files with 41 additions and 22 deletions

View file

@ -24,6 +24,7 @@
useMetalKit ? stdenv.isAarch64 && stdenv.isDarwin && !useOpenCL, useMetalKit ? stdenv.isAarch64 && stdenv.isDarwin && !useOpenCL,
useOpenCL ? false, useOpenCL ? false,
useRocm ? config.rocmSupport, useRocm ? config.rocmSupport,
llamaVersion ? "0.0.0", # Arbitrary version, substituted by the flake
}@inputs: }@inputs:
let let
@ -31,6 +32,7 @@ let
cmakeBool cmakeBool
cmakeFeature cmakeFeature
optionals optionals
strings
versionOlder versionOlder
; ;
@ -39,18 +41,19 @@ let
stdenv = throw "Use effectiveStdenv instead"; stdenv = throw "Use effectiveStdenv instead";
effectiveStdenv = if useCuda then cudaPackages.backendStdenv else inputs.stdenv; effectiveStdenv = if useCuda then cudaPackages.backendStdenv else inputs.stdenv;
# Give a little description difference between the flavors. suffices =
lib.optionals useOpenCL [ "OpenCL" ]
++ lib.optionals useCuda [ "CUDA" ]
++ lib.optionals useRocm [ "ROCm" ]
++ lib.optionals useMetalKit [ "MetalKit" ]
++ lib.optionals useBlas [ "BLAS" ];
pnameSuffix =
strings.optionalString (suffices != [ ])
"-${strings.concatMapStringsSep "-" strings.toLower suffices}";
descriptionSuffix = descriptionSuffix =
if useOpenCL then strings.optionalString (suffices != [ ])
" (OpenCL accelerated)" ", accelerated with ${strings.concatStringsSep ", " suffices}";
else if useCuda then
" (CUDA accelerated)"
else if useRocm then
" (ROCm accelerated)"
else if useMetalKit then
" (MetalKit accelerated)"
else
"";
# TODO: package the Python in this repository in a Nix-like way. # TODO: package the Python in this repository in a Nix-like way.
# It'd be nice to migrate to buildPythonPackage, as well as ensure this repo # It'd be nice to migrate to buildPythonPackage, as well as ensure this repo
@ -99,10 +102,12 @@ in
effectiveStdenv.mkDerivation ( effectiveStdenv.mkDerivation (
finalAttrs: { finalAttrs: {
name = "llama.cpp"; pname = "llama-cpp${pnameSuffix}";
version = llamaVersion;
src = ../../.; src = ../../.;
meta = { meta = {
description = "Inference of LLaMA model in pure C/C++${descriptionSuffix}"; description = "LLaMA model in pure C/C++${descriptionSuffix}";
mainProgram = "llama"; mainProgram = "llama";
}; };
@ -167,14 +172,14 @@ effectiveStdenv.mkDerivation (
# Define the shells here, but don't add in the inputsFrom to avoid recursion. # Define the shells here, but don't add in the inputsFrom to avoid recursion.
passthru = { passthru = {
shell = mkShell { shell = mkShell {
name = "default${descriptionSuffix}"; name = "shell-${finalAttrs.finalPackage.name}";
description = "contains numpy and sentencepiece"; description = "contains numpy and sentencepiece";
buildInputs = [ llama-python ]; buildInputs = [ llama-python ];
inputsFrom = [ finalAttrs.finalPackage ]; inputsFrom = [ finalAttrs.finalPackage ];
}; };
shell-extra = mkShell { shell-extra = mkShell {
name = "extra${descriptionSuffix}"; name = "shell-extra-${finalAttrs.finalPackage.name}";
description = "contains numpy, sentencepiece, torchWithoutCuda, and transformers"; description = "contains numpy, sentencepiece, torchWithoutCuda, and transformers";
buildInputs = [ llama-python-extra ]; buildInputs = [ llama-python-extra ];
inputsFrom = [ finalAttrs.finalPackage ]; inputsFrom = [ finalAttrs.finalPackage ];

View file

@ -1,3 +1,12 @@
{ lib, newScope }: {
lib,
newScope,
llamaVersion ? "0.0.0",
}:
lib.makeScope newScope (self: { llama-cpp = self.callPackage ./package.nix { }; }) lib.makeScope newScope (
self: {
inherit llamaVersion;
llama-cpp = self.callPackage ./package.nix { };
}
)

View file

@ -16,7 +16,10 @@
# { program = "/nix/store/00000000000000000000000000000000-llama.cpp/bin/quantize"; type = "app"; } # { program = "/nix/store/00000000000000000000000000000000-llama.cpp/bin/quantize"; type = "app"; }
# ``` # ```
outputs = outputs =
{ flake-parts, ... }@inputs: { self, flake-parts, ... }@inputs:
let
llamaVersion = self.dirtyShortRev;
in
flake-parts.lib.mkFlake { inherit inputs; } flake-parts.lib.mkFlake { inherit inputs; }
{ {
@ -48,7 +51,9 @@
# #
# Cf. https://nixos.org/manual/nix/unstable/command-ref/new-cli/nix3-flake.html?highlight=flake#flake-format # Cf. https://nixos.org/manual/nix/unstable/command-ref/new-cli/nix3-flake.html?highlight=flake#flake-format
flake.overlays.default = flake.overlays.default =
(final: prev: { llamaPackages = final.callPackage .devops/nix/scope.nix { }; }); (final: prev: {
llamaPackages = final.callPackage .devops/nix/scope.nix { inherit llamaVersion; };
});
systems = [ systems = [
"aarch64-darwin" "aarch64-darwin"
@ -69,10 +74,10 @@
# We don't use the overlay here so as to avoid making too many instances of nixpkgs, # We don't use the overlay here so as to avoid making too many instances of nixpkgs,
# cf. https://zimbatm.com/notes/1000-instances-of-nixpkgs # cf. https://zimbatm.com/notes/1000-instances-of-nixpkgs
packages = { packages = {
default = (pkgs.callPackage .devops/nix/scope.nix { }).llama-cpp; default = (pkgs.callPackage .devops/nix/scope.nix { inherit llamaVersion; }).llama-cpp;
opencl = config.packages.default.override { useOpenCL = true; }; opencl = config.packages.default.override { useOpenCL = true; };
cuda = (pkgsCuda.callPackage .devops/nix/scope.nix { }).llama-cpp; cuda = (pkgsCuda.callPackage .devops/nix/scope.nix { inherit llamaVersion; }).llama-cpp;
rocm = (pkgsRocm.callPackage .devops/nix/scope.nix { }).llama-cpp; rocm = (pkgsRocm.callPackage .devops/nix/scope.nix { inherit llamaVersion; }).llama-cpp;
}; };
}; };
}; };