From 2560bf63327529421c8b0aefbefeebbf44f91d37 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Thu, 4 Apr 2024 16:52:35 +0200
Subject: [PATCH 1/4] Fix check targets
Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
Makefile | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/Makefile b/Makefile
index c94c191..04edc8a 100644
--- a/Makefile
+++ b/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.
- 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:
From 791d093c62f8cdd48e1e18502371a800491112ed Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Tue, 26 Mar 2024 14:55:37 +0200
Subject: [PATCH 2/4] Refactor info API
This change adds a PropertyExtractor interface to encapsulate functions
that query a system for certain capabilities.
The IsTegraSystem has been renamed to HasTegraFiles function and marked
as Deprecated.
Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
pkg/nvlib/info/api.go | 34 ++++
pkg/nvlib/info/builder.go | 35 ++++
pkg/nvlib/info/options.go | 20 +-
.../info/{info.go => property-extractor.go} | 50 ++---
pkg/nvlib/info/property-extractor_mock.go | 178 ++++++++++++++++++
pkg/nvlib/info/root.go | 86 +++++++++
6 files changed, 353 insertions(+), 50 deletions(-)
create mode 100644 pkg/nvlib/info/api.go
create mode 100644 pkg/nvlib/info/builder.go
rename pkg/nvlib/info/{info.go => property-extractor.go} (64%)
create mode 100644 pkg/nvlib/info/property-extractor_mock.go
create mode 100644 pkg/nvlib/info/root.go
diff --git a/pkg/nvlib/info/api.go b/pkg/nvlib/info/api.go
new file mode 100644
index 0000000..41b8e22
--- /dev/null
+++ b/pkg/nvlib/info/api.go
@@ -0,0 +1,34 @@
+/**
+# 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 {
+ PropertyExtractor
+}
+
+// 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)
+}
diff --git a/pkg/nvlib/info/builder.go b/pkg/nvlib/info/builder.go
new file mode 100644
index 0000000..d9275ca
--- /dev/null
+++ b/pkg/nvlib/info/builder.go
@@ -0,0 +1,35 @@
+/**
+# 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 options struct {
+ root root
+}
+
+// New creates a new instance of the 'info' Interface.
+func New(opts ...Option) Interface {
+ o := &options{}
+ for _, opt := range opts {
+ opt(o)
+ }
+ if o.root == "" {
+ o.root = "/"
+ }
+ return &propertyExtractor{
+ root: o.root,
+ }
+}
diff --git a/pkg/nvlib/info/options.go b/pkg/nvlib/info/options.go
index ce72150..c4265d1 100644
--- a/pkg/nvlib/info/options.go
+++ b/pkg/nvlib/info/options.go
@@ -17,23 +17,11 @@
package info
// Option defines a function for passing options to the New() call.
-type Option func(*infolib)
-
-// New creates a new instance of the 'info' interface.
-func New(opts ...Option) Interface {
- i := &infolib{}
- for _, opt := range opts {
- opt(i)
- }
- if i.root == "" {
- i.root = "/"
- }
- return i
-}
+type Option func(*options)
// 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)
}
}
diff --git a/pkg/nvlib/info/info.go b/pkg/nvlib/info/property-extractor.go
similarity index 64%
rename from pkg/nvlib/info/info.go
rename to pkg/nvlib/info/property-extractor.go
index 677270c..9e41a54 100644
--- a/pkg/nvlib/info/info.go
+++ b/pkg/nvlib/info/property-extractor.go
@@ -19,31 +19,21 @@ 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
+type propertyExtractor struct {
+ root root
}
-var _ Interface = &infolib{}
+var _ Interface = &propertyExtractor{}
// HasDXCore returns true if DXCore is detected on the system.
-func (i *infolib) HasDXCore() (bool, string) {
+func (i *propertyExtractor) HasDXCore() (bool, string) {
const (
libraryName = "libdxcore.so"
)
- if err := assertHasLibrary(libraryName); err != nil {
+ if err := i.root.assertHasLibrary(libraryName); err != nil {
return false, fmt.Sprintf("could not load DXCore library: %v", err)
}
@@ -51,11 +41,11 @@ func (i *infolib) HasDXCore() (bool, string) {
}
// HasNvml returns true if NVML is detected on the system.
-func (i *infolib) HasNvml() (bool, string) {
+func (i *propertyExtractor) HasNvml() (bool, string) {
const (
libraryName = "libnvidia-ml.so.1"
)
- if err := assertHasLibrary(libraryName); err != nil {
+ if err := i.root.assertHasLibrary(libraryName); err != nil {
return false, fmt.Sprintf("could not load NVML library: %v", err)
}
@@ -63,9 +53,15 @@ func (i *infolib) HasNvml() (bool, string) {
}
// 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")
+// 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)
@@ -86,17 +82,3 @@ func (i *infolib) IsTegraSystem() (bool, string) {
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
-}
diff --git a/pkg/nvlib/info/property-extractor_mock.go b/pkg/nvlib/info/property-extractor_mock.go
new file mode 100644
index 0000000..570dfe3
--- /dev/null
+++ b/pkg/nvlib/info/property-extractor_mock.go
@@ -0,0 +1,178 @@
+// 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")
+// },
+// }
+//
+// // 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)
+
+ // 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 {
+ }
+ }
+ lockHasDXCore sync.RWMutex
+ lockHasNvml sync.RWMutex
+ lockHasTegraFiles sync.RWMutex
+ lockIsTegraSystem 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
+}
diff --git a/pkg/nvlib/info/root.go b/pkg/nvlib/info/root.go
new file mode 100644
index 0000000..d38dc73
--- /dev/null
+++ b/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
+}
From d1e08f17ea893eead136008184e0e94392918cc0 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Tue, 26 Mar 2024 10:46:42 +0200
Subject: [PATCH 3/4] Add UsesOnlyNVGPUModule check to PropertyExtractor
interface
Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
pkg/nvlib/info/api.go | 1 +
pkg/nvlib/info/builder.go | 22 +++++++-
pkg/nvlib/info/options.go | 20 ++++++++
pkg/nvlib/info/property-extractor.go | 61 ++++++++++++++++++++++-
pkg/nvlib/info/property-extractor_mock.go | 45 +++++++++++++++--
5 files changed, 142 insertions(+), 7 deletions(-)
diff --git a/pkg/nvlib/info/api.go b/pkg/nvlib/info/api.go
index 41b8e22..b466bcb 100644
--- a/pkg/nvlib/info/api.go
+++ b/pkg/nvlib/info/api.go
@@ -31,4 +31,5 @@ type PropertyExtractor interface {
HasTegraFiles() (bool, string)
// Deprecated: Use HasTegraFiles instead.
IsTegraSystem() (bool, string)
+ UsesOnlyNVGPUModule() (bool, string)
}
diff --git a/pkg/nvlib/info/builder.go b/pkg/nvlib/info/builder.go
index d9275ca..bf2dd89 100644
--- a/pkg/nvlib/info/builder.go
+++ b/pkg/nvlib/info/builder.go
@@ -16,8 +16,16 @@
package info
+import (
+ "github.com/NVIDIA/go-nvml/pkg/nvml"
+
+ "github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
+)
+
type options struct {
- root root
+ root root
+ nvmllib nvml.Interface
+ devicelib device.Interface
}
// New creates a new instance of the 'info' Interface.
@@ -29,7 +37,17 @@ func New(opts ...Option) Interface {
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))
+ }
return &propertyExtractor{
- root: o.root,
+ root: o.root,
+ nvmllib: o.nvmllib,
+ devicelib: o.devicelib,
}
}
diff --git a/pkg/nvlib/info/options.go b/pkg/nvlib/info/options.go
index c4265d1..f8b47aa 100644
--- a/pkg/nvlib/info/options.go
+++ b/pkg/nvlib/info/options.go
@@ -16,9 +16,29 @@
package info
+import (
+ "github.com/NVIDIA/go-nvml/pkg/nvml"
+
+ "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(l *options) {
+ l.devicelib = devicelib
+ }
+}
+
+// WithNvmlLib sets the nvml library for the library.
+func WithNvmlLib(nvmllib nvml.Interface) Option {
+ return func(l *options) {
+ l.nvmllib = nvmllib
+ }
+}
+
// WithRoot provides a Option to set the root of the 'info' interface.
func WithRoot(r string) Option {
return func(i *options) {
diff --git a/pkg/nvlib/info/property-extractor.go b/pkg/nvlib/info/property-extractor.go
index 9e41a54..43ec3b8 100644
--- a/pkg/nvlib/info/property-extractor.go
+++ b/pkg/nvlib/info/property-extractor.go
@@ -20,10 +20,16 @@ import (
"fmt"
"os"
"strings"
+
+ "github.com/NVIDIA/go-nvml/pkg/nvml"
+
+ "github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
)
type propertyExtractor struct {
- root root
+ root root
+ nvmllib nvml.Interface
+ devicelib device.Interface
}
var _ Interface = &propertyExtractor{}
@@ -82,3 +88,56 @@ func (i *propertyExtractor) HasTegraFiles() (bool, string) {
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"
+}
diff --git a/pkg/nvlib/info/property-extractor_mock.go b/pkg/nvlib/info/property-extractor_mock.go
index 570dfe3..f2b057e 100644
--- a/pkg/nvlib/info/property-extractor_mock.go
+++ b/pkg/nvlib/info/property-extractor_mock.go
@@ -29,6 +29,9 @@ var _ PropertyExtractor = &PropertyExtractorMock{}
// 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
@@ -48,6 +51,9 @@ type PropertyExtractorMock struct {
// 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.
@@ -62,11 +68,15 @@ type PropertyExtractorMock 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
+ lockHasDXCore sync.RWMutex
+ lockHasNvml sync.RWMutex
+ lockHasTegraFiles sync.RWMutex
+ lockIsTegraSystem sync.RWMutex
+ lockUsesOnlyNVGPUModule sync.RWMutex
}
// HasDXCore calls HasDXCoreFunc.
@@ -176,3 +186,30 @@ func (mock *PropertyExtractorMock) IsTegraSystemCalls() []struct {
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
+}
From 21c8f035ca66b29da378c3bb51e5c94dbcac4e73 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Tue, 26 Mar 2024 10:55:59 +0200
Subject: [PATCH 4/4] Add ResolvePlatform function to info package
Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
pkg/nvlib/info/api.go | 6 ++
pkg/nvlib/info/builder.go | 35 +++++++--
pkg/nvlib/info/logger.go | 28 +++++++
pkg/nvlib/info/options.go | 31 +++++++-
pkg/nvlib/info/property-extractor.go | 2 +-
pkg/nvlib/info/resolver.go | 64 ++++++++++++++++
pkg/nvlib/info/resolver_test.go | 110 +++++++++++++++++++++++++++
7 files changed, 266 insertions(+), 10 deletions(-)
create mode 100644 pkg/nvlib/info/logger.go
create mode 100644 pkg/nvlib/info/resolver.go
create mode 100644 pkg/nvlib/info/resolver_test.go
diff --git a/pkg/nvlib/info/api.go b/pkg/nvlib/info/api.go
index b466bcb..1c62d63 100644
--- a/pkg/nvlib/info/api.go
+++ b/pkg/nvlib/info/api.go
@@ -18,9 +18,15 @@ 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.
//
diff --git a/pkg/nvlib/info/builder.go b/pkg/nvlib/info/builder.go
index bf2dd89..87f20f0 100644
--- a/pkg/nvlib/info/builder.go
+++ b/pkg/nvlib/info/builder.go
@@ -22,18 +22,30 @@ import (
"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.
+// 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 = "/"
}
@@ -45,9 +57,22 @@ func New(opts ...Option) Interface {
if o.devicelib == nil {
o.devicelib = device.New(device.WithNvml(o.nvmllib))
}
- return &propertyExtractor{
- root: o.root,
- nvmllib: o.nvmllib,
- devicelib: o.devicelib,
+ 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,
}
}
diff --git a/pkg/nvlib/info/logger.go b/pkg/nvlib/info/logger.go
new file mode 100644
index 0000000..6a6f74e
--- /dev/null
+++ b/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{}) {}
diff --git a/pkg/nvlib/info/options.go b/pkg/nvlib/info/options.go
index f8b47aa..e05c2bf 100644
--- a/pkg/nvlib/info/options.go
+++ b/pkg/nvlib/info/options.go
@@ -27,15 +27,22 @@ type Option func(*options)
// WithDeviceLib sets the device library for the library.
func WithDeviceLib(devicelib device.Interface) Option {
- return func(l *options) {
- l.devicelib = devicelib
+ return func(i *options) {
+ i.devicelib = devicelib
+ }
+}
+
+// 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(l *options) {
- l.nvmllib = nvmllib
+ return func(i *options) {
+ i.nvmllib = nvmllib
}
}
@@ -45,3 +52,19 @@ func WithRoot(r string) Option {
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
+ }
+}
diff --git a/pkg/nvlib/info/property-extractor.go b/pkg/nvlib/info/property-extractor.go
index 43ec3b8..5d5d97c 100644
--- a/pkg/nvlib/info/property-extractor.go
+++ b/pkg/nvlib/info/property-extractor.go
@@ -32,7 +32,7 @@ type propertyExtractor struct {
devicelib device.Interface
}
-var _ Interface = &propertyExtractor{}
+var _ PropertyExtractor = &propertyExtractor{}
// HasDXCore returns true if DXCore is detected on the system.
func (i *propertyExtractor) HasDXCore() (bool, string) {
diff --git a/pkg/nvlib/info/resolver.go b/pkg/nvlib/info/resolver.go
new file mode 100644
index 0000000..1aeb04c
--- /dev/null
+++ b/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
+ }
+}
diff --git a/pkg/nvlib/info/resolver_test.go b/pkg/nvlib/info/resolver_test.go
new file mode 100644
index 0000000..611ec54
--- /dev/null
+++ b/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())
+ })
+ }
+}