From 01a7f7bb8e6c66c77a6338c69e1d8f1d26a872d4 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Thu, 22 Jun 2023 15:42:26 +0200
Subject: [PATCH] Explicitly generate CDI spec for CSV mode

This change explicitly generates a CDI specification from
the supplied CSV files when cdi mode is detected. This
ensures consistency between the behaviour on Tegra-based
systems.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 internal/modifier/csv.go | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/internal/modifier/csv.go b/internal/modifier/csv.go
index 0a834d0d..db89af1b 100644
--- a/internal/modifier/csv.go
+++ b/internal/modifier/csv.go
@@ -24,10 +24,11 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/cuda"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover/csv"
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover/tegra"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/modifier/cdi"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/oci"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/requirements"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi"
 )
 
 // csvMode represents the modifications as performed by the csv runtime mode
@@ -65,24 +66,33 @@ func NewCSVModifier(logger logger.Interface, cfg *config.Config, image image.CUD
 		csvFiles = csv.BaseFilesOnly(csvFiles)
 	}
 
-	d, err := tegra.New(
-		tegra.WithLogger(logger),
-		tegra.WithDriverRoot(cfg.NVIDIAContainerCLIConfig.Root),
-		tegra.WithNVIDIACTKPath(cfg.NVIDIACTKConfig.Path),
-		tegra.WithCSVFiles(csvFiles),
+	cdilib, err := nvcdi.New(
+		nvcdi.WithLogger(logger),
+		nvcdi.WithDriverRoot(cfg.NVIDIAContainerCLIConfig.Root),
+		nvcdi.WithNVIDIACTKPath(cfg.NVIDIACTKConfig.Path),
+		nvcdi.WithMode(nvcdi.ModeCSV),
+		nvcdi.WithCSVFiles(csvFiles),
 	)
 	if err != nil {
-		return nil, fmt.Errorf("failed to construct discoverer: %v", err)
+		return nil, fmt.Errorf("failed to construct CDI library: %v", err)
 	}
 
-	discoverModifier, err := NewModifierFromDiscoverer(logger, d)
+	spec, err := cdilib.GetSpec()
 	if err != nil {
-		return nil, fmt.Errorf("failed to construct modifier: %v", err)
+		return nil, fmt.Errorf("failed to get CDI spec: %v", err)
+	}
+
+	cdiModifier, err := cdi.New(
+		cdi.WithLogger(logger),
+		cdi.WithSpec(spec.Raw()),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to construct CDI modifier: %v", err)
 	}
 
 	modifiers := Merge(
 		nvidiaContainerRuntimeHookRemover{logger},
-		discoverModifier,
+		cdiModifier,
 	)
 
 	return modifiers, nil