Add functions for interacting with Events

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

pkg/nvml/consts.go

@@ -85,3 +85,10 @@ 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
+)

pkg/nvml/device.go

@@ -139,3 +139,14 @@ 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)
+}

pkg/nvml/device_mock.go

@@ -65,12 +65,18 @@ var _ Device = &DeviceMock{}
 // 			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")
 // 			},
@@ -129,12 +135,18 @@ type DeviceMock struct {
 	// 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)
 
@@ -196,12 +208,22 @@ type DeviceMock 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.
@@ -224,8 +246,10 @@ type DeviceMock struct {
 	lockGetMinorNumber                     sync.RWMutex
 	lockGetName                            sync.RWMutex
 	lockGetPciInfo                         sync.RWMutex
+	lockGetSupportedEventTypes             sync.RWMutex
 	lockGetUUID                            sync.RWMutex
 	lockIsMigDeviceHandle                  sync.RWMutex
+	lockRegisterEvents                     sync.RWMutex
 	lockSetMigMode                         sync.RWMutex
 }
 
@@ -665,6 +689,32 @@ func (mock *DeviceMock) GetPciInfoCalls() []struct {
 	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 {
@@ -717,6 +767,41 @@ func (mock *DeviceMock) IsMigDeviceHandleCalls() []struct {
 	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 {

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())
+}

pkg/nvml/nvml.go

@@ -102,3 +102,9 @@ func (n *nvmlLib) SystemGetCudaDriverVersion() (int, Return) {
 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)
+}

pkg/nvml/nvml_mock.go

@@ -29,6 +29,9 @@ var _ Interface = &InterfaceMock{}
 // 			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")
 // 			},
@@ -60,6 +63,9 @@ type InterfaceMock struct {
 	// 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
 
@@ -92,6 +98,9 @@ type InterfaceMock 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 {
 		}
@@ -109,6 +118,7 @@ type InterfaceMock struct {
 	lockDeviceGetHandleByIndex     sync.RWMutex
 	lockDeviceGetHandleByUUID      sync.RWMutex
 	lockErrorString                sync.RWMutex
+	lockEventSetCreate             sync.RWMutex
 	lockInit                       sync.RWMutex
 	lockShutdown                   sync.RWMutex
 	lockSystemGetCudaDriverVersion sync.RWMutex
@@ -234,6 +244,32 @@ func (mock *InterfaceMock) ErrorStringCalls() []struct {
 	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 {

pkg/nvml/types.go

@@ -28,6 +28,7 @@ type Interface interface {
 	DeviceGetHandleByIndex(Index int) (Device, Return)
 	DeviceGetHandleByUUID(UUID string) (Device, Return)
 	ErrorString(r Return) string
+	EventSetCreate() (EventSet, Return)
 	Init() Return
 	Shutdown() Return
 	SystemGetCudaDriverVersion() (int, Return)
@@ -54,8 +55,10 @@ type Device interface {
 	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)
 }
 
@@ -96,6 +99,18 @@ type ComputeInstanceInfo struct {
 	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