From ad20b6f89ce5de9d04d202acf9fd161890600fb6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Efe=20G=C3=B6kdemir?= Date: Tue, 27 Oct 2026 22:34:34 +0300 Subject: [PATCH] nvcdi: handle empty device selections without NVML MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Efe Gökdemir --- pkg/nvcdi/lib-nvml.go | 16 ++++++++++ pkg/nvcdi/lib-nvml_test.go | 63 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+) diff --git a/pkg/nvcdi/lib-nvml.go b/pkg/nvcdi/lib-nvml.go index d15d26498..a032427a2 100644 --- a/pkg/nvcdi/lib-nvml.go +++ b/pkg/nvcdi/lib-nvml.go @@ -17,6 +17,7 @@ package nvcdi import ( + "errors" "fmt" "strconv" "strings" @@ -50,7 +51,22 @@ func (l *nvmllib) GetCommonEdits() (*cdi.ContainerEdits, error) { // * a UUID of a GPU or MIG device // * the special ID 'all' func (l *nvmllib) DeviceSpecGenerators(ids ...string) (DeviceSpecGenerator, error) { + var specialID string + for _, id := range ids { + if id == "all" || id == "none" { + specialID = id + break + } + } + + if specialID == "none" { + return emptyDeviceSpecGenerator("none"), nil + } + if err := l.init(); err != nil { + if specialID == "all" && errors.Is(err, nvml.ERROR_DRIVER_NOT_LOADED) { + return emptyDeviceSpecGenerator("all"), nil + } return nil, err } defer l.tryShutdown() diff --git a/pkg/nvcdi/lib-nvml_test.go b/pkg/nvcdi/lib-nvml_test.go index 1e545a64e..29ebc0a85 100644 --- a/pkg/nvcdi/lib-nvml_test.go +++ b/pkg/nvcdi/lib-nvml_test.go @@ -143,6 +143,69 @@ func TestNvmllibGetDeviceSpecGeneratorsForIDs(t *testing.T) { } } +func TestNvmllibDeviceSpecGeneratorsWithoutNVIDIADevices(t *testing.T) { + testCases := []struct { + name string + ids []string + initResult nvml.Return + expectedDeviceName string + expectedInitCalls int + expectedError error + }{ + { + name: "all devices when the driver is not loaded", + ids: []string{"all"}, + initResult: nvml.ERROR_DRIVER_NOT_LOADED, + expectedDeviceName: "all", + expectedInitCalls: 1, + }, + { + name: "no devices when the driver is not loaded", + ids: []string{"none"}, + initResult: nvml.ERROR_DRIVER_NOT_LOADED, + expectedDeviceName: "none", + }, + { + name: "preserve other NVML initialization errors", + ids: []string{"all"}, + initResult: nvml.ERROR_NO_PERMISSION, + expectedInitCalls: 1, + expectedError: nvml.ERROR_NO_PERMISSION, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + mockNvml := dgxa100.New() + mockNvml.InitFunc = func() nvml.Return { + return tc.initResult + } + + l := &nvmllib{ + platformlibs: platformlibs{ + nvmllib: mockNvml, + devicelib: device.New(mockNvml), + }, + } + + generators, err := l.DeviceSpecGenerators(tc.ids...) + if tc.expectedError != nil { + require.ErrorIs(t, err, tc.expectedError) + require.Nil(t, generators) + } else { + require.NoError(t, err) + deviceSpecs, err := generators.GetDeviceSpecs() + require.NoError(t, err) + require.Len(t, deviceSpecs, 1) + require.Equal(t, tc.expectedDeviceName, deviceSpecs[0].Name) + require.Equal(t, []string{"NVCT_EMPTY_DEVICE="}, deviceSpecs[0].ContainerEdits.Env) + } + + require.Len(t, mockNvml.InitCalls(), tc.expectedInitCalls) + }) + } +} + // TODO: These need to be implemented in go-nvlib func mockOverrides(server *mockserver.Server) { for i, d := range server.Devices {