Update CDI dependency to v0.5.2

Signed-off-by: Evan Lezar <elezar@nvidia.com>

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/device/api.go

@@ -0,0 +1,63 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package device
+
+import (
+	"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml"
+)
+
+// Interface provides the API to the 'device' package
+type Interface interface {
+	GetDevices() ([]Device, error)
+	GetMigDevices() ([]MigDevice, error)
+	GetMigProfiles() ([]MigProfile, error)
+	NewDevice(d nvml.Device) (Device, error)
+	NewMigDevice(d nvml.Device) (MigDevice, error)
+	NewMigProfile(giProfileID, ciProfileID, ciEngProfileID int, migMemorySizeMB, deviceMemorySizeBytes uint64) (MigProfile, error)
+	ParseMigProfile(profile string) (MigProfile, error)
+	VisitDevices(func(i int, d Device) error) error
+	VisitMigDevices(func(i int, d Device, j int, m MigDevice) error) error
+	VisitMigProfiles(func(p MigProfile) error) error
+}
+
+type devicelib struct {
+	nvml nvml.Interface
+}
+
+var _ Interface = &devicelib{}
+
+// New creates a new instance of the 'device' interface
+func New(opts ...Option) Interface {
+	d := &devicelib{}
+	for _, opt := range opts {
+		opt(d)
+	}
+	if d.nvml == nil {
+		d.nvml = nvml.New()
+	}
+	return d
+}
+
+// WithNvml provides an Option to set the NVML library used by the 'device' interface
+func WithNvml(nvml nvml.Interface) Option {
+	return func(d *devicelib) {
+		d.nvml = nvml
+	}
+}
+
+// Option defines a function for passing options to the New() call
+type Option func(*devicelib)

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/device/device.go

@@ -0,0 +1,304 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package device
+
+import (
+	"fmt"
+
+	"github.com/NVIDIA/go-nvml/pkg/dl"
+	"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml"
+)
+
+// Device defines the set of extended functions associated with a device.Device
+type Device interface {
+	nvml.Device
+	GetMigDevices() ([]MigDevice, error)
+	GetMigProfiles() ([]MigProfile, error)
+	IsMigCapable() (bool, error)
+	IsMigEnabled() (bool, error)
+	VisitMigDevices(func(j int, m MigDevice) error) error
+	VisitMigProfiles(func(p MigProfile) error) error
+}
+
+type device struct {
+	nvml.Device
+	lib *devicelib
+}
+
+var _ Device = &device{}
+
+// NewDevice builds a new Device from an nvml.Device
+func (d *devicelib) NewDevice(dev nvml.Device) (Device, error) {
+	return &device{dev, d}, nil
+}
+
+// IsMigCapable checks if a device is capable of having MIG paprtitions created on it
+func (d *device) IsMigCapable() (bool, error) {
+	err := nvmlLookupSymbol("nvmlDeviceGetMigMode")
+	if err != nil {
+		return false, nil
+	}
+
+	_, _, ret := nvml.Device(d).GetMigMode()
+	if ret == nvml.ERROR_NOT_SUPPORTED {
+		return false, nil
+	}
+	if ret != nvml.SUCCESS {
+		return false, fmt.Errorf("error getting MIG mode: %v", ret)
+	}
+
+	return true, nil
+}
+
+// IsMigEnabled checks if a device has MIG mode currently enabled on it
+func (d *device) IsMigEnabled() (bool, error) {
+	err := nvmlLookupSymbol("nvmlDeviceGetMigMode")
+	if err != nil {
+		return false, nil
+	}
+
+	mode, _, ret := nvml.Device(d).GetMigMode()
+	if ret == nvml.ERROR_NOT_SUPPORTED {
+		return false, nil
+	}
+	if ret != nvml.SUCCESS {
+		return false, fmt.Errorf("error getting MIG mode: %v", ret)
+	}
+
+	return (mode == nvml.DEVICE_MIG_ENABLE), nil
+}
+
+// VisitMigDevices walks a top-level device and invokes a callback function for each MIG device configured on it
+func (d *device) VisitMigDevices(visit func(int, MigDevice) error) error {
+	count, ret := nvml.Device(d).GetMaxMigDeviceCount()
+	if ret != nvml.SUCCESS {
+		return fmt.Errorf("error getting max MIG device count: %v", ret)
+	}
+
+	for i := 0; i < count; i++ {
+		device, ret := nvml.Device(d).GetMigDeviceHandleByIndex(i)
+		if ret == nvml.ERROR_NOT_FOUND {
+			continue
+		}
+		if ret == nvml.ERROR_INVALID_ARGUMENT {
+			continue
+		}
+		if ret != nvml.SUCCESS {
+			return fmt.Errorf("error getting MIG device handle at index '%v': %v", i, ret)
+		}
+		mig, err := d.lib.NewMigDevice(device)
+		if err != nil {
+			return fmt.Errorf("error creating new MIG device wrapper: %v", err)
+		}
+		err = visit(i, mig)
+		if err != nil {
+			return fmt.Errorf("error visiting MIG device: %v", err)
+		}
+	}
+	return nil
+}
+
+// VisitMigProfiles walks a top-level device and invokes a callback function for each unique MIG Profile that can be configured on it
+func (d *device) VisitMigProfiles(visit func(MigProfile) error) error {
+	capable, err := d.IsMigCapable()
+	if err != nil {
+		return fmt.Errorf("error checking if GPU is MIG capable: %v", err)
+	}
+
+	if !capable {
+		return nil
+	}
+
+	memory, ret := d.GetMemoryInfo()
+	if ret != nvml.SUCCESS {
+		return fmt.Errorf("error getting device memory info: %v", ret)
+	}
+
+	for i := 0; i < nvml.GPU_INSTANCE_PROFILE_COUNT; i++ {
+		giProfileInfo, ret := d.GetGpuInstanceProfileInfo(i)
+		if ret != nvml.SUCCESS {
+			return fmt.Errorf("error getting GPU Instance profile info: %v", ret)
+		}
+
+		for j := 0; j < nvml.COMPUTE_INSTANCE_PROFILE_COUNT; j++ {
+			for k := 0; k < nvml.COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT; k++ {
+				p, err := d.lib.NewMigProfile(i, j, k, giProfileInfo.MemorySizeMB, memory.Total)
+				if err != nil {
+					return fmt.Errorf("error creating MIG profile: %v", err)
+				}
+
+				err = visit(p)
+				if err != nil {
+					return fmt.Errorf("error visiting MIG profile: %v", err)
+				}
+			}
+		}
+	}
+	return nil
+}
+
+// GetMigDevices gets the set of MIG devices associated with a top-level device
+func (d *device) GetMigDevices() ([]MigDevice, error) {
+	var migs []MigDevice
+	err := d.VisitMigDevices(func(j int, m MigDevice) error {
+		migs = append(migs, m)
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return migs, nil
+}
+
+// GetMigProfiles gets the set of unique MIG profiles associated with a top-level device
+func (d *device) GetMigProfiles() ([]MigProfile, error) {
+	var profiles []MigProfile
+	err := d.VisitMigProfiles(func(p MigProfile) error {
+		profiles = append(profiles, p)
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return profiles, nil
+}
+
+// VisitDevices visits each top-level device and invokes a callback function for it
+func (d *devicelib) VisitDevices(visit func(int, Device) error) error {
+	count, ret := d.nvml.DeviceGetCount()
+	if ret != nvml.SUCCESS {
+		return fmt.Errorf("error getting device count: %v", ret)
+	}
+
+	for i := 0; i < count; i++ {
+		device, ret := d.nvml.DeviceGetHandleByIndex(i)
+		if ret != nvml.SUCCESS {
+			return fmt.Errorf("error getting device handle for index '%v': %v", i, ret)
+		}
+		dev, err := d.NewDevice(device)
+		if err != nil {
+			return fmt.Errorf("error creating new device wrapper: %v", err)
+		}
+		err = visit(i, dev)
+		if err != nil {
+			return fmt.Errorf("error visiting device: %v", err)
+		}
+	}
+	return nil
+}
+
+// VisitMigDevices walks a top-level device and invokes a callback function for each MIG device configured on it
+func (d *devicelib) VisitMigDevices(visit func(int, Device, int, MigDevice) error) error {
+	err := d.VisitDevices(func(i int, dev Device) error {
+		err := dev.VisitMigDevices(func(j int, mig MigDevice) error {
+			err := visit(i, dev, j, mig)
+			if err != nil {
+				return fmt.Errorf("error visiting MIG device: %v", err)
+			}
+			return nil
+		})
+		if err != nil {
+			return fmt.Errorf("error visiting device: %v", err)
+		}
+		return nil
+	})
+	if err != nil {
+		return fmt.Errorf("error visiting devices: %v", err)
+	}
+	return nil
+}
+
+// VisitMigProfiles walks a top-level device and invokes a callback function for each unique MIG profile found on them
+func (d *devicelib) VisitMigProfiles(visit func(MigProfile) error) error {
+	visited := make(map[string]bool)
+	err := d.VisitDevices(func(i int, dev Device) error {
+		err := dev.VisitMigProfiles(func(p MigProfile) error {
+			if visited[p.String()] {
+				return nil
+			}
+
+			err := visit(p)
+			if err != nil {
+				return fmt.Errorf("error visiting MIG profile: %v", err)
+			}
+
+			visited[p.String()] = true
+			return nil
+		})
+		if err != nil {
+			return fmt.Errorf("error visiting device: %v", err)
+		}
+		return nil
+	})
+	if err != nil {
+		return fmt.Errorf("error visiting devices: %v", err)
+	}
+	return nil
+}
+
+// GetDevices gets the set of all top-level devices
+func (d *devicelib) GetDevices() ([]Device, error) {
+	var devs []Device
+	err := d.VisitDevices(func(i int, dev Device) error {
+		devs = append(devs, dev)
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return devs, nil
+}
+
+// GetMigDevices gets the set of MIG devices across all top-level devices
+func (d *devicelib) GetMigDevices() ([]MigDevice, error) {
+	var migs []MigDevice
+	err := d.VisitMigDevices(func(i int, dev Device, j int, m MigDevice) error {
+		migs = append(migs, m)
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return migs, nil
+}
+
+// GetMigProfiles gets the set of unique MIG profiles across all top-level devices
+func (d *devicelib) GetMigProfiles() ([]MigProfile, error) {
+	var profiles []MigProfile
+	err := d.VisitMigProfiles(func(p MigProfile) error {
+		profiles = append(profiles, p)
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return profiles, nil
+}
+
+// nvmlLookupSymbol checks to see if the given symbol is present in the NVML library
+func nvmlLookupSymbol(symbol string) error {
+	lib := dl.New("libnvidia-ml.so.1", dl.RTLD_LAZY|dl.RTLD_GLOBAL)
+	if lib == nil {
+		return fmt.Errorf("error instantiating DynamicLibrary for NVML")
+	}
+	err := lib.Open()
+	if err != nil {
+		return fmt.Errorf("error opening DynamicLibrary for NVML: %v", err)
+	}
+	defer lib.Close()
+	return lib.Lookup(symbol)
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/device/mig_device.go

@@ -0,0 +1,136 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package device
+
+import (
+	"fmt"
+
+	"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml"
+)
+
+// MigDevice defines the set of extended functions associated with a MIG device
+type MigDevice interface {
+	nvml.Device
+	GetProfile() (MigProfile, error)
+}
+
+type migdevice struct {
+	nvml.Device
+	lib     *devicelib
+	profile MigProfile
+}
+
+var _ MigDevice = &migdevice{}
+
+// NewMigDevice builds a new MigDevice from an nvml.Device
+func (d *devicelib) NewMigDevice(handle nvml.Device) (MigDevice, error) {
+	isMig, ret := handle.IsMigDeviceHandle()
+	if ret != nvml.SUCCESS {
+		return nil, fmt.Errorf("error checking if device is a MIG device: %v", ret)
+	}
+	if !isMig {
+		return nil, fmt.Errorf("not a MIG device")
+	}
+	return &migdevice{handle, d, nil}, nil
+}
+
+// GetProfile returns the MIG profile associated with a MIG device
+func (m *migdevice) GetProfile() (MigProfile, error) {
+	if m.profile != nil {
+		return m.profile, nil
+	}
+
+	parent, ret := m.Device.GetDeviceHandleFromMigDeviceHandle()
+	if ret != nvml.SUCCESS {
+		return nil, fmt.Errorf("error getting parent device handle: %v", ret)
+	}
+
+	parentMemoryInfo, ret := parent.GetMemoryInfo()
+	if ret != nvml.SUCCESS {
+		return nil, fmt.Errorf("error getting parent memory info: %v", ret)
+	}
+
+	attributes, ret := m.Device.GetAttributes()
+	if ret != nvml.SUCCESS {
+		return nil, fmt.Errorf("error getting MIG device attributes: %v", ret)
+	}
+
+	giID, ret := m.Device.GetGpuInstanceId()
+	if ret != nvml.SUCCESS {
+		return nil, fmt.Errorf("error getting MIG device GPU Instance ID: %v", ret)
+	}
+
+	ciID, ret := m.Device.GetComputeInstanceId()
+	if ret != nvml.SUCCESS {
+		return nil, fmt.Errorf("error getting MIG device Compute Instance ID: %v", ret)
+	}
+
+	gi, ret := parent.GetGpuInstanceById(giID)
+	if ret != nvml.SUCCESS {
+		return nil, fmt.Errorf("error getting GPU Instance: %v", ret)
+	}
+
+	ci, ret := gi.GetComputeInstanceById(ciID)
+	if ret != nvml.SUCCESS {
+		return nil, fmt.Errorf("error getting Compute Instance: %v", ret)
+	}
+
+	giInfo, ret := gi.GetInfo()
+	if ret != nvml.SUCCESS {
+		return nil, fmt.Errorf("error getting GPU Instance info: %v", ret)
+	}
+
+	ciInfo, ret := ci.GetInfo()
+	if ret != nvml.SUCCESS {
+		return nil, fmt.Errorf("error getting Compute Instance info: %v", ret)
+	}
+
+	for i := 0; i < nvml.GPU_INSTANCE_PROFILE_COUNT; i++ {
+		giProfileInfo, ret := parent.GetGpuInstanceProfileInfo(i)
+		if ret != nvml.SUCCESS {
+			return nil, fmt.Errorf("error getting GPU Instance profile info: %v", ret)
+		}
+
+		if giProfileInfo.Id != giInfo.ProfileId {
+			continue
+		}
+
+		for j := 0; j < nvml.COMPUTE_INSTANCE_PROFILE_COUNT; j++ {
+			for k := 0; k < nvml.COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT; k++ {
+				ciProfileInfo, ret := gi.GetComputeInstanceProfileInfo(j, k)
+				if ret != nvml.SUCCESS {
+					return nil, fmt.Errorf("error getting Compute Instance profile info: %v", ret)
+
+				}
+
+				if ciProfileInfo.Id != ciInfo.ProfileId {
+					continue
+				}
+
+				p, err := m.lib.NewMigProfile(i, j, k, attributes.MemorySizeMB, parentMemoryInfo.Total)
+				if err != nil {
+					return nil, fmt.Errorf("error creating MIG profile: %v", err)
+				}
+
+				m.profile = p
+				return p, nil
+			}
+		}
+	}
+
+	return nil, fmt.Errorf("no matching profile IDs found")
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/device/mig_profile.go

@@ -0,0 +1,332 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package device
+
+import (
+	"fmt"
+	"math"
+	"strconv"
+	"strings"
+
+	"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml"
+)
+
+const (
+	// AttributeMediaExtensions holds the string representation for the media extension MIG profile attribute.
+	AttributeMediaExtensions = "me"
+)
+
+// MigProfile represents a specific MIG profile.
+// Examples include "1g.5gb", "2g.10gb", "1c.2g.10gb", or "1c.1g.5gb+me", etc.
+type MigProfile interface {
+	String() string
+	GetInfo() MigProfileInfo
+	Equals(other MigProfile) bool
+}
+
+// MigProfileInfo holds all info associated with a specific MIG profile
+type MigProfileInfo struct {
+	C              int
+	G              int
+	GB             int
+	Attributes     []string
+	GIProfileID    int
+	CIProfileID    int
+	CIEngProfileID int
+}
+
+var _ MigProfile = &MigProfileInfo{}
+
+// NewProfile constructs a new Profile struct using info from the giProfiles and ciProfiles used to create it.
+func (d *devicelib) NewMigProfile(giProfileID, ciProfileID, ciEngProfileID int, migMemorySizeMB, deviceMemorySizeBytes uint64) (MigProfile, error) {
+	giSlices := 0
+	switch giProfileID {
+	case nvml.GPU_INSTANCE_PROFILE_1_SLICE:
+		giSlices = 1
+	case nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV1:
+		giSlices = 1
+	case nvml.GPU_INSTANCE_PROFILE_2_SLICE:
+		giSlices = 2
+	case nvml.GPU_INSTANCE_PROFILE_3_SLICE:
+		giSlices = 3
+	case nvml.GPU_INSTANCE_PROFILE_4_SLICE:
+		giSlices = 4
+	case nvml.GPU_INSTANCE_PROFILE_6_SLICE:
+		giSlices = 6
+	case nvml.GPU_INSTANCE_PROFILE_7_SLICE:
+		giSlices = 7
+	case nvml.GPU_INSTANCE_PROFILE_8_SLICE:
+		giSlices = 8
+	default:
+		return nil, fmt.Errorf("invalid GPU Instance Profile ID: %v", giProfileID)
+	}
+
+	ciSlices := 0
+	switch ciProfileID {
+	case nvml.COMPUTE_INSTANCE_PROFILE_1_SLICE:
+		ciSlices = 1
+	case nvml.COMPUTE_INSTANCE_PROFILE_2_SLICE:
+		ciSlices = 2
+	case nvml.COMPUTE_INSTANCE_PROFILE_3_SLICE:
+		ciSlices = 3
+	case nvml.COMPUTE_INSTANCE_PROFILE_4_SLICE:
+		ciSlices = 4
+	case nvml.COMPUTE_INSTANCE_PROFILE_6_SLICE:
+		ciSlices = 6
+	case nvml.COMPUTE_INSTANCE_PROFILE_7_SLICE:
+		ciSlices = 7
+	case nvml.COMPUTE_INSTANCE_PROFILE_8_SLICE:
+		ciSlices = 8
+	default:
+		return nil, fmt.Errorf("invalid Compute Instance Profile ID: %v", ciProfileID)
+	}
+
+	var attrs []string
+	switch giProfileID {
+	case nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV1:
+		attrs = append(attrs, AttributeMediaExtensions)
+	}
+
+	p := &MigProfileInfo{
+		C:              ciSlices,
+		G:              giSlices,
+		GB:             int(getMigMemorySizeGB(deviceMemorySizeBytes, migMemorySizeMB)),
+		Attributes:     attrs,
+		GIProfileID:    giProfileID,
+		CIProfileID:    ciProfileID,
+		CIEngProfileID: ciEngProfileID,
+	}
+
+	return p, nil
+}
+
+// ParseMigProfile converts a string representation of a MigProfile into an object
+func (d *devicelib) ParseMigProfile(profile string) (MigProfile, error) {
+	var err error
+	var c, g, gb int
+	var attrs []string
+
+	if len(profile) == 0 {
+		return nil, fmt.Errorf("empty Profile string")
+	}
+
+	split := strings.SplitN(profile, "+", 2)
+	if len(split) == 2 {
+		attrs, err = parseMigProfileAttributes(split[1])
+		if err != nil {
+			return nil, fmt.Errorf("error parsing attributes following '+' in Profile string: %v", err)
+		}
+	}
+
+	c, g, gb, err = parseMigProfileFields(split[0])
+	if err != nil {
+		return nil, fmt.Errorf("error parsing '.' separated fields in Profile string: %v", err)
+	}
+
+	p := &MigProfileInfo{
+		C:          c,
+		G:          g,
+		GB:         gb,
+		Attributes: attrs,
+	}
+
+	switch c {
+	case 1:
+		p.CIProfileID = nvml.COMPUTE_INSTANCE_PROFILE_1_SLICE
+	case 2:
+		p.CIProfileID = nvml.COMPUTE_INSTANCE_PROFILE_2_SLICE
+	case 3:
+		p.CIProfileID = nvml.COMPUTE_INSTANCE_PROFILE_3_SLICE
+	case 4:
+		p.CIProfileID = nvml.COMPUTE_INSTANCE_PROFILE_4_SLICE
+	case 6:
+		p.CIProfileID = nvml.COMPUTE_INSTANCE_PROFILE_6_SLICE
+	case 7:
+		p.CIProfileID = nvml.COMPUTE_INSTANCE_PROFILE_7_SLICE
+	case 8:
+		p.CIProfileID = nvml.COMPUTE_INSTANCE_PROFILE_8_SLICE
+	default:
+		return nil, fmt.Errorf("unknown Compute Instance slice size: %v", c)
+	}
+
+	switch g {
+	case 1:
+		p.GIProfileID = nvml.GPU_INSTANCE_PROFILE_1_SLICE
+	case 2:
+		p.GIProfileID = nvml.GPU_INSTANCE_PROFILE_2_SLICE
+	case 3:
+		p.GIProfileID = nvml.GPU_INSTANCE_PROFILE_3_SLICE
+	case 4:
+		p.GIProfileID = nvml.GPU_INSTANCE_PROFILE_4_SLICE
+	case 6:
+		p.GIProfileID = nvml.GPU_INSTANCE_PROFILE_6_SLICE
+	case 7:
+		p.GIProfileID = nvml.GPU_INSTANCE_PROFILE_7_SLICE
+	case 8:
+		p.GIProfileID = nvml.GPU_INSTANCE_PROFILE_8_SLICE
+	default:
+		return nil, fmt.Errorf("unknown GPU Instance slice size: %v", g)
+	}
+
+	p.CIEngProfileID = nvml.COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED
+
+	for _, a := range attrs {
+		switch a {
+		case AttributeMediaExtensions:
+			p.GIProfileID = nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV1
+		default:
+			return nil, fmt.Errorf("unknown Profile attribute: %v", a)
+		}
+	}
+
+	return p, nil
+}
+
+// String returns the string representation of a Profile
+func (p *MigProfileInfo) String() string {
+	var suffix string
+	if len(p.Attributes) > 0 {
+		suffix = "+" + strings.Join(p.Attributes, ",")
+	}
+	if p.C == p.G {
+		return fmt.Sprintf("%dg.%dgb%s", p.G, p.GB, suffix)
+	}
+	return fmt.Sprintf("%dc.%dg.%dgb%s", p.C, p.G, p.GB, suffix)
+}
+
+// GetInfo returns detailed info about a Profile
+func (p *MigProfileInfo) GetInfo() MigProfileInfo {
+	return *p
+}
+
+// Equals checks if two Profiles are identical or not
+func (p *MigProfileInfo) Equals(other MigProfile) bool {
+	switch o := other.(type) {
+	case *MigProfileInfo:
+		if p.C != o.C {
+			return false
+		}
+		if p.G != o.G {
+			return false
+		}
+		if p.GB != o.GB {
+			return false
+		}
+		if p.GIProfileID != o.GIProfileID {
+			return false
+		}
+		if p.CIProfileID != o.CIProfileID {
+			return false
+		}
+		if p.CIEngProfileID != o.CIEngProfileID {
+			return false
+		}
+		return true
+	}
+	return false
+}
+
+func parseMigProfileField(s string, field string) (int, error) {
+	if strings.TrimSpace(s) != s {
+		return -1, fmt.Errorf("leading or trailing spaces on '%%d%s'", field)
+	}
+
+	if !strings.HasSuffix(s, field) {
+		return -1, fmt.Errorf("missing '%s' from '%%d%s'", field, field)
+	}
+
+	v, err := strconv.Atoi(strings.TrimSuffix(s, field))
+	if err != nil {
+		return -1, fmt.Errorf("malformed number in '%%d%s'", field)
+	}
+
+	return v, nil
+}
+
+func parseMigProfileFields(s string) (int, int, int, error) {
+	var err error
+	var c, g, gb int
+
+	split := strings.SplitN(s, ".", 3)
+	if len(split) == 3 {
+		c, err = parseMigProfileField(split[0], "c")
+		if err != nil {
+			return -1, -1, -1, err
+		}
+		g, err = parseMigProfileField(split[1], "g")
+		if err != nil {
+			return -1, -1, -1, err
+		}
+		gb, err = parseMigProfileField(split[2], "gb")
+		if err != nil {
+			return -1, -1, -1, err
+		}
+		return c, g, gb, err
+	}
+	if len(split) == 2 {
+		g, err = parseMigProfileField(split[0], "g")
+		if err != nil {
+			return -1, -1, -1, err
+		}
+		gb, err = parseMigProfileField(split[1], "gb")
+		if err != nil {
+			return -1, -1, -1, err
+		}
+		return g, g, gb, nil
+	}
+
+	return -1, -1, -1, fmt.Errorf("parsed wrong number of fields, expected 2 or 3")
+}
+
+func parseMigProfileAttributes(s string) ([]string, error) {
+	attr := strings.Split(s, ",")
+	if len(attr) == 0 {
+		return nil, fmt.Errorf("empty attribute list")
+	}
+	unique := make(map[string]int)
+	for _, a := range attr {
+		if unique[a] > 0 {
+			return nil, fmt.Errorf("non unique attribute in list")
+		}
+		if a == "" {
+			return nil, fmt.Errorf("empty attribute in list")
+		}
+		if strings.TrimSpace(a) != a {
+			return nil, fmt.Errorf("leading or trailing spaces in attribute")
+		}
+		if a[0] >= '0' && a[0] <= '9' {
+			return nil, fmt.Errorf("attribute begins with a number")
+		}
+		for _, c := range a {
+			if (c < 'a' || c > 'z') && (c < 'A' || c > 'Z') && (c < '0' || c > '9') {
+				return nil, fmt.Errorf("non alpha-numeric character or digit in attribute")
+			}
+		}
+		unique[a]++
+	}
+	return attr, nil
+}
+
+func getMigMemorySizeGB(totalDeviceMemory, migMemorySizeMB uint64) uint64 {
+	const fracDenominator = 8
+	const oneMB = 1024 * 1024
+	const oneGB = 1024 * 1024 * 1024
+	fractionalGpuMem := (float64(migMemorySizeMB) * oneMB) / float64(totalDeviceMemory)
+	fractionalGpuMem = math.Ceil(fractionalGpuMem*fracDenominator) / fracDenominator
+	totalMemGB := float64((totalDeviceMemory + oneGB - 1) / oneGB)
+	return uint64(math.Round(fractionalGpuMem * totalMemGB))
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/ci.go

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package nvml
+
+import (
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+)
+
+type nvmlComputeInstance nvml.ComputeInstance
+
+var _ ComputeInstance = (*nvmlComputeInstance)(nil)
+
+// GetInfo() returns info about a Compute Instance
+func (ci nvmlComputeInstance) GetInfo() (ComputeInstanceInfo, Return) {
+	i, r := nvml.ComputeInstance(ci).GetInfo()
+	info := ComputeInstanceInfo{
+		Device:      nvmlDevice(i.Device),
+		GpuInstance: nvmlGpuInstance(i.GpuInstance),
+		Id:          i.Id,
+		ProfileId:   i.ProfileId,
+		Placement:   ComputeInstancePlacement(i.Placement),
+	}
+	return info, Return(r)
+}
+
+// Destroy() destroys a Compute Instance
+func (ci nvmlComputeInstance) Destroy() Return {
+	r := nvml.ComputeInstance(ci).Destroy()
+	return Return(r)
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/ci_mock.go

@@ -0,0 +1,102 @@
+// Code generated by moq; DO NOT EDIT.
+// github.com/matryer/moq
+
+package nvml
+
+import (
+	"sync"
+)
+
+// Ensure, that ComputeInstanceMock does implement ComputeInstance.
+// If this is not the case, regenerate this file with moq.
+var _ ComputeInstance = &ComputeInstanceMock{}
+
+// ComputeInstanceMock is a mock implementation of ComputeInstance.
+//
+// 	func TestSomethingThatUsesComputeInstance(t *testing.T) {
+//
+// 		// make and configure a mocked ComputeInstance
+// 		mockedComputeInstance := &ComputeInstanceMock{
+// 			DestroyFunc: func() Return {
+// 				panic("mock out the Destroy method")
+// 			},
+// 			GetInfoFunc: func() (ComputeInstanceInfo, Return) {
+// 				panic("mock out the GetInfo method")
+// 			},
+// 		}
+//
+// 		// use mockedComputeInstance in code that requires ComputeInstance
+// 		// and then make assertions.
+//
+// 	}
+type ComputeInstanceMock struct {
+	// DestroyFunc mocks the Destroy method.
+	DestroyFunc func() Return
+
+	// GetInfoFunc mocks the GetInfo method.
+	GetInfoFunc func() (ComputeInstanceInfo, Return)
+
+	// calls tracks calls to the methods.
+	calls struct {
+		// Destroy holds details about calls to the Destroy method.
+		Destroy []struct {
+		}
+		// GetInfo holds details about calls to the GetInfo method.
+		GetInfo []struct {
+		}
+	}
+	lockDestroy sync.RWMutex
+	lockGetInfo sync.RWMutex
+}
+
+// Destroy calls DestroyFunc.
+func (mock *ComputeInstanceMock) Destroy() Return {
+	if mock.DestroyFunc == nil {
+		panic("ComputeInstanceMock.DestroyFunc: method is nil but ComputeInstance.Destroy was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockDestroy.Lock()
+	mock.calls.Destroy = append(mock.calls.Destroy, callInfo)
+	mock.lockDestroy.Unlock()
+	return mock.DestroyFunc()
+}
+
+// DestroyCalls gets all the calls that were made to Destroy.
+// Check the length with:
+//     len(mockedComputeInstance.DestroyCalls())
+func (mock *ComputeInstanceMock) DestroyCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockDestroy.RLock()
+	calls = mock.calls.Destroy
+	mock.lockDestroy.RUnlock()
+	return calls
+}
+
+// GetInfo calls GetInfoFunc.
+func (mock *ComputeInstanceMock) GetInfo() (ComputeInstanceInfo, Return) {
+	if mock.GetInfoFunc == nil {
+		panic("ComputeInstanceMock.GetInfoFunc: method is nil but ComputeInstance.GetInfo was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetInfo.Lock()
+	mock.calls.GetInfo = append(mock.calls.GetInfo, callInfo)
+	mock.lockGetInfo.Unlock()
+	return mock.GetInfoFunc()
+}
+
+// GetInfoCalls gets all the calls that were made to GetInfo.
+// Check the length with:
+//     len(mockedComputeInstance.GetInfoCalls())
+func (mock *ComputeInstanceMock) GetInfoCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetInfo.RLock()
+	calls = mock.calls.GetInfo
+	mock.lockGetInfo.RUnlock()
+	return calls
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/consts.go

@@ -0,0 +1,94 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package nvml
+
+import (
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+)
+
+// Return constants
+const (
+	SUCCESS                       = Return(nvml.SUCCESS)
+	ERROR_UNINITIALIZED           = Return(nvml.ERROR_UNINITIALIZED)
+	ERROR_INVALID_ARGUMENT        = Return(nvml.ERROR_INVALID_ARGUMENT)
+	ERROR_NOT_SUPPORTED           = Return(nvml.ERROR_NOT_SUPPORTED)
+	ERROR_NO_PERMISSION           = Return(nvml.ERROR_NO_PERMISSION)
+	ERROR_ALREADY_INITIALIZED     = Return(nvml.ERROR_ALREADY_INITIALIZED)
+	ERROR_NOT_FOUND               = Return(nvml.ERROR_NOT_FOUND)
+	ERROR_INSUFFICIENT_SIZE       = Return(nvml.ERROR_INSUFFICIENT_SIZE)
+	ERROR_INSUFFICIENT_POWER      = Return(nvml.ERROR_INSUFFICIENT_POWER)
+	ERROR_DRIVER_NOT_LOADED       = Return(nvml.ERROR_DRIVER_NOT_LOADED)
+	ERROR_TIMEOUT                 = Return(nvml.ERROR_TIMEOUT)
+	ERROR_IRQ_ISSUE               = Return(nvml.ERROR_IRQ_ISSUE)
+	ERROR_LIBRARY_NOT_FOUND       = Return(nvml.ERROR_LIBRARY_NOT_FOUND)
+	ERROR_FUNCTION_NOT_FOUND      = Return(nvml.ERROR_FUNCTION_NOT_FOUND)
+	ERROR_CORRUPTED_INFOROM       = Return(nvml.ERROR_CORRUPTED_INFOROM)
+	ERROR_GPU_IS_LOST             = Return(nvml.ERROR_GPU_IS_LOST)
+	ERROR_RESET_REQUIRED          = Return(nvml.ERROR_RESET_REQUIRED)
+	ERROR_OPERATING_SYSTEM        = Return(nvml.ERROR_OPERATING_SYSTEM)
+	ERROR_LIB_RM_VERSION_MISMATCH = Return(nvml.ERROR_LIB_RM_VERSION_MISMATCH)
+	ERROR_IN_USE                  = Return(nvml.ERROR_IN_USE)
+	ERROR_MEMORY                  = Return(nvml.ERROR_MEMORY)
+	ERROR_NO_DATA                 = Return(nvml.ERROR_NO_DATA)
+	ERROR_VGPU_ECC_NOT_SUPPORTED  = Return(nvml.ERROR_VGPU_ECC_NOT_SUPPORTED)
+	ERROR_INSUFFICIENT_RESOURCES  = Return(nvml.ERROR_INSUFFICIENT_RESOURCES)
+	ERROR_UNKNOWN                 = Return(nvml.ERROR_UNKNOWN)
+)
+
+// MIG Mode constants
+const (
+	DEVICE_MIG_ENABLE  = nvml.DEVICE_MIG_ENABLE
+	DEVICE_MIG_DISABLE = nvml.DEVICE_MIG_DISABLE
+)
+
+// GPU Instance Profiles
+const (
+	GPU_INSTANCE_PROFILE_1_SLICE      = nvml.GPU_INSTANCE_PROFILE_1_SLICE
+	GPU_INSTANCE_PROFILE_2_SLICE      = nvml.GPU_INSTANCE_PROFILE_2_SLICE
+	GPU_INSTANCE_PROFILE_3_SLICE      = nvml.GPU_INSTANCE_PROFILE_3_SLICE
+	GPU_INSTANCE_PROFILE_4_SLICE      = nvml.GPU_INSTANCE_PROFILE_4_SLICE
+	GPU_INSTANCE_PROFILE_6_SLICE      = nvml.GPU_INSTANCE_PROFILE_6_SLICE
+	GPU_INSTANCE_PROFILE_7_SLICE      = nvml.GPU_INSTANCE_PROFILE_7_SLICE
+	GPU_INSTANCE_PROFILE_8_SLICE      = nvml.GPU_INSTANCE_PROFILE_8_SLICE
+	GPU_INSTANCE_PROFILE_1_SLICE_REV1 = nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV1
+	GPU_INSTANCE_PROFILE_COUNT        = nvml.GPU_INSTANCE_PROFILE_COUNT
+)
+
+// Compute Instance Profiles
+const (
+	COMPUTE_INSTANCE_PROFILE_1_SLICE = nvml.COMPUTE_INSTANCE_PROFILE_1_SLICE
+	COMPUTE_INSTANCE_PROFILE_2_SLICE = nvml.COMPUTE_INSTANCE_PROFILE_2_SLICE
+	COMPUTE_INSTANCE_PROFILE_3_SLICE = nvml.COMPUTE_INSTANCE_PROFILE_3_SLICE
+	COMPUTE_INSTANCE_PROFILE_4_SLICE = nvml.COMPUTE_INSTANCE_PROFILE_4_SLICE
+	COMPUTE_INSTANCE_PROFILE_6_SLICE = nvml.COMPUTE_INSTANCE_PROFILE_6_SLICE
+	COMPUTE_INSTANCE_PROFILE_7_SLICE = nvml.COMPUTE_INSTANCE_PROFILE_7_SLICE
+	COMPUTE_INSTANCE_PROFILE_8_SLICE = nvml.COMPUTE_INSTANCE_PROFILE_8_SLICE
+	COMPUTE_INSTANCE_PROFILE_COUNT   = nvml.COMPUTE_INSTANCE_PROFILE_COUNT
+)
+
+// Compute Instance Engine Profiles
+const (
+	COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED = nvml.COMPUTE_INSTANCE_ENGINE_PROFILE_SHARED
+	COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT  = nvml.COMPUTE_INSTANCE_ENGINE_PROFILE_COUNT
+)
+
+// Event Types
+const (
+	EventTypeXidCriticalError  = nvml.EventTypeXidCriticalError
+	EventTypeSingleBitEccError = nvml.EventTypeSingleBitEccError
+	EventTypeDoubleBitEccError = nvml.EventTypeDoubleBitEccError
+)

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/device.go

@@ -0,0 +1,152 @@
+/**
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+**/
+
+package nvml
+
+import "github.com/NVIDIA/go-nvml/pkg/nvml"
+
+type nvmlDevice nvml.Device
+
+var _ Device = (*nvmlDevice)(nil)
+
+// GetIndex returns the index of a Device
+func (d nvmlDevice) GetIndex() (int, Return) {
+	i, r := nvml.Device(d).GetIndex()
+	return i, Return(r)
+}
+
+// GetPciInfo returns the PCI info of a Device
+func (d nvmlDevice) GetPciInfo() (PciInfo, Return) {
+	p, r := nvml.Device(d).GetPciInfo()
+	return PciInfo(p), Return(r)
+}
+
+// GetMemoryInfo returns the memory info of a Device
+func (d nvmlDevice) GetMemoryInfo() (Memory, Return) {
+	p, r := nvml.Device(d).GetMemoryInfo()
+	return Memory(p), Return(r)
+}
+
+// GetUUID returns the UUID of a Device
+func (d nvmlDevice) GetUUID() (string, Return) {
+	u, r := nvml.Device(d).GetUUID()
+	return u, Return(r)
+}
+
+// GetMinorNumber returns the minor number of a Device
+func (d nvmlDevice) GetMinorNumber() (int, Return) {
+	m, r := nvml.Device(d).GetMinorNumber()
+	return m, Return(r)
+}
+
+// IsMigDeviceHandle returns whether a Device is a MIG device or not
+func (d nvmlDevice) IsMigDeviceHandle() (bool, Return) {
+	b, r := nvml.Device(d).IsMigDeviceHandle()
+	return b, Return(r)
+}
+
+// GetDeviceHandleFromMigDeviceHandle returns the parent Device of a MIG device
+func (d nvmlDevice) GetDeviceHandleFromMigDeviceHandle() (Device, Return) {
+	p, r := nvml.Device(d).GetDeviceHandleFromMigDeviceHandle()
+	return nvmlDevice(p), Return(r)
+}
+
+// SetMigMode sets the MIG mode of a Device
+func (d nvmlDevice) SetMigMode(mode int) (Return, Return) {
+	r1, r2 := nvml.Device(d).SetMigMode(mode)
+	return Return(r1), Return(r2)
+}
+
+// GetMigMode returns the MIG mode of a Device
+func (d nvmlDevice) GetMigMode() (int, int, Return) {
+	s1, s2, r := nvml.Device(d).GetMigMode()
+	return s1, s2, Return(r)
+}
+
+// GetGpuInstanceById returns the GPU Instance associated with a particular ID
+func (d nvmlDevice) GetGpuInstanceById(id int) (GpuInstance, Return) {
+	gi, r := nvml.Device(d).GetGpuInstanceById(id)
+	return nvmlGpuInstance(gi), Return(r)
+}
+
+// GetGpuInstanceProfileInfo returns the profile info of a GPU Instance
+func (d nvmlDevice) GetGpuInstanceProfileInfo(profile int) (GpuInstanceProfileInfo, Return) {
+	p, r := nvml.Device(d).GetGpuInstanceProfileInfo(profile)
+	return GpuInstanceProfileInfo(p), Return(r)
+}
+
+// GetGpuInstances returns the set of GPU Instances associated with a Device
+func (d nvmlDevice) GetGpuInstances(info *GpuInstanceProfileInfo) ([]GpuInstance, Return) {
+	nvmlGis, r := nvml.Device(d).GetGpuInstances((*nvml.GpuInstanceProfileInfo)(info))
+	var gis []GpuInstance
+	for _, gi := range nvmlGis {
+		gis = append(gis, nvmlGpuInstance(gi))
+	}
+	return gis, Return(r)
+}
+
+// GetMaxMigDeviceCount returns the maximum number of MIG devices that can be created on a Device
+func (d nvmlDevice) GetMaxMigDeviceCount() (int, Return) {
+	m, r := nvml.Device(d).GetMaxMigDeviceCount()
+	return m, Return(r)
+}
+
+// GetMigDeviceHandleByIndex returns the handle to a MIG device given its index
+func (d nvmlDevice) GetMigDeviceHandleByIndex(Index int) (Device, Return) {
+	h, r := nvml.Device(d).GetMigDeviceHandleByIndex(Index)
+	return nvmlDevice(h), Return(r)
+}
+
+// GetGpuInstanceId returns the GPU Instance ID of a MIG device
+func (d nvmlDevice) GetGpuInstanceId() (int, Return) {
+	gi, r := nvml.Device(d).GetGpuInstanceId()
+	return gi, Return(r)
+}
+
+// GetComputeInstanceId returns the Compute Instance ID of a MIG device
+func (d nvmlDevice) GetComputeInstanceId() (int, Return) {
+	ci, r := nvml.Device(d).GetComputeInstanceId()
+	return ci, Return(r)
+}
+
+// GetCudaComputeCapability returns the compute capability major and minor versions for a device
+func (d nvmlDevice) GetCudaComputeCapability() (int, int, Return) {
+	major, minor, r := nvml.Device(d).GetCudaComputeCapability()
+	return major, minor, Return(r)
+}
+
+// GetAttributes returns the device attributes for a MIG device
+func (d nvmlDevice) GetAttributes() (DeviceAttributes, Return) {
+	a, r := nvml.Device(d).GetAttributes()
+	return DeviceAttributes(a), Return(r)
+}
+
+// GetName returns the device attributes for a MIG device
+func (d nvmlDevice) GetName() (string, Return) {
+	n, r := nvml.Device(d).GetName()
+	return n, Return(r)
+}
+
+// RegisterEvents registers the specified event set and type with the device
+func (d nvmlDevice) RegisterEvents(EventTypes uint64, Set EventSet) Return {
+	return Return(nvml.Device(d).RegisterEvents(EventTypes, nvml.EventSet(Set)))
+}
+
+// GetSupportedEventTypes returns the events supported by the device
+func (d nvmlDevice) GetSupportedEventTypes() (uint64, Return) {
+	e, r := nvml.Device(d).GetSupportedEventTypes()
+	return e, Return(r)
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/device_mock.go

@@ -0,0 +1,834 @@
+// Code generated by moq; DO NOT EDIT.
+// github.com/matryer/moq
+
+package nvml
+
+import (
+	"sync"
+)
+
+// Ensure, that DeviceMock does implement Device.
+// If this is not the case, regenerate this file with moq.
+var _ Device = &DeviceMock{}
+
+// DeviceMock is a mock implementation of Device.
+//
+// 	func TestSomethingThatUsesDevice(t *testing.T) {
+//
+// 		// make and configure a mocked Device
+// 		mockedDevice := &DeviceMock{
+// 			GetAttributesFunc: func() (DeviceAttributes, Return) {
+// 				panic("mock out the GetAttributes method")
+// 			},
+// 			GetComputeInstanceIdFunc: func() (int, Return) {
+// 				panic("mock out the GetComputeInstanceId method")
+// 			},
+// 			GetCudaComputeCapabilityFunc: func() (int, int, Return) {
+// 				panic("mock out the GetCudaComputeCapability method")
+// 			},
+// 			GetDeviceHandleFromMigDeviceHandleFunc: func() (Device, Return) {
+// 				panic("mock out the GetDeviceHandleFromMigDeviceHandle method")
+// 			},
+// 			GetGpuInstanceByIdFunc: func(ID int) (GpuInstance, Return) {
+// 				panic("mock out the GetGpuInstanceById method")
+// 			},
+// 			GetGpuInstanceIdFunc: func() (int, Return) {
+// 				panic("mock out the GetGpuInstanceId method")
+// 			},
+// 			GetGpuInstanceProfileInfoFunc: func(Profile int) (GpuInstanceProfileInfo, Return) {
+// 				panic("mock out the GetGpuInstanceProfileInfo method")
+// 			},
+// 			GetGpuInstancesFunc: func(Info *GpuInstanceProfileInfo) ([]GpuInstance, Return) {
+// 				panic("mock out the GetGpuInstances method")
+// 			},
+// 			GetIndexFunc: func() (int, Return) {
+// 				panic("mock out the GetIndex method")
+// 			},
+// 			GetMaxMigDeviceCountFunc: func() (int, Return) {
+// 				panic("mock out the GetMaxMigDeviceCount method")
+// 			},
+// 			GetMemoryInfoFunc: func() (Memory, Return) {
+// 				panic("mock out the GetMemoryInfo method")
+// 			},
+// 			GetMigDeviceHandleByIndexFunc: func(Index int) (Device, Return) {
+// 				panic("mock out the GetMigDeviceHandleByIndex method")
+// 			},
+// 			GetMigModeFunc: func() (int, int, Return) {
+// 				panic("mock out the GetMigMode method")
+// 			},
+// 			GetMinorNumberFunc: func() (int, Return) {
+// 				panic("mock out the GetMinorNumber method")
+// 			},
+// 			GetNameFunc: func() (string, Return) {
+// 				panic("mock out the GetName method")
+// 			},
+// 			GetPciInfoFunc: func() (PciInfo, Return) {
+// 				panic("mock out the GetPciInfo method")
+// 			},
+// 			GetSupportedEventTypesFunc: func() (uint64, Return) {
+// 				panic("mock out the GetSupportedEventTypes method")
+// 			},
+// 			GetUUIDFunc: func() (string, Return) {
+// 				panic("mock out the GetUUID method")
+// 			},
+// 			IsMigDeviceHandleFunc: func() (bool, Return) {
+// 				panic("mock out the IsMigDeviceHandle method")
+// 			},
+// 			RegisterEventsFunc: func(v uint64, eventSet EventSet) Return {
+// 				panic("mock out the RegisterEvents method")
+// 			},
+// 			SetMigModeFunc: func(Mode int) (Return, Return) {
+// 				panic("mock out the SetMigMode method")
+// 			},
+// 		}
+//
+// 		// use mockedDevice in code that requires Device
+// 		// and then make assertions.
+//
+// 	}
+type DeviceMock struct {
+	// GetAttributesFunc mocks the GetAttributes method.
+	GetAttributesFunc func() (DeviceAttributes, Return)
+
+	// GetComputeInstanceIdFunc mocks the GetComputeInstanceId method.
+	GetComputeInstanceIdFunc func() (int, Return)
+
+	// GetCudaComputeCapabilityFunc mocks the GetCudaComputeCapability method.
+	GetCudaComputeCapabilityFunc func() (int, int, Return)
+
+	// GetDeviceHandleFromMigDeviceHandleFunc mocks the GetDeviceHandleFromMigDeviceHandle method.
+	GetDeviceHandleFromMigDeviceHandleFunc func() (Device, Return)
+
+	// GetGpuInstanceByIdFunc mocks the GetGpuInstanceById method.
+	GetGpuInstanceByIdFunc func(ID int) (GpuInstance, Return)
+
+	// GetGpuInstanceIdFunc mocks the GetGpuInstanceId method.
+	GetGpuInstanceIdFunc func() (int, Return)
+
+	// GetGpuInstanceProfileInfoFunc mocks the GetGpuInstanceProfileInfo method.
+	GetGpuInstanceProfileInfoFunc func(Profile int) (GpuInstanceProfileInfo, Return)
+
+	// GetGpuInstancesFunc mocks the GetGpuInstances method.
+	GetGpuInstancesFunc func(Info *GpuInstanceProfileInfo) ([]GpuInstance, Return)
+
+	// GetIndexFunc mocks the GetIndex method.
+	GetIndexFunc func() (int, Return)
+
+	// GetMaxMigDeviceCountFunc mocks the GetMaxMigDeviceCount method.
+	GetMaxMigDeviceCountFunc func() (int, Return)
+
+	// GetMemoryInfoFunc mocks the GetMemoryInfo method.
+	GetMemoryInfoFunc func() (Memory, Return)
+
+	// GetMigDeviceHandleByIndexFunc mocks the GetMigDeviceHandleByIndex method.
+	GetMigDeviceHandleByIndexFunc func(Index int) (Device, Return)
+
+	// GetMigModeFunc mocks the GetMigMode method.
+	GetMigModeFunc func() (int, int, Return)
+
+	// GetMinorNumberFunc mocks the GetMinorNumber method.
+	GetMinorNumberFunc func() (int, Return)
+
+	// GetNameFunc mocks the GetName method.
+	GetNameFunc func() (string, Return)
+
+	// GetPciInfoFunc mocks the GetPciInfo method.
+	GetPciInfoFunc func() (PciInfo, Return)
+
+	// GetSupportedEventTypesFunc mocks the GetSupportedEventTypes method.
+	GetSupportedEventTypesFunc func() (uint64, Return)
+
+	// GetUUIDFunc mocks the GetUUID method.
+	GetUUIDFunc func() (string, Return)
+
+	// IsMigDeviceHandleFunc mocks the IsMigDeviceHandle method.
+	IsMigDeviceHandleFunc func() (bool, Return)
+
+	// RegisterEventsFunc mocks the RegisterEvents method.
+	RegisterEventsFunc func(v uint64, eventSet EventSet) Return
+
+	// SetMigModeFunc mocks the SetMigMode method.
+	SetMigModeFunc func(Mode int) (Return, Return)
+
+	// calls tracks calls to the methods.
+	calls struct {
+		// GetAttributes holds details about calls to the GetAttributes method.
+		GetAttributes []struct {
+		}
+		// GetComputeInstanceId holds details about calls to the GetComputeInstanceId method.
+		GetComputeInstanceId []struct {
+		}
+		// GetCudaComputeCapability holds details about calls to the GetCudaComputeCapability method.
+		GetCudaComputeCapability []struct {
+		}
+		// GetDeviceHandleFromMigDeviceHandle holds details about calls to the GetDeviceHandleFromMigDeviceHandle method.
+		GetDeviceHandleFromMigDeviceHandle []struct {
+		}
+		// GetGpuInstanceById holds details about calls to the GetGpuInstanceById method.
+		GetGpuInstanceById []struct {
+			// ID is the ID argument value.
+			ID int
+		}
+		// GetGpuInstanceId holds details about calls to the GetGpuInstanceId method.
+		GetGpuInstanceId []struct {
+		}
+		// GetGpuInstanceProfileInfo holds details about calls to the GetGpuInstanceProfileInfo method.
+		GetGpuInstanceProfileInfo []struct {
+			// Profile is the Profile argument value.
+			Profile int
+		}
+		// GetGpuInstances holds details about calls to the GetGpuInstances method.
+		GetGpuInstances []struct {
+			// Info is the Info argument value.
+			Info *GpuInstanceProfileInfo
+		}
+		// GetIndex holds details about calls to the GetIndex method.
+		GetIndex []struct {
+		}
+		// GetMaxMigDeviceCount holds details about calls to the GetMaxMigDeviceCount method.
+		GetMaxMigDeviceCount []struct {
+		}
+		// GetMemoryInfo holds details about calls to the GetMemoryInfo method.
+		GetMemoryInfo []struct {
+		}
+		// GetMigDeviceHandleByIndex holds details about calls to the GetMigDeviceHandleByIndex method.
+		GetMigDeviceHandleByIndex []struct {
+			// Index is the Index argument value.
+			Index int
+		}
+		// GetMigMode holds details about calls to the GetMigMode method.
+		GetMigMode []struct {
+		}
+		// GetMinorNumber holds details about calls to the GetMinorNumber method.
+		GetMinorNumber []struct {
+		}
+		// GetName holds details about calls to the GetName method.
+		GetName []struct {
+		}
+		// GetPciInfo holds details about calls to the GetPciInfo method.
+		GetPciInfo []struct {
+		}
+		// GetSupportedEventTypes holds details about calls to the GetSupportedEventTypes method.
+		GetSupportedEventTypes []struct {
+		}
+		// GetUUID holds details about calls to the GetUUID method.
+		GetUUID []struct {
+		}
+		// IsMigDeviceHandle holds details about calls to the IsMigDeviceHandle method.
+		IsMigDeviceHandle []struct {
+		}
+		// RegisterEvents holds details about calls to the RegisterEvents method.
+		RegisterEvents []struct {
+			// V is the v argument value.
+			V uint64
+			// EventSet is the eventSet argument value.
+			EventSet EventSet
+		}
+		// SetMigMode holds details about calls to the SetMigMode method.
+		SetMigMode []struct {
+			// Mode is the Mode argument value.
+			Mode int
+		}
+	}
+	lockGetAttributes                      sync.RWMutex
+	lockGetComputeInstanceId               sync.RWMutex
+	lockGetCudaComputeCapability           sync.RWMutex
+	lockGetDeviceHandleFromMigDeviceHandle sync.RWMutex
+	lockGetGpuInstanceById                 sync.RWMutex
+	lockGetGpuInstanceId                   sync.RWMutex
+	lockGetGpuInstanceProfileInfo          sync.RWMutex
+	lockGetGpuInstances                    sync.RWMutex
+	lockGetIndex                           sync.RWMutex
+	lockGetMaxMigDeviceCount               sync.RWMutex
+	lockGetMemoryInfo                      sync.RWMutex
+	lockGetMigDeviceHandleByIndex          sync.RWMutex
+	lockGetMigMode                         sync.RWMutex
+	lockGetMinorNumber                     sync.RWMutex
+	lockGetName                            sync.RWMutex
+	lockGetPciInfo                         sync.RWMutex
+	lockGetSupportedEventTypes             sync.RWMutex
+	lockGetUUID                            sync.RWMutex
+	lockIsMigDeviceHandle                  sync.RWMutex
+	lockRegisterEvents                     sync.RWMutex
+	lockSetMigMode                         sync.RWMutex
+}
+
+// GetAttributes calls GetAttributesFunc.
+func (mock *DeviceMock) GetAttributes() (DeviceAttributes, Return) {
+	if mock.GetAttributesFunc == nil {
+		panic("DeviceMock.GetAttributesFunc: method is nil but Device.GetAttributes was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetAttributes.Lock()
+	mock.calls.GetAttributes = append(mock.calls.GetAttributes, callInfo)
+	mock.lockGetAttributes.Unlock()
+	return mock.GetAttributesFunc()
+}
+
+// GetAttributesCalls gets all the calls that were made to GetAttributes.
+// Check the length with:
+//     len(mockedDevice.GetAttributesCalls())
+func (mock *DeviceMock) GetAttributesCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetAttributes.RLock()
+	calls = mock.calls.GetAttributes
+	mock.lockGetAttributes.RUnlock()
+	return calls
+}
+
+// GetComputeInstanceId calls GetComputeInstanceIdFunc.
+func (mock *DeviceMock) GetComputeInstanceId() (int, Return) {
+	if mock.GetComputeInstanceIdFunc == nil {
+		panic("DeviceMock.GetComputeInstanceIdFunc: method is nil but Device.GetComputeInstanceId was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetComputeInstanceId.Lock()
+	mock.calls.GetComputeInstanceId = append(mock.calls.GetComputeInstanceId, callInfo)
+	mock.lockGetComputeInstanceId.Unlock()
+	return mock.GetComputeInstanceIdFunc()
+}
+
+// GetComputeInstanceIdCalls gets all the calls that were made to GetComputeInstanceId.
+// Check the length with:
+//     len(mockedDevice.GetComputeInstanceIdCalls())
+func (mock *DeviceMock) GetComputeInstanceIdCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetComputeInstanceId.RLock()
+	calls = mock.calls.GetComputeInstanceId
+	mock.lockGetComputeInstanceId.RUnlock()
+	return calls
+}
+
+// GetCudaComputeCapability calls GetCudaComputeCapabilityFunc.
+func (mock *DeviceMock) GetCudaComputeCapability() (int, int, Return) {
+	if mock.GetCudaComputeCapabilityFunc == nil {
+		panic("DeviceMock.GetCudaComputeCapabilityFunc: method is nil but Device.GetCudaComputeCapability was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetCudaComputeCapability.Lock()
+	mock.calls.GetCudaComputeCapability = append(mock.calls.GetCudaComputeCapability, callInfo)
+	mock.lockGetCudaComputeCapability.Unlock()
+	return mock.GetCudaComputeCapabilityFunc()
+}
+
+// GetCudaComputeCapabilityCalls gets all the calls that were made to GetCudaComputeCapability.
+// Check the length with:
+//     len(mockedDevice.GetCudaComputeCapabilityCalls())
+func (mock *DeviceMock) GetCudaComputeCapabilityCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetCudaComputeCapability.RLock()
+	calls = mock.calls.GetCudaComputeCapability
+	mock.lockGetCudaComputeCapability.RUnlock()
+	return calls
+}
+
+// GetDeviceHandleFromMigDeviceHandle calls GetDeviceHandleFromMigDeviceHandleFunc.
+func (mock *DeviceMock) GetDeviceHandleFromMigDeviceHandle() (Device, Return) {
+	if mock.GetDeviceHandleFromMigDeviceHandleFunc == nil {
+		panic("DeviceMock.GetDeviceHandleFromMigDeviceHandleFunc: method is nil but Device.GetDeviceHandleFromMigDeviceHandle was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetDeviceHandleFromMigDeviceHandle.Lock()
+	mock.calls.GetDeviceHandleFromMigDeviceHandle = append(mock.calls.GetDeviceHandleFromMigDeviceHandle, callInfo)
+	mock.lockGetDeviceHandleFromMigDeviceHandle.Unlock()
+	return mock.GetDeviceHandleFromMigDeviceHandleFunc()
+}
+
+// GetDeviceHandleFromMigDeviceHandleCalls gets all the calls that were made to GetDeviceHandleFromMigDeviceHandle.
+// Check the length with:
+//     len(mockedDevice.GetDeviceHandleFromMigDeviceHandleCalls())
+func (mock *DeviceMock) GetDeviceHandleFromMigDeviceHandleCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetDeviceHandleFromMigDeviceHandle.RLock()
+	calls = mock.calls.GetDeviceHandleFromMigDeviceHandle
+	mock.lockGetDeviceHandleFromMigDeviceHandle.RUnlock()
+	return calls
+}
+
+// GetGpuInstanceById calls GetGpuInstanceByIdFunc.
+func (mock *DeviceMock) GetGpuInstanceById(ID int) (GpuInstance, Return) {
+	if mock.GetGpuInstanceByIdFunc == nil {
+		panic("DeviceMock.GetGpuInstanceByIdFunc: method is nil but Device.GetGpuInstanceById was just called")
+	}
+	callInfo := struct {
+		ID int
+	}{
+		ID: ID,
+	}
+	mock.lockGetGpuInstanceById.Lock()
+	mock.calls.GetGpuInstanceById = append(mock.calls.GetGpuInstanceById, callInfo)
+	mock.lockGetGpuInstanceById.Unlock()
+	return mock.GetGpuInstanceByIdFunc(ID)
+}
+
+// GetGpuInstanceByIdCalls gets all the calls that were made to GetGpuInstanceById.
+// Check the length with:
+//     len(mockedDevice.GetGpuInstanceByIdCalls())
+func (mock *DeviceMock) GetGpuInstanceByIdCalls() []struct {
+	ID int
+} {
+	var calls []struct {
+		ID int
+	}
+	mock.lockGetGpuInstanceById.RLock()
+	calls = mock.calls.GetGpuInstanceById
+	mock.lockGetGpuInstanceById.RUnlock()
+	return calls
+}
+
+// GetGpuInstanceId calls GetGpuInstanceIdFunc.
+func (mock *DeviceMock) GetGpuInstanceId() (int, Return) {
+	if mock.GetGpuInstanceIdFunc == nil {
+		panic("DeviceMock.GetGpuInstanceIdFunc: method is nil but Device.GetGpuInstanceId was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetGpuInstanceId.Lock()
+	mock.calls.GetGpuInstanceId = append(mock.calls.GetGpuInstanceId, callInfo)
+	mock.lockGetGpuInstanceId.Unlock()
+	return mock.GetGpuInstanceIdFunc()
+}
+
+// GetGpuInstanceIdCalls gets all the calls that were made to GetGpuInstanceId.
+// Check the length with:
+//     len(mockedDevice.GetGpuInstanceIdCalls())
+func (mock *DeviceMock) GetGpuInstanceIdCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetGpuInstanceId.RLock()
+	calls = mock.calls.GetGpuInstanceId
+	mock.lockGetGpuInstanceId.RUnlock()
+	return calls
+}
+
+// GetGpuInstanceProfileInfo calls GetGpuInstanceProfileInfoFunc.
+func (mock *DeviceMock) GetGpuInstanceProfileInfo(Profile int) (GpuInstanceProfileInfo, Return) {
+	if mock.GetGpuInstanceProfileInfoFunc == nil {
+		panic("DeviceMock.GetGpuInstanceProfileInfoFunc: method is nil but Device.GetGpuInstanceProfileInfo was just called")
+	}
+	callInfo := struct {
+		Profile int
+	}{
+		Profile: Profile,
+	}
+	mock.lockGetGpuInstanceProfileInfo.Lock()
+	mock.calls.GetGpuInstanceProfileInfo = append(mock.calls.GetGpuInstanceProfileInfo, callInfo)
+	mock.lockGetGpuInstanceProfileInfo.Unlock()
+	return mock.GetGpuInstanceProfileInfoFunc(Profile)
+}
+
+// GetGpuInstanceProfileInfoCalls gets all the calls that were made to GetGpuInstanceProfileInfo.
+// Check the length with:
+//     len(mockedDevice.GetGpuInstanceProfileInfoCalls())
+func (mock *DeviceMock) GetGpuInstanceProfileInfoCalls() []struct {
+	Profile int
+} {
+	var calls []struct {
+		Profile int
+	}
+	mock.lockGetGpuInstanceProfileInfo.RLock()
+	calls = mock.calls.GetGpuInstanceProfileInfo
+	mock.lockGetGpuInstanceProfileInfo.RUnlock()
+	return calls
+}
+
+// GetGpuInstances calls GetGpuInstancesFunc.
+func (mock *DeviceMock) GetGpuInstances(Info *GpuInstanceProfileInfo) ([]GpuInstance, Return) {
+	if mock.GetGpuInstancesFunc == nil {
+		panic("DeviceMock.GetGpuInstancesFunc: method is nil but Device.GetGpuInstances was just called")
+	}
+	callInfo := struct {
+		Info *GpuInstanceProfileInfo
+	}{
+		Info: Info,
+	}
+	mock.lockGetGpuInstances.Lock()
+	mock.calls.GetGpuInstances = append(mock.calls.GetGpuInstances, callInfo)
+	mock.lockGetGpuInstances.Unlock()
+	return mock.GetGpuInstancesFunc(Info)
+}
+
+// GetGpuInstancesCalls gets all the calls that were made to GetGpuInstances.
+// Check the length with:
+//     len(mockedDevice.GetGpuInstancesCalls())
+func (mock *DeviceMock) GetGpuInstancesCalls() []struct {
+	Info *GpuInstanceProfileInfo
+} {
+	var calls []struct {
+		Info *GpuInstanceProfileInfo
+	}
+	mock.lockGetGpuInstances.RLock()
+	calls = mock.calls.GetGpuInstances
+	mock.lockGetGpuInstances.RUnlock()
+	return calls
+}
+
+// GetIndex calls GetIndexFunc.
+func (mock *DeviceMock) GetIndex() (int, Return) {
+	if mock.GetIndexFunc == nil {
+		panic("DeviceMock.GetIndexFunc: method is nil but Device.GetIndex was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetIndex.Lock()
+	mock.calls.GetIndex = append(mock.calls.GetIndex, callInfo)
+	mock.lockGetIndex.Unlock()
+	return mock.GetIndexFunc()
+}
+
+// GetIndexCalls gets all the calls that were made to GetIndex.
+// Check the length with:
+//     len(mockedDevice.GetIndexCalls())
+func (mock *DeviceMock) GetIndexCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetIndex.RLock()
+	calls = mock.calls.GetIndex
+	mock.lockGetIndex.RUnlock()
+	return calls
+}
+
+// GetMaxMigDeviceCount calls GetMaxMigDeviceCountFunc.
+func (mock *DeviceMock) GetMaxMigDeviceCount() (int, Return) {
+	if mock.GetMaxMigDeviceCountFunc == nil {
+		panic("DeviceMock.GetMaxMigDeviceCountFunc: method is nil but Device.GetMaxMigDeviceCount was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetMaxMigDeviceCount.Lock()
+	mock.calls.GetMaxMigDeviceCount = append(mock.calls.GetMaxMigDeviceCount, callInfo)
+	mock.lockGetMaxMigDeviceCount.Unlock()
+	return mock.GetMaxMigDeviceCountFunc()
+}
+
+// GetMaxMigDeviceCountCalls gets all the calls that were made to GetMaxMigDeviceCount.
+// Check the length with:
+//     len(mockedDevice.GetMaxMigDeviceCountCalls())
+func (mock *DeviceMock) GetMaxMigDeviceCountCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetMaxMigDeviceCount.RLock()
+	calls = mock.calls.GetMaxMigDeviceCount
+	mock.lockGetMaxMigDeviceCount.RUnlock()
+	return calls
+}
+
+// GetMemoryInfo calls GetMemoryInfoFunc.
+func (mock *DeviceMock) GetMemoryInfo() (Memory, Return) {
+	if mock.GetMemoryInfoFunc == nil {
+		panic("DeviceMock.GetMemoryInfoFunc: method is nil but Device.GetMemoryInfo was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetMemoryInfo.Lock()
+	mock.calls.GetMemoryInfo = append(mock.calls.GetMemoryInfo, callInfo)
+	mock.lockGetMemoryInfo.Unlock()
+	return mock.GetMemoryInfoFunc()
+}
+
+// GetMemoryInfoCalls gets all the calls that were made to GetMemoryInfo.
+// Check the length with:
+//     len(mockedDevice.GetMemoryInfoCalls())
+func (mock *DeviceMock) GetMemoryInfoCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetMemoryInfo.RLock()
+	calls = mock.calls.GetMemoryInfo
+	mock.lockGetMemoryInfo.RUnlock()
+	return calls
+}
+
+// GetMigDeviceHandleByIndex calls GetMigDeviceHandleByIndexFunc.
+func (mock *DeviceMock) GetMigDeviceHandleByIndex(Index int) (Device, Return) {
+	if mock.GetMigDeviceHandleByIndexFunc == nil {
+		panic("DeviceMock.GetMigDeviceHandleByIndexFunc: method is nil but Device.GetMigDeviceHandleByIndex was just called")
+	}
+	callInfo := struct {
+		Index int
+	}{
+		Index: Index,
+	}
+	mock.lockGetMigDeviceHandleByIndex.Lock()
+	mock.calls.GetMigDeviceHandleByIndex = append(mock.calls.GetMigDeviceHandleByIndex, callInfo)
+	mock.lockGetMigDeviceHandleByIndex.Unlock()
+	return mock.GetMigDeviceHandleByIndexFunc(Index)
+}
+
+// GetMigDeviceHandleByIndexCalls gets all the calls that were made to GetMigDeviceHandleByIndex.
+// Check the length with:
+//     len(mockedDevice.GetMigDeviceHandleByIndexCalls())
+func (mock *DeviceMock) GetMigDeviceHandleByIndexCalls() []struct {
+	Index int
+} {
+	var calls []struct {
+		Index int
+	}
+	mock.lockGetMigDeviceHandleByIndex.RLock()
+	calls = mock.calls.GetMigDeviceHandleByIndex
+	mock.lockGetMigDeviceHandleByIndex.RUnlock()
+	return calls
+}
+
+// GetMigMode calls GetMigModeFunc.
+func (mock *DeviceMock) GetMigMode() (int, int, Return) {
+	if mock.GetMigModeFunc == nil {
+		panic("DeviceMock.GetMigModeFunc: method is nil but Device.GetMigMode was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetMigMode.Lock()
+	mock.calls.GetMigMode = append(mock.calls.GetMigMode, callInfo)
+	mock.lockGetMigMode.Unlock()
+	return mock.GetMigModeFunc()
+}
+
+// GetMigModeCalls gets all the calls that were made to GetMigMode.
+// Check the length with:
+//     len(mockedDevice.GetMigModeCalls())
+func (mock *DeviceMock) GetMigModeCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetMigMode.RLock()
+	calls = mock.calls.GetMigMode
+	mock.lockGetMigMode.RUnlock()
+	return calls
+}
+
+// GetMinorNumber calls GetMinorNumberFunc.
+func (mock *DeviceMock) GetMinorNumber() (int, Return) {
+	if mock.GetMinorNumberFunc == nil {
+		panic("DeviceMock.GetMinorNumberFunc: method is nil but Device.GetMinorNumber was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetMinorNumber.Lock()
+	mock.calls.GetMinorNumber = append(mock.calls.GetMinorNumber, callInfo)
+	mock.lockGetMinorNumber.Unlock()
+	return mock.GetMinorNumberFunc()
+}
+
+// GetMinorNumberCalls gets all the calls that were made to GetMinorNumber.
+// Check the length with:
+//     len(mockedDevice.GetMinorNumberCalls())
+func (mock *DeviceMock) GetMinorNumberCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetMinorNumber.RLock()
+	calls = mock.calls.GetMinorNumber
+	mock.lockGetMinorNumber.RUnlock()
+	return calls
+}
+
+// GetName calls GetNameFunc.
+func (mock *DeviceMock) GetName() (string, Return) {
+	if mock.GetNameFunc == nil {
+		panic("DeviceMock.GetNameFunc: method is nil but Device.GetName was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetName.Lock()
+	mock.calls.GetName = append(mock.calls.GetName, callInfo)
+	mock.lockGetName.Unlock()
+	return mock.GetNameFunc()
+}
+
+// GetNameCalls gets all the calls that were made to GetName.
+// Check the length with:
+//     len(mockedDevice.GetNameCalls())
+func (mock *DeviceMock) GetNameCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetName.RLock()
+	calls = mock.calls.GetName
+	mock.lockGetName.RUnlock()
+	return calls
+}
+
+// GetPciInfo calls GetPciInfoFunc.
+func (mock *DeviceMock) GetPciInfo() (PciInfo, Return) {
+	if mock.GetPciInfoFunc == nil {
+		panic("DeviceMock.GetPciInfoFunc: method is nil but Device.GetPciInfo was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetPciInfo.Lock()
+	mock.calls.GetPciInfo = append(mock.calls.GetPciInfo, callInfo)
+	mock.lockGetPciInfo.Unlock()
+	return mock.GetPciInfoFunc()
+}
+
+// GetPciInfoCalls gets all the calls that were made to GetPciInfo.
+// Check the length with:
+//     len(mockedDevice.GetPciInfoCalls())
+func (mock *DeviceMock) GetPciInfoCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetPciInfo.RLock()
+	calls = mock.calls.GetPciInfo
+	mock.lockGetPciInfo.RUnlock()
+	return calls
+}
+
+// GetSupportedEventTypes calls GetSupportedEventTypesFunc.
+func (mock *DeviceMock) GetSupportedEventTypes() (uint64, Return) {
+	if mock.GetSupportedEventTypesFunc == nil {
+		panic("DeviceMock.GetSupportedEventTypesFunc: method is nil but Device.GetSupportedEventTypes was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetSupportedEventTypes.Lock()
+	mock.calls.GetSupportedEventTypes = append(mock.calls.GetSupportedEventTypes, callInfo)
+	mock.lockGetSupportedEventTypes.Unlock()
+	return mock.GetSupportedEventTypesFunc()
+}
+
+// GetSupportedEventTypesCalls gets all the calls that were made to GetSupportedEventTypes.
+// Check the length with:
+//     len(mockedDevice.GetSupportedEventTypesCalls())
+func (mock *DeviceMock) GetSupportedEventTypesCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetSupportedEventTypes.RLock()
+	calls = mock.calls.GetSupportedEventTypes
+	mock.lockGetSupportedEventTypes.RUnlock()
+	return calls
+}
+
+// GetUUID calls GetUUIDFunc.
+func (mock *DeviceMock) GetUUID() (string, Return) {
+	if mock.GetUUIDFunc == nil {
+		panic("DeviceMock.GetUUIDFunc: method is nil but Device.GetUUID was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetUUID.Lock()
+	mock.calls.GetUUID = append(mock.calls.GetUUID, callInfo)
+	mock.lockGetUUID.Unlock()
+	return mock.GetUUIDFunc()
+}
+
+// GetUUIDCalls gets all the calls that were made to GetUUID.
+// Check the length with:
+//     len(mockedDevice.GetUUIDCalls())
+func (mock *DeviceMock) GetUUIDCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetUUID.RLock()
+	calls = mock.calls.GetUUID
+	mock.lockGetUUID.RUnlock()
+	return calls
+}
+
+// IsMigDeviceHandle calls IsMigDeviceHandleFunc.
+func (mock *DeviceMock) IsMigDeviceHandle() (bool, Return) {
+	if mock.IsMigDeviceHandleFunc == nil {
+		panic("DeviceMock.IsMigDeviceHandleFunc: method is nil but Device.IsMigDeviceHandle was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockIsMigDeviceHandle.Lock()
+	mock.calls.IsMigDeviceHandle = append(mock.calls.IsMigDeviceHandle, callInfo)
+	mock.lockIsMigDeviceHandle.Unlock()
+	return mock.IsMigDeviceHandleFunc()
+}
+
+// IsMigDeviceHandleCalls gets all the calls that were made to IsMigDeviceHandle.
+// Check the length with:
+//     len(mockedDevice.IsMigDeviceHandleCalls())
+func (mock *DeviceMock) IsMigDeviceHandleCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockIsMigDeviceHandle.RLock()
+	calls = mock.calls.IsMigDeviceHandle
+	mock.lockIsMigDeviceHandle.RUnlock()
+	return calls
+}
+
+// RegisterEvents calls RegisterEventsFunc.
+func (mock *DeviceMock) RegisterEvents(v uint64, eventSet EventSet) Return {
+	if mock.RegisterEventsFunc == nil {
+		panic("DeviceMock.RegisterEventsFunc: method is nil but Device.RegisterEvents was just called")
+	}
+	callInfo := struct {
+		V        uint64
+		EventSet EventSet
+	}{
+		V:        v,
+		EventSet: eventSet,
+	}
+	mock.lockRegisterEvents.Lock()
+	mock.calls.RegisterEvents = append(mock.calls.RegisterEvents, callInfo)
+	mock.lockRegisterEvents.Unlock()
+	return mock.RegisterEventsFunc(v, eventSet)
+}
+
+// RegisterEventsCalls gets all the calls that were made to RegisterEvents.
+// Check the length with:
+//     len(mockedDevice.RegisterEventsCalls())
+func (mock *DeviceMock) RegisterEventsCalls() []struct {
+	V        uint64
+	EventSet EventSet
+} {
+	var calls []struct {
+		V        uint64
+		EventSet EventSet
+	}
+	mock.lockRegisterEvents.RLock()
+	calls = mock.calls.RegisterEvents
+	mock.lockRegisterEvents.RUnlock()
+	return calls
+}
+
+// SetMigMode calls SetMigModeFunc.
+func (mock *DeviceMock) SetMigMode(Mode int) (Return, Return) {
+	if mock.SetMigModeFunc == nil {
+		panic("DeviceMock.SetMigModeFunc: method is nil but Device.SetMigMode was just called")
+	}
+	callInfo := struct {
+		Mode int
+	}{
+		Mode: Mode,
+	}
+	mock.lockSetMigMode.Lock()
+	mock.calls.SetMigMode = append(mock.calls.SetMigMode, callInfo)
+	mock.lockSetMigMode.Unlock()
+	return mock.SetMigModeFunc(Mode)
+}
+
+// SetMigModeCalls gets all the calls that were made to SetMigMode.
+// Check the length with:
+//     len(mockedDevice.SetMigModeCalls())
+func (mock *DeviceMock) SetMigModeCalls() []struct {
+	Mode int
+} {
+	var calls []struct {
+		Mode int
+	}
+	mock.lockSetMigMode.RLock()
+	calls = mock.calls.SetMigMode
+	mock.lockSetMigMode.RUnlock()
+	return calls
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/event_set.go

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package nvml
+
+import (
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+)
+
+// Wait watches for an event with the specified timeout
+func (e EventSet) Wait(Timeoutms uint32) (EventData, Return) {
+	d, r := nvml.EventSet(e).Wait(Timeoutms)
+	eventData := EventData{
+		Device:            nvmlDevice(d.Device),
+		EventType:         d.EventType,
+		EventData:         d.EventData,
+		GpuInstanceId:     d.GpuInstanceId,
+		ComputeInstanceId: d.ComputeInstanceId,
+	}
+	return eventData, Return(r)
+}
+
+// Free deletes the event set
+func (e EventSet) Free() Return {
+	return Return(nvml.EventSet(e).Free())
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/gi.go

@@ -0,0 +1,71 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package nvml
+
+import (
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+)
+
+type nvmlGpuInstance nvml.GpuInstance
+
+var _ GpuInstance = (*nvmlGpuInstance)(nil)
+
+// GetInfo returns info about a GPU Intsance
+func (gi nvmlGpuInstance) GetInfo() (GpuInstanceInfo, Return) {
+	i, r := nvml.GpuInstance(gi).GetInfo()
+	info := GpuInstanceInfo{
+		Device:    nvmlDevice(i.Device),
+		Id:        i.Id,
+		ProfileId: i.ProfileId,
+		Placement: GpuInstancePlacement(i.Placement),
+	}
+	return info, Return(r)
+}
+
+// GetComputeInstanceById returns the Compute Instance associated with a particular ID.
+func (gi nvmlGpuInstance) GetComputeInstanceById(id int) (ComputeInstance, Return) {
+	ci, r := nvml.GpuInstance(gi).GetComputeInstanceById(id)
+	return nvmlComputeInstance(ci), Return(r)
+}
+
+// GetComputeInstanceProfileInfo returns info about a given Compute Instance profile
+func (gi nvmlGpuInstance) GetComputeInstanceProfileInfo(profile int, engProfile int) (ComputeInstanceProfileInfo, Return) {
+	p, r := nvml.GpuInstance(gi).GetComputeInstanceProfileInfo(profile, engProfile)
+	return ComputeInstanceProfileInfo(p), Return(r)
+}
+
+// CreateComputeInstance creates a Compute Instance within the GPU Instance
+func (gi nvmlGpuInstance) CreateComputeInstance(info *ComputeInstanceProfileInfo) (ComputeInstance, Return) {
+	ci, r := nvml.GpuInstance(gi).CreateComputeInstance((*nvml.ComputeInstanceProfileInfo)(info))
+	return nvmlComputeInstance(ci), Return(r)
+}
+
+// GetComputeInstances returns the set of Compute Instances associated with a GPU Instance
+func (gi nvmlGpuInstance) GetComputeInstances(info *ComputeInstanceProfileInfo) ([]ComputeInstance, Return) {
+	nvmlCis, r := nvml.GpuInstance(gi).GetComputeInstances((*nvml.ComputeInstanceProfileInfo)(info))
+	var cis []ComputeInstance
+	for _, ci := range nvmlCis {
+		cis = append(cis, nvmlComputeInstance(ci))
+	}
+	return cis, Return(r)
+}
+
+// Destroy destroys a GPU Instance
+func (gi nvmlGpuInstance) Destroy() Return {
+	r := nvml.GpuInstance(gi).Destroy()
+	return Return(r)
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/gi_mock.go

@@ -0,0 +1,280 @@
+// Code generated by moq; DO NOT EDIT.
+// github.com/matryer/moq
+
+package nvml
+
+import (
+	"sync"
+)
+
+// Ensure, that GpuInstanceMock does implement GpuInstance.
+// If this is not the case, regenerate this file with moq.
+var _ GpuInstance = &GpuInstanceMock{}
+
+// GpuInstanceMock is a mock implementation of GpuInstance.
+//
+// 	func TestSomethingThatUsesGpuInstance(t *testing.T) {
+//
+// 		// make and configure a mocked GpuInstance
+// 		mockedGpuInstance := &GpuInstanceMock{
+// 			CreateComputeInstanceFunc: func(Info *ComputeInstanceProfileInfo) (ComputeInstance, Return) {
+// 				panic("mock out the CreateComputeInstance method")
+// 			},
+// 			DestroyFunc: func() Return {
+// 				panic("mock out the Destroy method")
+// 			},
+// 			GetComputeInstanceByIdFunc: func(ID int) (ComputeInstance, Return) {
+// 				panic("mock out the GetComputeInstanceById method")
+// 			},
+// 			GetComputeInstanceProfileInfoFunc: func(Profile int, EngProfile int) (ComputeInstanceProfileInfo, Return) {
+// 				panic("mock out the GetComputeInstanceProfileInfo method")
+// 			},
+// 			GetComputeInstancesFunc: func(Info *ComputeInstanceProfileInfo) ([]ComputeInstance, Return) {
+// 				panic("mock out the GetComputeInstances method")
+// 			},
+// 			GetInfoFunc: func() (GpuInstanceInfo, Return) {
+// 				panic("mock out the GetInfo method")
+// 			},
+// 		}
+//
+// 		// use mockedGpuInstance in code that requires GpuInstance
+// 		// and then make assertions.
+//
+// 	}
+type GpuInstanceMock struct {
+	// CreateComputeInstanceFunc mocks the CreateComputeInstance method.
+	CreateComputeInstanceFunc func(Info *ComputeInstanceProfileInfo) (ComputeInstance, Return)
+
+	// DestroyFunc mocks the Destroy method.
+	DestroyFunc func() Return
+
+	// GetComputeInstanceByIdFunc mocks the GetComputeInstanceById method.
+	GetComputeInstanceByIdFunc func(ID int) (ComputeInstance, Return)
+
+	// GetComputeInstanceProfileInfoFunc mocks the GetComputeInstanceProfileInfo method.
+	GetComputeInstanceProfileInfoFunc func(Profile int, EngProfile int) (ComputeInstanceProfileInfo, Return)
+
+	// GetComputeInstancesFunc mocks the GetComputeInstances method.
+	GetComputeInstancesFunc func(Info *ComputeInstanceProfileInfo) ([]ComputeInstance, Return)
+
+	// GetInfoFunc mocks the GetInfo method.
+	GetInfoFunc func() (GpuInstanceInfo, Return)
+
+	// calls tracks calls to the methods.
+	calls struct {
+		// CreateComputeInstance holds details about calls to the CreateComputeInstance method.
+		CreateComputeInstance []struct {
+			// Info is the Info argument value.
+			Info *ComputeInstanceProfileInfo
+		}
+		// Destroy holds details about calls to the Destroy method.
+		Destroy []struct {
+		}
+		// GetComputeInstanceById holds details about calls to the GetComputeInstanceById method.
+		GetComputeInstanceById []struct {
+			// ID is the ID argument value.
+			ID int
+		}
+		// GetComputeInstanceProfileInfo holds details about calls to the GetComputeInstanceProfileInfo method.
+		GetComputeInstanceProfileInfo []struct {
+			// Profile is the Profile argument value.
+			Profile int
+			// EngProfile is the EngProfile argument value.
+			EngProfile int
+		}
+		// GetComputeInstances holds details about calls to the GetComputeInstances method.
+		GetComputeInstances []struct {
+			// Info is the Info argument value.
+			Info *ComputeInstanceProfileInfo
+		}
+		// GetInfo holds details about calls to the GetInfo method.
+		GetInfo []struct {
+		}
+	}
+	lockCreateComputeInstance         sync.RWMutex
+	lockDestroy                       sync.RWMutex
+	lockGetComputeInstanceById        sync.RWMutex
+	lockGetComputeInstanceProfileInfo sync.RWMutex
+	lockGetComputeInstances           sync.RWMutex
+	lockGetInfo                       sync.RWMutex
+}
+
+// CreateComputeInstance calls CreateComputeInstanceFunc.
+func (mock *GpuInstanceMock) CreateComputeInstance(Info *ComputeInstanceProfileInfo) (ComputeInstance, Return) {
+	if mock.CreateComputeInstanceFunc == nil {
+		panic("GpuInstanceMock.CreateComputeInstanceFunc: method is nil but GpuInstance.CreateComputeInstance was just called")
+	}
+	callInfo := struct {
+		Info *ComputeInstanceProfileInfo
+	}{
+		Info: Info,
+	}
+	mock.lockCreateComputeInstance.Lock()
+	mock.calls.CreateComputeInstance = append(mock.calls.CreateComputeInstance, callInfo)
+	mock.lockCreateComputeInstance.Unlock()
+	return mock.CreateComputeInstanceFunc(Info)
+}
+
+// CreateComputeInstanceCalls gets all the calls that were made to CreateComputeInstance.
+// Check the length with:
+//     len(mockedGpuInstance.CreateComputeInstanceCalls())
+func (mock *GpuInstanceMock) CreateComputeInstanceCalls() []struct {
+	Info *ComputeInstanceProfileInfo
+} {
+	var calls []struct {
+		Info *ComputeInstanceProfileInfo
+	}
+	mock.lockCreateComputeInstance.RLock()
+	calls = mock.calls.CreateComputeInstance
+	mock.lockCreateComputeInstance.RUnlock()
+	return calls
+}
+
+// Destroy calls DestroyFunc.
+func (mock *GpuInstanceMock) Destroy() Return {
+	if mock.DestroyFunc == nil {
+		panic("GpuInstanceMock.DestroyFunc: method is nil but GpuInstance.Destroy was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockDestroy.Lock()
+	mock.calls.Destroy = append(mock.calls.Destroy, callInfo)
+	mock.lockDestroy.Unlock()
+	return mock.DestroyFunc()
+}
+
+// DestroyCalls gets all the calls that were made to Destroy.
+// Check the length with:
+//     len(mockedGpuInstance.DestroyCalls())
+func (mock *GpuInstanceMock) DestroyCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockDestroy.RLock()
+	calls = mock.calls.Destroy
+	mock.lockDestroy.RUnlock()
+	return calls
+}
+
+// GetComputeInstanceById calls GetComputeInstanceByIdFunc.
+func (mock *GpuInstanceMock) GetComputeInstanceById(ID int) (ComputeInstance, Return) {
+	if mock.GetComputeInstanceByIdFunc == nil {
+		panic("GpuInstanceMock.GetComputeInstanceByIdFunc: method is nil but GpuInstance.GetComputeInstanceById was just called")
+	}
+	callInfo := struct {
+		ID int
+	}{
+		ID: ID,
+	}
+	mock.lockGetComputeInstanceById.Lock()
+	mock.calls.GetComputeInstanceById = append(mock.calls.GetComputeInstanceById, callInfo)
+	mock.lockGetComputeInstanceById.Unlock()
+	return mock.GetComputeInstanceByIdFunc(ID)
+}
+
+// GetComputeInstanceByIdCalls gets all the calls that were made to GetComputeInstanceById.
+// Check the length with:
+//     len(mockedGpuInstance.GetComputeInstanceByIdCalls())
+func (mock *GpuInstanceMock) GetComputeInstanceByIdCalls() []struct {
+	ID int
+} {
+	var calls []struct {
+		ID int
+	}
+	mock.lockGetComputeInstanceById.RLock()
+	calls = mock.calls.GetComputeInstanceById
+	mock.lockGetComputeInstanceById.RUnlock()
+	return calls
+}
+
+// GetComputeInstanceProfileInfo calls GetComputeInstanceProfileInfoFunc.
+func (mock *GpuInstanceMock) GetComputeInstanceProfileInfo(Profile int, EngProfile int) (ComputeInstanceProfileInfo, Return) {
+	if mock.GetComputeInstanceProfileInfoFunc == nil {
+		panic("GpuInstanceMock.GetComputeInstanceProfileInfoFunc: method is nil but GpuInstance.GetComputeInstanceProfileInfo was just called")
+	}
+	callInfo := struct {
+		Profile    int
+		EngProfile int
+	}{
+		Profile:    Profile,
+		EngProfile: EngProfile,
+	}
+	mock.lockGetComputeInstanceProfileInfo.Lock()
+	mock.calls.GetComputeInstanceProfileInfo = append(mock.calls.GetComputeInstanceProfileInfo, callInfo)
+	mock.lockGetComputeInstanceProfileInfo.Unlock()
+	return mock.GetComputeInstanceProfileInfoFunc(Profile, EngProfile)
+}
+
+// GetComputeInstanceProfileInfoCalls gets all the calls that were made to GetComputeInstanceProfileInfo.
+// Check the length with:
+//     len(mockedGpuInstance.GetComputeInstanceProfileInfoCalls())
+func (mock *GpuInstanceMock) GetComputeInstanceProfileInfoCalls() []struct {
+	Profile    int
+	EngProfile int
+} {
+	var calls []struct {
+		Profile    int
+		EngProfile int
+	}
+	mock.lockGetComputeInstanceProfileInfo.RLock()
+	calls = mock.calls.GetComputeInstanceProfileInfo
+	mock.lockGetComputeInstanceProfileInfo.RUnlock()
+	return calls
+}
+
+// GetComputeInstances calls GetComputeInstancesFunc.
+func (mock *GpuInstanceMock) GetComputeInstances(Info *ComputeInstanceProfileInfo) ([]ComputeInstance, Return) {
+	if mock.GetComputeInstancesFunc == nil {
+		panic("GpuInstanceMock.GetComputeInstancesFunc: method is nil but GpuInstance.GetComputeInstances was just called")
+	}
+	callInfo := struct {
+		Info *ComputeInstanceProfileInfo
+	}{
+		Info: Info,
+	}
+	mock.lockGetComputeInstances.Lock()
+	mock.calls.GetComputeInstances = append(mock.calls.GetComputeInstances, callInfo)
+	mock.lockGetComputeInstances.Unlock()
+	return mock.GetComputeInstancesFunc(Info)
+}
+
+// GetComputeInstancesCalls gets all the calls that were made to GetComputeInstances.
+// Check the length with:
+//     len(mockedGpuInstance.GetComputeInstancesCalls())
+func (mock *GpuInstanceMock) GetComputeInstancesCalls() []struct {
+	Info *ComputeInstanceProfileInfo
+} {
+	var calls []struct {
+		Info *ComputeInstanceProfileInfo
+	}
+	mock.lockGetComputeInstances.RLock()
+	calls = mock.calls.GetComputeInstances
+	mock.lockGetComputeInstances.RUnlock()
+	return calls
+}
+
+// GetInfo calls GetInfoFunc.
+func (mock *GpuInstanceMock) GetInfo() (GpuInstanceInfo, Return) {
+	if mock.GetInfoFunc == nil {
+		panic("GpuInstanceMock.GetInfoFunc: method is nil but GpuInstance.GetInfo was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockGetInfo.Lock()
+	mock.calls.GetInfo = append(mock.calls.GetInfo, callInfo)
+	mock.lockGetInfo.Unlock()
+	return mock.GetInfoFunc()
+}
+
+// GetInfoCalls gets all the calls that were made to GetInfo.
+// Check the length with:
+//     len(mockedGpuInstance.GetInfoCalls())
+func (mock *GpuInstanceMock) GetInfoCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockGetInfo.RLock()
+	calls = mock.calls.GetInfo
+	mock.lockGetInfo.RUnlock()
+	return calls
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/nvml.go

@@ -0,0 +1,110 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package nvml
+
+import (
+	"sync"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+)
+
+type nvmlLib struct {
+	sync.Mutex
+	refcount int
+}
+
+var _ Interface = (*nvmlLib)(nil)
+
+// New creates a new instance of the NVML Interface
+func New() Interface {
+	return &nvmlLib{}
+}
+
+// Init initializes an NVML Interface
+func (n *nvmlLib) Init() Return {
+	ret := nvml.Init()
+	if ret != nvml.SUCCESS {
+		return Return(ret)
+	}
+
+	n.Lock()
+	defer n.Unlock()
+	if n.refcount == 0 {
+		errorStringFunc = nvml.ErrorString
+	}
+	n.refcount++
+
+	return SUCCESS
+}
+
+// Shutdown shuts down an NVML Interface
+func (n *nvmlLib) Shutdown() Return {
+	ret := nvml.Shutdown()
+	if ret != nvml.SUCCESS {
+		return Return(ret)
+	}
+
+	n.Lock()
+	defer n.Unlock()
+	n.refcount--
+	if n.refcount == 0 {
+		errorStringFunc = defaultErrorStringFunc
+	}
+
+	return SUCCESS
+}
+
+// DeviceGetCount returns the total number of GPU Devices
+func (n *nvmlLib) DeviceGetCount() (int, Return) {
+	c, r := nvml.DeviceGetCount()
+	return c, Return(r)
+}
+
+// DeviceGetHandleByIndex returns a Device handle given its index
+func (n *nvmlLib) DeviceGetHandleByIndex(index int) (Device, Return) {
+	d, r := nvml.DeviceGetHandleByIndex(index)
+	return nvmlDevice(d), Return(r)
+}
+
+// DeviceGetHandleByUUID returns a Device handle given its UUID
+func (n *nvmlLib) DeviceGetHandleByUUID(uuid string) (Device, Return) {
+	d, r := nvml.DeviceGetHandleByUUID(uuid)
+	return nvmlDevice(d), Return(r)
+}
+
+// SystemGetDriverVersion returns the version of the installed NVIDIA driver
+func (n *nvmlLib) SystemGetDriverVersion() (string, Return) {
+	v, r := nvml.SystemGetDriverVersion()
+	return v, Return(r)
+}
+
+// SystemGetCudaDriverVersion returns the version of CUDA associated with the NVIDIA driver
+func (n *nvmlLib) SystemGetCudaDriverVersion() (int, Return) {
+	v, r := nvml.SystemGetCudaDriverVersion()
+	return v, Return(r)
+}
+
+// ErrorString returns the error string associated with a given return value
+func (n *nvmlLib) ErrorString(ret Return) string {
+	return nvml.ErrorString(nvml.Return(ret))
+}
+
+// EventSetCreate creates an event set
+func (n *nvmlLib) EventSetCreate() (EventSet, Return) {
+	e, r := nvml.EventSetCreate()
+	return EventSet(e), Return(r)
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/nvml_mock.go

@@ -0,0 +1,375 @@
+// Code generated by moq; DO NOT EDIT.
+// github.com/matryer/moq
+
+package nvml
+
+import (
+	"sync"
+)
+
+// Ensure, that InterfaceMock does implement Interface.
+// If this is not the case, regenerate this file with moq.
+var _ Interface = &InterfaceMock{}
+
+// InterfaceMock is a mock implementation of Interface.
+//
+// 	func TestSomethingThatUsesInterface(t *testing.T) {
+//
+// 		// make and configure a mocked Interface
+// 		mockedInterface := &InterfaceMock{
+// 			DeviceGetCountFunc: func() (int, Return) {
+// 				panic("mock out the DeviceGetCount method")
+// 			},
+// 			DeviceGetHandleByIndexFunc: func(Index int) (Device, Return) {
+// 				panic("mock out the DeviceGetHandleByIndex method")
+// 			},
+// 			DeviceGetHandleByUUIDFunc: func(UUID string) (Device, Return) {
+// 				panic("mock out the DeviceGetHandleByUUID method")
+// 			},
+// 			ErrorStringFunc: func(r Return) string {
+// 				panic("mock out the ErrorString method")
+// 			},
+// 			EventSetCreateFunc: func() (EventSet, Return) {
+// 				panic("mock out the EventSetCreate method")
+// 			},
+// 			InitFunc: func() Return {
+// 				panic("mock out the Init method")
+// 			},
+// 			ShutdownFunc: func() Return {
+// 				panic("mock out the Shutdown method")
+// 			},
+// 			SystemGetCudaDriverVersionFunc: func() (int, Return) {
+// 				panic("mock out the SystemGetCudaDriverVersion method")
+// 			},
+// 			SystemGetDriverVersionFunc: func() (string, Return) {
+// 				panic("mock out the SystemGetDriverVersion method")
+// 			},
+// 		}
+//
+// 		// use mockedInterface in code that requires Interface
+// 		// and then make assertions.
+//
+// 	}
+type InterfaceMock struct {
+	// DeviceGetCountFunc mocks the DeviceGetCount method.
+	DeviceGetCountFunc func() (int, Return)
+
+	// DeviceGetHandleByIndexFunc mocks the DeviceGetHandleByIndex method.
+	DeviceGetHandleByIndexFunc func(Index int) (Device, Return)
+
+	// DeviceGetHandleByUUIDFunc mocks the DeviceGetHandleByUUID method.
+	DeviceGetHandleByUUIDFunc func(UUID string) (Device, Return)
+
+	// ErrorStringFunc mocks the ErrorString method.
+	ErrorStringFunc func(r Return) string
+
+	// EventSetCreateFunc mocks the EventSetCreate method.
+	EventSetCreateFunc func() (EventSet, Return)
+
+	// InitFunc mocks the Init method.
+	InitFunc func() Return
+
+	// ShutdownFunc mocks the Shutdown method.
+	ShutdownFunc func() Return
+
+	// SystemGetCudaDriverVersionFunc mocks the SystemGetCudaDriverVersion method.
+	SystemGetCudaDriverVersionFunc func() (int, Return)
+
+	// SystemGetDriverVersionFunc mocks the SystemGetDriverVersion method.
+	SystemGetDriverVersionFunc func() (string, Return)
+
+	// calls tracks calls to the methods.
+	calls struct {
+		// DeviceGetCount holds details about calls to the DeviceGetCount method.
+		DeviceGetCount []struct {
+		}
+		// DeviceGetHandleByIndex holds details about calls to the DeviceGetHandleByIndex method.
+		DeviceGetHandleByIndex []struct {
+			// Index is the Index argument value.
+			Index int
+		}
+		// DeviceGetHandleByUUID holds details about calls to the DeviceGetHandleByUUID method.
+		DeviceGetHandleByUUID []struct {
+			// UUID is the UUID argument value.
+			UUID string
+		}
+		// ErrorString holds details about calls to the ErrorString method.
+		ErrorString []struct {
+			// R is the r argument value.
+			R Return
+		}
+		// EventSetCreate holds details about calls to the EventSetCreate method.
+		EventSetCreate []struct {
+		}
+		// Init holds details about calls to the Init method.
+		Init []struct {
+		}
+		// Shutdown holds details about calls to the Shutdown method.
+		Shutdown []struct {
+		}
+		// SystemGetCudaDriverVersion holds details about calls to the SystemGetCudaDriverVersion method.
+		SystemGetCudaDriverVersion []struct {
+		}
+		// SystemGetDriverVersion holds details about calls to the SystemGetDriverVersion method.
+		SystemGetDriverVersion []struct {
+		}
+	}
+	lockDeviceGetCount             sync.RWMutex
+	lockDeviceGetHandleByIndex     sync.RWMutex
+	lockDeviceGetHandleByUUID      sync.RWMutex
+	lockErrorString                sync.RWMutex
+	lockEventSetCreate             sync.RWMutex
+	lockInit                       sync.RWMutex
+	lockShutdown                   sync.RWMutex
+	lockSystemGetCudaDriverVersion sync.RWMutex
+	lockSystemGetDriverVersion     sync.RWMutex
+}
+
+// DeviceGetCount calls DeviceGetCountFunc.
+func (mock *InterfaceMock) DeviceGetCount() (int, Return) {
+	if mock.DeviceGetCountFunc == nil {
+		panic("InterfaceMock.DeviceGetCountFunc: method is nil but Interface.DeviceGetCount was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockDeviceGetCount.Lock()
+	mock.calls.DeviceGetCount = append(mock.calls.DeviceGetCount, callInfo)
+	mock.lockDeviceGetCount.Unlock()
+	return mock.DeviceGetCountFunc()
+}
+
+// DeviceGetCountCalls gets all the calls that were made to DeviceGetCount.
+// Check the length with:
+//     len(mockedInterface.DeviceGetCountCalls())
+func (mock *InterfaceMock) DeviceGetCountCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockDeviceGetCount.RLock()
+	calls = mock.calls.DeviceGetCount
+	mock.lockDeviceGetCount.RUnlock()
+	return calls
+}
+
+// DeviceGetHandleByIndex calls DeviceGetHandleByIndexFunc.
+func (mock *InterfaceMock) DeviceGetHandleByIndex(Index int) (Device, Return) {
+	if mock.DeviceGetHandleByIndexFunc == nil {
+		panic("InterfaceMock.DeviceGetHandleByIndexFunc: method is nil but Interface.DeviceGetHandleByIndex was just called")
+	}
+	callInfo := struct {
+		Index int
+	}{
+		Index: Index,
+	}
+	mock.lockDeviceGetHandleByIndex.Lock()
+	mock.calls.DeviceGetHandleByIndex = append(mock.calls.DeviceGetHandleByIndex, callInfo)
+	mock.lockDeviceGetHandleByIndex.Unlock()
+	return mock.DeviceGetHandleByIndexFunc(Index)
+}
+
+// DeviceGetHandleByIndexCalls gets all the calls that were made to DeviceGetHandleByIndex.
+// Check the length with:
+//     len(mockedInterface.DeviceGetHandleByIndexCalls())
+func (mock *InterfaceMock) DeviceGetHandleByIndexCalls() []struct {
+	Index int
+} {
+	var calls []struct {
+		Index int
+	}
+	mock.lockDeviceGetHandleByIndex.RLock()
+	calls = mock.calls.DeviceGetHandleByIndex
+	mock.lockDeviceGetHandleByIndex.RUnlock()
+	return calls
+}
+
+// DeviceGetHandleByUUID calls DeviceGetHandleByUUIDFunc.
+func (mock *InterfaceMock) DeviceGetHandleByUUID(UUID string) (Device, Return) {
+	if mock.DeviceGetHandleByUUIDFunc == nil {
+		panic("InterfaceMock.DeviceGetHandleByUUIDFunc: method is nil but Interface.DeviceGetHandleByUUID was just called")
+	}
+	callInfo := struct {
+		UUID string
+	}{
+		UUID: UUID,
+	}
+	mock.lockDeviceGetHandleByUUID.Lock()
+	mock.calls.DeviceGetHandleByUUID = append(mock.calls.DeviceGetHandleByUUID, callInfo)
+	mock.lockDeviceGetHandleByUUID.Unlock()
+	return mock.DeviceGetHandleByUUIDFunc(UUID)
+}
+
+// DeviceGetHandleByUUIDCalls gets all the calls that were made to DeviceGetHandleByUUID.
+// Check the length with:
+//     len(mockedInterface.DeviceGetHandleByUUIDCalls())
+func (mock *InterfaceMock) DeviceGetHandleByUUIDCalls() []struct {
+	UUID string
+} {
+	var calls []struct {
+		UUID string
+	}
+	mock.lockDeviceGetHandleByUUID.RLock()
+	calls = mock.calls.DeviceGetHandleByUUID
+	mock.lockDeviceGetHandleByUUID.RUnlock()
+	return calls
+}
+
+// ErrorString calls ErrorStringFunc.
+func (mock *InterfaceMock) ErrorString(r Return) string {
+	if mock.ErrorStringFunc == nil {
+		panic("InterfaceMock.ErrorStringFunc: method is nil but Interface.ErrorString was just called")
+	}
+	callInfo := struct {
+		R Return
+	}{
+		R: r,
+	}
+	mock.lockErrorString.Lock()
+	mock.calls.ErrorString = append(mock.calls.ErrorString, callInfo)
+	mock.lockErrorString.Unlock()
+	return mock.ErrorStringFunc(r)
+}
+
+// ErrorStringCalls gets all the calls that were made to ErrorString.
+// Check the length with:
+//     len(mockedInterface.ErrorStringCalls())
+func (mock *InterfaceMock) ErrorStringCalls() []struct {
+	R Return
+} {
+	var calls []struct {
+		R Return
+	}
+	mock.lockErrorString.RLock()
+	calls = mock.calls.ErrorString
+	mock.lockErrorString.RUnlock()
+	return calls
+}
+
+// EventSetCreate calls EventSetCreateFunc.
+func (mock *InterfaceMock) EventSetCreate() (EventSet, Return) {
+	if mock.EventSetCreateFunc == nil {
+		panic("InterfaceMock.EventSetCreateFunc: method is nil but Interface.EventSetCreate was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockEventSetCreate.Lock()
+	mock.calls.EventSetCreate = append(mock.calls.EventSetCreate, callInfo)
+	mock.lockEventSetCreate.Unlock()
+	return mock.EventSetCreateFunc()
+}
+
+// EventSetCreateCalls gets all the calls that were made to EventSetCreate.
+// Check the length with:
+//     len(mockedInterface.EventSetCreateCalls())
+func (mock *InterfaceMock) EventSetCreateCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockEventSetCreate.RLock()
+	calls = mock.calls.EventSetCreate
+	mock.lockEventSetCreate.RUnlock()
+	return calls
+}
+
+// Init calls InitFunc.
+func (mock *InterfaceMock) Init() Return {
+	if mock.InitFunc == nil {
+		panic("InterfaceMock.InitFunc: method is nil but Interface.Init was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockInit.Lock()
+	mock.calls.Init = append(mock.calls.Init, callInfo)
+	mock.lockInit.Unlock()
+	return mock.InitFunc()
+}
+
+// InitCalls gets all the calls that were made to Init.
+// Check the length with:
+//     len(mockedInterface.InitCalls())
+func (mock *InterfaceMock) InitCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockInit.RLock()
+	calls = mock.calls.Init
+	mock.lockInit.RUnlock()
+	return calls
+}
+
+// Shutdown calls ShutdownFunc.
+func (mock *InterfaceMock) Shutdown() Return {
+	if mock.ShutdownFunc == nil {
+		panic("InterfaceMock.ShutdownFunc: method is nil but Interface.Shutdown was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockShutdown.Lock()
+	mock.calls.Shutdown = append(mock.calls.Shutdown, callInfo)
+	mock.lockShutdown.Unlock()
+	return mock.ShutdownFunc()
+}
+
+// ShutdownCalls gets all the calls that were made to Shutdown.
+// Check the length with:
+//     len(mockedInterface.ShutdownCalls())
+func (mock *InterfaceMock) ShutdownCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockShutdown.RLock()
+	calls = mock.calls.Shutdown
+	mock.lockShutdown.RUnlock()
+	return calls
+}
+
+// SystemGetCudaDriverVersion calls SystemGetCudaDriverVersionFunc.
+func (mock *InterfaceMock) SystemGetCudaDriverVersion() (int, Return) {
+	if mock.SystemGetCudaDriverVersionFunc == nil {
+		panic("InterfaceMock.SystemGetCudaDriverVersionFunc: method is nil but Interface.SystemGetCudaDriverVersion was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockSystemGetCudaDriverVersion.Lock()
+	mock.calls.SystemGetCudaDriverVersion = append(mock.calls.SystemGetCudaDriverVersion, callInfo)
+	mock.lockSystemGetCudaDriverVersion.Unlock()
+	return mock.SystemGetCudaDriverVersionFunc()
+}
+
+// SystemGetCudaDriverVersionCalls gets all the calls that were made to SystemGetCudaDriverVersion.
+// Check the length with:
+//     len(mockedInterface.SystemGetCudaDriverVersionCalls())
+func (mock *InterfaceMock) SystemGetCudaDriverVersionCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockSystemGetCudaDriverVersion.RLock()
+	calls = mock.calls.SystemGetCudaDriverVersion
+	mock.lockSystemGetCudaDriverVersion.RUnlock()
+	return calls
+}
+
+// SystemGetDriverVersion calls SystemGetDriverVersionFunc.
+func (mock *InterfaceMock) SystemGetDriverVersion() (string, Return) {
+	if mock.SystemGetDriverVersionFunc == nil {
+		panic("InterfaceMock.SystemGetDriverVersionFunc: method is nil but Interface.SystemGetDriverVersion was just called")
+	}
+	callInfo := struct {
+	}{}
+	mock.lockSystemGetDriverVersion.Lock()
+	mock.calls.SystemGetDriverVersion = append(mock.calls.SystemGetDriverVersion, callInfo)
+	mock.lockSystemGetDriverVersion.Unlock()
+	return mock.SystemGetDriverVersionFunc()
+}
+
+// SystemGetDriverVersionCalls gets all the calls that were made to SystemGetDriverVersion.
+// Check the length with:
+//     len(mockedInterface.SystemGetDriverVersionCalls())
+func (mock *InterfaceMock) SystemGetDriverVersionCalls() []struct {
+} {
+	var calls []struct {
+	}
+	mock.lockSystemGetDriverVersion.RLock()
+	calls = mock.calls.SystemGetDriverVersion
+	mock.lockSystemGetDriverVersion.RUnlock()
+	return calls
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/return.go

@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package nvml
+
+import (
+	"fmt"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+)
+
+// String returns the string representation of a Return
+func (r Return) String() string {
+	return errorStringFunc(nvml.Return(r))
+}
+
+// Error returns the string representation of a Return
+func (r Return) Error() string {
+	return errorStringFunc(nvml.Return(r))
+}
+
+// Assigned to nvml.ErrorString if the system nvml library is in use
+var errorStringFunc = defaultErrorStringFunc
+
+var defaultErrorStringFunc = func(r nvml.Return) string {
+	switch Return(r) {
+	case SUCCESS:
+		return "SUCCESS"
+	case ERROR_UNINITIALIZED:
+		return "ERROR_UNINITIALIZED"
+	case ERROR_INVALID_ARGUMENT:
+		return "ERROR_INVALID_ARGUMENT"
+	case ERROR_NOT_SUPPORTED:
+		return "ERROR_NOT_SUPPORTED"
+	case ERROR_NO_PERMISSION:
+		return "ERROR_NO_PERMISSION"
+	case ERROR_ALREADY_INITIALIZED:
+		return "ERROR_ALREADY_INITIALIZED"
+	case ERROR_NOT_FOUND:
+		return "ERROR_NOT_FOUND"
+	case ERROR_INSUFFICIENT_SIZE:
+		return "ERROR_INSUFFICIENT_SIZE"
+	case ERROR_INSUFFICIENT_POWER:
+		return "ERROR_INSUFFICIENT_POWER"
+	case ERROR_DRIVER_NOT_LOADED:
+		return "ERROR_DRIVER_NOT_LOADED"
+	case ERROR_TIMEOUT:
+		return "ERROR_TIMEOUT"
+	case ERROR_IRQ_ISSUE:
+		return "ERROR_IRQ_ISSUE"
+	case ERROR_LIBRARY_NOT_FOUND:
+		return "ERROR_LIBRARY_NOT_FOUND"
+	case ERROR_FUNCTION_NOT_FOUND:
+		return "ERROR_FUNCTION_NOT_FOUND"
+	case ERROR_CORRUPTED_INFOROM:
+		return "ERROR_CORRUPTED_INFOROM"
+	case ERROR_GPU_IS_LOST:
+		return "ERROR_GPU_IS_LOST"
+	case ERROR_RESET_REQUIRED:
+		return "ERROR_RESET_REQUIRED"
+	case ERROR_OPERATING_SYSTEM:
+		return "ERROR_OPERATING_SYSTEM"
+	case ERROR_LIB_RM_VERSION_MISMATCH:
+		return "ERROR_LIB_RM_VERSION_MISMATCH"
+	case ERROR_IN_USE:
+		return "ERROR_IN_USE"
+	case ERROR_MEMORY:
+		return "ERROR_MEMORY"
+	case ERROR_NO_DATA:
+		return "ERROR_NO_DATA"
+	case ERROR_VGPU_ECC_NOT_SUPPORTED:
+		return "ERROR_VGPU_ECC_NOT_SUPPORTED"
+	case ERROR_INSUFFICIENT_RESOURCES:
+		return "ERROR_INSUFFICIENT_RESOURCES"
+	case ERROR_UNKNOWN:
+		return "ERROR_UNKNOWN"
+	default:
+		return fmt.Sprintf("Unknown return value: %d", r)
+	}
+}

vendor/gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml/types.go

@@ -0,0 +1,136 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package nvml
+
+import (
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+)
+
+// Interface defines the functions implemented by an NVML library
+//
+//go:generate moq -out nvml_mock.go . Interface
+type Interface interface {
+	DeviceGetCount() (int, Return)
+	DeviceGetHandleByIndex(Index int) (Device, Return)
+	DeviceGetHandleByUUID(UUID string) (Device, Return)
+	ErrorString(r Return) string
+	EventSetCreate() (EventSet, Return)
+	Init() Return
+	Shutdown() Return
+	SystemGetCudaDriverVersion() (int, Return)
+	SystemGetDriverVersion() (string, Return)
+}
+
+// Device defines the functions implemented by an NVML device
+//
+//go:generate moq -out device_mock.go . Device
+type Device interface {
+	GetAttributes() (DeviceAttributes, Return)
+	GetComputeInstanceId() (int, Return)
+	GetCudaComputeCapability() (int, int, Return)
+	GetDeviceHandleFromMigDeviceHandle() (Device, Return)
+	GetGpuInstanceById(ID int) (GpuInstance, Return)
+	GetGpuInstanceId() (int, Return)
+	GetGpuInstanceProfileInfo(Profile int) (GpuInstanceProfileInfo, Return)
+	GetGpuInstances(Info *GpuInstanceProfileInfo) ([]GpuInstance, Return)
+	GetIndex() (int, Return)
+	GetMaxMigDeviceCount() (int, Return)
+	GetMemoryInfo() (Memory, Return)
+	GetMigDeviceHandleByIndex(Index int) (Device, Return)
+	GetMigMode() (int, int, Return)
+	GetMinorNumber() (int, Return)
+	GetName() (string, Return)
+	GetPciInfo() (PciInfo, Return)
+	GetSupportedEventTypes() (uint64, Return)
+	GetUUID() (string, Return)
+	IsMigDeviceHandle() (bool, Return)
+	RegisterEvents(uint64, EventSet) Return
+	SetMigMode(Mode int) (Return, Return)
+}
+
+// GpuInstance defines the functions implemented by a GpuInstance
+//
+//go:generate moq -out gi_mock.go . GpuInstance
+type GpuInstance interface {
+	CreateComputeInstance(Info *ComputeInstanceProfileInfo) (ComputeInstance, Return)
+	Destroy() Return
+	GetComputeInstanceById(ID int) (ComputeInstance, Return)
+	GetComputeInstanceProfileInfo(Profile int, EngProfile int) (ComputeInstanceProfileInfo, Return)
+	GetComputeInstances(Info *ComputeInstanceProfileInfo) ([]ComputeInstance, Return)
+	GetInfo() (GpuInstanceInfo, Return)
+}
+
+// ComputeInstance defines the functions implemented by a ComputeInstance
+//
+//go:generate moq -out ci_mock.go . ComputeInstance
+type ComputeInstance interface {
+	Destroy() Return
+	GetInfo() (ComputeInstanceInfo, Return)
+}
+
+// GpuInstanceInfo holds info about a GPU Instance
+type GpuInstanceInfo struct {
+	Device    Device
+	Id        uint32
+	ProfileId uint32
+	Placement GpuInstancePlacement
+}
+
+// ComputeInstanceInfo holds info about a Compute Instance
+type ComputeInstanceInfo struct {
+	Device      Device
+	GpuInstance GpuInstance
+	Id          uint32
+	ProfileId   uint32
+	Placement   ComputeInstancePlacement
+}
+
+// EventData defines NVML event Data
+type EventData struct {
+	Device            Device
+	EventType         uint64
+	EventData         uint64
+	GpuInstanceId     uint32
+	ComputeInstanceId uint32
+}
+
+// EventSet defines NVML event Data
+type EventSet nvml.EventSet
+
+// Return defines an NVML return type
+type Return nvml.Return
+
+// Memory holds info about GPU device memory
+type Memory nvml.Memory
+
+// PciInfo holds info about the PCI connections of a GPU dvice
+type PciInfo nvml.PciInfo
+
+// GpuInstanceProfileInfo holds info about a GPU Instance Profile
+type GpuInstanceProfileInfo nvml.GpuInstanceProfileInfo
+
+// GpuInstancePlacement holds placement info about a GPU Instance
+type GpuInstancePlacement nvml.GpuInstancePlacement
+
+// ComputeInstanceProfileInfo holds info about a Compute Instance Profile
+type ComputeInstanceProfileInfo nvml.ComputeInstanceProfileInfo
+
+// ComputeInstancePlacement holds placement info about a Compute Instance
+type ComputeInstancePlacement nvml.ComputeInstancePlacement
+
+// DeviceAttributes stores information about MIG devices
+type DeviceAttributes nvml.DeviceAttributes