Преглед изворни кода

Finish unwinding idempotent payload logic

The recent ROCm change partially removed idempotent
payloads, but the ggml-metal.metal file for mac was still
idempotent.  This finishes switching to always extract
the payloads, and now that idempotentcy is gone, the
version directory is no longer useful.
Daniel Hiltgen пре 1 година
родитељ
комит
4a5c9b8035
6 измењених фајлова са 71 додато и 87 уклоњено
  1. 5 0
      docs/linux.md
  2. 13 21
      gpu/amd_linux.go
  3. 7 17
      gpu/amd_windows.go
  4. 25 22
      gpu/assets.go
  5. 20 27
      llm/payload_common.go
  6. 1 0
      server/routes.go

+ 5 - 0
docs/linux.md

@@ -72,6 +72,11 @@ Verify that the drivers are installed by running the following command, which sh
 nvidia-smi
 ```
 
+### Install ROCm (optional - for Radeon GPUs)
+[Download and Install](https://rocm.docs.amd.com/projects/install-on-linux/en/latest/tutorial/quick-start.html)
+
+Make sure to install ROCm v6
+
 ### Start Ollama
 
 Start Ollama using `systemd`:

+ 13 - 21
gpu/amd_linux.go

@@ -11,14 +11,11 @@ import (
 	"slices"
 	"strconv"
 	"strings"
-
-	"github.com/jmorganca/ollama/version"
 )
 
 // Discovery logic for AMD/ROCm GPUs
 
 const (
-	curlMsg               = "curl -fsSL https://github.com/ollama/ollama/releases/download/v%s/rocm-amd64-deps.tgz | tar -zxf - -C %s"
 	DriverVersionFile     = "/sys/module/amdgpu/version"
 	AMDNodesSysfsDir      = "/sys/class/kfd/kfd/topology/nodes/"
 	GPUPropertiesFileGlob = AMDNodesSysfsDir + "*/properties"
@@ -278,22 +275,22 @@ func setupLink(source, target string) error {
 func AMDValidateLibDir() (string, error) {
 	// We rely on the rpath compiled into our library to find rocm
 	// so we establish a symlink to wherever we find it on the system
-	// to $AssetsDir/rocm
-
-	// If we already have a rocm dependency wired, nothing more to do
-	assetsDir, err := AssetsDir()
+	// to <payloads>/rocm
+	payloadsDir, err := PayloadsDir()
 	if err != nil {
-		return "", fmt.Errorf("unable to lookup lib dir: %w", err)
+		return "", err
 	}
-	// Versioned directory
-	rocmTargetDir := filepath.Join(assetsDir, "rocm")
+
+	// If we already have a rocm dependency wired, nothing more to do
+	rocmTargetDir := filepath.Join(payloadsDir, "rocm")
 	if rocmLibUsable(rocmTargetDir) {
 		return rocmTargetDir, nil
 	}
-	// Parent dir (unversioned)
-	commonRocmDir := filepath.Join(filepath.Dir(assetsDir), "rocm")
-	if rocmLibUsable(commonRocmDir) {
-		return rocmTargetDir, setupLink(commonRocmDir, rocmTargetDir)
+
+	// Well known ollama installer path
+	installedRocmDir := "/usr/share/ollama/lib/rocm"
+	if rocmLibUsable(installedRocmDir) {
+		return rocmTargetDir, setupLink(installedRocmDir, rocmTargetDir)
 	}
 
 	// Prefer explicit HIP env var
@@ -322,14 +319,9 @@ func AMDValidateLibDir() (string, error) {
 	if rocmLibUsable("/opt/rocm/lib") {
 		return rocmTargetDir, setupLink("/opt/rocm/lib", rocmTargetDir)
 	}
-	err = os.MkdirAll(rocmTargetDir, 0755)
-	if err != nil {
-		return "", fmt.Errorf("failed to create empty rocm dir %s %w", rocmTargetDir, err)
-	}
 
-	// If we still haven't found a usable rocm, the user will have to download it on their own
-	slog.Warn("amdgpu detected, but no compatible rocm library found.  Either install rocm v6, or run the following")
-	slog.Warn(fmt.Sprintf(curlMsg, version.Version, rocmTargetDir))
+	// If we still haven't found a usable rocm, the user will have to install it on their own
+	slog.Warn("amdgpu detected, but no compatible rocm library found.  Either install rocm v6, or follow manual install instructions at https://github.com/ollama/ollama/blob/main/docs/linux.md#manual-install")
 	return "", fmt.Errorf("no suitable rocm found, falling back to CPU")
 }
 

+ 7 - 17
gpu/amd_windows.go

@@ -140,7 +140,7 @@ func AMDValidateLibDir() (string, error) {
 	// $LibDir/rocm, we instead rely on setting PATH to point
 	// to the location of the ROCm library
 
-	// Installer payload location
+	// Installer payload location if we're running the installed binary
 	exe, err := os.Executable()
 	if err == nil {
 		rocmTargetDir := filepath.Join(filepath.Dir(exe), "rocm")
@@ -150,13 +150,12 @@ func AMDValidateLibDir() (string, error) {
 		}
 	}
 
-	// If we already have a rocm dependency wired, nothing more to do
-	libDir, err := AssetsDir()
-	if err != nil {
-		return "", fmt.Errorf("unable to lookup lib dir: %w", err)
-	}
-	rocmTargetDir := filepath.Join(libDir, "rocm")
+	// Installer payload (if we're running from some other location)
+	localAppData := os.Getenv("LOCALAPPDATA")
+	appDir := filepath.Join(localAppData, "Programs", "Ollama")
+	rocmTargetDir := filepath.Join(appDir, "rocm")
 	if rocmLibUsable(rocmTargetDir) {
+		slog.Debug("detected ollama installed ROCm at " + rocmTargetDir)
 		return rocmTargetDir, nil
 	}
 
@@ -175,16 +174,7 @@ func AMDValidateLibDir() (string, error) {
 		return RocmStandardLocation, nil
 	}
 
-	// Installer payload (if we're running from some other location)
-	localAppData := os.Getenv("LOCALAPPDATA")
-	appDir := filepath.Join(localAppData, "Programs", "Ollama")
-	rocmTargetDir = filepath.Join(appDir, "rocm")
-	if rocmLibUsable(rocmTargetDir) {
-		slog.Debug("detected ollama installed ROCm at " + rocmTargetDir)
-		return rocmTargetDir, nil
-	}
-
 	// Should not happen on windows since we include it in the installer, but stand-alone binary might hit this
-	slog.Warn("amdgpu detected, but no compatible rocm library found.  Please install ROCm v6")
+	slog.Warn("amdgpu detected, but no compatible rocm library found.  Please install ROCm")
 	return "", fmt.Errorf("no suitable rocm found, falling back to CPU")
 }

+ 25 - 22
gpu/assets.go

@@ -7,34 +7,37 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"sync"
+)
 
-	"github.com/jmorganca/ollama/version"
+var (
+	lock        sync.Mutex
+	payloadsDir = ""
 )
 
-func AssetsDir() (string, error) {
-	home, err := os.UserHomeDir()
-	if err != nil {
-		return "", err
+func PayloadsDir() (string, error) {
+	lock.Lock()
+	defer lock.Unlock()
+	if payloadsDir == "" {
+		tmpDir, err := os.MkdirTemp("", "ollama")
+		if err != nil {
+			return "", fmt.Errorf("failed to generate tmp dir: %w", err)
+		}
+		payloadsDir = tmpDir
 	}
-	baseDir := filepath.Join(home, ".ollama", "assets")
-	libDirs, err := os.ReadDir(baseDir)
-	if err == nil {
-		for _, d := range libDirs {
-			if d.Name() == version.Version {
-				continue
-			}
-			// Special case the rocm dependencies, which are handled by the installer
-			if d.Name() == "rocm" {
-				continue
-			}
-			slog.Debug("stale lib detected, cleaning up " + d.Name())
-			err = os.RemoveAll(filepath.Join(baseDir, d.Name()))
-			if err != nil {
-				slog.Warn(fmt.Sprintf("unable to clean up stale library %s: %s", filepath.Join(baseDir, d.Name()), err))
-			}
+	return payloadsDir, nil
+}
+
+func Cleanup() {
+	lock.Lock()
+	defer lock.Unlock()
+	if payloadsDir != "" {
+		slog.Debug("cleaning up payloads dir " + payloadsDir)
+		err := os.RemoveAll(payloadsDir)
+		if err != nil {
+			slog.Warn(fmt.Sprintf("failed to cleanup tmp dir: %s", err))
 		}
 	}
-	return filepath.Join(baseDir, version.Version), nil
 }
 
 func UpdatePath(dir string) {

+ 20 - 27
llm/payload_common.go

@@ -104,13 +104,14 @@ func rocmDynLibPresent() bool {
 }
 
 func nativeInit() error {
-	slog.Info("Extracting dynamic libraries...")
-	assetsDir, err := gpu.AssetsDir()
+	payloadsDir, err := gpu.PayloadsDir()
 	if err != nil {
 		return err
 	}
+	slog.Info(fmt.Sprintf("Extracting dynamic libraries to %s ...", payloadsDir))
+
 	if runtime.GOOS == "darwin" {
-		err := extractPayloadFiles(assetsDir, "llama.cpp/ggml-metal.metal")
+		err := extractPayloadFiles(payloadsDir, "llama.cpp/ggml-metal.metal")
 		if err != nil {
 			if err == payloadMissing {
 				// TODO perhaps consider this a hard failure on arm macs?
@@ -119,10 +120,10 @@ func nativeInit() error {
 			}
 			return err
 		}
-		os.Setenv("GGML_METAL_PATH_RESOURCES", assetsDir)
+		os.Setenv("GGML_METAL_PATH_RESOURCES", payloadsDir)
 	}
 
-	libs, err := extractDynamicLibs(assetsDir, "llama.cpp/build/*/*/*/lib/*")
+	libs, err := extractDynamicLibs(payloadsDir, "llama.cpp/build/*/*/*/lib/*")
 	if err != nil {
 		if err == payloadMissing {
 			slog.Info(fmt.Sprintf("%s", payloadMissing))
@@ -153,7 +154,7 @@ func nativeInit() error {
 	return nil
 }
 
-func extractDynamicLibs(assetsDir, glob string) ([]string, error) {
+func extractDynamicLibs(payloadsDir, glob string) ([]string, error) {
 	files, err := fs.Glob(libEmbed, glob)
 	if err != nil || len(files) == 0 {
 		return nil, payloadMissing
@@ -172,14 +173,14 @@ func extractDynamicLibs(assetsDir, glob string) ([]string, error) {
 		g.Go(func() error {
 			// llama.cpp/build/$OS/$GOARCH/$VARIANT/lib/$LIBRARY
 			// Include the variant in the path to avoid conflicts between multiple server libs
-			targetDir := filepath.Join(assetsDir, pathComps[pathComponentCount-3])
+			targetDir := filepath.Join(payloadsDir, pathComps[pathComponentCount-3])
 			srcFile, err := libEmbed.Open(file)
 			if err != nil {
 				return fmt.Errorf("read payload %s: %v", file, err)
 			}
 			defer srcFile.Close()
 			if err := os.MkdirAll(targetDir, 0o755); err != nil {
-				return fmt.Errorf("create payload lib dir %s: %v", assetsDir, err)
+				return fmt.Errorf("create payload lib dir %s: %v", payloadsDir, err)
 			}
 			src := io.Reader(srcFile)
 			filename := file
@@ -210,7 +211,7 @@ func extractDynamicLibs(assetsDir, glob string) ([]string, error) {
 	return libs, g.Wait()
 }
 
-func extractPayloadFiles(assetsDir, glob string) error {
+func extractPayloadFiles(payloadsDir, glob string) error {
 	files, err := fs.Glob(libEmbed, glob)
 	if err != nil || len(files) == 0 {
 		return payloadMissing
@@ -222,8 +223,8 @@ func extractPayloadFiles(assetsDir, glob string) error {
 			return fmt.Errorf("read payload %s: %v", file, err)
 		}
 		defer srcFile.Close()
-		if err := os.MkdirAll(assetsDir, 0o755); err != nil {
-			return fmt.Errorf("create payload lib dir %s: %v", assetsDir, err)
+		if err := os.MkdirAll(payloadsDir, 0o755); err != nil {
+			return fmt.Errorf("create payload lib dir %s: %v", payloadsDir, err)
 		}
 		src := io.Reader(srcFile)
 		filename := file
@@ -235,22 +236,14 @@ func extractPayloadFiles(assetsDir, glob string) error {
 			filename = strings.TrimSuffix(filename, ".gz")
 		}
 
-		destFile := filepath.Join(assetsDir, filepath.Base(filename))
-		_, err = os.Stat(destFile)
-		switch {
-		case errors.Is(err, os.ErrNotExist):
-			destFp, err := os.OpenFile(destFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o755)
-			if err != nil {
-				return fmt.Errorf("write payload %s: %v", file, err)
-			}
-			defer destFp.Close()
-			if _, err := io.Copy(destFp, src); err != nil {
-				return fmt.Errorf("copy payload %s: %v", file, err)
-			}
-		case err != nil:
-			return fmt.Errorf("stat payload %s: %v", file, err)
-		case err == nil:
-			slog.Debug("payload already exists: " + destFile)
+		destFile := filepath.Join(payloadsDir, filepath.Base(filename))
+		destFp, err := os.OpenFile(destFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o755)
+		if err != nil {
+			return fmt.Errorf("write payload %s: %v", file, err)
+		}
+		defer destFp.Close()
+		if _, err := io.Copy(destFp, src); err != nil {
+			return fmt.Errorf("copy payload %s: %v", file, err)
 		}
 	}
 	return nil

+ 1 - 0
server/routes.go

@@ -1092,6 +1092,7 @@ func Serve(ln net.Listener) error {
 		if loaded.runner != nil {
 			loaded.runner.Close()
 		}
+		gpu.Cleanup()
 		os.Exit(0)
 	}()