Add and resolve auto discovery mode for cdi generation

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2023-02-13 16:04:30 +01:00
parent f36c775d50
commit 4ccb0b9a53
2 changed files with 28 additions and 4 deletions

View File

@ -36,6 +36,7 @@ import (
) )
const ( const (
discoveryModeAuto = "auto"
discoveryModeNVML = "nvml" discoveryModeNVML = "nvml"
discoveryModeWSL = "wsl" discoveryModeWSL = "wsl"
@ -96,8 +97,8 @@ func (m command) build() *cli.Command {
}, },
&cli.StringFlag{ &cli.StringFlag{
Name: "discovery-mode", Name: "discovery-mode",
Usage: "The mode to use when discovering the available entities. One of [nvml | wsl]", Usage: "The mode to use when discovering the available entities. One of [auto | nvml | wsl]. I mode is set to 'auto' the mode will be determined based on the system configuration.",
Value: discoveryModeNVML, Value: discoveryModeAuto,
Destination: &cfg.discoveryMode, Destination: &cfg.discoveryMode,
}, },
&cli.StringFlag{ &cli.StringFlag{
@ -132,6 +133,7 @@ func (m command) validateFlags(r *cli.Context, cfg *config) error {
cfg.discoveryMode = strings.ToLower(cfg.discoveryMode) cfg.discoveryMode = strings.ToLower(cfg.discoveryMode)
switch cfg.discoveryMode { switch cfg.discoveryMode {
case discoveryModeAuto:
case discoveryModeNVML: case discoveryModeNVML:
case discoveryModeWSL: case discoveryModeWSL:
default: default:

View File

@ -19,6 +19,7 @@ package nvcdi
import ( import (
"github.com/sirupsen/logrus" "github.com/sirupsen/logrus"
"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/device" "gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/device"
"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/info"
"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml" "gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml"
) )
@ -39,7 +40,7 @@ func New(opts ...Option) Interface {
opt(l) opt(l)
} }
if l.mode == "" { if l.mode == "" {
l.mode = "nvml" l.mode = "auto"
} }
if l.logger == nil { if l.logger == nil {
l.logger = logrus.StandardLogger() l.logger = logrus.StandardLogger()
@ -54,7 +55,7 @@ func New(opts ...Option) Interface {
l.nvidiaCTKPath = "/usr/bin/nvidia-ctk" l.nvidiaCTKPath = "/usr/bin/nvidia-ctk"
} }
switch l.mode { switch l.resolveMode() {
case "nvml": case "nvml":
if l.nvmllib == nil { if l.nvmllib == nil {
l.nvmllib = nvml.New() l.nvmllib = nvml.New()
@ -71,3 +72,24 @@ func New(opts ...Option) Interface {
// TODO: We want an error here. // TODO: We want an error here.
return nil return nil
} }
// resolveMode resolves the mode for CDI spec generation based on the current system.
func (l *nvcdilib) resolveMode() (rmode string) {
if l.mode != "auto" {
return l.mode
}
defer func() {
l.logger.Infof("Auto-detected mode as %q", rmode)
}()
nvinfo := info.New()
isWSL, reason := nvinfo.HasDXCore()
l.logger.Debugf("Is WSL-based system? %v: %v", isWSL, reason)
if isWSL {
return "wsl"
}
return "nvml"
}