Skip to content

Commit 21c8f03

Browse files
elezarelezar
committed
Add ResolvePlatform function to info package
Signed-off-by: Evan Lezar <[email protected]>
1 parent d1e08f1 commit 21c8f03

File tree

7 files changed

+266
-10
lines changed

7 files changed

+266
-10
lines changed

‎pkg/nvlib/info/api.go

+6
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,15 @@ package info
1818

1919
// Interface provides the API to the info package.
2020
type Interface interface {
21+
PlatformResolver
2122
PropertyExtractor
2223
}
2324

25+
// PlatformResolver defines a function to resolve the current platform.
26+
type PlatformResolver interface {
27+
ResolvePlatform() Platform
28+
}
29+
2430
// PropertyExtractor provides a set of functions to query capabilities of the
2531
// system.
2632
//

‎pkg/nvlib/info/builder.go

+30-5
Original file line numberDiff line numberDiff line change
@@ -22,18 +22,30 @@ import (
2222
"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
2323
)
2424

25+
type infolib struct {
26+
PropertyExtractor
27+
PlatformResolver
28+
}
29+
2530
type options struct {
31+
logger basicLogger
2632
root root
2733
nvmllib nvml.Interface
2834
devicelib device.Interface
35+
36+
platform Platform
37+
propertyExtractor PropertyExtractor
2938
}
3039

31-
// New creates a new instance of the 'info' Interface.
40+
// New creates a new instance of the 'info' interface.
3241
func New(opts ...Option) Interface {
3342
o := &options{}
3443
for _, opt := range opts {
3544
opt(o)
3645
}
46+
if o.logger == nil {
47+
o.logger = &nullLogger{}
48+
}
3749
if o.root == "" {
3850
o.root = "/"
3951
}
@@ -45,9 +57,22 @@ func New(opts ...Option) Interface {
4557
if o.devicelib == nil {
4658
o.devicelib = device.New(device.WithNvml(o.nvmllib))
4759
}
48-
return &propertyExtractor{
49-
root: o.root,
50-
nvmllib: o.nvmllib,
51-
devicelib: o.devicelib,
60+
if o.platform == "" {
61+
o.platform = PlatformAuto
62+
}
63+
if o.propertyExtractor == nil {
64+
o.propertyExtractor = &propertyExtractor{
65+
root: o.root,
66+
nvmllib: o.nvmllib,
67+
devicelib: o.devicelib,
68+
}
69+
}
70+
return &infolib{
71+
PlatformResolver: &platformResolver{
72+
logger: o.logger,
73+
platform: o.platform,
74+
propertyExtractor: o.propertyExtractor,
75+
},
76+
PropertyExtractor: o.propertyExtractor,
5277
}
5378
}

‎pkg/nvlib/info/logger.go

+28
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
/**
2+
# Copyright 2024 NVIDIA CORPORATION
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
**/
16+
17+
package info
18+
19+
type basicLogger interface {
20+
Debugf(string, ...interface{})
21+
Infof(string, ...interface{})
22+
}
23+
24+
type nullLogger struct{}
25+
26+
func (n *nullLogger) Debugf(string, ...interface{}) {}
27+
28+
func (n *nullLogger) Infof(string, ...interface{}) {}

‎pkg/nvlib/info/options.go

+27-4
Original file line numberDiff line numberDiff line change
@@ -27,15 +27,22 @@ type Option func(*options)
2727

2828
// WithDeviceLib sets the device library for the library.
2929
func WithDeviceLib(devicelib device.Interface) Option {
30-
return func(l *options) {
31-
l.devicelib = devicelib
30+
return func(i *options) {
31+
i.devicelib = devicelib
32+
}
33+
}
34+
35+
// WithLogger sets the logger for the library.
36+
func WithLogger(logger basicLogger) Option {
37+
return func(i *options) {
38+
i.logger = logger
3239
}
3340
}
3441

3542
// WithNvmlLib sets the nvml library for the library.
3643
func WithNvmlLib(nvmllib nvml.Interface) Option {
37-
return func(l *options) {
38-
l.nvmllib = nvmllib
44+
return func(i *options) {
45+
i.nvmllib = nvmllib
3946
}
4047
}
4148

@@ -45,3 +52,19 @@ func WithRoot(r string) Option {
4552
i.root = root(r)
4653
}
4754
}
55+
56+
// WithPropertyExtractor provides an Option to set the PropertyExtractor
57+
// interface implementation.
58+
// This is predominantly used for testing.
59+
func WithPropertyExtractor(propertyExtractor PropertyExtractor) Option {
60+
return func(i *options) {
61+
i.propertyExtractor = propertyExtractor
62+
}
63+
}
64+
65+
// WithPlatform provides an option to set the platform explicitly.
66+
func WithPlatform(platform Platform) Option {
67+
return func(i *options) {
68+
i.platform = platform
69+
}
70+
}

‎pkg/nvlib/info/property-extractor.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ type propertyExtractor struct {
3232
devicelib device.Interface
3333
}
3434

35-
var _ Interface = &propertyExtractor{}
35+
var _ PropertyExtractor = &propertyExtractor{}
3636

3737
// HasDXCore returns true if DXCore is detected on the system.
3838
func (i *propertyExtractor) HasDXCore() (bool, string) {

‎pkg/nvlib/info/resolver.go

+64
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/**
2+
# Copyright 2024 NVIDIA CORPORATION
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
**/
16+
17+
package info
18+
19+
// Platform represents a supported plaform.
20+
type Platform string
21+
22+
const (
23+
PlatformAuto = Platform("auto")
24+
PlatformNVML = Platform("nvml")
25+
PlatformTegra = Platform("tegra")
26+
PlatformWSL = Platform("wsl")
27+
PlatformUnknown = Platform("unknown")
28+
)
29+
30+
type platformResolver struct {
31+
logger basicLogger
32+
platform Platform
33+
propertyExtractor PropertyExtractor
34+
}
35+
36+
func (p platformResolver) ResolvePlatform() Platform {
37+
if p.platform != PlatformAuto {
38+
p.logger.Infof("Using requested platform '%s'", p.platform)
39+
return p.platform
40+
}
41+
42+
hasDXCore, reason := p.propertyExtractor.HasDXCore()
43+
p.logger.Debugf("Is WSL-based system? %v: %v", hasDXCore, reason)
44+
45+
hasTegraFiles, reason := p.propertyExtractor.HasTegraFiles()
46+
p.logger.Debugf("Is Tegra-based system? %v: %v", hasTegraFiles, reason)
47+
48+
hasNVML, reason := p.propertyExtractor.HasNvml()
49+
p.logger.Debugf("Is NVML-based system? %v: %v", hasNVML, reason)
50+
51+
usesOnlyNVGPUModule, reason := p.propertyExtractor.UsesOnlyNVGPUModule()
52+
p.logger.Debugf("Uses nvgpu kernel module? %v: %v", usesOnlyNVGPUModule, reason)
53+
54+
switch {
55+
case hasDXCore:
56+
return PlatformWSL
57+
case (hasTegraFiles && !hasNVML), usesOnlyNVGPUModule:
58+
return PlatformTegra
59+
case hasNVML:
60+
return PlatformNVML
61+
default:
62+
return PlatformUnknown
63+
}
64+
}

‎pkg/nvlib/info/resolver_test.go

+110
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
/**
2+
# Copyright (c) NVIDIA CORPORATION. All rights reserved.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
**/
16+
17+
package info
18+
19+
import (
20+
"fmt"
21+
"testing"
22+
23+
"github.com/stretchr/testify/require"
24+
)
25+
26+
func TestResolvePlatform(t *testing.T) {
27+
testCases := []struct {
28+
platform string
29+
hasTegraFiles bool
30+
hasDXCore bool
31+
hasNVML bool
32+
usesOnlyNVGPUModule bool
33+
expected string
34+
}{
35+
{
36+
platform: "auto",
37+
hasDXCore: true,
38+
expected: "wsl",
39+
},
40+
{
41+
platform: "auto",
42+
hasDXCore: false,
43+
hasTegraFiles: true,
44+
hasNVML: false,
45+
expected: "tegra",
46+
},
47+
{
48+
platform: "auto",
49+
hasDXCore: false,
50+
hasTegraFiles: false,
51+
hasNVML: false,
52+
expected: "unknown",
53+
},
54+
{
55+
platform: "auto",
56+
hasDXCore: false,
57+
hasTegraFiles: true,
58+
hasNVML: true,
59+
expected: "nvml",
60+
},
61+
{
62+
platform: "auto",
63+
hasDXCore: false,
64+
hasTegraFiles: true,
65+
hasNVML: true,
66+
usesOnlyNVGPUModule: true,
67+
expected: "tegra",
68+
},
69+
{
70+
platform: "nvml",
71+
hasDXCore: true,
72+
hasTegraFiles: true,
73+
expected: "nvml",
74+
},
75+
{
76+
platform: "wsl",
77+
hasDXCore: false,
78+
expected: "wsl",
79+
},
80+
{
81+
platform: "not-auto",
82+
hasDXCore: true,
83+
expected: "not-auto",
84+
},
85+
}
86+
87+
for i, tc := range testCases {
88+
t.Run(fmt.Sprintf("test case %d", i), func(t *testing.T) {
89+
l := New(
90+
WithPropertyExtractor(&PropertyExtractorMock{
91+
HasDXCoreFunc: func() (bool, string) {
92+
return tc.hasDXCore, ""
93+
},
94+
HasNvmlFunc: func() (bool, string) {
95+
return tc.hasNVML, ""
96+
},
97+
HasTegraFilesFunc: func() (bool, string) {
98+
return tc.hasTegraFiles, ""
99+
},
100+
UsesOnlyNVGPUModuleFunc: func() (bool, string) {
101+
return tc.usesOnlyNVGPUModule, ""
102+
},
103+
}),
104+
WithPlatform(Platform(tc.platform)),
105+
)
106+
107+
require.Equal(t, Platform(tc.expected), l.ResolvePlatform())
108+
})
109+
}
110+
}

0 commit comments

Comments
 (0)