ClearML/go-nvlib
1
0
Fork 0
mirror of https://github.com/clearml/go-nvlib synced 2025-02-11 23:23:41 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add Placement related calls for GPUInstances in nvml wrapper

Browse Source 
Signed-off-by: Kevin Klues <kklues@nvidia.com>
This commit is contained in:
Kevin Klues 2022-12-08 14:47:10 +00:00
parent 49dcad67c4
commit 6a4886e49e
2 changed files with 18 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
pkg/nvml/device.go
View File

@ -88,6 +88,16 @@ func (d nvmlDevice) GetGpuInstanceProfileInfo(profile int) (GpuInstanceProfileIn
return GpuInstanceProfileInfo(p), Return(r) return GpuInstanceProfileInfo(p), Return(r)
} }
// GetGpuInstancePossiblePlacements returns the possible placements of a GPU Instance
func (d nvmlDevice) GetGpuInstancePossiblePlacements(info *GpuInstanceProfileInfo) ([]GpuInstancePlacement, Return) {
nvmlPlacements, r := nvml.Device(d).GetGpuInstancePossiblePlacements((*nvml.GpuInstanceProfileInfo)(info))
var placements []GpuInstancePlacement
for _, p := range nvmlPlacements {
placements = append(placements, GpuInstancePlacement(p))
}
return placements, Return(r)
}
// GetGpuInstances returns the set of GPU Instances associated with a Device // GetGpuInstances returns the set of GPU Instances associated with a Device
func (d nvmlDevice) GetGpuInstances(info *GpuInstanceProfileInfo) ([]GpuInstance, Return) { func (d nvmlDevice) GetGpuInstances(info *GpuInstanceProfileInfo) ([]GpuInstance, Return) {
nvmlGis, r := nvml.Device(d).GetGpuInstances((*nvml.GpuInstanceProfileInfo)(info)) nvmlGis, r := nvml.Device(d).GetGpuInstances((*nvml.GpuInstanceProfileInfo)(info))
@ -98,6 +108,12 @@ func (d nvmlDevice) GetGpuInstances(info *GpuInstanceProfileInfo) ([]GpuInstance
return gis, Return(r) return gis, Return(r)
} }
// CreateGpuInstanceWithPlacement creates a GPU Instance with a specific placement
func (d nvmlDevice) CreateGpuInstanceWithPlacement(info *GpuInstanceProfileInfo, placement *GpuInstancePlacement) (GpuInstance, Return) {
gi, r := nvml.Device(d).CreateGpuInstanceWithPlacement((*nvml.GpuInstanceProfileInfo)(info), (*nvml.GpuInstancePlacement)(placement))
return nvmlGpuInstance(gi), Return(r)
}
// GetMaxMigDeviceCount returns the maximum number of MIG devices that can be created on a Device // GetMaxMigDeviceCount returns the maximum number of MIG devices that can be created on a Device
func (d nvmlDevice) GetMaxMigDeviceCount() (int, Return) { func (d nvmlDevice) GetMaxMigDeviceCount() (int, Return) {
m, r := nvml.Device(d).GetMaxMigDeviceCount() m, r := nvml.Device(d).GetMaxMigDeviceCount()

2
pkg/nvml/types.go
View File

@ -39,12 +39,14 @@ type Interface interface {
// //
//go:generate moq -out device_mock.go . Device //go:generate moq -out device_mock.go . Device
type Device interface { type Device interface {
CreateGpuInstanceWithPlacement(*GpuInstanceProfileInfo, *GpuInstancePlacement) (GpuInstance, Return)
GetAttributes() (DeviceAttributes, Return) GetAttributes() (DeviceAttributes, Return)
GetComputeInstanceId() (int, Return) GetComputeInstanceId() (int, Return)
GetCudaComputeCapability() (int, int, Return) GetCudaComputeCapability() (int, int, Return)
GetDeviceHandleFromMigDeviceHandle() (Device, Return) GetDeviceHandleFromMigDeviceHandle() (Device, Return)
GetGpuInstanceById(ID int) (GpuInstance, Return) GetGpuInstanceById(ID int) (GpuInstance, Return)
GetGpuInstanceId() (int, Return) GetGpuInstanceId() (int, Return)
GetGpuInstancePossiblePlacements(*GpuInstanceProfileInfo) ([]GpuInstancePlacement, Return)
GetGpuInstanceProfileInfo(Profile int) (GpuInstanceProfileInfo, Return) GetGpuInstanceProfileInfo(Profile int) (GpuInstanceProfileInfo, Return)
GetGpuInstances(Info *GpuInstanceProfileInfo) ([]GpuInstance, Return) GetGpuInstances(Info *GpuInstanceProfileInfo) ([]GpuInstance, Return)
GetIndex() (int, Return) GetIndex() (int, Return)