From 38a40b08a22bed86672ebbcc139ec5c503f68d6f Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Tue, 11 Feb 2025 15:37:41 +0100
Subject: [PATCH] [no-relnote] Add RuntimeMode type

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 internal/info/auto.go               | 25 +++++++++++++++++--------
 internal/runtime/runtime_factory.go | 18 +++++++++---------
 2 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/internal/info/auto.go b/internal/info/auto.go
index 5246477c..14cabb81 100644
--- a/internal/info/auto.go
+++ b/internal/info/auto.go
@@ -23,8 +23,17 @@ 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")
+)
+
 type RuntimeModeResolver interface {
-	ResolveRuntimeMode(string) string
+	ResolveRuntimeMode(string) RuntimeMode
 }
 
 type modeResolver struct {
@@ -67,7 +76,7 @@ func NewRuntimeModeResolver(opts ...Option) RuntimeModeResolver {
 }
 
 // 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) {
 	r := modeResolver{
 		logger:            logger,
 		image:             &image,
@@ -76,17 +85,17 @@ func ResolveAutoMode(logger logger.Interface, mode string, image image.CUDA) (rm
 	return r.ResolveRuntimeMode(mode)
 }
 
-func (m *modeResolver) ResolveRuntimeMode(mode string) (rmode string) {
+func (m *modeResolver) ResolveRuntimeMode(mode string) (rmode RuntimeMode) {
 	if mode != "auto" {
 		m.logger.Infof("Using requested mode '%s'", mode)
-		return mode
+		return RuntimeMode(mode)
 	}
 	defer func() {
 		m.logger.Infof("Auto-detected mode as '%v'", rmode)
 	}()
 
 	if m.image.OnlyFullyQualifiedCDIDevices() {
-		return "cdi"
+		return RuntimeModeCDI
 	}
 
 	nvinfo := info.New(
@@ -96,9 +105,9 @@ func (m *modeResolver) ResolveRuntimeMode(mode string) (rmode string) {
 
 	switch nvinfo.ResolvePlatform() {
 	case info.PlatformNVML, info.PlatformWSL:
-		return "legacy"
+		return RuntimeModeLegacy
 	case info.PlatformTegra:
-		return "csv"
+		return RuntimeModeCSV
 	}
-	return "legacy"
+	return RuntimeModeLegacy
 }
diff --git a/internal/runtime/runtime_factory.go b/internal/runtime/runtime_factory.go
index 8bb050a5..adaaa917 100644
--- a/internal/runtime/runtime_factory.go
+++ b/internal/runtime/runtime_factory.go
@@ -101,13 +101,13 @@ func newSpecModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Sp
 	return modifiers, nil
 }
 
-func newModeModifier(logger logger.Interface, mode string, cfg *config.Config, image image.CUDA) (oci.SpecModifier, error) {
+func newModeModifier(logger logger.Interface, mode info.RuntimeMode, cfg *config.Config, 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, image)
 	}
 
@@ -119,7 +119,7 @@ func newModeModifier(logger logger.Interface, mode string, cfg *config.Config, i
 // The image is also used to determine the runtime mode to apply.
 // If a non-CDI mode is detected we ensure that the image does not process
 // annotation devices.
-func initRuntimeModeAndImage(logger logger.Interface, cfg *config.Config, ociSpec oci.Spec) (string, *image.CUDA, error) {
+func initRuntimeModeAndImage(logger logger.Interface, cfg *config.Config, ociSpec oci.Spec) (info.RuntimeMode, *image.CUDA, error) {
 	rawSpec, err := ociSpec.Load()
 	if err != nil {
 		return "", nil, fmt.Errorf("failed to load OCI spec: %v", err)
@@ -142,7 +142,7 @@ func initRuntimeModeAndImage(logger logger.Interface, cfg *config.Config, ociSpe
 	)
 	mode := modeResolver.ResolveRuntimeMode(cfg.NVIDIAContainerRuntimeConfig.Mode)
 	// We update the mode here so that we can continue passing just the config to other functions.
-	cfg.NVIDIAContainerRuntimeConfig.Mode = mode
+	cfg.NVIDIAContainerRuntimeConfig.Mode = string(mode)
 
 	if mode == "cdi" || len(cfg.NVIDIAContainerRuntimeConfig.Modes.CDI.AnnotationPrefixes) == 0 {
 		return mode, &image, nil
@@ -158,12 +158,12 @@ func initRuntimeModeAndImage(logger logger.Interface, cfg *config.Config, ociSpe
 }
 
 // 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", "feature-gated", "mode"}
 	default: