From 1ce1243e2acb0cc594494bb99e984cb6652d9ced Mon Sep 17 00:00:00 2001
From: Dmitry Meyer <me@undef.im>
Date: Mon, 23 Sep 2024 08:38:32 +0000
Subject: [PATCH] [shim] Request all Nvidia driver capabilities (#1714)

Closes: https://github.com/dstackai/dstack/issues/1713
---
 runner/internal/shim/docker.go | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go
index 730159f63..6f1844d12 100644
--- a/runner/internal/shim/docker.go
+++ b/runner/internal/shim/docker.go
@@ -595,7 +595,13 @@ func requestGpuIfAvailable(ctx context.Context, client docker.APIClient) ([]cont
 	for runtime := range info.Runtimes {
 		if runtime == consts.NVIDIA_RUNTIME {
 			return []container.DeviceRequest{
-				{Capabilities: [][]string{{"gpu"}}, Count: -1}, // --gpus=all
+				{
+					// Request all capabilities to maximize compatibility with all sorts of GPU workloads.
+					// Default capabilities: utility, compute.
+					// https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/1.16.0/docker-specialized.html
+					Capabilities: [][]string{{"gpu", "utility", "compute", "graphics", "video", "display", "compat32"}},
+					Count:        -1, // --gpus=all
+				},
 			}, nil
 		}
 	}