Add ResolvePlatform function to info package

Signed-off-by: Evan Lezar <[email protected]>

Author: elezar

pkg/nvlib/info/api.go

@@ -18,9 +18,15 @@ package info
 
 // Interface provides the API to the info package.
 type Interface interface {
+	PlatformResolver
 	Properties
 }
 
+// PlatformResolver defines a function to resolve a mode.
+type PlatformResolver interface {
+	ResolvePlatform() Platform
+}
+
 // Properties provides a set of functions to query capabilities of the system.
 //
 //go:generate moq -rm -stub -out properties_mock.go . Properties

pkg/nvlib/info/builder.go

@@ -23,17 +23,24 @@ import (
 )
 
 type options struct {
+	logger    basicLogger
 	root      root
 	nvmllib   nvml.Interface
 	devicelib device.Interface
+
+	platform   Platform
+	properties Properties
 }
 
-// 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 +52,23 @@ func New(opts ...Option) Interface {
 	if o.devicelib == nil {
 		o.devicelib = device.New(device.WithNvml(o.nvmllib))
 	}
+	if o.platform == "" {
+		o.platform = PlatformAuto
+	}
+	if o.properties == nil {
+		o.properties = &info{
+			root:      o.root,
+			nvmllib:   o.nvmllib,
+			devicelib: o.devicelib,
+		}
+	}
 	return &infolib{
-		root:      o.root,
-		nvmllib:   o.nvmllib,
-		devicelib: o.devicelib,
+		logger: o.logger,
+		PlatformResolver: &platformResolver{
+			logger:     o.logger,
+			platform:   o.platform,
+			Properties: o.properties,
+		},
+		Properties: o.properties,
 	}
 }

pkg/nvlib/info/info.go

@@ -27,15 +27,22 @@ import (
 )
 
 type infolib struct {
+	logger basicLogger
+	Properties
+	PlatformResolver
+}
+
+type info struct {
 	root      root
 	nvmllib   nvml.Interface
 	devicelib device.Interface
 }
 
+var _ Properties = &info{}
 var _ Interface = &infolib{}
 
 // HasDXCore returns true if DXCore is detected on the system.
-func (i *infolib) HasDXCore() (bool, string) {
+func (i *info) HasDXCore() (bool, string) {
 	const (
 		libraryName = "libdxcore.so"
 	)
@@ -47,7 +54,7 @@ func (i *infolib) HasDXCore() (bool, string) {
 }
 
 // HasNvml returns true if NVML is detected on the system.
-func (i *infolib) HasNvml() (bool, string) {
+func (i *info) HasNvml() (bool, string) {
 	const (
 		libraryName = "libnvidia-ml.so.1"
 	)
@@ -59,7 +66,7 @@ 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) {
+func (i *info) IsTegraSystem() (bool, string) {
 	tegraReleaseFile := i.root.join("/etc/nv_tegra_release")
 	tegraFamilyFile := i.root.join("/sys/devices/soc0/family")
 
@@ -93,7 +100,7 @@ func (i *infolib) IsTegraSystem() (bool, string) {
 //	GPU 0: Orin (nvgpu) (UUID: 54d0709b-558d-5a59-9c65-0c5fc14a21a4)
 //
 // This function returns true if ALL devices use the nvgpu module.
-func (i *infolib) UsesOnlyNVGPUModule() (uses bool, reason string) {
+func (i *info) UsesOnlyNVGPUModule() (uses bool, reason string) {
 	// We ensure that this function never panics
 	defer func() {
 		if err := recover(); err != 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

@@ -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,18 @@ func WithRoot(r string) Option {
 		i.root = root(r)
 	}
 }
+
+// WithProperties provides an Option to set the Properties interface implementation.
+// This is predominantly used for testing.
+func WithProperties(properties Properties) Option {
+	return func(i *options) {
+		i.properties = properties
+	}
+}
+
+// WithPlatform provides an option to set the platform explicitly.
+func WithPlatform(platform Platform) Option {
+	return func(i *options) {
+		i.platform = platform
+	}
+}

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
+	Properties
+}
+
+func (p platformResolver) ResolvePlatform() Platform {
+	if p.platform != PlatformAuto {
+		p.logger.Infof("Using requested platform '%s'", p.platform)
+		return p.platform
+	}
+
+	isWSL, reason := p.HasDXCore()
+	p.logger.Debugf("Is WSL-based system? %v: %v", isWSL, reason)
+
+	isTegra, reason := p.IsTegraSystem()
+	p.logger.Debugf("Is Tegra-based system? %v: %v", isTegra, reason)
+
+	hasNVML, reason := p.HasNvml()
+	p.logger.Debugf("Is NVML-based system? %v: %v", hasNVML, reason)
+
+	usesNVGPUModule, reason := p.UsesOnlyNVGPUModule()
+	p.logger.Debugf("Uses nvgpu kernel module? %v: %v", usesNVGPUModule, reason)
+
+	switch {
+	case isWSL:
+		return PlatformWSL
+	case (isTegra && !hasNVML), usesNVGPUModule:
+		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
+		isTegra             bool
+		hasDXCore           bool
+		hasNVML             bool
+		usesOnlyNVGPUModule bool
+		expected            string
+	}{
+		{
+			platform:  "auto",
+			hasDXCore: true,
+			expected:  "wsl",
+		},
+		{
+			platform:  "auto",
+			hasDXCore: false,
+			isTegra:   true,
+			hasNVML:   false,
+			expected:  "tegra",
+		},
+		{
+			platform:  "auto",
+			hasDXCore: false,
+			isTegra:   false,
+			hasNVML:   false,
+			expected:  "unknown",
+		},
+		{
+			platform:  "auto",
+			hasDXCore: false,
+			isTegra:   true,
+			hasNVML:   true,
+			expected:  "nvml",
+		},
+		{
+			platform:            "auto",
+			hasDXCore:           false,
+			isTegra:             true,
+			hasNVML:             true,
+			usesOnlyNVGPUModule: true,
+			expected:            "tegra",
+		},
+		{
+			platform:  "nvml",
+			hasDXCore: true,
+			isTegra:   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(
+				WithProperties(&PropertiesMock{
+					HasDXCoreFunc: func() (bool, string) {
+						return tc.hasDXCore, ""
+					},
+					HasNvmlFunc: func() (bool, string) {
+						return tc.hasNVML, ""
+					},
+					IsTegraSystemFunc: func() (bool, string) {
+						return tc.isTegra, ""
+					},
+					UsesOnlyNVGPUModuleFunc: func() (bool, string) {
+						return tc.usesOnlyNVGPUModule, ""
+					},
+				}),
+				WithPlatform(Platform(tc.platform)),
+			)
+
+			require.Equal(t, Platform(tc.expected), l.ResolvePlatform())
+		})
+	}
+}