Merge branch 'NVIDIA:main' into main

This commit is contained in:
Seungmin Kim 2024-06-04 16:29:35 +09:00 committed by GitHub
commit 155fe66575
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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)