gpu.go 3.65 KB
Newer Older
1
package discover
2
3

import (
4
	"context"
5
	"fmt"
6
	"log/slog"
7
8
	"os"
	"path/filepath"
9
	"runtime"
10
	"strings"
Michael Yang's avatar
Michael Yang committed
11

12
	"github.com/ollama/ollama/format"
13
	"github.com/ollama/ollama/ml"
14
15
)

16
17
18
// Jetson devices have JETSON_JETPACK="x.y.z" factory set to the Jetpack version installed.
// Included to drive logic for reducing Ollama-allocated overhead on L4T/Jetson devices.
var CudaTegra string = os.Getenv("JETSON_JETPACK")
19

20
21
22
23
func GetCPUInfo() GpuInfo {
	mem, err := GetCPUMem()
	if err != nil {
		slog.Warn("error looking up system memory", "error", err)
24
25
	}

26
27
28
29
30
31
	return GpuInfo{
		memInfo: mem,
		DeviceID: ml.DeviceID{
			Library: "cpu",
			ID:      "0",
		},
32
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
33
34
}

35
36
37
func GetGPUInfo(ctx context.Context, runners []FilteredRunnerDiscovery) GpuInfoList {
	devs := GPUDevices(ctx, runners)
	return devInfoToInfoList(devs)
38
39
}

40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
func devInfoToInfoList(devs []ml.DeviceInfo) GpuInfoList {
	resp := []GpuInfo{}
	// Our current packaging model places ggml-hip in the main directory
	// but keeps rocm in an isolated directory.  We have to add it to
	// the [LD_LIBRARY_]PATH so ggml-hip will load properly
	rocmDir := filepath.Join(LibOllamaPath, "rocm")
	if _, err := os.Stat(rocmDir); err != nil {
		rocmDir = ""
	}

	for _, dev := range devs {
		info := GpuInfo{
			DeviceID: dev.DeviceID,
			filterID: dev.FilteredID,
			Name:     dev.Description,
			memInfo: memInfo{
				TotalMemory: dev.TotalMemory,
				FreeMemory:  dev.FreeMemory,
58
			},
59
60
61
62
			// TODO can we avoid variant
			DependencyPath: dev.LibraryPath,
			DriverMajor:    dev.DriverMajor,
			DriverMinor:    dev.DriverMinor,
Michael Yang's avatar
lint  
Michael Yang committed
63
		}
64
65
		if dev.Library == "CUDA" || dev.Library == "ROCm" {
			info.MinimumMemory = 457 * format.MebiByte
Daniel Hiltgen's avatar
Daniel Hiltgen committed
66
		}
67
68
69
70
		if dev.Library == "ROCm" {
			info.Compute = fmt.Sprintf("gfx%x%02x", dev.ComputeMajor, dev.ComputeMinor)
			if rocmDir != "" {
				info.DependencyPath = append(info.DependencyPath, rocmDir)
71
			}
72
73
		} else {
			info.Compute = fmt.Sprintf("%d.%d", dev.ComputeMajor, dev.ComputeMinor)
74
		}
75
		resp = append(resp, info)
76
	}
77
	if len(resp) == 0 {
78
79
80
		mem, err := GetCPUMem()
		if err != nil {
			slog.Warn("error looking up system memory", "error", err)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
81
82
		}

83
84
85
86
87
88
89
		resp = append(resp, GpuInfo{
			memInfo: mem,
			DeviceID: ml.DeviceID{
				Library: "cpu",
				ID:      "0",
			},
		})
90
91
92
93
	}
	return resp
}

94
95
96
97
98
99
100
// Given the list of GPUs this instantiation is targeted for,
// figure out the visible devices environment variable
//
// If different libraries are detected, the first one is what we use
func (l GpuInfoList) GetVisibleDevicesEnv() []string {
	if len(l) == 0 {
		return nil
101
	}
102
	return []string{rocmGetVisibleDevicesEnv(l)}
103
104
}

105
106
107
108
109
func rocmGetVisibleDevicesEnv(gpuInfo []GpuInfo) string {
	ids := []string{}
	for _, info := range gpuInfo {
		if info.Library != "ROCm" {
			continue
110
		}
111
112
113
		// If the devices requires a numeric ID, for filtering purposes, we use the unfiltered ID number
		if info.filterID != "" {
			ids = append(ids, info.filterID)
Wang,Zhe's avatar
Wang,Zhe committed
114
		} else {
115
			ids = append(ids, info.ID)
Wang,Zhe's avatar
Wang,Zhe committed
116
117
		}
	}
118
119
	if len(ids) == 0 {
		return ""
120
	}
121
122
123
	envVar := "ROCR_VISIBLE_DEVICES="
	if runtime.GOOS != "linux" {
		envVar = "HIP_VISIBLE_DEVICES="
Daniel Hiltgen's avatar
Daniel Hiltgen committed
124
	}
125
126
127
128
129
	// There are 3 potential env vars to use to select GPUs.
	// ROCR_VISIBLE_DEVICES supports UUID or numeric but does not work on Windows
	// HIP_VISIBLE_DEVICES supports numeric IDs only
	// GPU_DEVICE_ORDINAL supports numeric IDs only
	return envVar + strings.Join(ids, ",")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
130
}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
131

132
// GetSystemInfo returns the last cached state of the GPUs on the system
133
func GetSystemInfo() SystemInfo {
134
135
136
	deviceMu.Lock()
	defer deviceMu.Unlock()
	gpus := devInfoToInfoList(devices)
137
138
139
140
141
	if len(gpus) == 1 && gpus[0].Library == "cpu" {
		gpus = []GpuInfo{}
	}

	return SystemInfo{
142
143
144
145
146
		System: CPUInfo{
			CPUs:    GetCPUDetails(),
			GpuInfo: GetCPUInfo(),
		},
		GPUs: gpus,
147
148
	}
}