From 08ef242afb47ee6ba4aa9fd54eea6dc1ebe965e6 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Mon, 13 Mar 2023 18:18:54 +0200
Subject: [PATCH] Add --cdi-enabled option to control generating CDI spec

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 test/container/containerd_test.sh  |  1 -
 test/container/docker_test.sh      |  1 -
 test/container/toolkit_test.sh     |  2 +-
 tools/container/toolkit/toolkit.go | 11 +++++++++++
 4 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/test/container/containerd_test.sh b/test/container/containerd_test.sh
index 48b00347..25e1fede 100755
--- a/test/container/containerd_test.sh
+++ b/test/container/containerd_test.sh
@@ -54,7 +54,6 @@ testing::containerd::toolkit::run() {
 		--pid "container:${containerd_dind_ctr}" \
 		-e RUNTIME="containerd" \
 		-e RUNTIME_ARGS="--config=${containerd_dind_containerd_dir}/containerd.toml --socket=${containerd_dind_containerd_dir}/containerd.sock" \
-		-e CDI_OUTPUT_DIR="" \
 		--name "${containerd_test_ctr}" \
 		"${toolkit_container_image}" "/usr/local/nvidia" "--no-daemon"
 
diff --git a/test/container/docker_test.sh b/test/container/docker_test.sh
index b314c8c0..decc198f 100755
--- a/test/container/docker_test.sh
+++ b/test/container/docker_test.sh
@@ -39,7 +39,6 @@ testing::docker::toolkit::run() {
 		--volumes-from "${docker_dind_ctr}" \
 		--pid "container:${docker_dind_ctr}" \
 		-e RUNTIME_ARGS="--socket ${docker_dind_socket}" \
-		-e CDI_OUTPUT_DIR="" \
 		--name "${docker_test_ctr}" \
 		"${toolkit_container_image}" "/usr/local/nvidia" "--no-daemon"
 
diff --git a/test/container/toolkit_test.sh b/test/container/toolkit_test.sh
index 6ae724ab..d11b0a58 100644
--- a/test/container/toolkit_test.sh
+++ b/test/container/toolkit_test.sh
@@ -23,7 +23,7 @@ testing::toolkit::install() {
 		READLINK="greadlink"
 	fi
 
-	testing::docker_run::toolkit::shell 'toolkit install --toolkit-root=/usr/local/nvidia/toolkit --cdi-output-dir=""'
+	testing::docker_run::toolkit::shell 'toolkit install --toolkit-root=/usr/local/nvidia/toolkit'
 	docker run --rm -v "${shared_dir}:/work" alpine sh -c "chown -R ${uid}:${gid} /work/"
 
 	# Ensure toolkit dir is correctly setup
diff --git a/tools/container/toolkit/toolkit.go b/tools/container/toolkit/toolkit.go
index 8f93db7d..ef66f88b 100644
--- a/tools/container/toolkit/toolkit.go
+++ b/tools/container/toolkit/toolkit.go
@@ -56,6 +56,7 @@ type options struct {
 	ContainerCLIDebug string
 	toolkitRoot       string
 
+	cdiEnabled   bool
 	cdiOutputDir string
 	cdiKind      string
 	cdiVendor    string
@@ -171,6 +172,13 @@ func main() {
 			Destination: &opts.toolkitRoot,
 			EnvVars:     []string{"TOOLKIT_ROOT"},
 		},
+		&cli.BoolFlag{
+			Name:        "cdi-enabled",
+			Aliases:     []string{"enable-cdi"},
+			Usage:       "enable the generation of a CDI specification",
+			Destination: &opts.cdiEnabled,
+			EnvVars:     []string{"CDI_ENABLED", "ENABLE_CDI"},
+		},
 		&cli.StringFlag{
 			Name:        "cdi-output-dir",
 			Usage:       "the directory where the CDI output files are to be written. If this is set to '', no CDI specification is generated.",
@@ -592,6 +600,9 @@ func createDirectories(dir ...string) error {
 
 // generateCDISpec generates a CDI spec for use in managemnt containers
 func generateCDISpec(opts *options, nvidiaCTKPath string) error {
+	if !opts.cdiEnabled {
+		return nil
+	}
 	if opts.cdiOutputDir == "" {
 		log.Info("Skipping CDI spec generation (no output directory specified)")
 		return nil