From 4cfb7e139d8980e0c197cee0f74de1b89d2780ef Mon Sep 17 00:00:00 2001 From: Evan Lezar Date: Tue, 11 Feb 2025 15:37:41 +0100 Subject: [PATCH] [no-relnote] Add RuntimeMode type Signed-off-by: Evan Lezar --- cmd/nvidia-container-runtime-hook/main.go | 2 +- internal/info/auto.go | 23 ++++++++++++++++------- internal/runtime/runtime_factory.go | 14 +++++++------- 3 files changed, 24 insertions(+), 15 deletions(-) diff --git a/cmd/nvidia-container-runtime-hook/main.go b/cmd/nvidia-container-runtime-hook/main.go index e864a51d..4d6c0029 100644 --- a/cmd/nvidia-container-runtime-hook/main.go +++ b/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.") } diff --git a/internal/info/auto.go b/internal/info/auto.go index c6800da1..27e5132e 100644 --- a/internal/info/auto.go +++ b/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 } diff --git a/internal/runtime/runtime_factory.go b/internal/runtime/runtime_factory.go index 2b5cd9c6..90e004d4 100644 --- a/internal/runtime/runtime_factory.go +++ b/internal/runtime/runtime_factory.go @@ -105,13 +105,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, 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) } @@ -119,12 +119,12 @@ func newModeModifier(logger logger.Interface, mode string, cfg *config.Config, o } // 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: