From 1cb5426db8bc311fa7a650bd13f8a2164bb65610 Mon Sep 17 00:00:00 2001 From: Evan Lezar Date: Wed, 21 Sep 2022 15:10:06 +0200 Subject: [PATCH] Add functions for interacting with Events Signed-off-by: Evan Lezar --- pkg/nvml/consts.go | 7 ++++ pkg/nvml/device.go | 11 ++++++ pkg/nvml/device_mock.go | 85 +++++++++++++++++++++++++++++++++++++++++ pkg/nvml/event_set.go | 39 +++++++++++++++++++ pkg/nvml/nvml.go | 6 +++ pkg/nvml/nvml_mock.go | 36 +++++++++++++++++ pkg/nvml/types.go | 15 ++++++++ 7 files changed, 199 insertions(+) create mode 100644 pkg/nvml/event_set.go diff --git a/pkg/nvml/consts.go b/pkg/nvml/consts.go index 2ded48b..b353ef9 100644 --- a/pkg/nvml/consts.go +++ b/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 +) diff --git a/pkg/nvml/device.go b/pkg/nvml/device.go index 37ba0b1..90c3e67 100644 --- a/pkg/nvml/device.go +++ b/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) +} diff --git a/pkg/nvml/device_mock.go b/pkg/nvml/device_mock.go index be397ae..a216803 100644 --- a/pkg/nvml/device_mock.go +++ b/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 { diff --git a/pkg/nvml/event_set.go b/pkg/nvml/event_set.go new file mode 100644 index 0000000..2e0b9e2 --- /dev/null +++ b/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()) +} diff --git a/pkg/nvml/nvml.go b/pkg/nvml/nvml.go index 9619168..4e3eaa6 100644 --- a/pkg/nvml/nvml.go +++ b/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) +} diff --git a/pkg/nvml/nvml_mock.go b/pkg/nvml/nvml_mock.go index 56b50a9..7570ff4 100644 --- a/pkg/nvml/nvml_mock.go +++ b/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 { diff --git a/pkg/nvml/types.go b/pkg/nvml/types.go index 38ae3f2..4379c68 100644 --- a/pkg/nvml/types.go +++ b/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