From 226c54613e88f32fd8820d48c7dabb7d27347fc5 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Wed, 22 Mar 2023 14:04:12 +0200
Subject: [PATCH] Also return an error from nvcdi.New

This change allows nvcdi.New to return an error in addition to the
constructed library instead of panicing.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 CHANGELOG.md                            | 3 ++-
 cmd/nvidia-ctk/cdi/generate/generate.go | 5 ++++-
 pkg/nvcdi/lib.go                        | 7 +++----
 tools/container/toolkit/toolkit.go      | 5 ++++-
 4 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 095cd76f..4f38cc7b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,11 +2,12 @@
 
 ## v1.13.0-rc.3
 
+* Also return an error from the nvcdi.New constructor instead of panicing.
 * Generate a simplified CDI specification by default. This means that entities in the common edits in a spec are not included in device definitions.
 * Add transformers to deduplicate and simplify CDI specifications.
 * Fix the generation of CDI specifications for management containers when the driver libraries are not in the LDCache.
 * Prefer /run over /var/run when locating nvidia-persistenced and nvidia-fabricmanager sockets.
-* Only initialize NVML for modes that require it when runing `nvidia-ctk cdi generate`
+* Only initialize NVML for modes that require it when runing `nvidia-ctk cdi generate`.
 * [libnvidia-container] Fix segmentation fault when RPC initialization fails.
 * [libnvidia-container] Build centos variants of the NVIDIA Container Library with static libtirpc v1.3.2.
 * [libnvidia-container] Remove make targets for fedora35 as the centos8 packages are compatible.
diff --git a/cmd/nvidia-ctk/cdi/generate/generate.go b/cmd/nvidia-ctk/cdi/generate/generate.go
index 12caca6a..b67bcb72 100644
--- a/cmd/nvidia-ctk/cdi/generate/generate.go
+++ b/cmd/nvidia-ctk/cdi/generate/generate.go
@@ -188,13 +188,16 @@ func (m command) generateSpec(cfg *config) (spec.Interface, error) {
 		return nil, fmt.Errorf("failed to create device namer: %v", err)
 	}
 
-	cdilib := nvcdi.New(
+	cdilib, err := nvcdi.New(
 		nvcdi.WithLogger(m.logger),
 		nvcdi.WithDriverRoot(cfg.driverRoot),
 		nvcdi.WithNVIDIACTKPath(cfg.nvidiaCTKPath),
 		nvcdi.WithDeviceNamer(deviceNamer),
 		nvcdi.WithMode(string(cfg.mode)),
 	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create CDI library: %v", err)
+	}
 
 	deviceSpecs, err := cdilib.GetAllDeviceSpecs()
 	if err != nil {
diff --git a/pkg/nvcdi/lib.go b/pkg/nvcdi/lib.go
index 646b06bd..d692b9ff 100644
--- a/pkg/nvcdi/lib.go
+++ b/pkg/nvcdi/lib.go
@@ -49,7 +49,7 @@ type nvcdilib struct {
 }
 
 // New creates a new nvcdi library
-func New(opts ...Option) Interface {
+func New(opts ...Option) (Interface, error) {
 	l := &nvcdilib{}
 	for _, opt := range opts {
 		opt(l)
@@ -102,8 +102,7 @@ func New(opts ...Option) Interface {
 		}
 		lib = (*mofedlib)(l)
 	default:
-		// TODO: We would like to return an error here instead of panicking
-		panic("Unknown mode")
+		return nil, fmt.Errorf("unknown mode %q", l.mode)
 	}
 
 	w := wrapper{
@@ -111,7 +110,7 @@ func New(opts ...Option) Interface {
 		vendor:    l.vendor,
 		class:     l.class,
 	}
-	return &w
+	return &w, nil
 }
 
 // GetSpec combines the device specs and common edits from the wrapped Interface to a single spec.Interface.
diff --git a/tools/container/toolkit/toolkit.go b/tools/container/toolkit/toolkit.go
index ef66f88b..f07d48b9 100644
--- a/tools/container/toolkit/toolkit.go
+++ b/tools/container/toolkit/toolkit.go
@@ -608,13 +608,16 @@ func generateCDISpec(opts *options, nvidiaCTKPath string) error {
 		return nil
 	}
 
-	cdilib := nvcdi.New(
+	cdilib, err := nvcdi.New(
 		nvcdi.WithMode(nvcdi.ModeManagement),
 		nvcdi.WithDriverRoot(opts.DriverRootCtrPath),
 		nvcdi.WithNVIDIACTKPath(nvidiaCTKPath),
 		nvcdi.WithVendor(opts.cdiVendor),
 		nvcdi.WithClass(opts.cdiClass),
 	)
+	if err != nil {
+		return fmt.Errorf("failed to create CDI library for management containers: %v", err)
+	}
 
 	spec, err := cdilib.GetSpec()
 	if err != nil {