Refine default thread selection for NUMA systems (#7322)

Until we have full NUMA support, this adjusts the default thread selection
algorithm to count up the number of performance cores across all sockets.
This commit is contained in:
Daniel Hiltgen 2024-10-30 15:05:45 -07:00 committed by GitHub
parent c826e57475
commit 16f4eabe2d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 2121 additions and 6 deletions

View file

@ -3,9 +3,11 @@ package discover
import (
"bufio"
"fmt"
"io"
"os"
"reflect"
"regexp"
"sort"
"strings"
"github.com/ollama/ollama/format"
@ -109,6 +111,10 @@ func GetCPUDetails() ([]CPU, error) {
if err != nil {
return nil, err
}
return linuxCPUDetails(file)
}
func linuxCPUDetails(file io.Reader) ([]CPU, error) {
reColumns := regexp.MustCompile("\t+: ")
scanner := bufio.NewScanner(file)
cpuInfos := []linuxCpuInfo{}
@ -131,6 +137,9 @@ func GetCPUDetails() ([]CPU, error) {
cpu = &linuxCpuInfo{}
}
}
if cpu.ID != "" {
cpuInfos = append(cpuInfos, *cpu)
}
// Process the sockets/cores/threads
socketByID := map[string]*CPU{}
@ -177,10 +186,14 @@ func GetCPUDetails() ([]CPU, error) {
s.EfficiencyCoreCount = efficiencyCoreCount
}
}
result := []CPU{}
for _, c := range socketByID {
result = append(result, *c)
keys := make([]string, 0, len(socketByID))
result := make([]CPU, 0, len(socketByID))
for k := range socketByID {
keys = append(keys, k)
}
sort.Strings(keys)
for _, k := range keys {
result = append(result, *socketByID[k])
}
return result, nil
}

2097
discover/gpu_linux_test.go Normal file

File diff suppressed because it is too large Load diff

View file

@ -175,6 +175,11 @@ func (si SystemInfo) GetOptimalThreadCount() int {
if len(si.System.CPUs) == 0 {
return 0
}
// Allocate thread count matching the performance cores on a single socket
return si.System.CPUs[0].CoreCount - si.System.CPUs[0].EfficiencyCoreCount
coreCount := 0
for _, c := range si.System.CPUs {
coreCount += c.CoreCount - c.EfficiencyCoreCount
}
return coreCount
}