From 7ff23999e87db57a38165e0e234982b734e92b8c Mon Sep 17 00:00:00 2001
From: Christopher Desiniotis <cdesiniotis@nvidia.com>
Date: Mon, 20 Nov 2023 14:55:21 -0800
Subject: [PATCH] Add option to nvidia-ctk to enable CDI in docker

Signed-off-by: Christopher Desiniotis <cdesiniotis@nvidia.com>
Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 CHANGELOG.md                                  |  1 +
 cmd/nvidia-ctk/runtime/configure/configure.go | 26 ++++++++++++++-----
 pkg/config/engine/docker/docker.go            |  5 ++--
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6212856c..7d182aa8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,7 @@
 * Fix bug in determining default nvidia-container-runtime.user config value on SUSE-based systems.
 * Add `crun` to the list of configured low-level runtimes.
 * Add `--cdi.enabled` option to `nvidia-ctk runtime configure` command to enable CDI in containerd.
+* Added support for `nvidia-ctk runtime configure --enable-cdi` for the `docker` runtime. Note that this requires Docker >= 25.
 
 * [toolkit-container] Bump CUDA base image version to 12.3.1.
 * [libnvidia-container] Added detection of libnvdxgdmal.so.1 on WSL2.
diff --git a/cmd/nvidia-ctk/runtime/configure/configure.go b/cmd/nvidia-ctk/runtime/configure/configure.go
index 11f37387..145fef0e 100644
--- a/cmd/nvidia-ctk/runtime/configure/configure.go
+++ b/cmd/nvidia-ctk/runtime/configure/configure.go
@@ -20,13 +20,14 @@ import (
 	"fmt"
 	"path/filepath"
 
+	"github.com/urfave/cli/v2"
+
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/crio"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/docker"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/ocihook"
-	"github.com/urfave/cli/v2"
 )
 
 const (
@@ -185,7 +186,7 @@ func (m command) validateFlags(c *cli.Context, config *config) error {
 		}
 	}
 
-	if config.runtime != "containerd" {
+	if config.runtime != "containerd" && config.runtime != "docker" {
 		if config.cdi.enabled {
 			m.logger.Warningf("Ignoring cdi.enabled flag for %v", config.runtime)
 		}
@@ -244,10 +245,9 @@ func (m command) configureConfigFile(c *cli.Context, config *config) error {
 		return fmt.Errorf("unable to update config: %v", err)
 	}
 
-	if config.cdi.enabled {
-		if err := cfg.Set("enable_cdi", true); err != nil {
-			return fmt.Errorf("failed enable CDI in containerd: %w", err)
-		}
+	err = enableCDI(config, cfg)
+	if err != nil {
+		return fmt.Errorf("failed to enable CDI in %s: %w", config.runtime, err)
 	}
 
 	outputPath := config.getOuputConfigPath()
@@ -300,3 +300,17 @@ func (m *command) configureOCIHook(c *cli.Context, config *config) error {
 	}
 	return nil
 }
+
+// enableCDI enables the use of CDI in the corresponding container engine
+func enableCDI(config *config, cfg engine.Interface) error {
+	if !config.cdi.enabled {
+		return nil
+	}
+	switch config.runtime {
+	case "containerd":
+		return cfg.Set("enable_cdi", true)
+	case "docker":
+		return cfg.Set("experimental", true)
+	}
+	return fmt.Errorf("enabling CDI in %s is not supported", config.runtime)
+}
diff --git a/pkg/config/engine/docker/docker.go b/pkg/config/engine/docker/docker.go
index 965dc2e4..a6828528 100644
--- a/pkg/config/engine/docker/docker.go
+++ b/pkg/config/engine/docker/docker.go
@@ -114,9 +114,10 @@ func (c *Config) RemoveRuntime(name string) error {
 	return nil
 }
 
-// Set is not supported for docker configs.
+// Set sets the specified docker option
 func (c *Config) Set(key string, value interface{}) error {
-	return fmt.Errorf("Set is not supported for crio configs")
+	(*c)[key] = value
+	return nil
 }
 
 // Save writes the config to the specified path