diff --git a/cmd/nvidia-ctk/cdi/generate/generate.go b/cmd/nvidia-ctk/cdi/generate/generate.go index 7aef47b8..caee034c 100644 --- a/cmd/nvidia-ctk/cdi/generate/generate.go +++ b/cmd/nvidia-ctk/cdi/generate/generate.go @@ -36,10 +36,6 @@ import ( ) const ( - discoveryModeAuto = "auto" - discoveryModeNVML = "nvml" - discoveryModeWSL = "wsl" - formatJSON = "json" formatYAML = "yaml" @@ -97,8 +93,8 @@ func (m command) build() *cli.Command { }, &cli.StringFlag{ 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.", - Value: discoveryModeAuto, + 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: nvcdi.ModeAuto, Destination: &cfg.discoveryMode, }, &cli.StringFlag{ @@ -133,9 +129,9 @@ func (m command) validateFlags(r *cli.Context, cfg *config) error { cfg.discoveryMode = strings.ToLower(cfg.discoveryMode) switch cfg.discoveryMode { - case discoveryModeAuto: - case discoveryModeNVML: - case discoveryModeWSL: + case nvcdi.ModeAuto: + case nvcdi.ModeNvml: + case nvcdi.ModeWsl: default: return fmt.Errorf("invalid discovery mode: %v", cfg.discoveryMode) } diff --git a/pkg/nvcdi/api.go b/pkg/nvcdi/api.go index 0fd07e72..267010e9 100644 --- a/pkg/nvcdi/api.go +++ b/pkg/nvcdi/api.go @@ -22,6 +22,15 @@ import ( "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 type Interface interface { GetCommonEdits() (*cdi.ContainerEdits, error) diff --git a/pkg/nvcdi/lib.go b/pkg/nvcdi/lib.go index 18d69d84..be5554eb 100644 --- a/pkg/nvcdi/lib.go +++ b/pkg/nvcdi/lib.go @@ -42,7 +42,7 @@ func New(opts ...Option) Interface { opt(l) } if l.mode == "" { - l.mode = "auto" + l.mode = ModeAuto } if l.logger == nil { l.logger = logrus.StandardLogger() @@ -61,7 +61,7 @@ func New(opts ...Option) Interface { } switch l.resolveMode() { - case "nvml": + case ModeNvml: if l.nvmllib == nil { l.nvmllib = nvml.New() } @@ -70,7 +70,7 @@ func New(opts ...Option) Interface { } return (*nvmllib)(l) - case "wsl": + case ModeWsl: 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. func (l *nvcdilib) resolveMode() (rmode string) { - if l.mode != "auto" { + if l.mode != ModeAuto { return l.mode } defer func() { @@ -91,8 +91,8 @@ func (l *nvcdilib) resolveMode() (rmode string) { l.logger.Debugf("Is WSL-based system? %v: %v", isWSL, reason) if isWSL { - return "wsl" + return ModeWsl } - return "nvml" + return ModeNvml }