ollama: add libs to wrapper's LD_LIBRARY_PATH to fix `ollama-cuda` (#342685) (d1e6ee24) · Commits · nix / nixpkgs

pkgs/by-name/ol/ollama/package.nix

+116 −112

Original line number	Diff line number	Diff line
		@@ -79,13 +79,16 @@ let
		paths = rocmLibs ++ [ rocmClang ];
		};

		cudaLibs = [
		cudaPackages.cuda_cudart
		cudaPackages.libcublas
		cudaPackages.cuda_cccl
		];
		cudaToolkit = buildEnv {
		name = "cuda-merged";
		paths = [
		(lib.getBin (cudaPackages.cuda_nvcc.__spliced.buildHost or cudaPackages.cuda_nvcc))
		(lib.getLib cudaPackages.cuda_cudart)
		paths = map lib.getLib cudaLibs ++ [
		(lib.getOutput "static" cudaPackages.cuda_cudart)
		(lib.getLib cudaPackages.libcublas)
		(lib.getBin (cudaPackages.cuda_nvcc.__spliced.buildHost or cudaPackages.cuda_nvcc))
		];
		};

		@@ -107,6 +110,9 @@ let
		++ lib.optionals enableRocm [
		"--suffix LD_LIBRARY_PATH : '${rocmPath}/lib'"
		"--set-default HIP_PATH '${rocmPath}'"
		]
		++ lib.optionals enableCuda [
		"--suffix LD_LIBRARY_PATH : '${lib.makeLibraryPath (map lib.getLib cudaLibs)}'"
		];
		wrapperArgs = builtins.concatStringsSep " " wrapperOptions;

		@@ -114,13 +120,7 @@ let
		if enableCuda then buildGoModule.override { stdenv = overrideCC stdenv gcc12; } else buildGoModule;
		inherit (lib) licenses platforms maintainers;
		in
		goBuild (
		(lib.optionalAttrs enableRocm {
		ROCM_PATH = rocmPath;
		CLBlast_DIR = "${clblast}/lib/cmake/CLBlast";
		})
		// (lib.optionalAttrs enableCuda { CUDA_LIB_DIR = "${cudaToolkit}/lib"; })
		// {
		goBuild {
		inherit
		pname
		version
		@@ -128,6 +128,13 @@ goBuild (
		vendorHash
		;

		env =
		lib.optionalAttrs enableRocm {
		ROCM_PATH = rocmPath;
		CLBlast_DIR = "${clblast}/lib/cmake/CLBlast";
		}
		// lib.optionalAttrs enableCuda { CUDA_LIB_DIR = "${cudaToolkit}/lib"; };

		nativeBuildInputs =
		[ cmake ]
		++ lib.optionals enableRocm [ rocmPackages.llvm.bintools ]
		@@ -140,12 +147,7 @@ goBuild (

		buildInputs =
		lib.optionals enableRocm (rocmLibs ++ [ libdrm ])
		++ lib.optionals enableCuda [
		cudaToolkit
		cudaPackages.cuda_cudart
		cudaPackages.cuda_cccl
		cudaPackages.libcublas
		]
		++ lib.optionals enableCuda cudaLibs
		++ lib.optionals stdenv.isDarwin metalFrameworks;

		patches = [
		@@ -160,8 +162,10 @@ goBuild (
		# replace inaccurate version number with actual release version
		substituteInPlace version/version.go --replace-fail 0.0.0 '${version}'

		# apply llama.cpp patches
		for cur in llm/patches/*; do patch -p1 -d llm/llama.cpp < $cur; done
		# apply ollama's patches to `llama.cpp` submodule
		for diff in llm/patches/*; do
		patch -p1 -d llm/llama.cpp < $diff
		done
		'';

		overrideModAttrs = (
		@@ -177,6 +181,7 @@ goBuild (
		# build llama.cpp libraries for ollama
		go generate ./...
		'';

		postFixup =
		''
		# the app doesn't appear functional at the moment, so hide it
		@@ -227,4 +232,3 @@ goBuild (
		];
		};
		}
		)