ClearML/go-nvlib
1
0
Fork 0
mirror of https://github.com/clearml/go-nvlib synced 2025-04-25 08:35:21 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request from NVIDIA/fix-igpu-nvml-auto

Browse Source 
Add platform detection logic from nvidia-container-toolkit
This commit is contained in:
Evan Lezar 2024-05-21 11:42:32 +02:00 committed by GitHub
commit 7604335102
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 812 additions and 118 deletions

6
Makefile
View File

@ -18,6 +18,7 @@ DOCKER ?= docker
PCI_IDS_URL ?= https://pci-ids.ucw.cz/v2.2/pci.ids
CHECK_TARGETS := lint
TARGETS := binary build all check fmt assert-fmt generate lint vet test coverage
DOCKER_TARGETS := $(patsubst %,docker-%, $(TARGETS))
.PHONY: $(TARGETS) $(DOCKER_TARGETS) vendor check-vendor
@ -28,7 +29,7 @@ build:
GOOS=$(GOOS) go build ./...
all: check build binary
check: assert-fmt lint vet
check: $(CHECK_TARGETS)
vendor:
go mod tidy
@ -59,8 +60,7 @@ generate:
go generate $(MODULE)/...
lint:
# We use `go list -f '{{.Dir}}' $(MODULE)/...` to skip the `vendor` folder.
go list -f '{{.Dir}}' $(MODULE)/... | grep -v pkg/nvml | xargs golint -set_exit_status
golangci-lint run ./...
## goimports: Apply goimports -local to the codebase
goimports:

41
pkg/nvlib/info/api.go Normal file
View File

@ -0,0 +1,41 @@
/**
# Copyright 2024 NVIDIA CORPORATION
#
# 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 info
// Interface provides the API to the info package.
type Interface interface {
PlatformResolver
PropertyExtractor
}
// PlatformResolver defines a function to resolve the current platform.
type PlatformResolver interface {
ResolvePlatform() Platform
}
// PropertyExtractor provides a set of functions to query capabilities of the
// system.
//
//go:generate moq -rm -out property-extractor_mock.go . PropertyExtractor
type PropertyExtractor interface {
HasDXCore() (bool, string)
HasNvml() (bool, string)
HasTegraFiles() (bool, string)
// Deprecated: Use HasTegraFiles instead.
IsTegraSystem() (bool, string)
UsesOnlyNVGPUModule() (bool, string)
}

78
pkg/nvlib/info/builder.go Normal file
View File

@ -0,0 +1,78 @@
/**
# Copyright 2024 NVIDIA CORPORATION
#
# 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 info
import (
"github.com/NVIDIA/go-nvml/pkg/nvml"
"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
)
type infolib struct {
PropertyExtractor
PlatformResolver
}
type options struct {
logger basicLogger
root root
nvmllib nvml.Interface
devicelib device.Interface
platform Platform
propertyExtractor PropertyExtractor
}
// New creates a new instance of the 'info' interface.
func New(opts ...Option) Interface {
o := &options{}
for _, opt := range opts {
opt(o)
}
if o.logger == nil {
o.logger = &nullLogger{}
}
if o.root == "" {
o.root = "/"
}
if o.nvmllib == nil {
o.nvmllib = nvml.New(
nvml.WithLibraryPath(o.root.tryResolveLibrary("libnvidia-ml.so.1")),
)
}
if o.devicelib == nil {
o.devicelib = device.New(device.WithNvml(o.nvmllib))
}
if o.platform == "" {
o.platform = PlatformAuto
}
if o.propertyExtractor == nil {
o.propertyExtractor = &propertyExtractor{
root: o.root,
nvmllib: o.nvmllib,
devicelib: o.devicelib,
}
}
return &infolib{
PlatformResolver: &platformResolver{
logger: o.logger,
platform: o.platform,
propertyExtractor: o.propertyExtractor,
},
PropertyExtractor: o.propertyExtractor,
}
}

102
pkg/nvlib/info/info.go
View File

@ -1,102 +0,0 @@
/**
# 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 info
import (
"fmt"
"os"
"path/filepath"
"strings"
"github.com/NVIDIA/go-nvml/pkg/dl"
)
// Interface provides the API to the info package.
type Interface interface {
HasDXCore() (bool, string)
HasNvml() (bool, string)
IsTegraSystem() (bool, string)
}
type infolib struct {
root string
}
var _ Interface = &infolib{}
// HasDXCore returns true if DXCore is detected on the system.
func (i *infolib) HasDXCore() (bool, string) {
const (
libraryName = "libdxcore.so"
)
if err := assertHasLibrary(libraryName); err != nil {
return false, fmt.Sprintf("could not load DXCore library: %v", err)
}
return true, "found DXCore library"
}
// HasNvml returns true if NVML is detected on the system.
func (i *infolib) HasNvml() (bool, string) {
const (
libraryName = "libnvidia-ml.so.1"
)
if err := assertHasLibrary(libraryName); err != nil {
return false, fmt.Sprintf("could not load NVML library: %v", err)
}
return true, "found NVML library"
}
// IsTegraSystem returns true if the system is detected as a Tegra-based system.
func (i *infolib) IsTegraSystem() (bool, string) {
tegraReleaseFile := filepath.Join(i.root, "/etc/nv_tegra_release")
tegraFamilyFile := filepath.Join(i.root, "/sys/devices/soc0/family")
if info, err := os.Stat(tegraReleaseFile); err == nil && !info.IsDir() {
return true, fmt.Sprintf("%v found", tegraReleaseFile)
}
if info, err := os.Stat(tegraFamilyFile); err != nil || info.IsDir() {
return false, fmt.Sprintf("%v file not found", tegraFamilyFile)
}
contents, err := os.ReadFile(tegraFamilyFile)
if err != nil {
return false, fmt.Sprintf("could not read %v", tegraFamilyFile)
}
if strings.HasPrefix(strings.ToLower(string(contents)), "tegra") {
return true, fmt.Sprintf("%v has 'tegra' prefix", tegraFamilyFile)
}
return false, fmt.Sprintf("%v has no 'tegra' prefix", tegraFamilyFile)
}
// assertHasLibrary returns an error if the specified library cannot be loaded.
func assertHasLibrary(libraryName string) error {
const (
libraryLoadFlags = dl.RTLD_LAZY
)
lib := dl.New(libraryName, libraryLoadFlags)
if err := lib.Open(); err != nil {
return err
}
defer lib.Close()
return nil
}

28
pkg/nvlib/info/logger.go Normal file
View File

@ -0,0 +1,28 @@
/**
# Copyright 2024 NVIDIA CORPORATION
#
# 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 info
type basicLogger interface {
Debugf(string, ...interface{})
Infof(string, ...interface{})
}
type nullLogger struct{}
func (n *nullLogger) Debugf(string, ...interface{}) {}
func (n *nullLogger) Infof(string, ...interface{}) {}

57
pkg/nvlib/info/options.go
View File

@ -16,24 +16,55 @@
package info
// Option defines a function for passing options to the New() call.
type Option func(*infolib)
import (
"github.com/NVIDIA/go-nvml/pkg/nvml"
// New creates a new instance of the 'info' interface.
func New(opts ...Option) Interface {
i := &infolib{}
for _, opt := range opts {
opt(i)
"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
)
// Option defines a function for passing options to the New() call.
type Option func(*options)
// WithDeviceLib sets the device library for the library.
func WithDeviceLib(devicelib device.Interface) Option {
return func(i *options) {
i.devicelib = devicelib
}
if i.root == "" {
i.root = "/"
}
// WithLogger sets the logger for the library.
func WithLogger(logger basicLogger) Option {
return func(i *options) {
i.logger = logger
}
}
// WithNvmlLib sets the nvml library for the library.
func WithNvmlLib(nvmllib nvml.Interface) Option {
return func(i *options) {
i.nvmllib = nvmllib
}
return i
}
// WithRoot provides a Option to set the root of the 'info' interface.
func WithRoot(root string) Option {
return func(i *infolib) {
i.root = root
func WithRoot(r string) Option {
return func(i *options) {
i.root = root(r)
}
}
// WithPropertyExtractor provides an Option to set the PropertyExtractor
// interface implementation.
// This is predominantly used for testing.
func WithPropertyExtractor(propertyExtractor PropertyExtractor) Option {
return func(i *options) {
i.propertyExtractor = propertyExtractor
}
}
// WithPlatform provides an option to set the platform explicitly.
func WithPlatform(platform Platform) Option {
return func(i *options) {
i.platform = platform
}
}

143
pkg/nvlib/info/property-extractor.go Normal file
View File

@ -0,0 +1,143 @@
/**
# 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 info
import (
"fmt"
"os"
"strings"
"github.com/NVIDIA/go-nvml/pkg/nvml"
"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
)
type propertyExtractor struct {
root root
nvmllib nvml.Interface
devicelib device.Interface
}
var _ PropertyExtractor = &propertyExtractor{}
// HasDXCore returns true if DXCore is detected on the system.
func (i *propertyExtractor) HasDXCore() (bool, string) {
const (
libraryName = "libdxcore.so"
)
if err := i.root.assertHasLibrary(libraryName); err != nil {
return false, fmt.Sprintf("could not load DXCore library: %v", err)
}
return true, "found DXCore library"
}
// HasNvml returns true if NVML is detected on the system.
func (i *propertyExtractor) HasNvml() (bool, string) {
const (
libraryName = "libnvidia-ml.so.1"
)
if err := i.root.assertHasLibrary(libraryName); err != nil {
return false, fmt.Sprintf("could not load NVML library: %v", err)
}
return true, "found NVML library"
}
// IsTegraSystem returns true if the system is detected as a Tegra-based system.
// Deprecated: Use HasTegraFiles instead.
func (i *propertyExtractor) IsTegraSystem() (bool, string) {
return i.HasTegraFiles()
}
// HasTegraFiles returns true if tegra-based files are detected on the system.
func (i *propertyExtractor) HasTegraFiles() (bool, string) {
tegraReleaseFile := i.root.join("/etc/nv_tegra_release")
tegraFamilyFile := i.root.join("/sys/devices/soc0/family")
if info, err := os.Stat(tegraReleaseFile); err == nil && !info.IsDir() {
return true, fmt.Sprintf("%v found", tegraReleaseFile)
}
if info, err := os.Stat(tegraFamilyFile); err != nil || info.IsDir() {
return false, fmt.Sprintf("%v file not found", tegraFamilyFile)
}
contents, err := os.ReadFile(tegraFamilyFile)
if err != nil {
return false, fmt.Sprintf("could not read %v", tegraFamilyFile)
}
if strings.HasPrefix(strings.ToLower(string(contents)), "tegra") {
return true, fmt.Sprintf("%v has 'tegra' prefix", tegraFamilyFile)
}
return false, fmt.Sprintf("%v has no 'tegra' prefix", tegraFamilyFile)
}
// UsesOnlyNVGPUModule checks whether the only the nvgpu module is used.
// This kernel module is used on Tegra-based systems when using the iGPU.
// Since some of these systems also support NVML, we use the device name
// reported by NVML to determine whether the system is an iGPU system.
//
// Devices that use the nvgpu module have their device names as:
//
// GPU 0: Orin (nvgpu) (UUID: 54d0709b-558d-5a59-9c65-0c5fc14a21a4)
//
// This function returns true if ALL devices use the nvgpu module.
func (i *propertyExtractor) UsesOnlyNVGPUModule() (uses bool, reason string) {
// We ensure that this function never panics
defer func() {
if err := recover(); err != nil {
uses = false
reason = fmt.Sprintf("panic: %v", err)
}
}()
ret := i.nvmllib.Init()
if ret != nvml.SUCCESS {
return false, fmt.Sprintf("failed to initialize nvml: %v", ret)
}
defer func() {
_ = i.nvmllib.Shutdown()
}()
var names []string
err := i.devicelib.VisitDevices(func(i int, d device.Device) error {
name, ret := d.GetName()
if ret != nvml.SUCCESS {
return fmt.Errorf("device %v: %v", i, ret)
}
names = append(names, name)
return nil
})
if err != nil {
return false, fmt.Sprintf("failed to get device names: %v", err)
}
if len(names) == 0 {
return false, "no devices found"
}
for _, name := range names {
if !strings.Contains(name, "(nvgpu)") {
return false, fmt.Sprintf("device %q does not use nvgpu module", name)
}
}
return true, "all devices use nvgpu module"
}

215
pkg/nvlib/info/property-extractor_mock.go Normal file
View File

@ -0,0 +1,215 @@
// Code generated by moq; DO NOT EDIT.
// github.com/matryer/moq
package info
import (
"sync"
)
// Ensure, that PropertyExtractorMock does implement PropertyExtractor.
// If this is not the case, regenerate this file with moq.
var _ PropertyExtractor = &PropertyExtractorMock{}
// PropertyExtractorMock is a mock implementation of PropertyExtractor.
//
// func TestSomethingThatUsesPropertyExtractor(t *testing.T) {
//
// // make and configure a mocked PropertyExtractor
// mockedPropertyExtractor := &PropertyExtractorMock{
// HasDXCoreFunc: func() (bool, string) {
// panic("mock out the HasDXCore method")
// },
// HasNvmlFunc: func() (bool, string) {
// panic("mock out the HasNvml method")
// },
// HasTegraFilesFunc: func() (bool, string) {
// panic("mock out the HasTegraFiles method")
// },
// IsTegraSystemFunc: func() (bool, string) {
// panic("mock out the IsTegraSystem method")
// },
// UsesOnlyNVGPUModuleFunc: func() (bool, string) {
// panic("mock out the UsesOnlyNVGPUModule method")
// },
// }
//
// // use mockedPropertyExtractor in code that requires PropertyExtractor
// // and then make assertions.
//
// }
type PropertyExtractorMock struct {
// HasDXCoreFunc mocks the HasDXCore method.
HasDXCoreFunc func() (bool, string)
// HasNvmlFunc mocks the HasNvml method.
HasNvmlFunc func() (bool, string)
// HasTegraFilesFunc mocks the HasTegraFiles method.
HasTegraFilesFunc func() (bool, string)
// IsTegraSystemFunc mocks the IsTegraSystem method.
IsTegraSystemFunc func() (bool, string)
// UsesOnlyNVGPUModuleFunc mocks the UsesOnlyNVGPUModule method.
UsesOnlyNVGPUModuleFunc func() (bool, string)
// calls tracks calls to the methods.
calls struct {
// HasDXCore holds details about calls to the HasDXCore method.
HasDXCore []struct {
}
// HasNvml holds details about calls to the HasNvml method.
HasNvml []struct {
}
// HasTegraFiles holds details about calls to the HasTegraFiles method.
HasTegraFiles []struct {
}
// IsTegraSystem holds details about calls to the IsTegraSystem method.
IsTegraSystem []struct {
}
// UsesOnlyNVGPUModule holds details about calls to the UsesOnlyNVGPUModule method.
UsesOnlyNVGPUModule []struct {
}
}
lockHasDXCore sync.RWMutex
lockHasNvml sync.RWMutex
lockHasTegraFiles sync.RWMutex
lockIsTegraSystem sync.RWMutex
lockUsesOnlyNVGPUModule sync.RWMutex
}
// HasDXCore calls HasDXCoreFunc.
func (mock *PropertyExtractorMock) HasDXCore() (bool, string) {
if mock.HasDXCoreFunc == nil {
panic("PropertyExtractorMock.HasDXCoreFunc: method is nil but PropertyExtractor.HasDXCore was just called")
}
callInfo := struct {
}{}
mock.lockHasDXCore.Lock()
mock.calls.HasDXCore = append(mock.calls.HasDXCore, callInfo)
mock.lockHasDXCore.Unlock()
return mock.HasDXCoreFunc()
}
// HasDXCoreCalls gets all the calls that were made to HasDXCore.
// Check the length with:
//
// len(mockedPropertyExtractor.HasDXCoreCalls())
func (mock *PropertyExtractorMock) HasDXCoreCalls() []struct {
} {
var calls []struct {
}
mock.lockHasDXCore.RLock()
calls = mock.calls.HasDXCore
mock.lockHasDXCore.RUnlock()
return calls
}
// HasNvml calls HasNvmlFunc.
func (mock *PropertyExtractorMock) HasNvml() (bool, string) {
if mock.HasNvmlFunc == nil {
panic("PropertyExtractorMock.HasNvmlFunc: method is nil but PropertyExtractor.HasNvml was just called")
}
callInfo := struct {
}{}
mock.lockHasNvml.Lock()
mock.calls.HasNvml = append(mock.calls.HasNvml, callInfo)
mock.lockHasNvml.Unlock()
return mock.HasNvmlFunc()
}
// HasNvmlCalls gets all the calls that were made to HasNvml.
// Check the length with:
//
// len(mockedPropertyExtractor.HasNvmlCalls())
func (mock *PropertyExtractorMock) HasNvmlCalls() []struct {
} {
var calls []struct {
}
mock.lockHasNvml.RLock()
calls = mock.calls.HasNvml
mock.lockHasNvml.RUnlock()
return calls
}
// HasTegraFiles calls HasTegraFilesFunc.
func (mock *PropertyExtractorMock) HasTegraFiles() (bool, string) {
if mock.HasTegraFilesFunc == nil {
panic("PropertyExtractorMock.HasTegraFilesFunc: method is nil but PropertyExtractor.HasTegraFiles was just called")
}
callInfo := struct {
}{}
mock.lockHasTegraFiles.Lock()
mock.calls.HasTegraFiles = append(mock.calls.HasTegraFiles, callInfo)
mock.lockHasTegraFiles.Unlock()
return mock.HasTegraFilesFunc()
}
// HasTegraFilesCalls gets all the calls that were made to HasTegraFiles.
// Check the length with:
//
// len(mockedPropertyExtractor.HasTegraFilesCalls())
func (mock *PropertyExtractorMock) HasTegraFilesCalls() []struct {
} {
var calls []struct {
}
mock.lockHasTegraFiles.RLock()
calls = mock.calls.HasTegraFiles
mock.lockHasTegraFiles.RUnlock()
return calls
}
// IsTegraSystem calls IsTegraSystemFunc.
func (mock *PropertyExtractorMock) IsTegraSystem() (bool, string) {
if mock.IsTegraSystemFunc == nil {
panic("PropertyExtractorMock.IsTegraSystemFunc: method is nil but PropertyExtractor.IsTegraSystem was just called")
}
callInfo := struct {
}{}
mock.lockIsTegraSystem.Lock()
mock.calls.IsTegraSystem = append(mock.calls.IsTegraSystem, callInfo)
mock.lockIsTegraSystem.Unlock()
return mock.IsTegraSystemFunc()
}
// IsTegraSystemCalls gets all the calls that were made to IsTegraSystem.
// Check the length with:
//
// len(mockedPropertyExtractor.IsTegraSystemCalls())
func (mock *PropertyExtractorMock) IsTegraSystemCalls() []struct {
} {
var calls []struct {
}
mock.lockIsTegraSystem.RLock()
calls = mock.calls.IsTegraSystem
mock.lockIsTegraSystem.RUnlock()
return calls
}
// UsesOnlyNVGPUModule calls UsesOnlyNVGPUModuleFunc.
func (mock *PropertyExtractorMock) UsesOnlyNVGPUModule() (bool, string) {
if mock.UsesOnlyNVGPUModuleFunc == nil {
panic("PropertyExtractorMock.UsesOnlyNVGPUModuleFunc: method is nil but PropertyExtractor.UsesOnlyNVGPUModule was just called")
}
callInfo := struct {
}{}
mock.lockUsesOnlyNVGPUModule.Lock()
mock.calls.UsesOnlyNVGPUModule = append(mock.calls.UsesOnlyNVGPUModule, callInfo)
mock.lockUsesOnlyNVGPUModule.Unlock()
return mock.UsesOnlyNVGPUModuleFunc()
}
// UsesOnlyNVGPUModuleCalls gets all the calls that were made to UsesOnlyNVGPUModule.
// Check the length with:
//
// len(mockedPropertyExtractor.UsesOnlyNVGPUModuleCalls())
func (mock *PropertyExtractorMock) UsesOnlyNVGPUModuleCalls() []struct {
} {
var calls []struct {
}
mock.lockUsesOnlyNVGPUModule.RLock()
calls = mock.calls.UsesOnlyNVGPUModule
mock.lockUsesOnlyNVGPUModule.RUnlock()
return calls
}

64
pkg/nvlib/info/resolver.go Normal file
View File

@ -0,0 +1,64 @@
/**
# Copyright 2024 NVIDIA CORPORATION
#
# 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 info
// Platform represents a supported plaform.
type Platform string
const (
PlatformAuto = Platform("auto")
PlatformNVML = Platform("nvml")
PlatformTegra = Platform("tegra")
PlatformWSL = Platform("wsl")
PlatformUnknown = Platform("unknown")
)
type platformResolver struct {
logger basicLogger
platform Platform
propertyExtractor PropertyExtractor
}
func (p platformResolver) ResolvePlatform() Platform {
if p.platform != PlatformAuto {
p.logger.Infof("Using requested platform '%s'", p.platform)
return p.platform
}
hasDXCore, reason := p.propertyExtractor.HasDXCore()
p.logger.Debugf("Is WSL-based system? %v: %v", hasDXCore, reason)
hasTegraFiles, reason := p.propertyExtractor.HasTegraFiles()
p.logger.Debugf("Is Tegra-based system? %v: %v", hasTegraFiles, reason)
hasNVML, reason := p.propertyExtractor.HasNvml()
p.logger.Debugf("Is NVML-based system? %v: %v", hasNVML, reason)
usesOnlyNVGPUModule, reason := p.propertyExtractor.UsesOnlyNVGPUModule()
p.logger.Debugf("Uses nvgpu kernel module? %v: %v", usesOnlyNVGPUModule, reason)
switch {
case hasDXCore:
return PlatformWSL
case (hasTegraFiles && !hasNVML), usesOnlyNVGPUModule:
return PlatformTegra
case hasNVML:
return PlatformNVML
default:
return PlatformUnknown
}
}

110
pkg/nvlib/info/resolver_test.go Normal file
View File

@ -0,0 +1,110 @@
/**
# Copyright (c) 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 info
import (
"fmt"
"testing"
"github.com/stretchr/testify/require"
)
func TestResolvePlatform(t *testing.T) {
testCases := []struct {
platform string
hasTegraFiles bool
hasDXCore bool
hasNVML bool
usesOnlyNVGPUModule bool
expected string
}{
{
platform: "auto",
hasDXCore: true,
expected: "wsl",
},
{
platform: "auto",
hasDXCore: false,
hasTegraFiles: true,
hasNVML: false,
expected: "tegra",
},
{
platform: "auto",
hasDXCore: false,
hasTegraFiles: false,
hasNVML: false,
expected: "unknown",
},
{
platform: "auto",
hasDXCore: false,
hasTegraFiles: true,
hasNVML: true,
expected: "nvml",
},
{
platform: "auto",
hasDXCore: false,
hasTegraFiles: true,
hasNVML: true,
usesOnlyNVGPUModule: true,
expected: "tegra",
},
{
platform: "nvml",
hasDXCore: true,
hasTegraFiles: true,
expected: "nvml",
},
{
platform: "wsl",
hasDXCore: false,
expected: "wsl",
},
{
platform: "not-auto",
hasDXCore: true,
expected: "not-auto",
},
}
for i, tc := range testCases {
t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
l := New(
WithPropertyExtractor(&PropertyExtractorMock{
HasDXCoreFunc: func() (bool, string) {
return tc.hasDXCore, ""
},
HasNvmlFunc: func() (bool, string) {
return tc.hasNVML, ""
},
HasTegraFilesFunc: func() (bool, string) {
return tc.hasTegraFiles, ""
},
UsesOnlyNVGPUModuleFunc: func() (bool, string) {
return tc.usesOnlyNVGPUModule, ""
},
}),
WithPlatform(Platform(tc.platform)),
)
require.Equal(t, Platform(tc.expected), l.ResolvePlatform())
})
}
}

86
pkg/nvlib/info/root.go Normal file
View File

@ -0,0 +1,86 @@
/**
# Copyright 2024 NVIDIA CORPORATION
#
# 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 info
import (
"fmt"
"path/filepath"
"github.com/NVIDIA/go-nvml/pkg/dl"
)
// root represents a directory on the filesystem relative to which libraries
// such as the NVIDIA driver libraries can be found.
type root string
func (r root) join(parts ...string) string {
return filepath.Join(append([]string{string(r)}, parts...)...)
}
// assertHasLibrary returns an error if the specified library cannot be loaded.
func (r root) assertHasLibrary(libraryName string) error {
const (
libraryLoadFlags = dl.RTLD_LAZY
)
lib := dl.New(r.tryResolveLibrary(libraryName), libraryLoadFlags)
if err := lib.Open(); err != nil {
return err
}
defer lib.Close()
return nil
}
// tryResolveLibrary attempts to locate the specified library in the root.
// If the root is not specified, is "/", or the library cannot be found in the
// set of predefined paths, the input is returned as is.
func (r root) tryResolveLibrary(libraryName string) string {
if r == "" || r == "/" {
return libraryName
}
librarySearchPaths := []string{
"/usr/lib64",
"/usr/lib/x86_64-linux-gnu",
"/usr/lib/aarch64-linux-gnu",
"/lib64",
"/lib/x86_64-linux-gnu",
"/lib/aarch64-linux-gnu",
}
for _, d := range librarySearchPaths {
l := r.join(d, libraryName)
resolved, err := resolveLink(l)
if err != nil {
continue
}
return resolved
}
return libraryName
}
// resolveLink finds the target of a symlink or the file itself in the
// case of a regular file.
// This is equivalent to running `readlink -f ${l}`.
func resolveLink(l string) (string, error) {
resolved, err := filepath.EvalSymlinks(l)
if err != nil {
return "", fmt.Errorf("error resolving link '%v': %w", l, err)
}
return resolved, nil
}