16a7d7da64
Checking the existence of /dev/nvidia-uvm had two limitations: - If kernel module nvidia_uvm is unloaded, /dev/nvidia-uvm is not removed, so nvidia-modprobe was not called when it should have. - Even if kernel module nvidia_uvm is loaded and /dev/nvidia-uvm exists, /dev/nvidia-uvm-tools might not exist. Using nvidia-modprobe fixes this issue. nvidia-modprobe parses /proc/modules to check which modules are already loaded, so there should be no impact on performance.
62 lines
1.2 KiB
Go
62 lines
1.2 KiB
Go
// Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved.
|
|
|
|
package nvidia
|
|
|
|
import (
|
|
"errors"
|
|
"os"
|
|
"os/exec"
|
|
|
|
"github.com/NVIDIA/nvidia-docker/src/cuda"
|
|
"github.com/NVIDIA/nvidia-docker/src/nvml"
|
|
)
|
|
|
|
const (
|
|
DockerPlugin = "nvidia-docker"
|
|
DeviceCtl = "/dev/nvidiactl"
|
|
DeviceUVM = "/dev/nvidia-uvm"
|
|
DeviceUVMTools = "/dev/nvidia-uvm-tools"
|
|
)
|
|
|
|
func Init() error {
|
|
if err := os.Setenv("CUDA_DISABLE_UNIFIED_MEMORY", "1"); err != nil {
|
|
return err
|
|
}
|
|
if err := os.Setenv("CUDA_CACHE_DISABLE", "1"); err != nil {
|
|
return err
|
|
}
|
|
if err := os.Unsetenv("CUDA_VISIBLE_DEVICES"); err != nil {
|
|
return err
|
|
}
|
|
return nvml.Init()
|
|
}
|
|
|
|
func Shutdown() error {
|
|
return nvml.Shutdown()
|
|
}
|
|
|
|
func LoadUVM() error {
|
|
if exec.Command("nvidia-modprobe", "-u", "-c=0").Run() != nil {
|
|
return errors.New("Could not load UVM kernel module. Is nvidia-modprobe installed?")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func GetDriverVersion() (string, error) {
|
|
return nvml.GetDriverVersion()
|
|
}
|
|
|
|
func GetCUDAVersion() (string, error) {
|
|
return cuda.GetDriverVersion()
|
|
}
|
|
|
|
func GetControlDevicePaths() ([]string, error) {
|
|
devs := []string{DeviceCtl, DeviceUVM}
|
|
|
|
_, err := os.Stat(DeviceUVMTools)
|
|
if os.IsNotExist(err) {
|
|
return devs, nil
|
|
}
|
|
return append(devs, DeviceUVMTools), err
|
|
}
|