diff --git a/go.mod b/go.mod index 71f0c68..845a90c 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,7 @@ module github.com/NVIDIA/go-gpuallocator go 1.20 require ( - github.com/NVIDIA/go-nvlib v0.5.0 + github.com/NVIDIA/go-nvlib v0.7.3 github.com/NVIDIA/go-nvml v0.12.9-0 github.com/stretchr/testify v1.10.0 ) diff --git a/go.sum b/go.sum index 71cea1f..f0e9e1d 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,5 @@ -github.com/NVIDIA/go-nvlib v0.5.0 h1:951KGrfr+p3cs89alO9z/ZxPPWKxwht9tx9rxiADoLI= -github.com/NVIDIA/go-nvlib v0.5.0/go.mod h1:87z49ULPr4GWPSGfSIp3taU4XENRYN/enIg88MzcL4k= +github.com/NVIDIA/go-nvlib v0.7.3 h1:kXc8PkWUlrwedSpM4fR8xT/DAq1NKy8HqhpgteFcGAw= +github.com/NVIDIA/go-nvlib v0.7.3/go.mod h1:i95Je7GinMy/+BDs++DAdbPmT2TubjNP8i8joC7DD7I= github.com/NVIDIA/go-nvml v0.12.9-0 h1:e344UK8ZkeMeeLkdQtRhmXRxNf+u532LDZPGMtkdus0= github.com/NVIDIA/go-nvml v0.12.9-0/go.mod h1:+KNA7c7gIBH7SKSJ1ntlwkfN80zdx8ovl4hrK3LmPt4= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= diff --git a/vendor/github.com/NVIDIA/go-nvlib/pkg/nvlib/device/device.go b/vendor/github.com/NVIDIA/go-nvlib/pkg/nvlib/device/device.go index 5e1510c..f00d1ac 100644 --- a/vendor/github.com/NVIDIA/go-nvlib/pkg/nvlib/device/device.go +++ b/vendor/github.com/NVIDIA/go-nvlib/pkg/nvlib/device/device.go @@ -18,6 +18,7 @@ package device import ( "fmt" + "strings" "github.com/NVIDIA/go-nvml/pkg/nvml" ) @@ -30,6 +31,8 @@ type Device interface { GetCudaComputeCapabilityAsString() (string, error) GetMigDevices() ([]MigDevice, error) GetMigProfiles() ([]MigProfile, error) + GetPCIBusID() (string, error) + IsFabricAttached() (bool, error) IsMigCapable() (bool, error) IsMigEnabled() (bool, error) VisitMigDevices(func(j int, m MigDevice) error) error @@ -83,9 +86,13 @@ func (d *device) GetArchitectureAsString() (string, error) { case nvml.DEVICE_ARCH_AMPERE: return "Ampere", nil case nvml.DEVICE_ARCH_ADA: - return "Ada", nil + return "Ada Lovelace", nil case nvml.DEVICE_ARCH_HOPPER: return "Hopper", nil + case nvml.DEVICE_ARCH_BLACKWELL: + return "Blackwell", nil + case nvml.DEVICE_ARCH_T23X: + return "Orin", nil case nvml.DEVICE_ARCH_UNKNOWN: return "Unknown", nil } @@ -122,7 +129,7 @@ func (d *device) GetBrandAsString() (string, error) { case nvml.BRAND_NVIDIA_VWS: return "NvidiaVWS", nil // Deprecated in favor of nvml.BRAND_NVIDIA_CLOUD_GAMING - //case nvml.BRAND_NVIDIA_VGAMING: + // case nvml.BRAND_NVIDIA_VGAMING: // return "VGaming", nil case nvml.BRAND_NVIDIA_CLOUD_GAMING: return "NvidiaCloudGaming", nil @@ -140,6 +147,29 @@ func (d *device) GetBrandAsString() (string, error) { return "", fmt.Errorf("error interpreting device brand as string: %v", brand) } +// GetPCIBusID returns the string representation of the bus ID. +func (d *device) GetPCIBusID() (string, error) { + info, ret := d.GetPciInfo() + if ret != nvml.SUCCESS { + return "", fmt.Errorf("error getting PCI info: %w", ret) + } + + var bytes []byte + for _, b := range info.BusId { + if byte(b) == '\x00' { + break + } + bytes = append(bytes, byte(b)) + } + id := strings.ToLower(string(bytes)) + + if id != "0000" { + id = strings.TrimPrefix(id, "0000") + } + + return id, nil +} + // GetCudaComputeCapabilityAsString returns the Device's CUDA compute capability as a version string. func (d *device) GetCudaComputeCapabilityAsString() (string, error) { major, minor, ret := d.GetCudaComputeCapability() @@ -183,6 +213,53 @@ func (d *device) IsMigEnabled() (bool, error) { return (mode == nvml.DEVICE_MIG_ENABLE), nil } +// IsFabricAttached checks if a device is attached to a GPU fabric. +func (d *device) IsFabricAttached() (bool, error) { + if d.lib.hasSymbol("nvmlDeviceGetGpuFabricInfo") { + info, ret := d.GetGpuFabricInfo() + if ret == nvml.ERROR_NOT_SUPPORTED { + return false, nil + } + if ret != nvml.SUCCESS { + return false, fmt.Errorf("error getting GPU Fabric Info: %v", ret) + } + if info.State != nvml.GPU_FABRIC_STATE_COMPLETED { + return false, nil + } + if info.ClusterUuid == [16]uint8{} { + return false, nil + } + if nvml.Return(info.Status) != nvml.SUCCESS { + return false, nil + } + + return true, nil + } + + if d.lib.hasSymbol("nvmlDeviceGetGpuFabricInfoV") { + info, ret := d.GetGpuFabricInfoV().V2() + if ret == nvml.ERROR_NOT_SUPPORTED { + return false, nil + } + if ret != nvml.SUCCESS { + return false, fmt.Errorf("error getting GPU Fabric Info: %v", ret) + } + if info.State != nvml.GPU_FABRIC_STATE_COMPLETED { + return false, nil + } + if info.ClusterUuid == [16]uint8{} { + return false, nil + } + if nvml.Return(info.Status) != nvml.SUCCESS { + return false, nil + } + + return true, nil + } + + return false, nil +} + // VisitMigDevices walks a top-level device and invokes a callback function for each MIG device configured on it. func (d *device) VisitMigDevices(visit func(int, MigDevice) error) error { capable, err := d.IsMigCapable() diff --git a/vendor/github.com/NVIDIA/go-nvlib/pkg/nvlib/device/mig_device.go b/vendor/github.com/NVIDIA/go-nvlib/pkg/nvlib/device/mig_device.go index 7145a06..cf485a4 100644 --- a/vendor/github.com/NVIDIA/go-nvlib/pkg/nvlib/device/mig_device.go +++ b/vendor/github.com/NVIDIA/go-nvlib/pkg/nvlib/device/mig_device.go @@ -63,7 +63,7 @@ func (m *migdevice) GetProfile() (MigProfile, error) { return m.profile, nil } - parent, ret := m.Device.GetDeviceHandleFromMigDeviceHandle() + parent, ret := m.GetDeviceHandleFromMigDeviceHandle() if ret != nvml.SUCCESS { return nil, fmt.Errorf("error getting parent device handle: %v", ret) } @@ -73,17 +73,17 @@ func (m *migdevice) GetProfile() (MigProfile, error) { return nil, fmt.Errorf("error getting parent memory info: %v", ret) } - attributes, ret := m.Device.GetAttributes() + attributes, ret := m.GetAttributes() if ret != nvml.SUCCESS { return nil, fmt.Errorf("error getting MIG device attributes: %v", ret) } - giID, ret := m.Device.GetGpuInstanceId() + giID, ret := m.GetGpuInstanceId() if ret != nvml.SUCCESS { return nil, fmt.Errorf("error getting MIG device GPU Instance ID: %v", ret) } - ciID, ret := m.Device.GetComputeInstanceId() + ciID, ret := m.GetComputeInstanceId() if ret != nvml.SUCCESS { return nil, fmt.Errorf("error getting MIG device Compute Instance ID: %v", ret) } diff --git a/vendor/modules.txt b/vendor/modules.txt index 8e97acc..b4bf369 100644 --- a/vendor/modules.txt +++ b/vendor/modules.txt @@ -1,4 +1,4 @@ -# github.com/NVIDIA/go-nvlib v0.5.0 +# github.com/NVIDIA/go-nvlib v0.7.3 ## explicit; go 1.20 github.com/NVIDIA/go-nvlib/pkg/nvlib/device # github.com/NVIDIA/go-nvml v0.12.9-0