mirror of
https://github.com/likelovewant/ollama-for-amd.git
synced 2025-12-21 22:33:56 +00:00
discover: Support cgroups cores and memory limitations (#10292)
* Add supports for cgroups cores and memory limitations * fix compile error and add logs * remove cpu info log
This commit is contained in:
@@ -2,6 +2,7 @@ package discover
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"bufio"
|
"bufio"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
@@ -10,12 +11,21 @@ import (
|
|||||||
"reflect"
|
"reflect"
|
||||||
"regexp"
|
"regexp"
|
||||||
"sort"
|
"sort"
|
||||||
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"github.com/ollama/ollama/format"
|
"github.com/ollama/ollama/format"
|
||||||
)
|
)
|
||||||
|
|
||||||
func GetCPUMem() (memInfo, error) {
|
func GetCPUMem() (memInfo, error) {
|
||||||
|
mem, err := getCPUMem()
|
||||||
|
if err != nil {
|
||||||
|
return memInfo{}, err
|
||||||
|
}
|
||||||
|
return getCPUMemByCgroups(mem), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func getCPUMem() (memInfo, error) {
|
||||||
var mem memInfo
|
var mem memInfo
|
||||||
var total, available, free, buffers, cached, freeSwap uint64
|
var total, available, free, buffers, cached, freeSwap uint64
|
||||||
f, err := os.Open("/proc/meminfo")
|
f, err := os.Open("/proc/meminfo")
|
||||||
@@ -56,6 +66,32 @@ func GetCPUMem() (memInfo, error) {
|
|||||||
return mem, nil
|
return mem, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func getCPUMemByCgroups(mem memInfo) memInfo {
|
||||||
|
total, err := getUint64ValueFromFile("/sys/fs/cgroup/memory.max")
|
||||||
|
if err == nil {
|
||||||
|
mem.TotalMemory = total
|
||||||
|
}
|
||||||
|
used, err := getUint64ValueFromFile("/sys/fs/cgroup/memory.current")
|
||||||
|
if err == nil {
|
||||||
|
mem.FreeMemory = mem.TotalMemory - used
|
||||||
|
}
|
||||||
|
return mem
|
||||||
|
}
|
||||||
|
|
||||||
|
func getUint64ValueFromFile(path string) (uint64, error) {
|
||||||
|
f, err := os.Open(path)
|
||||||
|
if err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
defer f.Close()
|
||||||
|
s := bufio.NewScanner(f)
|
||||||
|
for s.Scan() {
|
||||||
|
line := s.Text()
|
||||||
|
return strconv.ParseUint(line, 10, 64)
|
||||||
|
}
|
||||||
|
return 0, errors.New("empty file content")
|
||||||
|
}
|
||||||
|
|
||||||
const CpuInfoFilename = "/proc/cpuinfo"
|
const CpuInfoFilename = "/proc/cpuinfo"
|
||||||
|
|
||||||
type linuxCpuInfo struct {
|
type linuxCpuInfo struct {
|
||||||
@@ -74,7 +110,41 @@ func GetCPUDetails() []CPU {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
defer file.Close()
|
defer file.Close()
|
||||||
return linuxCPUDetails(file)
|
cpus := linuxCPUDetails(file)
|
||||||
|
return overwriteThreadCountByLinuxCgroups(cpus)
|
||||||
|
}
|
||||||
|
|
||||||
|
func overwriteThreadCountByLinuxCgroups(cpus []CPU) []CPU {
|
||||||
|
file, err := os.Open("/sys/fs/cgroup/cpu.max")
|
||||||
|
if err != nil {
|
||||||
|
return cpus
|
||||||
|
}
|
||||||
|
defer file.Close()
|
||||||
|
|
||||||
|
scanner := bufio.NewScanner(file)
|
||||||
|
for scanner.Scan() {
|
||||||
|
line := scanner.Text()
|
||||||
|
if sl := strings.Split(line, " "); len(sl) == 2 {
|
||||||
|
allowdUs, err := strconv.ParseInt(sl[0], 10, 64)
|
||||||
|
if err != nil {
|
||||||
|
slog.Warn("failed to parse CPU allowed micro secs", "error", err)
|
||||||
|
return cpus
|
||||||
|
}
|
||||||
|
unitUs, err := strconv.ParseInt(sl[1], 10, 64)
|
||||||
|
if err != nil {
|
||||||
|
slog.Warn("failed to parse CPU unit micro secs", "error", err)
|
||||||
|
return cpus
|
||||||
|
}
|
||||||
|
|
||||||
|
threads := int(max(allowdUs/unitUs, 1))
|
||||||
|
|
||||||
|
cpu := cpus[0]
|
||||||
|
cpu.CoreCount = threads
|
||||||
|
cpu.ThreadCount = threads
|
||||||
|
return []CPU{cpu}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return cpus
|
||||||
}
|
}
|
||||||
|
|
||||||
func linuxCPUDetails(file io.Reader) []CPU {
|
func linuxCPUDetails(file io.Reader) []CPU {
|
||||||
|
|||||||
Reference in New Issue
Block a user