UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2025-06-16 11:30:20 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[no-relnote] Add RuntimeMode type

Browse Source 
Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2025-02-11 15:37:41 +01:00
parent 7fe0aad96b
commit 38a40b08a2
No known key found for this signature in database
2 changed files with 26 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
internal/info/auto.go
View File

@ -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
}

18
internal/runtime/runtime_factory.go
View File

@ -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: