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.
-	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:

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/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
 
+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(*infolib)
+type Option func(*options)
 
-// New creates a new instance of the 'info' interface.
-func New(opts ...Option) Interface {
-	i := &infolib{}
-	for _, opt := range opts {
-		opt(i)
+// 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
 	}
 }