Merge pull request #725 from elezar/fix-nvsandboxutils

Ensure that nvsandboxutils is available for version
NVIDIA · Oct 3, 2024 · 4604e3b · 4604e3b
2 parents 19482da + 7cd2aef
commit 4604e3b
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/pkg/nvcdi/lib.go b/pkg/nvcdi/lib.go
@@ -263,6 +263,10 @@ func (l *nvcdilib) getCudaVersionNvml() (string, error) {
 }
 
 func (l *nvcdilib) getCudaVersionNvsandboxutils() (string, error) {
+	if l.nvsandboxutilslib == nil {
+		return "", fmt.Errorf("libnvsandboxutils is not available")
+	}
+
 	// Sandboxutils initialization should happen before this function is called
 	version, ret := l.nvsandboxutilslib.GetDriverVersion()
 	if ret != nvsandboxutils.SUCCESS {

diff --git a/pkg/nvcdi/management.go b/pkg/nvcdi/management.go
@@ -28,6 +28,7 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/edits"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/cuda"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/nvsandboxutils"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
 )
 
@@ -61,6 +62,19 @@ func (m *managementlib) GetAllDeviceSpecs() ([]specs.Device, error) {
 
 // GetCommonEdits returns the common edits for use in managementlib containers.
 func (m *managementlib) GetCommonEdits() (*cdi.ContainerEdits, error) {
+	if m.nvsandboxutilslib != nil {
+		if r := m.nvsandboxutilslib.Init(m.driverRoot); r != nvsandboxutils.SUCCESS {
+			m.logger.Warningf("Failed to init nvsandboxutils: %v; ignoring", r)
+			m.nvsandboxutilslib = nil
+		}
+		defer func() {
+			if m.nvsandboxutilslib == nil {
+				return
+			}
+			_ = m.nvsandboxutilslib.Shutdown()
+		}()
+	}
+
 	version, err := m.getCudaVersion()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get CUDA version: %v", err)