152 lines
3.8 KiB
Go
152 lines
3.8 KiB
Go
// Package sysinfo 提供系统资源检测和 worker 数量自动计算功能。
|
||
package sysinfo
|
||
|
||
import (
|
||
"os/exec"
|
||
"runtime"
|
||
"strconv"
|
||
"strings"
|
||
)
|
||
|
||
// SystemInfo 系统资源信息
|
||
type SystemInfo struct {
|
||
CPUCores int // CPU 逻辑核心数
|
||
MemoryGB float64 // 内存大小（GB）
|
||
MemoryMB int // 内存大小（MB）
|
||
}
|
||
|
||
// GetSystemInfo 获取系统资源信息
|
||
func GetSystemInfo() SystemInfo {
|
||
info := SystemInfo{
|
||
CPUCores: runtime.NumCPU(),
|
||
}
|
||
|
||
// 获取内存信息
|
||
info.MemoryMB = getMemoryMB()
|
||
info.MemoryGB = float64(info.MemoryMB) / 1024.0
|
||
|
||
return info
|
||
}
|
||
|
||
// getMemoryMB 获取系统内存大小（MB）
|
||
func getMemoryMB() int {
|
||
switch runtime.GOOS {
|
||
case "windows":
|
||
return getMemoryWindows()
|
||
case "linux":
|
||
return getMemoryLinux()
|
||
case "darwin":
|
||
return getMemoryDarwin()
|
||
default:
|
||
return 0
|
||
}
|
||
}
|
||
|
||
// getMemoryWindows 使用 wmic 获取 Windows 内存
|
||
func getMemoryWindows() int {
|
||
cmd := exec.Command("wmic", "computersystem", "get", "totalphysicalmemory", "/value")
|
||
output, err := cmd.Output()
|
||
if err != nil {
|
||
return 0
|
||
}
|
||
|
||
// 解析输出：TotalPhysicalMemory=1234567890
|
||
for _, line := range strings.Split(string(output), "\n") {
|
||
line = strings.TrimSpace(line)
|
||
if strings.HasPrefix(line, "TotalPhysicalMemory=") {
|
||
value := strings.TrimPrefix(line, "TotalPhysicalMemory=")
|
||
if bytes, err := strconv.ParseInt(value, 10, 64); err == nil {
|
||
return int(bytes / 1024 / 1024)
|
||
}
|
||
}
|
||
}
|
||
return 0
|
||
}
|
||
|
||
// getMemoryLinux 从 /proc/meminfo 获取 Linux 内存
|
||
func getMemoryLinux() int {
|
||
cmd := exec.Command("cat", "/proc/meminfo")
|
||
output, err := cmd.Output()
|
||
if err != nil {
|
||
return 0
|
||
}
|
||
|
||
// 解析 MemTotal: 12345678 kB
|
||
for _, line := range strings.Split(string(output), "\n") {
|
||
if strings.HasPrefix(line, "MemTotal:") {
|
||
fields := strings.Fields(line)
|
||
if len(fields) >= 2 {
|
||
if kb, err := strconv.ParseInt(fields[1], 10, 64); err == nil {
|
||
return int(kb / 1024)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
return 0
|
||
}
|
||
|
||
// getMemoryDarwin 使用 sysctl 获取 macOS 内存
|
||
func getMemoryDarwin() int {
|
||
cmd := exec.Command("sysctl", "-n", "hw.memsize")
|
||
output, err := cmd.Output()
|
||
if err != nil {
|
||
return 0
|
||
}
|
||
|
||
value := strings.TrimSpace(string(output))
|
||
if bytes, err := strconv.ParseInt(value, 10, 64); err == nil {
|
||
return int(bytes / 1024 / 1024)
|
||
}
|
||
return 0
|
||
}
|
||
|
||
// WorkerConfig worker 配置
|
||
type WorkerConfig struct {
|
||
GoWorkers int // Go 端 worker 数量（探测/MIME/移动）
|
||
NsfwWorkers int // NudeNet Python 推理 worker 数量
|
||
AnimeNsfwWorkers int // WD Tagger Python 推理 worker 数量
|
||
}
|
||
|
||
// CalculateWorkerConfig 根据系统资源自动计算 worker 配置
|
||
// 规则：
|
||
// - Go worker: CPU核心数 - 2（留2核给Python推理和系统）
|
||
// - NudeNet worker: 因为模型小（~300MB），可以开更多，上限为 CPU核心数
|
||
// - WD Tagger worker: 因为模型大（~1.5GB），受内存限制，每个worker约需2GB
|
||
func CalculateWorkerConfig(info SystemInfo) WorkerConfig {
|
||
config := WorkerConfig{}
|
||
|
||
// Go worker: CPU核心数 - 2，最少为1
|
||
goWorkers := info.CPUCores - 2
|
||
if goWorkers < 1 {
|
||
goWorkers = 1
|
||
}
|
||
config.GoWorkers = goWorkers
|
||
|
||
// NudeNet worker: 模型小，可以开到 CPU核心数，但留一些给系统
|
||
nsfwWorkers := info.CPUCores
|
||
if nsfwWorkers > 12 {
|
||
nsfwWorkers = 12 // 上限12，避免过度并发
|
||
}
|
||
config.NsfwWorkers = nsfwWorkers
|
||
|
||
// WD Tagger worker: 受内存限制，每个worker约需2GB
|
||
// 预留4GB给系统和其他进程
|
||
availableMemoryGB := info.MemoryGB - 4
|
||
if availableMemoryGB < 0 {
|
||
availableMemoryGB = 0
|
||
}
|
||
animeWorkers := int(availableMemoryGB / 2)
|
||
if animeWorkers < 1 {
|
||
animeWorkers = 1
|
||
}
|
||
if animeWorkers > info.CPUCores {
|
||
animeWorkers = info.CPUCores // 不超过CPU核心数
|
||
}
|
||
if animeWorkers > 12 {
|
||
animeWorkers = 12 // 上限12
|
||
}
|
||
config.AnimeNsfwWorkers = animeWorkers
|
||
|
||
return config
|
||
}
|