From 2680c458117fa6dc38d9009859c033b45f90fda9 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Mon, 20 Feb 2023 16:27:34 +0200
Subject: [PATCH] Add mode constants to nvcdi

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 cmd/nvidia-ctk/cdi/generate/generate.go | 14 +++++---------
 pkg/nvcdi/api.go                        |  9 +++++++++
 pkg/nvcdi/lib.go                        | 12 ++++++------
 3 files changed, 20 insertions(+), 15 deletions(-)

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
 }