From f44e31f7afecee1aa89e9f182dd82af205e9bde6 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 sockets in CDI specification

This change adds with-fabricmanager and with-persistenced
flags to the nvidia-ctk cdi generate command that ensures that
a generated specification includes these sockets 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 | 15 +++++++++++++++
 pkg/nvcdi/lib.go                        |  4 ++++
 2 files changed, 19 insertions(+)

diff --git a/cmd/nvidia-ctk/cdi/generate/generate.go b/cmd/nvidia-ctk/cdi/generate/generate.go
index 9f9e994b2..d6eb7680e 100644
--- a/cmd/nvidia-ctk/cdi/generate/generate.go
+++ b/cmd/nvidia-ctk/cdi/generate/generate.go
@@ -60,6 +60,9 @@ type options struct {
 		files          cli.StringSlice
 		ignorePatterns cli.StringSlice
 	}
+
+	withFabricmanager bool
+	withPersistenced  bool
 }
 
 // NewCommand constructs a generate-cdi command with the specified logger
@@ -169,6 +172,16 @@ func (m command) build() *cli.Command {
 			Usage:       "Specify a pattern the CSV mount specifications.",
 			Destination: &opts.csv.ignorePatterns,
 		},
+		&cli.BoolFlag{
+			Name:        "with-fabricmanager",
+			Usage:       "Include the nvidia-fabricmanager socket in the generated CDI specification.",
+			Destination: &opts.withFabricmanager,
+		},
+		&cli.BoolFlag{
+			Name:        "with-persistenced",
+			Usage:       "Include the nvidia-persistenced socket in the generated CDI specification.",
+			Destination: &opts.withPersistenced,
+		},
 	}
 
 	return &c
@@ -273,6 +286,8 @@ 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("allow-fabricmanager", opts.withFabricmanager),
+		nvcdi.WithOptInFeature("allow-persistenced", opts.withPersistenced),
 	)
 	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 e14b29716..4be46b79f 100644
--- a/pkg/nvcdi/lib.go
+++ b/pkg/nvcdi/lib.go
@@ -132,6 +132,10 @@ 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("allow-fabricmanager", true)(l)
+		WithOptInFeature("allow-persistenced", true)(l)
 		lib = (*managementlib)(l)
 	case ModeNvml:
 		lib = (*nvmllib)(l)