github.com/sentienttechnologies/studio-go-runner@v0.0.0-20201118202441-6d21f2ced8ee/internal/runner/cpu.go (about) 1 // Copyright 2018-2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License. 2 3 package runner 4 5 // This file contains functions and data structures used to handle CPU based hardware information, along 6 // with CPU and memory resource accounting 7 8 import ( 9 "fmt" 10 "sync" 11 12 "github.com/shirou/gopsutil/cpu" 13 memory "github.com/shirou/gopsutil/mem" 14 15 "github.com/dustin/go-humanize" 16 17 "github.com/go-stack/stack" 18 "github.com/jjeffery/kv" // MIT License 19 ) 20 21 type cpuTracker struct { 22 cpuInfo []cpu.InfoStat // CPU Information is static so cache it for later reference 23 24 AllocCores uint // The number of cores currently consumed and allocated 25 AllocMem uint64 // The amount of memory currently allocated 26 27 HardMaxCores uint // The number of cores that the hardware has provisioned 28 HardMaxMem uint64 // The amount of RAM the system has provisioned 29 30 SoftMaxCores uint // User specified limit on the number of cores to permit to be used in allocations 31 SoftMaxMem uint64 // User specified memory that is available for allocation 32 33 InitErr kv.Error // Any error that might have been recorded during initialization, if set this package may produce unexpected results 34 35 sync.Mutex 36 } 37 38 var ( 39 cpuTrack = &cpuTracker{} 40 ) 41 42 func init() { 43 cpuTrack.cpuInfo, _ = cpu.Info() 44 45 cpuTrack.HardMaxCores = uint(len(cpuTrack.cpuInfo)) 46 mem, err := memory.VirtualMemory() 47 if err != nil { 48 cpuTrack.InitErr = kv.Wrap(err).With("stack", stack.Trace().TrimRuntime()) 49 return 50 } 51 cpuTrack.HardMaxMem = mem.Available 52 53 cpuTrack.SoftMaxCores = cpuTrack.HardMaxCores 54 cpuTrack.SoftMaxMem = cpuTrack.HardMaxMem 55 } 56 57 // CPUAllocated is used to track an individual allocation of CPU 58 // resources that will be returned at a later time 59 // 60 type CPUAllocated struct { 61 cores uint 62 mem uint64 63 } 64 65 // CPUFree is used to retrieve information about the currently available CPU resources 66 // 67 func CPUFree() (cores uint, mem uint64) { 68 cpuTrack.Lock() 69 defer cpuTrack.Unlock() 70 71 return cpuTrack.SoftMaxCores - cpuTrack.AllocCores, 72 cpuTrack.SoftMaxMem - cpuTrack.AllocMem 73 } 74 75 // SetCPULimits is used to set the soft limits for the CPU that is premitted to be allocated to 76 // callers 77 // 78 func SetCPULimits(maxCores uint, maxMem uint64) (err kv.Error) { 79 80 cpuTrack.Lock() 81 defer cpuTrack.Unlock() 82 83 if cpuTrack.InitErr != nil { 84 return cpuTrack.InitErr 85 } 86 87 if maxCores > cpuTrack.HardMaxCores { 88 msg := fmt.Sprintf("new soft cores limit %d, violated hard limit %d", maxCores, cpuTrack.HardMaxCores) 89 return kv.NewError(msg).With("stack", stack.Trace().TrimRuntime()) 90 } 91 if maxMem > cpuTrack.HardMaxMem { 92 msg := fmt.Sprintf("new soft memory limit %d, violated hard limit %d", maxMem, cpuTrack.HardMaxMem) 93 return kv.NewError(msg).With("stack", stack.Trace().TrimRuntime()) 94 } 95 96 if maxCores == 0 { 97 cpuTrack.SoftMaxCores = cpuTrack.HardMaxCores 98 } else { 99 cpuTrack.SoftMaxCores = maxCores 100 } 101 102 if maxMem == 0 { 103 cpuTrack.SoftMaxMem = cpuTrack.HardMaxMem 104 } else { 105 cpuTrack.SoftMaxMem = maxMem 106 } 107 108 return nil 109 } 110 111 // AllocCPU is used by callers to attempt to allocate a CPU resource from the system, CPU affinity is not implemented 112 // and so this is soft accounting 113 // 114 // live can be used to test the capacity is sufficient for the request without making the request itself 115 // 116 func AllocCPU(maxCores uint, maxMem uint64, live bool) (alloc *CPUAllocated, err kv.Error) { 117 118 cpuTrack.Lock() 119 defer cpuTrack.Unlock() 120 121 if cpuTrack.InitErr != nil { 122 return nil, cpuTrack.InitErr 123 } 124 125 if maxCores+cpuTrack.AllocCores > cpuTrack.SoftMaxCores { 126 return nil, kv.NewError("insufficient CPU").With("cores_wanted", maxCores).With("cores_available", cpuTrack.SoftMaxCores-cpuTrack.AllocCores).With("stack", stack.Trace().TrimRuntime()) 127 } 128 if maxMem+cpuTrack.AllocMem > cpuTrack.SoftMaxMem { 129 msg := fmt.Sprintf("insufficient memory %s requested from pool of %s", humanize.Bytes(maxMem), humanize.Bytes(cpuTrack.SoftMaxMem)) 130 return nil, kv.NewError(msg).With("stack", stack.Trace().TrimRuntime()) 131 } 132 133 if !live { 134 return nil, nil 135 } 136 137 cpuTrack.AllocCores += maxCores 138 cpuTrack.AllocMem += maxMem 139 140 return &CPUAllocated{ 141 cores: maxCores, 142 mem: maxMem, 143 }, nil 144 } 145 146 // Release is used to return a soft allocation to the system accounting 147 // 148 func (cpu *CPUAllocated) Release() { 149 150 cpuTrack.Lock() 151 defer cpuTrack.Unlock() 152 153 if cpuTrack.InitErr != nil { 154 return 155 } 156 157 cpuTrack.AllocCores -= cpu.cores 158 cpuTrack.AllocMem -= cpu.mem 159 }