amd_linux.go 16.4 KB
Newer Older
1
package discover
Daniel Hiltgen's avatar
Daniel Hiltgen committed
2
3
4
5
6
7

import (
	"bufio"
	"errors"
	"fmt"
	"io"
8
	"io/fs"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
9
10
11
	"log/slog"
	"os"
	"path/filepath"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
12
	"regexp"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
13
	"slices"
14
	"sort"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
15
16
	"strconv"
	"strings"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
17

18
	"github.com/ollama/ollama/envconfig"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
19
	"github.com/ollama/ollama/format"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
20
21
22
23
24
25
26
27
28
29
30
)

// Discovery logic for AMD/ROCm GPUs

const (
	DriverVersionFile     = "/sys/module/amdgpu/version"
	AMDNodesSysfsDir      = "/sys/class/kfd/kfd/topology/nodes/"
	GPUPropertiesFileGlob = AMDNodesSysfsDir + "*/properties"

	// Prefix with the node dir
	GPUTotalMemoryFileGlob = "mem_banks/*/properties" // size_in_bytes line
31
32

	// Direct Rendering Manager sysfs location
33
	DRMDeviceDirGlob   = "/sys/class/drm/card*/device"
34
35
36
37
38
39
40
	DRMTotalMemoryFile = "mem_info_vram_total"
	DRMUsedMemoryFile  = "mem_info_vram_used"

	// In hex; properties file is in decimal
	DRMUniqueIDFile = "unique_id"
	DRMVendorFile   = "vendor"
	DRMDeviceFile   = "device"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
41
42
43
44
)

var (
	// Used to validate if the given ROCm lib is usable
45
46
	ROCmLibGlobs          = []string{"libhipblas.so.2*", "rocblas"} // TODO - probably include more coverage of files here...
	RocmStandardLocations = []string{"/opt/rocm/lib", "/usr/lib64"}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
47
48
49
)

// Gather GPU information from the amdgpu driver if any supported GPUs are detected
50
51
// Only called once during bootstrap
func AMDGetGPUInfo() ([]RocmGPUInfo, error) {
52
	resp := []RocmGPUInfo{}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
53
	if !AMDDetected() {
54
		return resp, fmt.Errorf("AMD GPUs not detected")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
55
56
57
	}

	// Opportunistic logging of driver version to aid in troubleshooting
Daniel Hiltgen's avatar
Daniel Hiltgen committed
58
59
	driverMajor, driverMinor, err := AMDDriverVersion()
	if err != nil {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
60
		// TODO - if we see users crash and burn with the upstreamed kernel this can be adjusted to hard-fail rocm support and fallback to CPU
Daniel Hiltgen's avatar
Daniel Hiltgen committed
61
		slog.Warn("ollama recommends running the https://www.amd.com/en/support/linux-drivers", "error", err)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
62
63
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
64
65
	// Determine if the user has already pre-selected which GPUs to look at, then ignore the others
	var visibleDevices []string
Michael Yang's avatar
string  
Michael Yang committed
66
67
68
	hipVD := envconfig.HipVisibleDevices()   // zero based index only
	rocrVD := envconfig.RocrVisibleDevices() // zero based index or UUID, but consumer cards seem to not support UUID
	gpuDO := envconfig.GpuDeviceOrdinal()    // zero based index
Daniel Hiltgen's avatar
Daniel Hiltgen committed
69
70
71
72
73
74
75
76
77
78
	switch {
	// TODO is this priorty order right?
	case hipVD != "":
		visibleDevices = strings.Split(hipVD, ",")
	case rocrVD != "":
		visibleDevices = strings.Split(rocrVD, ",")
		// TODO - since we don't yet support UUIDs, consider detecting and reporting here
		// all our test systems show GPU-XX indicating UUID is not supported
	case gpuDO != "":
		visibleDevices = strings.Split(gpuDO, ",")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
79
80
	}

Michael Yang's avatar
string  
Michael Yang committed
81
	gfxOverride := envconfig.HsaOverrideGfxVersion()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
82
83
84
85
86
87
	var supported []string
	libDir := ""

	// The amdgpu driver always exposes the host CPU(s) first, but we have to skip them and subtract
	// from the other IDs to get alignment with the HIP libraries expectations (zero is the first GPU, not the CPU)
	matches, _ := filepath.Glob(GPUPropertiesFileGlob)
88
89
90
91
92
93
94
95
96
97
98
99
100
101
	sort.Slice(matches, func(i, j int) bool {
		// /sys/class/kfd/kfd/topology/nodes/<number>/properties
		a, err := strconv.ParseInt(filepath.Base(filepath.Dir(matches[i])), 10, 64)
		if err != nil {
			slog.Debug("parse err", "error", err, "match", matches[i])
			return false
		}
		b, err := strconv.ParseInt(filepath.Base(filepath.Dir(matches[j])), 10, 64)
		if err != nil {
			slog.Debug("parse err", "error", err, "match", matches[i])
			return false
		}
		return a < b
	})
Daniel Hiltgen's avatar
Daniel Hiltgen committed
102
103
104
105
106
107
	cpuCount := 0
	for _, match := range matches {
		slog.Debug("evaluating amdgpu node " + match)
		fp, err := os.Open(match)
		if err != nil {
			slog.Debug("failed to open sysfs node", "file", match, "error", err)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
108
109
			continue
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
110
111
112
113
114
		defer fp.Close()
		nodeID, err := strconv.Atoi(filepath.Base(filepath.Dir(match)))
		if err != nil {
			slog.Debug("failed to parse node ID", "error", err)
			continue
Daniel Hiltgen's avatar
Daniel Hiltgen committed
115
116
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
117
118
119
		scanner := bufio.NewScanner(fp)
		isCPU := false
		var major, minor, patch uint64
120
		var vendor, device, uniqueID uint64
Daniel Hiltgen's avatar
Daniel Hiltgen committed
121
122
123
124
125
		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			// Note: we could also use "cpu_cores_count X" where X is greater than zero to detect CPUs
			if strings.HasPrefix(line, "gfx_target_version") {
				ver := strings.Fields(line)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
126

Daniel Hiltgen's avatar
Daniel Hiltgen committed
127
128
129
130
131
132
				// Detect CPUs
				if len(ver) == 2 && ver[1] == "0" {
					slog.Debug("detected CPU " + match)
					isCPU = true
					break
				}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
133

Daniel Hiltgen's avatar
Daniel Hiltgen committed
134
135
136
137
138
139
140
141
142
143
144
145
146
147
				if len(ver) != 2 || len(ver[1]) < 5 {
					slog.Warn("malformed "+match, "gfx_target_version", line)
					// If this winds up being a CPU, our offsets may be wrong
					continue
				}
				l := len(ver[1])
				var err1, err2, err3 error
				patch, err1 = strconv.ParseUint(ver[1][l-2:l], 10, 32)
				minor, err2 = strconv.ParseUint(ver[1][l-4:l-2], 10, 32)
				major, err3 = strconv.ParseUint(ver[1][:l-4], 10, 32)
				if err1 != nil || err2 != nil || err3 != nil {
					slog.Debug("malformed int " + line)
					continue
				}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
148
149
150
			} else if strings.HasPrefix(line, "vendor_id") {
				ver := strings.Fields(line)
				if len(ver) != 2 {
151
					slog.Debug("malformed", "vendor_id", line)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
152
153
					continue
				}
154
				vendor, err = strconv.ParseUint(ver[1], 10, 64)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
155
				if err != nil {
156
					slog.Debug("malformed", "vendor_id", line, "error", err)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
157
158
159
160
				}
			} else if strings.HasPrefix(line, "device_id") {
				ver := strings.Fields(line)
				if len(ver) != 2 {
161
					slog.Debug("malformed", "device_id", line)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
162
163
					continue
				}
164
				device, err = strconv.ParseUint(ver[1], 10, 64)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
165
				if err != nil {
166
167
168
169
170
171
172
173
174
175
176
					slog.Debug("malformed", "device_id", line, "error", err)
				}
			} else if strings.HasPrefix(line, "unique_id") {
				ver := strings.Fields(line)
				if len(ver) != 2 {
					slog.Debug("malformed", "unique_id", line)
					continue
				}
				uniqueID, err = strconv.ParseUint(ver[1], 10, 64)
				if err != nil {
					slog.Debug("malformed", "unique_id", line, "error", err)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
177
				}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
178
179
180
181
			}
			// TODO - any other properties we want to extract and record?
			// vendor_id + device_id -> pci lookup for "Name"
			// Other metrics that may help us understand relative performance between multiple GPUs
Daniel Hiltgen's avatar
Daniel Hiltgen committed
182
183
		}

184
185
186
187
		// Note: while ./mem_banks/*/used_memory exists, it doesn't appear to take other VRAM consumers
		// into consideration, so we instead map the device over to the DRM driver sysfs nodes which
		// do reliably report VRAM usage.

Daniel Hiltgen's avatar
Daniel Hiltgen committed
188
189
190
		if isCPU {
			cpuCount++
			continue
Daniel Hiltgen's avatar
Daniel Hiltgen committed
191
192
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
193
194
		// CPUs are always first in the list
		gpuID := nodeID - cpuCount
Daniel Hiltgen's avatar
Daniel Hiltgen committed
195

Daniel Hiltgen's avatar
Daniel Hiltgen committed
196
197
		// Shouldn't happen, but just in case...
		if gpuID < 0 {
198
199
200
			err := fmt.Errorf("unexpected amdgpu sysfs data resulted in negative GPU ID, please set OLLAMA_DEBUG=1 and report an issue")
			slog.Error(err.Error())
			return nil, err
Daniel Hiltgen's avatar
Daniel Hiltgen committed
201
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
202
203

		// Look up the memory for the current node
Daniel Hiltgen's avatar
Daniel Hiltgen committed
204
205
		totalMemory := uint64(0)
		usedMemory := uint64(0)
206
		var usedFile string
207
208
209
210
211
212
213
		mapping := []struct {
			id       uint64
			filename string
		}{
			{vendor, DRMVendorFile},
			{device, DRMDeviceFile},
			{uniqueID, DRMUniqueIDFile}, // Not all devices will report this
Daniel Hiltgen's avatar
Daniel Hiltgen committed
214
		}
215
216
217
218
219
220
221
		slog.Debug("mapping amdgpu to drm sysfs nodes", "amdgpu", match, "vendor", vendor, "device", device, "unique_id", uniqueID)
		// Map over to DRM location to find the total/free memory
		drmMatches, _ := filepath.Glob(DRMDeviceDirGlob)
		for _, devDir := range drmMatches {
			matched := true
			for _, m := range mapping {
				if m.id == 0 {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
222
					// Null ID means it didn't populate, so we can't use it to match
223
224
225
					continue
				}
				filename := filepath.Join(devDir, m.filename)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
226
				buf, err := os.ReadFile(filename)
227
228
229
230
231
				if err != nil {
					slog.Debug("failed to read sysfs node", "file", filename, "error", err)
					matched = false
					break
				}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
232
				// values here are in hex, strip off the lead 0x and parse so we can compare the numeric (decimal) values in amdgpu
233
234
235
236
237
238
239
240
241
242
243
244
				cmp, err := strconv.ParseUint(strings.TrimPrefix(strings.TrimSpace(string(buf)), "0x"), 16, 64)
				if err != nil {
					slog.Debug("failed to parse sysfs node", "file", filename, "error", err)
					matched = false
					break
				}
				if cmp != m.id {
					matched = false
					break
				}
			}
			if !matched {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
245
246
				continue
			}
247
248
249
250

			// Found the matching DRM directory
			slog.Debug("matched", "amdgpu", match, "drm", devDir)
			totalFile := filepath.Join(devDir, DRMTotalMemoryFile)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
251
			buf, err := os.ReadFile(totalFile)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
252
			if err != nil {
253
254
				slog.Debug("failed to read sysfs node", "file", totalFile, "error", err)
				break
Daniel Hiltgen's avatar
Daniel Hiltgen committed
255
			}
256
			totalMemory, err = strconv.ParseUint(strings.TrimSpace(string(buf)), 10, 64)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
257
			if err != nil {
258
259
260
261
				slog.Debug("failed to parse sysfs node", "file", totalFile, "error", err)
				break
			}

262
263
			usedFile = filepath.Join(devDir, DRMUsedMemoryFile)
			usedMemory, err = getFreeMemory(usedFile)
264
			if err != nil {
265
				slog.Debug("failed to update used memory", "error", err)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
266
			}
267
			break
Daniel Hiltgen's avatar
Daniel Hiltgen committed
268
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
269

Daniel Hiltgen's avatar
Daniel Hiltgen committed
270
271
272
273
274
		var name string
		// TODO - PCI ID lookup
		if vendor > 0 && device > 0 {
			name = fmt.Sprintf("%04x:%04x", vendor, device)
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
275

276
277
278
279
280
281
282
		gpuInfo := RocmGPUInfo{
			GpuInfo: GpuInfo{
				Library: "rocm",
				memInfo: memInfo{
					TotalMemory: totalMemory,
					FreeMemory:  (totalMemory - usedMemory),
				},
Daniel Hiltgen's avatar
Daniel Hiltgen committed
283
				ID:            strconv.Itoa(gpuID),
284
285
286
287
288
				Name:          name,
				Compute:       fmt.Sprintf("gfx%d%x%x", major, minor, patch),
				MinimumMemory: rocmMinimumMemory,
				DriverMajor:   driverMajor,
				DriverMinor:   driverMinor,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
289
			},
290
			usedFilepath: usedFile,
Daniel Hiltgen's avatar
Daniel Hiltgen committed
291
292
		}

293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
		// iGPU detection, remove this check once we can support an iGPU variant of the rocm library
		if totalMemory < IGPUMemLimit {
			reason := "unsupported Radeon iGPU detected skipping"
			slog.Info(reason, "id", gpuID, "total", format.HumanBytes2(totalMemory))
			unsupportedGPUs = append(unsupportedGPUs, UnsupportedGPUInfo{
				GpuInfo: gpuInfo.GpuInfo,
				Reason:  reason,
			})
			continue
		}

		if int(major) < RocmComputeMin {
			reason := fmt.Sprintf("amdgpu too old gfx%d%x%x", major, minor, patch)
			slog.Warn(reason, "gpu", gpuID)
			unsupportedGPUs = append(unsupportedGPUs, UnsupportedGPUInfo{
				GpuInfo: gpuInfo.GpuInfo,
				Reason:  reason,
			})

			continue
		}

		slog.Debug("amdgpu memory", "gpu", gpuID, "total", format.HumanBytes2(totalMemory))
		slog.Debug("amdgpu memory", "gpu", gpuID, "available", format.HumanBytes2(totalMemory-usedMemory))

Daniel Hiltgen's avatar
Daniel Hiltgen committed
318
319
320
321
322
323
324
325
326
327
		// If the user wants to filter to a subset of devices, filter out if we aren't a match
		if len(visibleDevices) > 0 {
			include := false
			for _, visible := range visibleDevices {
				if visible == gpuInfo.ID {
					include = true
					break
				}
			}
			if !include {
328
329
330
331
332
333
334
				reason := "filtering out device per user request"
				slog.Info(reason, "id", gpuInfo.ID, "visible_devices", visibleDevices)
				unsupportedGPUs = append(unsupportedGPUs, UnsupportedGPUInfo{
					GpuInfo: gpuInfo.GpuInfo,
					Reason:  reason,
				})

Daniel Hiltgen's avatar
Daniel Hiltgen committed
335
336
337
338
339
340
341
342
343
				continue
			}
		}

		// Final validation is gfx compatibility - load the library if we haven't already loaded it
		// even if the user overrides, we still need to validate the library
		if libDir == "" {
			libDir, err = AMDValidateLibDir()
			if err != nil {
344
345
346
347
348
349
350
				err = fmt.Errorf("unable to verify rocm library: %w", err)
				slog.Warn(err.Error())
				unsupportedGPUs = append(unsupportedGPUs, UnsupportedGPUInfo{
					GpuInfo: gpuInfo.GpuInfo,
					Reason:  err.Error(),
				})
				return nil, err
Daniel Hiltgen's avatar
Daniel Hiltgen committed
351
352
353
354
355
356
357
358
359
			}
		}
		gpuInfo.DependencyPath = libDir

		if gfxOverride == "" {
			// Only load supported list once
			if len(supported) == 0 {
				supported, err = GetSupportedGFX(libDir)
				if err != nil {
360
361
362
363
364
365
366
					err = fmt.Errorf("failed to lookup supported GFX types: %w", err)
					slog.Warn(err.Error())
					unsupportedGPUs = append(unsupportedGPUs, UnsupportedGPUInfo{
						GpuInfo: gpuInfo.GpuInfo,
						Reason:  err.Error(),
					})
					return nil, err
Daniel Hiltgen's avatar
Daniel Hiltgen committed
367
368
369
				}
				slog.Debug("rocm supported GPUs", "types", supported)
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
370
			gfx := gpuInfo.Compute
Daniel Hiltgen's avatar
Daniel Hiltgen committed
371
			if !slices.Contains[[]string, string](supported, gfx) {
372
373
374
375
376
377
378
				reason := fmt.Sprintf("amdgpu is not supported (supported types:%s)", supported)
				slog.Warn(reason, "gpu_type", gfx, "gpu", gpuInfo.ID, "library", libDir)
				unsupportedGPUs = append(unsupportedGPUs, UnsupportedGPUInfo{
					GpuInfo: gpuInfo.GpuInfo,
					Reason:  reason,
				})

Daniel Hiltgen's avatar
Daniel Hiltgen committed
379
380
381
382
383
384
385
				// TODO - consider discrete markdown just for ROCM troubleshooting?
				slog.Warn("See https://github.com/ollama/ollama/blob/main/docs/gpu.md#overrides for HSA_OVERRIDE_GFX_VERSION usage")
				continue
			} else {
				slog.Info("amdgpu is supported", "gpu", gpuInfo.ID, "gpu_type", gfx)
			}
		} else {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
386
			slog.Info("skipping rocm gfx compatibility check", "HSA_OVERRIDE_GFX_VERSION", gfxOverride)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
387
388
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
389
390
391
392
393
		// Check for env var workarounds
		if name == "1002:687f" { // Vega RX 56
			gpuInfo.EnvWorkarounds = append(gpuInfo.EnvWorkarounds, [2]string{"HSA_ENABLE_SDMA", "0"})
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
394
395
		// The GPU has passed all the verification steps and is supported
		resp = append(resp, gpuInfo)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
396
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
397
	if len(resp) == 0 {
398
399
400
		err := fmt.Errorf("no compatible amdgpu devices detected")
		slog.Info(err.Error())
		return nil, err
Daniel Hiltgen's avatar
Daniel Hiltgen committed
401
	}
402
	if err := verifyKFDDriverAccess(); err != nil {
403
404
405
		err = fmt.Errorf("amdgpu devices detected but permission problems block access: %w", err)
		slog.Error(err.Error())
		return nil, err
406
	}
407
	return resp, nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
408
409
410
411
412
413
414
415
416
417
418
}

// Quick check for AMD driver so we can skip amdgpu discovery if not present
func AMDDetected() bool {
	// Some driver versions (older?) don't have a version file, so just lookup the parent dir
	sysfsDir := filepath.Dir(DriverVersionFile)
	_, err := os.Stat(sysfsDir)
	if errors.Is(err, os.ErrNotExist) {
		slog.Debug("amdgpu driver not detected " + sysfsDir)
		return false
	} else if err != nil {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
419
		slog.Debug("error looking up amd driver", "path", sysfsDir, "error", err)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
420
421
422
423
424
425
426
427
		return false
	}
	return true
}

// Prefer to use host installed ROCm, as long as it meets our minimum requirements
// failing that, tell the user how to download it on their own
func AMDValidateLibDir() (string, error) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
428
	libDir, err := commonAMDValidateLibDir()
429
	if err == nil {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
430
		return libDir, nil
431
432
	}

433
434
435
	// Well known ollama installer path
	installedRocmDir := "/usr/share/ollama/lib/rocm"
	if rocmLibUsable(installedRocmDir) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
436
		return installedRocmDir, nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
437
438
	}

439
440
	// 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")
Michael Yang's avatar
lint  
Michael Yang committed
441
	return "", errors.New("no suitable rocm found, falling back to CPU")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
442
443
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
444
445
func AMDDriverVersion() (driverMajor, driverMinor int, err error) {
	_, err = os.Stat(DriverVersionFile)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
446
	if err != nil {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
447
		return 0, 0, fmt.Errorf("amdgpu version file missing: %s %w", DriverVersionFile, err)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
448
449
450
	}
	fp, err := os.Open(DriverVersionFile)
	if err != nil {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
451
		return 0, 0, err
Daniel Hiltgen's avatar
Daniel Hiltgen committed
452
453
454
455
	}
	defer fp.Close()
	verString, err := io.ReadAll(fp)
	if err != nil {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
		return 0, 0, err
	}

	pattern := `\A(\d+)\.(\d+).*`
	regex := regexp.MustCompile(pattern)
	match := regex.FindStringSubmatch(string(verString))
	if len(match) < 2 {
		return 0, 0, fmt.Errorf("malformed version string %s", string(verString))
	}
	driverMajor, err = strconv.Atoi(match[1])
	if err != nil {
		return 0, 0, err
	}
	driverMinor, err = strconv.Atoi(match[2])
	if err != nil {
		return 0, 0, err
Daniel Hiltgen's avatar
Daniel Hiltgen committed
472
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
473
	return driverMajor, driverMinor, nil
Daniel Hiltgen's avatar
Daniel Hiltgen committed
474
}
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491

func (gpus RocmGPUInfoList) RefreshFreeMemory() error {
	if len(gpus) == 0 {
		return nil
	}
	for i := range gpus {
		usedMemory, err := getFreeMemory(gpus[i].usedFilepath)
		if err != nil {
			return err
		}
		slog.Debug("updating rocm free memory", "gpu", gpus[i].ID, "name", gpus[i].Name, "before", format.HumanBytes2(gpus[i].FreeMemory), "now", format.HumanBytes2(gpus[i].TotalMemory-usedMemory))
		gpus[i].FreeMemory = gpus[i].TotalMemory - usedMemory
	}
	return nil
}

func getFreeMemory(usedFile string) (uint64, error) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
492
	buf, err := os.ReadFile(usedFile)
493
494
495
496
497
498
499
500
501
502
	if err != nil {
		return 0, fmt.Errorf("failed to read sysfs node %s %w", usedFile, err)
	}
	usedMemory, err := strconv.ParseUint(strings.TrimSpace(string(buf)), 10, 64)
	if err != nil {
		slog.Debug("failed to parse sysfs node", "file", usedFile, "error", err)
		return 0, fmt.Errorf("failed to parse sysfs node %s %w", usedFile, err)
	}
	return usedMemory, nil
}
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518

func verifyKFDDriverAccess() error {
	// 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, 0o666)
	if err != nil {
		if errors.Is(err, fs.ErrPermission) {
			return fmt.Errorf("permissions not set up properly.  Either run ollama as root, or add you user account to the render group. %w", err)
		} else if errors.Is(err, fs.ErrNotExist) {
			// Container runtime failure?
			return fmt.Errorf("kfd driver not loaded.  If running in a container, remember to include '--device /dev/kfd --device /dev/dri'")
		}
		return fmt.Errorf("failed to check permission on /dev/kfd: %w", err)
	}
	fd.Close()
	return nil
}