Skip to content

Commit f794cfd

Browse files
committed
flake.nix: use flake-parts
1 parent 8403bf2 commit f794cfd

File tree

5 files changed

+172
-119
lines changed

5 files changed

+172
-119
lines changed

.devops/nix/apps.nix

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,22 @@
1-
{ package, binaries }:
2-
3-
let
4-
default = builtins.elemAt binaries 0;
5-
mkApp = name: {
6-
${name} = {
7-
type = "app";
8-
program = "${package}/bin/${name}";
1+
{
2+
perSystem =
3+
{ config, lib, ... }:
4+
{
5+
apps =
6+
let
7+
inherit (config.packages) default;
8+
binaries = [
9+
"llama"
10+
"llama-embedding"
11+
"llama-server"
12+
"quantize"
13+
"train-text-from-scratch"
14+
];
15+
mkApp = name: {
16+
type = "app";
17+
program = "${default}/bin/${name}";
18+
};
19+
in
20+
lib.genAttrs binaries mkApp;
921
};
10-
};
11-
result = builtins.foldl' (acc: name: (mkApp name) // acc) { } binaries;
12-
in
13-
14-
result // { default = result.${default}; }
22+
}

.devops/nix/devshells.nix

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,13 @@
1-
{ concatMapAttrs, packages }:
2-
3-
concatMapAttrs
4-
(name: package: {
5-
${name} = package.passthru.shell;
6-
${name + "-extra"} = package.passthru.shell-extra;
7-
})
8-
packages
1+
{
2+
perSystem =
3+
{ config, lib, ... }:
4+
{
5+
devShells =
6+
lib.concatMapAttrs
7+
(name: package: {
8+
${name} = package.passthru.shell;
9+
${name + "-extra"} = package.passthru.shell-extra;
10+
})
11+
config.packages;
12+
};
13+
}

.devops/nix/nixpkgs-instances.nix

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
{ inputs, ... }:
2+
{
3+
# The _module.args definitions are passed on to modules as arguments. E.g.
4+
# the module `{ pkgs ... }: { /* config */ }` implicitly uses
5+
# `_module.args.pkgs` (defined in this case by flake-parts).
6+
perSystem =
7+
{ system, ... }:
8+
{
9+
_module.args = {
10+
pkgsCuda = import inputs.nixpkgs {
11+
inherit system;
12+
# Ensure dependencies use CUDA consistently (e.g. that openmpi, ucc,
13+
# and ucx are built with CUDA support)
14+
config.cudaSupport = true;
15+
config.allowUnfreePredicate =
16+
p:
17+
builtins.all
18+
(
19+
license:
20+
license.free
21+
|| builtins.elem license.shortName [
22+
"CUDA EULA"
23+
"cuDNN EULA"
24+
]
25+
)
26+
(p.meta.licenses or [ p.meta.license ]);
27+
};
28+
# Ensure dependencies use ROCm consistently
29+
pkgsRocm = import inputs.nixpkgs {
30+
inherit system;
31+
config.rocmSupport = true;
32+
};
33+
};
34+
};
35+
}

flake.lock

Lines changed: 37 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

flake.nix

Lines changed: 66 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -1,111 +1,79 @@
11
{
2+
description = "Port of Facebook's LLaMA model in C/C++";
3+
24
inputs = {
35
nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
6+
flake-parts.url = "github:hercules-ci/flake-parts";
47
};
58

9+
# For inspection, use `nix flake show github:ggerganov/llama.cpp` or the nix repl:
10+
#
11+
# ```bash
12+
# ❯ nix repl
13+
# nix-repl> :lf github:ggerganov/llama.cpp
14+
# Added 13 variables.
15+
# nix-repl> outputs.apps.x86_64-linux.quantize
16+
# { program = "/nix/store/00000000000000000000000000000000-llama.cpp/bin/quantize"; type = "app"; }
17+
# ```
618
outputs =
7-
{ self, nixpkgs }:
8-
9-
let
10-
systems = [
11-
"aarch64-darwin"
12-
"aarch64-linux"
13-
"x86_64-darwin" # x86_64-darwin isn't tested (and likely isn't relevant)
14-
"x86_64-linux"
15-
];
16-
eachSystem = f: nixpkgs.lib.genAttrs systems (system: f system);
17-
in
19+
{ flake-parts, ... }@inputs:
20+
flake-parts.lib.mkFlake { inherit inputs; }
1821

19-
{
20-
# An overlay can be used to have a more granular control over llama-cpp's
21-
# dependencies and configuration, than that offered by the `.override`
22-
# mechanism. Cf. https://nixos.org/manual/nixpkgs/stable/#chap-overlays.
23-
#
24-
# E.g. in a flake:
25-
# ```
26-
# { nixpkgs, llama-cpp, ... }:
27-
# let pkgs = import nixpkgs {
28-
# overlays = [ (llama-cpp.overlays.default) ];
29-
# system = "aarch64-linux";
30-
# config.allowUnfree = true;
31-
# config.cudaSupport = true;
32-
# config.cudaCapabilities = [ "7.2" ];
33-
# config.cudaEnableForwardCompat = false;
34-
# }; in {
35-
# packages.aarch64-linux.llamaJetsonXavier = pkgs.llamaPackages.llama-cpp;
36-
# }
37-
# ```
38-
#
39-
# Cf. https://nixos.org/manual/nix/unstable/command-ref/new-cli/nix3-flake.html?highlight=flake#flake-format
40-
overlays.default = (final: prev: { llamaPackages = final.callPackage .devops/nix/scope.nix { }; });
22+
{
4123

42-
# These use the package definition from `./.devops/nix/package.nix`.
43-
# There's one per backend that llama-cpp uses. Add more as needed!
44-
packages = eachSystem (
45-
system:
46-
let
47-
# Avoid re-evaluation for the nixpkgs instance,
48-
# cf. https://zimbatm.com/notes/1000-instances-of-nixpkgs
49-
pkgs = nixpkgs.legacyPackages.${system};
24+
imports = [
25+
.devops/nix/nixpkgs-instances.nix
26+
.devops/nix/apps.nix
27+
.devops/nix/devshells.nix
28+
];
5029

51-
# Ensure dependencies use CUDA consistently (e.g. that openmpi, ucc,
52-
# and ucx are built with CUDA support)
53-
pkgsCuda = import nixpkgs {
54-
inherit system;
30+
# An overlay can be used to have a more granular control over llama-cpp's
31+
# dependencies and configuration, than that offered by the `.override`
32+
# mechanism. Cf. https://nixos.org/manual/nixpkgs/stable/#chap-overlays.
33+
#
34+
# E.g. in a flake:
35+
# ```
36+
# { nixpkgs, llama-cpp, ... }:
37+
# let pkgs = import nixpkgs {
38+
# overlays = [ (llama-cpp.overlays.default) ];
39+
# system = "aarch64-linux";
40+
# config.allowUnfree = true;
41+
# config.cudaSupport = true;
42+
# config.cudaCapabilities = [ "7.2" ];
43+
# config.cudaEnableForwardCompat = false;
44+
# }; in {
45+
# packages.aarch64-linux.llamaJetsonXavier = pkgs.llamaPackages.llama-cpp;
46+
# }
47+
# ```
48+
#
49+
# Cf. https://nixos.org/manual/nix/unstable/command-ref/new-cli/nix3-flake.html?highlight=flake#flake-format
50+
flake.overlays.default =
51+
(final: prev: { llamaPackages = final.callPackage .devops/nix/scope.nix { }; });
5552

56-
config.cudaSupport = true;
57-
config.allowUnfreePredicate =
58-
p:
59-
builtins.all
60-
(
61-
license:
62-
license.free
63-
|| builtins.elem license.shortName [
64-
"CUDA EULA"
65-
"cuDNN EULA"
66-
]
67-
)
68-
(p.meta.licenses or [ p.meta.license ]);
69-
};
53+
systems = [
54+
"aarch64-darwin"
55+
"aarch64-linux"
56+
"x86_64-darwin" # x86_64-darwin isn't tested (and likely isn't relevant)
57+
"x86_64-linux"
58+
];
7059

71-
# Ensure dependencies use ROCm consistently
72-
pkgsRocm = import nixpkgs {
73-
inherit system;
74-
config.rocmSupport = true;
60+
perSystem =
61+
{
62+
config,
63+
pkgs,
64+
pkgsCuda,
65+
pkgsRocm,
66+
...
67+
}:
68+
{
69+
# We don't use the overlay here so as to avoid making too many instances of nixpkgs,
70+
# cf. https://zimbatm.com/notes/1000-instances-of-nixpkgs
71+
packages = {
72+
default = (pkgs.callPackage .devops/nix/scope.nix { }).llama-cpp;
73+
opencl = config.packages.default.override { useOpenCL = true; };
74+
cuda = (pkgsCuda.callPackage .devops/nix/scope.nix { }).llama-cpp;
75+
rocm = (pkgsRocm.callPackage .devops/nix/scope.nix { }).llama-cpp;
76+
};
7577
};
76-
in
77-
{
78-
default = (pkgs.callPackage .devops/nix/scope.nix { }).llama-cpp;
79-
opencl = self.packages.${system}.default.override { useOpenCL = true; };
80-
cuda = (pkgsCuda.callPackage .devops/nix/scope.nix { }).llama-cpp;
81-
rocm = (pkgsRocm.callPackage .devops/nix/scope.nix { }).llama-cpp;
82-
}
83-
);
84-
85-
# These use the definition of llama-cpp from `./.devops/nix/package.nix`
86-
# and expose various binaries as apps with `nix run .#app-name`.
87-
# Note that none of these apps use anything other than the default backend.
88-
apps = eachSystem (
89-
system:
90-
import ./.devops/nix/apps.nix {
91-
package = self.packages.${system}.default;
92-
binaries = [
93-
"llama"
94-
"llama-embedding"
95-
"llama-server"
96-
"quantize"
97-
"train-text-from-scratch"
98-
];
99-
}
100-
);
101-
102-
# These expose a build environment for either a "default" or an "extra" set of dependencies.
103-
devShells = eachSystem (
104-
system:
105-
import ./.devops/nix/devshells.nix {
106-
concatMapAttrs = nixpkgs.lib.concatMapAttrs;
107-
packages = self.packages.${system};
108-
}
109-
);
110-
};
78+
};
11179
}

0 commit comments

Comments
 (0)