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

Add platform detection logic from nvidia-container-toolkit

Makefile

@@ -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.
+	golangci-lint run ./...
-	go list -f '{{.Dir}}' $(MODULE)/... | grep -v pkg/nvml | xargs golint -set_exit_status
 
 ## goimports: Apply goimports -local to the codebase
 goimports:

pkg/nvlib/info/api.go

@@ -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)
+}

pkg/nvlib/info/builder.go

@@ -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,
+	}
+}

pkg/nvlib/info/info.go

@@ -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
-}

pkg/nvlib/info/logger.go

@@ -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{}) {}

pkg/nvlib/info/options.go

@@ -16,24 +16,55 @@
 
 package info
 
-// Option defines a function for passing options to the New() call.
+import (
-type Option func(*infolib)
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
 
-// New creates a new instance of the 'info' interface.
+	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
-func New(opts ...Option) Interface {
+)
-	i := &infolib{}
+
-	for _, opt := range opts {
+// Option defines a function for passing options to the New() call.
-		opt(i)
+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 {
+func WithRoot(r string) Option {
-	return func(i *infolib) {
+	return func(i *options) {
-		i.root = root
+		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
 	}
 }

pkg/nvlib/info/property-extractor.go

@@ -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"
+}

pkg/nvlib/info/property-extractor_mock.go

@@ -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
+}

pkg/nvlib/info/resolver.go

@@ -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
+	}
+}

pkg/nvlib/info/resolver_test.go

@@ -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())
+		})
+	}
+}

pkg/nvlib/info/root.go

@@ -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
+}