Use jit-cdi mode for CSV systems

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

Author: elezar

internal/info/auto.go

@@ -53,9 +53,10 @@ type RuntimeModeResolver interface {
 type modeResolver struct {
 	logger logger.Interface
 	// TODO: This only needs to consider the requested devices.
-	image             *image.CUDA
-	propertyExtractor info.PropertyExtractor
-	defaultMode       RuntimeMode
+	image                       *image.CUDA
+	propertyExtractor           info.PropertyExtractor
+	defaultMode                 RuntimeMode
+	forceCSVModeForTegraSystems bool
 }
 
 type Option func(*modeResolver)
@@ -66,6 +67,12 @@ func WithDefaultMode(defaultMode RuntimeMode) Option {
 	}
 }
 
+func WithForceCSVModeForTegraSystems(forceCSVModeForTegraSystems bool) Option {
+	return func(mr *modeResolver) {
+		mr.forceCSVModeForTegraSystems = forceCSVModeForTegraSystems
+	}
+}
+
 func WithLogger(logger logger.Interface) Option {
 	return func(mr *modeResolver) {
 		mr.logger = logger
@@ -130,7 +137,10 @@ func (m *modeResolver) ResolveRuntimeMode(mode string) (rmode RuntimeMode) {
 	case info.PlatformNVML, info.PlatformWSL:
 		return m.defaultMode
 	case info.PlatformTegra:
-		return CSVRuntimeMode
+		if m.forceCSVModeForTegraSystems {
+			return CSVRuntimeMode
+		}
+		return JitCDIRuntimeMode
 	}
 	return m.defaultMode
 }

internal/modifier/cdi.go

@@ -65,7 +65,7 @@ func NewCDIModifier(logger logger.Interface, cfg *config.Config, image image.CUD
 		automaticDevices = append(automaticDevices, withUniqueDevices(gatedDevices(image)).DeviceRequests()...)
 		automaticDevices = append(automaticDevices, withUniqueDevices(imexDevices(image)).DeviceRequests()...)
 
-		automaticModifier, err := newAutomaticCDISpecModifier(logger, cfg, automaticDevices)
+		automaticModifier, err := newAutomaticCDISpecModifier(logger, cfg, image, automaticDevices)
 		if err == nil {
 			return automaticModifier, nil
 		}
@@ -163,9 +163,10 @@ func filterAutomaticDevices(devices []string) []string {
 	return automatic
 }
 
-func newAutomaticCDISpecModifier(logger logger.Interface, cfg *config.Config, devices []string) (oci.SpecModifier, error) {
+func newAutomaticCDISpecModifier(logger logger.Interface, cfg *config.Config, image image.CUDA, devices []string) (oci.SpecModifier, error) {
 	logger.Debugf("Generating in-memory CDI specs for devices %v", devices)
 
+	csvFileList := getCSVFileList(cfg, image)
 	cdiModeIdentifiers := cdiModeIdentfiersFromDevices(devices...)
 
 	logger.Debugf("Per-mode identifiers: %v", cdiModeIdentifiers)
@@ -179,6 +180,7 @@ func newAutomaticCDISpecModifier(logger logger.Interface, cfg *config.Config, de
 			nvcdi.WithClass(cdiModeIdentifiers.deviceClassByMode[mode]),
 			nvcdi.WithMode(mode),
 			nvcdi.WithFeatureFlags(cfg.NVIDIAContainerRuntimeConfig.Modes.JitCDI.NVCDIFeatureFlags...),
+			nvcdi.WithCSVFiles(csvFileList),
 		)
 		if err != nil {
 			return nil, fmt.Errorf("failed to construct CDI library for mode %q: %w", mode, err)

internal/modifier/csv.go

@@ -44,14 +44,7 @@ func NewCSVModifier(logger logger.Interface, cfg *config.Config, container image
 		return nil, fmt.Errorf("requirements not met: %v", err)
 	}
 
-	csvFiles, err := csv.GetFileList(cfg.NVIDIAContainerRuntimeConfig.Modes.CSV.MountSpecPath)
-	if err != nil {
-		return nil, fmt.Errorf("failed to get list of CSV files: %v", err)
-	}
-
-	if container.Getenv(image.EnvVarNvidiaRequireJetpack) != "csv-mounts=all" {
-		csvFiles = csv.BaseFilesOnly(csvFiles)
-	}
+	csvFiles := getCSVFileList(cfg, container)
 
 	cdilib, err := nvcdi.New(
 		nvcdi.WithLogger(logger),
@@ -106,3 +99,14 @@ func checkRequirements(logger logger.Interface, image image.CUDA) error {
 
 	return r.Assert()
 }
+
+func getCSVFileList(cfg *config.Config, container image.CUDA) []string {
+	csvFiles, err := csv.GetFileList(cfg.NVIDIAContainerRuntimeConfig.Modes.CSV.MountSpecPath)
+	if err != nil {
+		return nil
+	}
+	if container.Getenv(image.EnvVarNvidiaRequireJetpack) != "csv-mounts=all" {
+		csvFiles = csv.BaseFilesOnly(csvFiles)
+	}
+	return csvFiles
+}

internal/platform-support/tegra/csv/csv.go

@@ -34,7 +34,7 @@ const (
 )
 
 // DefaultFileList returns the list of CSV files that are used by default.
-func DefaultFileList() []string {
+func DefaultFileList(inFolder string) []string {
 	files := []string{
 		"devices.csv",
 		"drivers.csv",

internal/runtime/runtime_factory.go

@@ -141,6 +141,8 @@ func initRuntimeModeAndImage(logger logger.Interface, cfg *config.Config, ociSpe
 	modeResolver := info.NewRuntimeModeResolver(
 		info.WithLogger(logger),
 		info.WithImage(&image),
+		// TODO: Add a feature flag.
+		info.WithForceCSVModeForTegraSystems(false),
 	)
 	mode := modeResolver.ResolveRuntimeMode(cfg.NVIDIAContainerRuntimeConfig.Mode)
 	// We update the mode here so that we can continue passing just the config to other functions.