UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2024-11-22 00:08:11 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add mode constants to nvcdi

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2023-02-20 16:27:34 +02:00
parent b76808dbd5
commit 2680c45811
3 changed files with 20 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
cmd/nvidia-ctk/cdi/generate/generate.go
View File

@ -36,10 +36,6 @@ import (
) )
const ( const (
discoveryModeAuto = "auto"
discoveryModeNVML = "nvml"
discoveryModeWSL = "wsl"
formatJSON = "json" formatJSON = "json"
formatYAML = "yaml" formatYAML = "yaml"
@ -97,8 +93,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 [auto | nvml | wsl]. I mode is set to 'auto' the mode will be determined based on the system configuration.", Usage: "The mode to use when discovering the available entities. One of [auto | nvml | wsl]. If mode is set to 'auto' the mode will be determined based on the system configuration.",
Value: discoveryModeAuto, Value: nvcdi.ModeAuto,
Destination: &cfg.discoveryMode, Destination: &cfg.discoveryMode,
}, },
&cli.StringFlag{ &cli.StringFlag{
@ -133,9 +129,9 @@ 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 nvcdi.ModeAuto:
case discoveryModeNVML: case nvcdi.ModeNvml:
case discoveryModeWSL: case nvcdi.ModeWsl:
default: default:
return fmt.Errorf("invalid discovery mode: %v", cfg.discoveryMode) return fmt.Errorf("invalid discovery mode: %v", cfg.discoveryMode)
} }

9
pkg/nvcdi/api.go
View File

@ -22,6 +22,15 @@ import (
"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/device" "gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvlib/device"
) )
const (
// ModeAuto configures the CDI spec generator to automatically detect the system configuration
ModeAuto = "auto"
// ModeNvml configures the CDI spec generator to use the NVML library.
ModeNvml = "nvml"
// ModeWsl configures the CDI spec generator to generate a WSL spec.
ModeWsl = "wsl"
)
// Interface defines the API for the nvcdi package // Interface defines the API for the nvcdi package
type Interface interface { type Interface interface {
GetCommonEdits() (*cdi.ContainerEdits, error) GetCommonEdits() (*cdi.ContainerEdits, error)

12
pkg/nvcdi/lib.go
View File

@ -42,7 +42,7 @@ func New(opts ...Option) Interface {
opt(l) opt(l)
} }
if l.mode == "" { if l.mode == "" {
l.mode = "auto" l.mode = ModeAuto
} }
if l.logger == nil { if l.logger == nil {
l.logger = logrus.StandardLogger() l.logger = logrus.StandardLogger()
@ -61,7 +61,7 @@ func New(opts ...Option) Interface {
} }
switch l.resolveMode() { switch l.resolveMode() {
case "nvml": case ModeNvml:
if l.nvmllib == nil { if l.nvmllib == nil {
l.nvmllib = nvml.New() l.nvmllib = nvml.New()
} }
@ -70,7 +70,7 @@ func New(opts ...Option) Interface {
} }
return (*nvmllib)(l) return (*nvmllib)(l)
case "wsl": case ModeWsl:
return (*wsllib)(l) return (*wsllib)(l)
} }
@ -80,7 +80,7 @@ func New(opts ...Option) Interface {
// resolveMode resolves the mode for CDI spec generation based on the current system. // resolveMode resolves the mode for CDI spec generation based on the current system.
func (l *nvcdilib) resolveMode() (rmode string) { func (l *nvcdilib) resolveMode() (rmode string) {
if l.mode != "auto" { if l.mode != ModeAuto {
return l.mode return l.mode
} }
defer func() { defer func() {
@ -91,8 +91,8 @@ func (l *nvcdilib) resolveMode() (rmode string) {
l.logger.Debugf("Is WSL-based system? %v: %v", isWSL, reason) l.logger.Debugf("Is WSL-based system? %v: %v", isWSL, reason)
if isWSL { if isWSL {
return "wsl" return ModeWsl
} }
return "nvml" return ModeNvml
} }