@@ -19,18 +19,30 @@ const (
1919
2020// ProbeGPUSupport determines whether or not the Docker engine has GPU support.
2121func ProbeGPUSupport (ctx context.Context , dockerClient client.SystemAPIClient ) (GPUSupport , error ) {
22- info , err := dockerClient .Info (ctx )
23- if err != nil {
24- return GPUSupportNone , err
25- }
26- if _ , hasNvidia := info .Runtimes ["nvidia" ]; hasNvidia {
22+ // First search for nvidia-container-runtime on PATH
23+ if _ , err := exec .LookPath ("nvidia-container-runtime" ); err == nil {
2724 return GPUSupportCUDA , nil
2825 }
2926
30- // If nvidia runtime is not listed, try searching for nvidia-container-runtime on PATH
31- if _ , err := exec .LookPath ("nvidia-container-runtime" ); err == nil {
27+ // Next look for explicitly configured nvidia runtime. This is not required in Docker 19.03+ but
28+ // may be configured on some systems
29+ hasNvidia , err := HasNVIDIARuntime (ctx , dockerClient )
30+ if err != nil {
31+ return GPUSupportNone , err
32+ }
33+ if hasNvidia {
3234 return GPUSupportCUDA , nil
3335 }
3436
3537 return GPUSupportNone , nil
3638}
39+
40+ // HasNVIDIARuntime determines whether there is an nvidia runtime available
41+ func HasNVIDIARuntime (ctx context.Context , dockerClient client.SystemAPIClient ) (bool , error ) {
42+ info , err := dockerClient .Info (ctx )
43+ if err != nil {
44+ return false , err
45+ }
46+ _ , hasNvidia := info .Runtimes ["nvidia" ]
47+ return hasNvidia , nil
48+ }
0 commit comments