llama.cpp's start-up output is read as a shard becomes ready: the devices it used, its backend, and the layers it offloaded. That goes up with READY in a new ShardStatus.compute, is kept on the shard, and is shown by the admin API, ctl replicas (ON) and the console (COMPUTE). A shard given GPUs whose engine says it found no usable GPU or offloaded no layer is stopped and failed, even where nvidia-smi cannot be asked. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
80 lines
2.3 KiB
Go
80 lines
2.3 KiB
Go
package engines
|
|
|
|
import (
|
|
"regexp"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
// Compute is where an engine runs its model, as it said while loading it.
|
|
type Compute struct {
|
|
// Backend is CUDA, Vulkan, Metal, ROCm or CPU.
|
|
Backend string
|
|
// Devices are those it named: "CUDA0: Quadro RTX 6000 (7.5)".
|
|
Devices []string
|
|
// Offloaded of Total layers went to the devices; 0 and 0 when it did
|
|
// not say.
|
|
Offloaded, Total int
|
|
// NoGPU is the engine saying it found no device it could use.
|
|
NoGPU bool
|
|
}
|
|
|
|
var (
|
|
offloadedRE = regexp.MustCompile(`offloaded (\d+)/(\d+) layers to GPU`)
|
|
usingRE = regexp.MustCompile(`using device ([A-Za-z]+\d+) \(([^)]*)\)`)
|
|
ccRE = regexp.MustCompile(`Device (\d+): ([^,]+), compute capability ([0-9.]+)`)
|
|
noDevicesRE = regexp.MustCompile(`found 0 (CUDA|Vulkan|ROCm) devices|no usable GPU found`)
|
|
)
|
|
|
|
// ReadLlamaCompute reads llama.cpp's start-up output for where it put the
|
|
// model, reporting false when the output says nothing of it.
|
|
func ReadLlamaCompute(lines []string) (Compute, bool) {
|
|
var c Compute
|
|
known := false
|
|
cc := map[string]string{} // device index → compute capability
|
|
for _, l := range lines {
|
|
if m := ccRE.FindStringSubmatch(l); m != nil {
|
|
cc[m[1]] = m[3]
|
|
}
|
|
}
|
|
seen := map[string]bool{}
|
|
for _, l := range lines {
|
|
if m := offloadedRE.FindStringSubmatch(l); m != nil {
|
|
c.Offloaded, _ = strconv.Atoi(m[1])
|
|
c.Total, _ = strconv.Atoi(m[2])
|
|
known = true
|
|
}
|
|
if m := usingRE.FindStringSubmatch(l); m != nil && !seen[m[1]] {
|
|
seen[m[1]] = true
|
|
d := m[1] + ": " + strings.TrimSpace(m[2])
|
|
idx := strings.TrimLeft(m[1], "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")
|
|
if v, ok := cc[idx]; ok && strings.HasPrefix(m[1], "CUDA") {
|
|
d += " (" + v + ")"
|
|
}
|
|
c.Devices = append(c.Devices, d)
|
|
if c.Backend == "" {
|
|
c.Backend = strings.TrimRight(m[1], "0123456789")
|
|
}
|
|
known = true
|
|
}
|
|
if noDevicesRE.MatchString(l) {
|
|
c.NoGPU, known = true, true
|
|
}
|
|
}
|
|
switch {
|
|
case c.OnTheCPU():
|
|
c.Backend = "CPU"
|
|
case c.Backend == "" && c.Offloaded > 0:
|
|
c.Backend = "GPU"
|
|
case c.Backend == "":
|
|
c.Backend = "CPU"
|
|
}
|
|
return c, known
|
|
}
|
|
|
|
// OnTheCPU is whether a model the engine was given GPUs for is computed on
|
|
// the CPU, by its own account: no device found, or no layer offloaded.
|
|
func (c Compute) OnTheCPU() bool {
|
|
return c.NoGPU || (c.Total > 0 && c.Offloaded == 0)
|
|
}
|