From 70da6cfa50fb934d8da321db39254ffbb688c7ff Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Mon, 16 Sep 2024 11:27:57 +0200
Subject: [PATCH] Allow inclusion of persistenced socket in CDI specification

This change adds an include-persistenced-socket flag to the
nvidia-ctk cdi generate command that ensures that a generated
specification includes the nvidia-persistenced socket if present on
the host.

Note that for mangement mode, these sockets are always included
if detected.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 cmd/nvidia-ctk/cdi/generate/generate.go | 8 ++++++++
 pkg/nvcdi/lib.go                        | 3 +++
 2 files changed, 11 insertions(+)

diff --git a/cmd/nvidia-ctk/cdi/generate/generate.go b/cmd/nvidia-ctk/cdi/generate/generate.go
index 9f9e994b..9c2da04f 100644
--- a/cmd/nvidia-ctk/cdi/generate/generate.go
+++ b/cmd/nvidia-ctk/cdi/generate/generate.go
@@ -60,6 +60,8 @@ type options struct {
 		files          cli.StringSlice
 		ignorePatterns cli.StringSlice
 	}
+
+	includePersistencedSocket bool
 }
 
 // NewCommand constructs a generate-cdi command with the specified logger
@@ -169,6 +171,11 @@ func (m command) build() *cli.Command {
 			Usage:       "Specify a pattern the CSV mount specifications.",
 			Destination: &opts.csv.ignorePatterns,
 		},
+		&cli.BoolFlag{
+			Name:        "include-persistenced-socket",
+			Usage:       "Include the nvidia-persistenced socket in the generated CDI specification.",
+			Destination: &opts.includePersistencedSocket,
+		},
 	}
 
 	return &c
@@ -273,6 +280,7 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 		nvcdi.WithLibrarySearchPaths(opts.librarySearchPaths.Value()),
 		nvcdi.WithCSVFiles(opts.csv.files.Value()),
 		nvcdi.WithCSVIgnorePatterns(opts.csv.ignorePatterns.Value()),
+		nvcdi.WithOptInFeature("include-persistenced-socket", opts.includePersistencedSocket),
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create CDI library: %v", err)
diff --git a/pkg/nvcdi/lib.go b/pkg/nvcdi/lib.go
index e14b2971..e535763b 100644
--- a/pkg/nvcdi/lib.go
+++ b/pkg/nvcdi/lib.go
@@ -132,6 +132,9 @@ func New(opts ...Option) (Interface, error) {
 		if l.vendor == "" {
 			l.vendor = "management.nvidia.com"
 		}
+		// For management specifications we always allow the fabricmanager and
+		// persistenced sockets.
+		WithOptInFeature("include-persistenced-socket", true)(l)
 		lib = (*managementlib)(l)
 	case ModeNvml:
 		lib = (*nvmllib)(l)