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

nvpci: Add GetGPUByIndex()

Browse Source 
Signed-off-by: Christopher Desiniotis <cdesiniotis@nvidia.com>
This commit is contained in:
Christopher Desiniotis 2022-08-23 15:24:26 -07:00
parent 09ae86c8e0
commit 6ff7845b92
2 changed files with 22 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
pkg/nvpci/nvpci.go
View File

@ -46,6 +46,7 @@ type Interface interface {
GetVGAControllers() ([]*NvidiaPCIDevice, error) GetVGAControllers() ([]*NvidiaPCIDevice, error)
GetNVSwitches() ([]*NvidiaPCIDevice, error) GetNVSwitches() ([]*NvidiaPCIDevice, error)
GetGPUs() ([]*NvidiaPCIDevice, error) GetGPUs() ([]*NvidiaPCIDevice, error)
GetGPUByIndex(int) (*NvidiaPCIDevice, error)
} }
// MemoryResources a more human readable handle // MemoryResources a more human readable handle
@ -353,3 +354,17 @@ func (p *nvpci) GetGPUs() ([]*NvidiaPCIDevice, error) {
return filtered, nil return filtered, nil
} }
// GetGPUByIndex returns an NVIDIA GPU device at a particular index
func (p *nvpci) GetGPUByIndex(i int) (*NvidiaPCIDevice, error) {
gpus, err := p.GetGPUs()
if err != nil {
return nil, fmt.Errorf("error getting all gpus: %v", err)
}
if i < 0 || i >= len(gpus) {
return nil, fmt.Errorf("invalid index '%d'", i)
}
return gpus[i], nil
}

7
pkg/nvpci/nvpci_test.go
View File

@ -66,6 +66,13 @@ func TestNvpci(t *testing.T) {
require.Equal(t, ga100PmcID, bar0.Read32(0)) require.Equal(t, ga100PmcID, bar0.Read32(0))
require.Equal(t, devices[0].IsVF, false, "Device incorrectly identified as a VF") require.Equal(t, devices[0].IsVF, false, "Device incorrectly identified as a VF")
device, err := nvpci.GetGPUByIndex(0)
require.Nil(t, err, "Error getting GPU at index 0")
require.Equal(t, "0000:80:05.1", device.Address, "Wrong Address found for device")
device, err = nvpci.GetGPUByIndex(1)
require.Error(t, err, "No error returned when getting GPU at invalid index")
} }
func TestNvpciNUMANode(t *testing.T) { func TestNvpciNUMANode(t *testing.T) {