From c5eda7af8e0d6efc4f4e9811a3b207ab5b5fa132 Mon Sep 17 00:00:00 2001 From: Evan Lezar Date: Thu, 4 Apr 2024 15:00:21 +0200 Subject: [PATCH] Ensure that libnvidia-ml.so.1 is found in driver root This change ensures that the driver root is used to locate libnvidia-ml.so.1 if required. Signed-off-by: Evan Lezar --- pkg/nvcdi/lib.go | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/pkg/nvcdi/lib.go b/pkg/nvcdi/lib.go index 249bd311..d2db3b6c 100644 --- a/pkg/nvcdi/lib.go +++ b/pkg/nvcdi/lib.go @@ -96,7 +96,16 @@ func New(opts ...Option) (Interface, error) { root.WithLibrarySearchPaths(l.librarySearchPaths...), ) if l.nvmllib == nil { - l.nvmllib = nvml.New() + var nvmlOpts []nvml.LibraryOption + candidates, err := l.driver.Libraries().Locate("libnvidia-ml.so.1") + if err != nil { + l.logger.Warningf("Ignoring error in locating libnvidia-ml.so.1: %v", err) + } else { + libNvidiaMlPath := candidates[0] + l.logger.Infof("Using %v", libNvidiaMlPath) + nvmlOpts = append(nvmlOpts, nvml.WithLibraryPath(libNvidiaMlPath)) + } + l.nvmllib = nvml.New(nvmlOpts...) } if l.devicelib == nil { l.devicelib = device.New(l.nvmllib)