Source of nvml-mock.go - NVIDIA GPU

Raw file
Source viewDiff to previous
5.25 KB
// GetUtilizationRates returns the GPU and memory utilization in that order (GPU, Memory).
 
package nvmlmock
​
import (
    "golang.zabbix.com/plugin/nvidia/pkg/nvml"
)
​
var (
    _ nvml.Device = (*MockDevice)(nil)
    _ nvml.Runner = (*MockRunner)(nil)
)
​
type MockRunner struct {
    IsInit         bool
    DriverVersion  string
    NvmlVersion    string
    DevicesByIndex []*nvml.NVMLDevice
    DevicesByUUID  map[string]*nvml.NVMLDevice
    WantedErr      error
}
​
type MockDevice struct {
    Temperature    int
    Name           string
    MemoryInfo     *nvml.MemoryInfo
    MemoryInfoV2   *nvml.MemoryInfoV2
    PcieThroughput uint
    FanSpeed       uint
    PowerUsage     uint
    WantedErr      error
}
​
func (m *MockRunner) Init() error {
    if m.WantedErr != nil {
        return m.WantedErr
    }
​
    if m.IsInit {
        return nvml.ErrAlreadyInitialized
    }
​
    m.IsInit = true
    return nil
}
​
func (m *MockRunner) InitV2() error {
    if m.WantedErr != nil {
        return m.WantedErr
    }
​
    if m.IsInit {
        return nvml.ErrAlreadyInitialized
    }
​
    m.IsInit = true
    return nil
}
​
// Initialize NVML and any necessary resources
func (m *MockRunner) GetDeviceCount() (uint, error) {
    return uint(len(m.DevicesByIndex)), nil
}
​
// Initialize NVML and any necessary resources