Skip to content
This repository was archived by the owner on Oct 6, 2025. It is now read-only.

Improve CUDA detection #156

Merged
ericcurtin merged 2 commits into from
Sep 26, 2025
Merged

Improve CUDA detection #156

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cb46972
Improve CUDA detection
ekcasey Sep 25, 2025
6dfffc5
Only specify nvidia runtime if one exists
ekcasey Sep 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 21 additions & 2 deletions pkg/gpu/gpu.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package gpu

import (
"context"
"os/exec"

"github.com/docker/docker/client"
)
Expand All @@ -18,12 +19,30 @@ const (

// ProbeGPUSupport determines whether or not the Docker engine has GPU support.
func ProbeGPUSupport(ctx context.Context, dockerClient client.SystemAPIClient) (GPUSupport, error) {
info, err := dockerClient.Info(ctx)
// First search for nvidia-container-runtime on PATH
if _, err := exec.LookPath("nvidia-container-runtime"); err == nil {
return GPUSupportCUDA, nil
}

// Next look for explicitly configured nvidia runtime. This is not required in Docker 19.03+ but
// may be configured on some systems
hasNvidia, err := HasNVIDIARuntime(ctx, dockerClient)
if err != nil {
return GPUSupportNone, err
}
if _, hasNvidia := info.Runtimes["nvidia"]; hasNvidia {
if hasNvidia {
return GPUSupportCUDA, nil
}

return GPUSupportNone, nil
}

// HasNVIDIARuntime determines whether there is an nvidia runtime available
func HasNVIDIARuntime(ctx context.Context, dockerClient client.SystemAPIClient) (bool, error) {
info, err := dockerClient.Info(ctx)
if err != nil {
return false, err
}
_, hasNvidia := info.Runtimes["nvidia"]
return hasNvidia, nil
}
4 changes: 3 additions & 1 deletion pkg/standalone/containers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,7 +268,9 @@ func CreateControllerContainer(ctx context.Context, dockerClient *client.Client,
nat.Port(portStr + "/tcp"): portBindings,
}
if gpu == gpupkg.GPUSupportCUDA {
hostConfig.Runtime = "nvidia"
if ok, err := gpupkg.HasNVIDIARuntime(ctx, dockerClient); err == nil && ok {
hostConfig.Runtime = "nvidia"
}
Comment on lines +271 to +273
Copy link

@gemini-code-assist gemini-code-assist bot Sep 25, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The current implementation silently ignores any error returned from gpupkg.HasNVIDIARuntime. While the logic correctly avoids setting the nvidia runtime in case of an error, it would be better to log a warning to inform the user that the check failed. This could help in debugging issues where GPU support is expected but not working correctly.

		if ok, err := gpupkg.HasNVIDIARuntime(ctx, dockerClient); err != nil {
			printer.Printf("Warning: failed to check for nvidia runtime: %v\n", err)
		} else if ok {
			hostConfig.Runtime = "nvidia"
		}

hostConfig.DeviceRequests = []container.DeviceRequest{{Count: -1, Capabilities: [][]string{{"gpu"}}}}
}

Expand Down
Loading