payload_common.go 7.6 KB
Newer Older
1
2
3
package llm

import (
4
	"compress/bzip2"
5
6
7
8
9
10
11
12
13
14
	"errors"
	"fmt"
	"io"
	"io/fs"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strings"

15
16
17
	"golang.org/x/exp/slices"
	"golang.org/x/sync/errgroup"

18
19
20
	"github.com/jmorganca/ollama/gpu"
)

21
// Libraries names may contain an optional variant separated by '_'
22
// For example, "rocm_v6" and "rocm_v5" or "cpu" and "cpu_avx2"
23
// Any library without a variant is the lowest common denominator
24
var availableDynLibs = map[string]string{}
25

26
const pathComponentCount = 7
27

28
29
// getDynLibs returns an ordered list of LLM libraries to try, starting with the best
func getDynLibs(gpuInfo gpu.GpuInfo) []string {
30
31
32
33
	// Short circuit if we know we're using the default built-in (darwin only)
	if gpuInfo.Library == "default" {
		return []string{"default"}
	}
34
35
36
37
38
39
40
	// TODO - temporary until we have multiple CPU variations for Darwin
	// Short circuit on darwin with metal only
	if len(availableDynLibs) == 1 {
		if _, onlyMetal := availableDynLibs["metal"]; onlyMetal {
			return []string{availableDynLibs["metal"]}
		}
	}
41

42
	exactMatch := ""
43
44
	dynLibs := []string{}
	altDynLibs := []string{}
45
46
47
48
	requested := gpuInfo.Library
	if gpuInfo.Variant != "" {
		requested += "_" + gpuInfo.Variant
	}
49
	// Try to find an exact match
50
	for cmp := range availableDynLibs {
51
52
		if requested == cmp {
			exactMatch = cmp
53
			dynLibs = []string{availableDynLibs[cmp]}
54
55
56
			break
		}
	}
57
	// Then for GPUs load alternates and sort the list for consistent load ordering
58
	if gpuInfo.Library != "cpu" {
59
		for cmp := range availableDynLibs {
60
			if gpuInfo.Library == strings.Split(cmp, "_")[0] && cmp != exactMatch {
61
				altDynLibs = append(altDynLibs, cmp)
62
63
			}
		}
64
65
66
		slices.Sort(altDynLibs)
		for _, altDynLib := range altDynLibs {
			dynLibs = append(dynLibs, availableDynLibs[altDynLib])
67
68
		}
	}
69
70
71
72
73
74
75
76
77

	// Load up the best CPU variant if not primary requested
	if gpuInfo.Library != "cpu" {
		variant := gpu.GetCPUVariant()
		// If no variant, then we fall back to default
		// If we have a variant, try that if we find an exact match
		// Attempting to run the wrong CPU instructions will panic the
		// process
		if variant != "" {
78
			for cmp := range availableDynLibs {
79
				if cmp == "cpu_"+variant {
80
					dynLibs = append(dynLibs, availableDynLibs[cmp])
81
82
83
84
					break
				}
			}
		} else {
85
			dynLibs = append(dynLibs, availableDynLibs["cpu"])
86
87
88
		}
	}

Michael Yang's avatar
Michael Yang committed
89
	// Finally, if we didn't find any matches, LCD CPU FTW
90
91
	if len(dynLibs) == 0 {
		dynLibs = []string{availableDynLibs["cpu"]}
92
	}
93
	return dynLibs
94
95
}

96
97
98
func rocmDynLibPresent() bool {
	for dynLibName := range availableDynLibs {
		if strings.HasPrefix(dynLibName, "rocm") {
99
100
101
102
103
104
105
			return true
		}
	}
	return false
}

func nativeInit(workdir string) error {
106
	log.Printf("Extracting dynamic libraries...")
107
108
109
110
111
112
113
114
115
116
117
118
119
	if runtime.GOOS == "darwin" {
		err := extractPayloadFiles(workdir, "llama.cpp/ggml-metal.metal")
		if err != nil {
			if err == payloadMissing {
				// TODO perhaps consider this a hard failure on arm macs?
				log.Printf("ggml-meta.metal payload missing")
				return nil
			}
			return err
		}
		os.Setenv("GGML_METAL_PATH_RESOURCES", workdir)
	}

120
	libs, err := extractDynamicLibs(workdir, "llama.cpp/build/*/*/*/lib/*")
121
122
123
124
125
126
127
128
129
130
	if err != nil {
		if err == payloadMissing {
			log.Printf("%s", payloadMissing)
			return nil
		}
		return err
	}
	for _, lib := range libs {
		// The last dir component is the variant name
		variant := filepath.Base(filepath.Dir(lib))
131
		availableDynLibs[variant] = lib
132
133
134
135
136
137
138
	}

	if err := verifyDriverAccess(); err != nil {
		return err
	}

	// Report which dynamic libraries we have loaded to assist troubleshooting
139
	variants := make([]string, len(availableDynLibs))
140
	i := 0
141
	for variant := range availableDynLibs {
142
143
144
		variants[i] = variant
		i++
	}
145
146
	log.Printf("Dynamic LLM libraries %v", variants)
	log.Printf("Override detection logic by setting OLLAMA_LLM_LIBRARY")
147
148
149
150
151
152
153
154
155
156
157

	return nil
}

func extractDynamicLibs(workDir, glob string) ([]string, error) {
	files, err := fs.Glob(libEmbed, glob)
	if err != nil || len(files) == 0 {
		return nil, payloadMissing
	}
	libs := []string{}

158
159
160
161
162
	// TODO consider making this idempotent with some sort of persistent directory (where we store models probably)
	// and tracking by version so we don't reexpand the files every time
	// Also maybe consider lazy loading only what is needed

	g := new(errgroup.Group)
163
164
165
166
167
168
169
	for _, file := range files {
		pathComps := strings.Split(file, "/")
		if len(pathComps) != pathComponentCount {
			log.Printf("unexpected payload components: %v", pathComps)
			continue
		}

170
171
172
173
174
175
		file := file
		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(workDir, pathComps[pathComponentCount-3])
			srcFile, err := libEmbed.Open(file)
176
			if err != nil {
177
				return fmt.Errorf("read payload %s: %v", file, err)
178
			}
179
180
181
			defer srcFile.Close()
			if err := os.MkdirAll(targetDir, 0o755); err != nil {
				return fmt.Errorf("create payload temp dir %s: %v", workDir, err)
182
			}
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
			src := io.Reader(srcFile)
			filename := file
			if strings.HasSuffix(file, ".bz2") {
				src = bzip2.NewReader(src)
				filename = strings.TrimSuffix(filename, ".bz2")
			}

			destFile := filepath.Join(targetDir, filepath.Base(filename))
			if strings.Contains(destFile, "server") {
				libs = append(libs, destFile)
			}

			_, err = os.Stat(destFile)
			switch {
			case errors.Is(err, os.ErrNotExist):
				destFile, 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 destFile.Close()
				if _, err := io.Copy(destFile, src); err != nil {
					return fmt.Errorf("copy payload %s: %v", file, err)
				}
			case err != nil:
				return fmt.Errorf("stat payload %s: %v", file, err)
			}
			return nil
		})
211
	}
212
	return libs, g.Wait()
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
}

func extractPayloadFiles(workDir, glob string) error {
	files, err := fs.Glob(libEmbed, glob)
	if err != nil || len(files) == 0 {
		return payloadMissing
	}

	for _, file := range files {
		srcFile, err := libEmbed.Open(file)
		if err != nil {
			return fmt.Errorf("read payload %s: %v", file, err)
		}
		defer srcFile.Close()
		if err := os.MkdirAll(workDir, 0o755); err != nil {
			return fmt.Errorf("create payload temp dir %s: %v", workDir, err)
		}
230
231
232
233
234
235
		src := io.Reader(srcFile)
		filename := file
		if strings.HasSuffix(file, ".bz2") {
			src = bzip2.NewReader(src)
			filename = strings.TrimSuffix(filename, ".bz2")
		}
236

237
		destFile := filepath.Join(workDir, filepath.Base(filename))
238
239
240
241
242
243
244
245
		_, err = os.Stat(destFile)
		switch {
		case errors.Is(err, os.ErrNotExist):
			destFile, 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 destFile.Close()
246
			if _, err := io.Copy(destFile, src); err != nil {
247
248
249
250
251
252
253
254
255
256
257
258
259
260
				return fmt.Errorf("copy payload %s: %v", file, err)
			}
		case err != nil:
			return fmt.Errorf("stat payload %s: %v", file, err)
		}
	}
	return nil
}

func verifyDriverAccess() error {
	if runtime.GOOS != "linux" {
		return nil
	}
	// Only check ROCm access if we have the dynamic lib loaded
261
	if rocmDynLibPresent() {
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
		// Verify we have permissions - either running as root, or we have group access to the driver
		fd, err := os.OpenFile("/dev/kfd", os.O_RDWR, 0666)
		if err != nil {
			if errors.Is(err, fs.ErrPermission) {
				return fmt.Errorf("Radeon card detected, but permissions not set up properly.  Either run ollama as root, or add you user account to the render group.")
			} else if errors.Is(err, fs.ErrNotExist) {
				// expected behavior without a radeon card
				return nil
			}

			return fmt.Errorf("failed to check permission on /dev/kfd: %w", err)
		}
		fd.Close()
	}
	return nil
}