Improve CUDA detection

ekcasey · ekcasey · commit cb46972ddefb · 2025-09-25T15:41:21.000-06:00
In Docker 19.03+ NVIDIA Container Toolkit can integrates via an OCI prestart hook
instead of requiring  to be listed as a separate runtime.

Signed-off-by: Emily Casey &lt;emily.casey@docker.com&gt;
diff --git a/pkg/gpu/gpu.go b/pkg/gpu/gpu.go
@@ -2,6 +2,7 @@ package gpu
 
 import (
 	"context"
+	"os/exec"
 
 	"github.com/docker/docker/client"
 )
@@ -25,5 +26,11 @@ func ProbeGPUSupport(ctx context.Context, dockerClient client.SystemAPIClient) (
 	if _, hasNvidia := info.Runtimes["nvidia"]; hasNvidia {
 		return GPUSupportCUDA, nil
 	}
+
+	// If nvidia runtime is not listed, try searching for nvidia-container-runtime on PATH
+	if _, err := exec.LookPath("nvidia-container-runtime"); err == nil {
+		return GPUSupportCUDA, nil
+	}
+
 	return GPUSupportNone, nil
 }

Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,7 @@ package gpu`
`2`	`2`
`3`	`3`	`import (`
`4`	`4`	`"context"`
	`5`	`+ "os/exec"`
`5`	`6`
`6`	`7`	`"github.com/docker/docker/client"`
`7`	`8`	`)`
`@@ -25,5 +26,11 @@ func ProbeGPUSupport(ctx context.Context, dockerClient client.SystemAPIClient) (`
`25`	`26`	`if _, hasNvidia := info.Runtimes["nvidia"]; hasNvidia {`
`26`	`27`	`return GPUSupportCUDA, nil`
`27`	`28`	`}`
	`29`	`+`
	`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 {`
	`32`	`+ return GPUSupportCUDA, nil`
	`33`	`+ }`
	`34`	`+`
`28`	`35`	`return GPUSupportNone, nil`
`29`	`36`	`}`