From e96d9c58f131faf3d49813d0a7c9e7ff56cbe036 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Tue, 15 Nov 2022 14:29:32 +0100
Subject: [PATCH] Add GetGPUByPciBusID to nvpci.Interface

This change adds a GetGPUByPciBusID method to the nvpci Interface.
The exising NewDevice function is moved to nvmdev where it is used.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 pkg/nvmdev/nvmdev.go | 14 +++++++++++---
 pkg/nvpci/nvpci.go   | 13 +++++++------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/pkg/nvmdev/nvmdev.go b/pkg/nvmdev/nvmdev.go
index 9f7b5f1..fa2c3fc 100644
--- a/pkg/nvmdev/nvmdev.go
+++ b/pkg/nvmdev/nvmdev.go
@@ -18,13 +18,14 @@ package nvmdev
 
 import (
 	"fmt"
-	"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvpci"
 	"os"
 	"path"
 	"path/filepath"
 	"sort"
 	"strconv"
 	"strings"
+
+	"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvpci"
 )
 
 const (
@@ -241,7 +242,7 @@ func (m mdev) iommuGroup() (int, error) {
 
 // NewParentDevice constructs a ParentDevice
 func NewParentDevice(devicePath string) (*ParentDevice, error) {
-	nvdevice, err := nvpci.NewDevice(devicePath)
+	nvdevice, err := newNvidiaPCIDeviceFromPath(devicePath)
 	if err != nil {
 		return nil, fmt.Errorf("failed to construct NVIDIA PCI device: %v", err)
 	}
@@ -330,7 +331,7 @@ func (p *ParentDevice) GetPhysicalFunction() (*nvpci.NvidiaPCIDevice, error) {
 		return nil, fmt.Errorf("unable to resolve %s: %v", path.Join(p.Path, "physfn"), err)
 	}
 
-	return nvpci.NewDevice(physfnPath)
+	return newNvidiaPCIDeviceFromPath(physfnPath)
 }
 
 // GetPhysicalFunction gets the physical PCI device that a vGPU is created on
@@ -374,3 +375,10 @@ func (p *ParentDevice) GetAvailableMDEVInstances(mdevType string) (int, error) {
 
 	return availableInstances, nil
 }
+
+// newNvidiaPCIDeviceFromPath constructs an NvidiaPCIDevice for the specified device path.
+func newNvidiaPCIDeviceFromPath(devicePath string) (*nvpci.NvidiaPCIDevice, error) {
+	root := filepath.Dir(devicePath)
+	address := filepath.Base(devicePath)
+	return nvpci.NewFrom(root).GetGPUByPciBusID(address)
+}
diff --git a/pkg/nvpci/nvpci.go b/pkg/nvpci/nvpci.go
index ae5bcd2..e6b23b6 100644
--- a/pkg/nvpci/nvpci.go
+++ b/pkg/nvpci/nvpci.go
@@ -49,6 +49,7 @@ type Interface interface {
 	GetNVSwitches() ([]*NvidiaPCIDevice, error)
 	GetGPUs() ([]*NvidiaPCIDevice, error)
 	GetGPUByIndex(int) (*NvidiaPCIDevice, error)
+	GetGPUByPciBusID(string) (*NvidiaPCIDevice, error)
 	GetNetworkControllers() ([]*NvidiaPCIDevice, error)
 	GetPciBridges() ([]*NvidiaPCIDevice, error)
 	GetDPUs() ([]*NvidiaPCIDevice, error)
@@ -143,10 +144,10 @@ func (p *nvpci) GetAllDevices() ([]*NvidiaPCIDevice, error) {
 
 	var nvdevices []*NvidiaPCIDevice
 	for _, deviceDir := range deviceDirs {
-		devicePath := path.Join(p.pciDevicesRoot, deviceDir.Name())
-		nvdevice, err := NewDevice(devicePath)
+		deviceAddress := deviceDir.Name()
+		nvdevice, err := p.GetGPUByPciBusID(deviceAddress)
 		if err != nil {
-			return nil, fmt.Errorf("error constructing NVIDIA PCI device %s: %v", deviceDir.Name(), err)
+			return nil, fmt.Errorf("error constructing NVIDIA PCI device %s: %v", deviceAddress, err)
 		}
 		if nvdevice == nil {
 			continue
@@ -168,9 +169,9 @@ func (p *nvpci) GetAllDevices() ([]*NvidiaPCIDevice, error) {
 	return nvdevices, nil
 }
 
-// NewDevice constructs an NvidiaPCIDevice
-func NewDevice(devicePath string) (*NvidiaPCIDevice, error) {
-	address := path.Base(devicePath)
+// GetGPUByPciBusID constructs an NvidiaPCIDevice for the specified address (PCI Bus ID)
+func (p *nvpci) GetGPUByPciBusID(address string) (*NvidiaPCIDevice, error) {
+	devicePath := filepath.Join(p.pciDevicesRoot, address)
 
 	vendor, err := os.ReadFile(path.Join(devicePath, "vendor"))
 	if err != nil {