413 lines
13 KiB
Nix
413 lines
13 KiB
Nix
{
|
|
inputs = {
|
|
nixpkgs.url = "github:nixos/nixpkgs?ref=nixos-unstable";
|
|
|
|
pyproject-nix = {
|
|
url = "github:pyproject-nix/pyproject.nix";
|
|
inputs.nixpkgs.follows = "nixpkgs";
|
|
};
|
|
|
|
uv2nix = {
|
|
url = "github:pyproject-nix/uv2nix";
|
|
inputs.nixpkgs.follows = "nixpkgs";
|
|
inputs.pyproject-nix.follows = "pyproject-nix";
|
|
};
|
|
|
|
pyproject-build-systems = {
|
|
url = "github:pyproject-nix/build-system-pkgs";
|
|
inputs.pyproject-nix.follows = "pyproject-nix";
|
|
inputs.uv2nix.follows = "uv2nix";
|
|
inputs.nixpkgs.follows = "nixpkgs";
|
|
};
|
|
|
|
treefmt-nix.url = "github:numtide/treefmt-nix";
|
|
};
|
|
|
|
outputs =
|
|
{
|
|
self,
|
|
nixpkgs,
|
|
pyproject-nix,
|
|
uv2nix,
|
|
pyproject-build-systems,
|
|
treefmt-nix,
|
|
}:
|
|
let
|
|
supportedSystems = [ "x86_64-linux" ];
|
|
|
|
pyproject = builtins.fromTOML (builtins.readFile ./pyproject.toml);
|
|
|
|
forEachSupportedSystem =
|
|
f:
|
|
nixpkgs.lib.genAttrs supportedSystems (
|
|
system:
|
|
let
|
|
pkgs = import nixpkgs {
|
|
inherit system;
|
|
config = {
|
|
allowUnfree = true;
|
|
};
|
|
};
|
|
|
|
python = pkgs.python313;
|
|
|
|
workspace = uv2nix.lib.workspace.loadWorkspace { workspaceRoot = ./.; };
|
|
|
|
overlay = workspace.mkPyprojectOverlay {
|
|
sourcePreference = "wheel";
|
|
};
|
|
|
|
# some packages don't have their dependencies declared properly, or need native libraries from other packages
|
|
fixupLibsOverlay =
|
|
final: prev:
|
|
let
|
|
packages = {
|
|
llama-cpp-binaries = {
|
|
buildInputs = [
|
|
pkgs.openssl
|
|
];
|
|
|
|
postInstall = ''
|
|
pushd "$out/${python.sitePackages}/llama_cpp_binaries/bin"
|
|
|
|
for f in ./*.so.0*; do
|
|
dest=$(echo $f | sed -E 's/(.*\.so\.[0-9])[.0-9]+/\1/g')
|
|
|
|
if [ ! -f $dest ]; then
|
|
echo "linking $f to $dest"
|
|
ln -s $f $dest
|
|
fi
|
|
done
|
|
|
|
popd
|
|
'';
|
|
|
|
preFixup = ''
|
|
addAutoPatchelfSearchPath $out/${python.sitePackages}/llama_cpp_binaries/bin
|
|
'';
|
|
};
|
|
flash-attn = {
|
|
buildInputs = [
|
|
final.torch
|
|
final.nvidia-cuda-runtime-cu12
|
|
];
|
|
};
|
|
exllamav2 = {
|
|
buildInputs = [
|
|
final.torch
|
|
final.nvidia-cuda-runtime-cu12
|
|
];
|
|
};
|
|
exllamav3 = {
|
|
buildInputs = [
|
|
final.torch
|
|
final.nvidia-cuda-runtime-cu12
|
|
];
|
|
};
|
|
markupsafe = {
|
|
buildInputs = [
|
|
final.setuptools
|
|
];
|
|
|
|
dontAutoPatchelf = true;
|
|
};
|
|
torch = {
|
|
buildInputs = [
|
|
final.nvidia-cusparse-cu12
|
|
final.nvidia-cusparselt-cu12
|
|
final.nvidia-cublas-cu12
|
|
final.nvidia-cufile-cu12
|
|
final.nvidia-cusolver-cu12
|
|
final.nvidia-cufft-cu12
|
|
final.nvidia-cuda-runtime-cu12
|
|
final.nvidia-cudnn-cu12
|
|
final.nvidia-cuda-nvrtc-cu12
|
|
final.nvidia-cuda-cupti-cu12
|
|
final.nvidia-curand-cu12
|
|
final.nvidia-nccl-cu12
|
|
];
|
|
};
|
|
bitsandbytes = {
|
|
buildInputs = [
|
|
final.nvidia-cublas-cu12
|
|
final.nvidia-cusparse-cu12
|
|
final.nvidia-nvjitlink-cu12
|
|
final.nvidia-cuda-runtime-cu12
|
|
];
|
|
ignoreDeps = [
|
|
# amd
|
|
"libhip*"
|
|
|
|
# unnecessary cuda versions
|
|
"libcu*.so.11"
|
|
"libcu*.so.13"
|
|
"libnvJitLink*.11"
|
|
"libnvJitLink*.13"
|
|
|
|
# intel
|
|
"libimf*"
|
|
"libintlc*"
|
|
"libirng*"
|
|
"libsvml*"
|
|
"libsycl*"
|
|
];
|
|
};
|
|
torchao = {
|
|
buildInputs = [
|
|
final.torch
|
|
final.nvidia-cuda-runtime-cu12
|
|
];
|
|
};
|
|
nvidia-cufile-cu12 = { };
|
|
nvidia-cusolver-cu12 = {
|
|
buildInputs = [
|
|
final.nvidia-cublas-cu12
|
|
final.nvidia-cusparse-cu12
|
|
final.nvidia-nvjitlink-cu12
|
|
];
|
|
};
|
|
nvidia-cusparse-cu12 = {
|
|
buildInputs = [
|
|
final.nvidia-nvjitlink-cu12
|
|
];
|
|
};
|
|
nvidia-nvshmem-cu12 = {
|
|
nativeBuildInputs = [
|
|
pkgs.mpi
|
|
];
|
|
};
|
|
};
|
|
|
|
ignoreDeps = [
|
|
"libmlx5.so.1"
|
|
"librdmacm.so.1"
|
|
"libibverbs.so.1"
|
|
|
|
# don't know
|
|
"libcuda.so.1"
|
|
|
|
# old cuda versions
|
|
"libcudart.so.11.0"
|
|
"libcudart.so.11"
|
|
"libcublas.so.11"
|
|
"libcublasLt.so.11"
|
|
"libcusparse.so.11"
|
|
];
|
|
in
|
|
builtins.listToAttrs (
|
|
map (name: {
|
|
inherit name;
|
|
value = prev.${name}.overrideAttrs (old: {
|
|
buildInputs = (old.buildInputs or [ ]) ++ (packages.${name}.buildInputs or [ ]);
|
|
propagatedBuildInputs =
|
|
(old.propagatedBuildInputs or [ ]) ++ (packages.${name}.propagatedBuildInputs or [ ]);
|
|
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ (packages.${name}.nativeBuildInputs or [ ]);
|
|
|
|
autoPatchelfIgnoreMissingDeps = ignoreDeps ++ (packages.${name}.ignoreDeps or [ ]);
|
|
|
|
postInstall = (old.postInstall or "") + (packages.${name}.postInstall or "");
|
|
|
|
preFixup =
|
|
(old.preFixup or "")
|
|
+ (packages.${name}.preFixup or "")
|
|
+ (pkgs.lib.optionalString
|
|
(!(packages.${name} ? dontAutoPatchelf) || (packages.${name}.dontAutoPatchelf == false))
|
|
(
|
|
builtins.concatStringsSep "\n" (
|
|
map (e: "addAutoPatchelfSearchPath \"${e}\"") (packages.${name}.buildInputs or [ ])
|
|
)
|
|
)
|
|
);
|
|
});
|
|
}) (builtins.attrNames packages)
|
|
);
|
|
|
|
pythonSet = (pkgs.callPackage pyproject-nix.build.packages { inherit python; }).overrideScope (
|
|
pkgs.lib.composeManyExtensions [
|
|
pyproject-build-systems.overlays.default
|
|
overlay
|
|
fixupLibsOverlay
|
|
(final: prev: {
|
|
text-generation-webui = prev.text-generation-webui.overrideAttrs (old: {
|
|
postInstall = (old.postInstall or "") + ''
|
|
cp -r ${./css} $out/${python.sitePackages}/css
|
|
cp -r ${./js} $out/${python.sitePackages}/js
|
|
'';
|
|
|
|
meta = {
|
|
mainProgram = "server.py";
|
|
};
|
|
});
|
|
})
|
|
]
|
|
);
|
|
|
|
treefmtEval = treefmt-nix.lib.evalModule pkgs ./treefmt.nix;
|
|
|
|
treefmt = treefmtEval.config.build.wrapper;
|
|
in
|
|
f {
|
|
inherit
|
|
fixupLibsOverlay
|
|
overlay
|
|
pkgs
|
|
python
|
|
pythonSet
|
|
system
|
|
treefmt
|
|
treefmtEval
|
|
workspace
|
|
;
|
|
}
|
|
);
|
|
in
|
|
{
|
|
formatter = forEachSupportedSystem ({ treefmt, ... }: treefmt);
|
|
|
|
packages = forEachSupportedSystem (
|
|
{
|
|
pkgs,
|
|
pythonSet,
|
|
workspace,
|
|
...
|
|
}:
|
|
{
|
|
default =
|
|
let
|
|
inherit (pkgs.callPackages pyproject-nix.build.util { }) mkApplication;
|
|
in
|
|
mkApplication {
|
|
venv = pythonSet.mkVirtualEnv "text-generation-webui-env" workspace.deps.default;
|
|
package = pythonSet.text-generation-webui;
|
|
};
|
|
}
|
|
);
|
|
|
|
devShells = forEachSupportedSystem (
|
|
{
|
|
fixupLibsOverlay,
|
|
pkgs,
|
|
python,
|
|
pythonSet,
|
|
system,
|
|
treefmt,
|
|
workspace,
|
|
...
|
|
}:
|
|
{
|
|
default = self.devShells.${system}.impure;
|
|
|
|
# shell for managing uv / python deps imperatively
|
|
impure = pkgs.mkShell {
|
|
packages = [
|
|
pkgs.uv
|
|
python
|
|
treefmt
|
|
];
|
|
|
|
env = {
|
|
# Let Nix manage Python
|
|
UV_PYTHON_DOWNLOADS = "never";
|
|
UV_PYTHON = python.interpreter;
|
|
|
|
# Python libraries often load native shared objects using dlopen(3).
|
|
# Setting LD_LIBRARY_PATH makes the dynamic library loader aware of libraries without using RPATH for lookup.
|
|
LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath pkgs.pythonManylinuxPackages.manylinux1;
|
|
};
|
|
|
|
shellHook = ''
|
|
# Needed for Pytorch to find cuda libraries
|
|
export LD_LIBRARY_PATH=/run/opengl-driver/lib:$LD_LIBRARY_PATH
|
|
|
|
# Undo dependency propagation by nixpkgs.
|
|
unset PYTHONPATH
|
|
|
|
if [ -f ./.venv/bin/activate ]; then
|
|
source ./.venv/bin/activate
|
|
fi
|
|
'';
|
|
};
|
|
|
|
# fully declarative development shell with all deps managed by nix
|
|
uv2nix =
|
|
let
|
|
editableOverlay = workspace.mkEditablePyprojectOverlay {
|
|
root = "$REPO_ROOT";
|
|
};
|
|
|
|
editablePythonSet = pythonSet.overrideScope (
|
|
pkgs.lib.composeManyExtensions [
|
|
editableOverlay
|
|
fixupLibsOverlay
|
|
(final: prev: {
|
|
${pyproject.project.name} = prev.${pyproject.project.name}.overrideAttrs (old: {
|
|
src = pkgs.lib.fileset.toSource {
|
|
root = old.src;
|
|
fileset = pkgs.lib.fileset.unions [
|
|
(old.src + "/pyproject.toml")
|
|
];
|
|
};
|
|
|
|
nativeBuildInputs = old.nativeBuildInputs ++ (final.resolveBuildSystem { editables = [ ]; });
|
|
});
|
|
})
|
|
]
|
|
);
|
|
|
|
virtualenv = editablePythonSet.mkVirtualEnv "${pyproject.project.name}-dev-env" workspace.deps.all;
|
|
in
|
|
pkgs.mkShell {
|
|
packages = [
|
|
pkgs.uv
|
|
treefmt
|
|
virtualenv
|
|
];
|
|
|
|
env = {
|
|
# Let Nix manage all things Python
|
|
UV_NO_SYNC = "1";
|
|
UV_PYTHON = python.interpreter;
|
|
UV_PYTHON_DOWNLOADS = "never";
|
|
TRITON_LIBCUDA_PATH = "${pkgs.lib.getLib pkgs.cudaPackages.cudatoolkit}/lib/stubs";
|
|
CUDA_HOME = "${pkgs.cudaPackages.cudatoolkit}";
|
|
CUDA_DIR = "${pkgs.cudaPackages.cudatoolkit}";
|
|
CUDATookit_ROOT = "${pkgs.cudaPackages.cudatoolkit}";
|
|
|
|
LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath (
|
|
pkgs.pythonManylinuxPackages.manylinux1
|
|
++ [
|
|
pkgs.cudaPackages.cudatoolkit
|
|
]
|
|
);
|
|
};
|
|
|
|
shellHook = ''
|
|
# Undo dependency propagation by nixpkgs.
|
|
unset PYTHONPATH
|
|
|
|
# Get repository root using git. This is expanded at runtime by the editable `.pth` machinery.
|
|
export REPO_ROOT=$(git rev-parse --show-toplevel)
|
|
|
|
# make sure all libraries are findable by python
|
|
for f in $(find -L "${virtualenv}/${python.sitePackages}" -type d -name "lib"); do
|
|
export LD_LIBRARY_PATH="$f:$LD_LIBRARY_PATH"
|
|
done
|
|
|
|
# Needed for Pytorch to find cuda libraries
|
|
export LD_LIBRARY_PATH=/run/opengl-driver/lib:$LD_LIBRARY_PATH
|
|
'';
|
|
};
|
|
}
|
|
);
|
|
|
|
checks = forEachSupportedSystem (
|
|
{ treefmtEval, system, ... }:
|
|
{
|
|
treefmt = treefmtEval.config.build.check self;
|
|
|
|
uv2nixShell = self.devShells.${system}.uv2nix;
|
|
}
|
|
);
|
|
};
|
|
}
|