[no-relnote] Add RuntimeMode type

Signed-off-by: Evan Lezar <[email protected]>

Author: elezar

cmd/nvidia-container-runtime-hook/main.go

@@ -82,7 +82,7 @@ func doPrestart() {
 		return
 	}
 
-	if !hook.NVIDIAContainerRuntimeHookConfig.SkipModeDetection && info.ResolveAutoMode(&logInterceptor{}, hook.NVIDIAContainerRuntimeConfig.Mode, container.Image) != "legacy" {
+	if !hook.NVIDIAContainerRuntimeHookConfig.SkipModeDetection && info.ResolveAutoMode(&logInterceptor{}, hook.NVIDIAContainerRuntimeConfig.Mode, container.Image) != info.RuntimeModeLegacy {
 		log.Panicln("invoking the NVIDIA Container Runtime Hook directly (e.g. specifying the docker --gpus flag) is not supported. Please use the NVIDIA Container Runtime (e.g. specify the --runtime=nvidia flag) instead.")
 	}
 

internal/info/auto.go

@@ -23,22 +23,31 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 )
 
+// A RuntimeMode is used to select a specific mode of operation for the NVIDIA Container Runtime.
+type RuntimeMode string
+
+const (
+	RuntimeModeLegacy = RuntimeMode("legacy")
+	RuntimeModeCSV    = RuntimeMode("csv")
+	RuntimeModeCDI    = RuntimeMode("cdi")
+)
+
 // ResolveAutoMode determines the correct mode for the platform if set to "auto"
-func ResolveAutoMode(logger logger.Interface, mode string, image image.CUDA) (rmode string) {
+func ResolveAutoMode(logger logger.Interface, mode string, image image.CUDA) (rmode RuntimeMode) {
 	return resolveMode(logger, mode, image, nil)
 }
 
-func resolveMode(logger logger.Interface, mode string, image image.CUDA, propertyExtractor info.PropertyExtractor) (rmode string) {
+func resolveMode(logger logger.Interface, mode string, image image.CUDA, propertyExtractor info.PropertyExtractor) (rmode RuntimeMode) {
 	if mode != "auto" {
 		logger.Infof("Using requested mode '%s'", mode)
-		return mode
+		return RuntimeMode(mode)
 	}
 	defer func() {
 		logger.Infof("Auto-detected mode as '%v'", rmode)
 	}()
 
 	if image.OnlyFullyQualifiedCDIDevices() {
-		return "cdi"
+		return RuntimeModeCDI
 	}
 
 	nvinfo := info.New(
@@ -48,9 +57,9 @@ func resolveMode(logger logger.Interface, mode string, image image.CUDA, propert
 
 	switch nvinfo.ResolvePlatform() {
 	case info.PlatformNVML, info.PlatformWSL:
-		return "legacy"
+		return RuntimeModeLegacy
 	case info.PlatformTegra:
-		return "csv"
+		return RuntimeModeCSV
 	}
-	return "legacy"
+	return RuntimeModeLegacy
 }

internal/runtime/runtime_factory.go

@@ -105,26 +105,26 @@ func newSpecModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Sp
 	return modifiers, nil
 }
 
-func newModeModifier(logger logger.Interface, mode string, cfg *config.Config, ociSpec oci.Spec, image image.CUDA) (oci.SpecModifier, error) {
+func newModeModifier(logger logger.Interface, mode info.RuntimeMode, cfg *config.Config, ociSpec oci.Spec, image image.CUDA) (oci.SpecModifier, error) {
 	switch mode {
-	case "legacy":
+	case info.RuntimeModeLegacy:
 		return modifier.NewStableRuntimeModifier(logger, cfg.NVIDIAContainerRuntimeHookConfig.Path), nil
-	case "csv":
+	case info.RuntimeModeCSV:
 		return modifier.NewCSVModifier(logger, cfg, image)
-	case "cdi":
+	case info.RuntimeModeCDI:
 		return modifier.NewCDIModifier(logger, cfg, ociSpec)
 	}
 
 	return nil, fmt.Errorf("invalid runtime mode: %v", cfg.NVIDIAContainerRuntimeConfig.Mode)
 }
 
 // supportedModifierTypes returns the modifiers supported for a specific runtime mode.
-func supportedModifierTypes(mode string) []string {
+func supportedModifierTypes(mode info.RuntimeMode) []string {
 	switch mode {
-	case "cdi":
+	case info.RuntimeModeCDI:
 		// For CDI mode we make no additional modifications.
 		return []string{"nvidia-hook-remover", "mode"}
-	case "csv":
+	case info.RuntimeModeCSV:
 		// For CSV mode we support mode and feature-gated modification.
 		return []string{"nvidia-hook-remover", "mode", "feature-gated"}
 	default: