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  }