Skip to content

Commit a6ddb3e

Browse files
committed
refactor(gpu): simplify runtime detection with ordered runtime list
Signed-off-by: pnkcaht <samzoovsk19@gmail.com>
1 parent 7bcfe40 commit a6ddb3e

File tree

1 file changed

+19
-23
lines changed

1 file changed

+19
-23
lines changed

cmd/cli/pkg/gpu/gpu.go

Lines changed: 19 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -29,35 +29,31 @@ func ProbeGPUSupport(ctx context.Context, dockerClient client.SystemAPIClient) (
2929
// Docker Info is the source of truth for which runtimes are actually usable.
3030
info, err := dockerClient.Info(ctx)
3131
if err != nil {
32-
return GPUSupportNone, err
32+
// Preserve best-effort behavior: if Docker Info is unavailable (e.g. in
33+
// restricted or degraded environments), do not treat this as a hard failure.
34+
// Instead, assume no GPU support and allow callers to continue.
35+
return GPUSupportNone, nil
3336
}
3437

35-
// 1. CUDA (NVIDIA)
36-
// NVIDIA remains the highest priority due to its wide adoption and
37-
// first-class support in Docker (>= 19.03).
38-
if _, ok := info.Runtimes["nvidia"]; ok {
39-
return GPUSupportCUDA, nil
40-
}
41-
42-
// 2. ROCm (AMD)
43-
// Explicit ROCm runtime configured in the Docker Engine.
44-
if _, ok := info.Runtimes["rocm"]; ok {
45-
return GPUSupportROCm, nil
46-
}
47-
48-
// 3. MUSA (MThreads)
49-
// Used primarily on specific accelerator platforms.
50-
if _, ok := info.Runtimes["mthreads"]; ok {
51-
return GPUSupportMUSA, nil
38+
// Runtimes are checked in priority order, from highest to lowest.
39+
// The first matching runtime determines the selected GPU support.
40+
supportedRuntimes := []struct {
41+
name string
42+
support GPUSupport
43+
}{
44+
{"nvidia", GPUSupportCUDA}, // 1. CUDA (NVIDIA)
45+
{"rocm", GPUSupportROCm}, // 2. ROCm (AMD)
46+
{"mthreads", GPUSupportMUSA}, // 3. MUSA (MThreads)
47+
{"cann", GPUSupportCANN}, // 4. Ascend CANN (Huawei)
5248
}
5349

54-
// 4. Ascend CANN (Huawei)
55-
// Ascend NPU runtime registered in Docker.
56-
if _, ok := info.Runtimes["cann"]; ok {
57-
return GPUSupportCANN, nil
50+
for _, r := range supportedRuntimes {
51+
if _, ok := info.Runtimes[r.name]; ok {
52+
return r.support, nil
53+
}
5854
}
5955

60-
// 5. Legacy fallback
56+
// Legacy fallback
6157
// Older Docker setups may not register the NVIDIA runtime explicitly,
6258
// but still have the legacy nvidia-container-runtime available on PATH.
6359
if _, err := exec.LookPath("nvidia-container-runtime"); err == nil {

0 commit comments

Comments
 (0)