Show host GPU driver on fleet instances

runner/docs/shim.openapi.yaml

@@ -479,6 +479,21 @@ components:
           type: string
           examples:
             - 0.18.34
+        gpu_vendor:
+          description: >
+            (since [0.20.28](https://github.com/dstackai/dstack/releases/tag/0.20.28))
+            Host GPU vendor. Omitted on hosts without GPUs.
+          type: string
+          examples:
+            - nvidia
+        gpu_driver_version:
+          description: >
+            (since [0.20.28](https://github.com/dstackai/dstack/releases/tag/0.20.28))
+            Host GPU driver version. Omitted on hosts without GPUs
+            or if detection failed.
+          type: string
+          examples:
+            - 570.86.15
       required:
         - service
         - version

runner/internal/shim/api/api_test.go

@@ -5,10 +5,12 @@ import (
 	"sync"
 
 	"github.com/dstackai/dstack/runner/internal/shim"
+	"github.com/dstackai/dstack/runner/internal/shim/host"
 )
 
 type DummyRunner struct {
 	tasks map[string]bool
+	gpus  []host.GpuInfo
 	mu    sync.Mutex
 }
 
@@ -46,6 +48,10 @@ func (ds *DummyRunner) Resources(context.Context) shim.Resources {
 	return shim.Resources{}
 }
 
+func (ds *DummyRunner) Gpus(context.Context) []host.GpuInfo {
+	return ds.gpus
+}
+
 func NewDummyRunner() *DummyRunner {
 	return &DummyRunner{
 		tasks: map[string]bool{},

runner/internal/shim/api/handlers.go

@@ -16,10 +16,15 @@ func (s *ShimServer) HealthcheckHandler(w http.ResponseWriter, r *http.Request)
 	s.mu.RLock()
 	defer s.mu.RUnlock()
 
-	return &HealthcheckResponse{
+	response := &HealthcheckResponse{
 		Service: "dstack-shim",
 		Version: s.version,
-	}, nil
+	}
+	if gpus := s.runner.Gpus(r.Context()); len(gpus) > 0 {
+		response.GpuVendor = string(gpus[0].Vendor)
+		response.GpuDriverVersion = gpus[0].DriverVersion
+	}
+	return response, nil
 }
 
 func (s *ShimServer) ShutdownHandler(w http.ResponseWriter, r *http.Request) (interface{}, error) {

runner/internal/shim/api/handlers_test.go

@@ -7,6 +7,8 @@ import (
 	"testing"
 
 	commonapi "github.com/dstackai/dstack/runner/internal/common/api"
+	"github.com/dstackai/dstack/runner/internal/common/gpu"
+	"github.com/dstackai/dstack/runner/internal/shim/host"
 )
 
 func TestHealthcheck(t *testing.T) {
@@ -29,6 +31,30 @@ func TestHealthcheck(t *testing.T) {
 	}
 }
 
+func TestHealthcheckWithGpus(t *testing.T) {
+	request := httptest.NewRequest("GET", "/api/healthcheck", nil)
+	responseRecorder := httptest.NewRecorder()
+
+	runner := NewDummyRunner()
+	runner.gpus = []host.GpuInfo{
+		{Vendor: gpu.GpuVendorNvidia, Name: "T4", Vram: 16384, DriverVersion: "570.86.15"},
+	}
+	server := NewShimServer(context.Background(), ":12346", "0.0.1.dev2", runner, nil, nil, nil, nil)
+
+	f := commonapi.JSONResponseHandler(server.HealthcheckHandler)
+	f(responseRecorder, request)
+
+	if responseRecorder.Code != 200 {
+		t.Errorf("Want status '%d', got '%d'", 200, responseRecorder.Code)
+	}
+
+	expected := `{"service":"dstack-shim","version":"0.0.1.dev2","gpu_vendor":"nvidia","gpu_driver_version":"570.86.15"}`
+
+	if strings.TrimSpace(responseRecorder.Body.String()) != expected {
+		t.Errorf("Want '%s', got '%s'", expected, responseRecorder.Body.String())
+	}
+}
+
 func TestTaskSubmit(t *testing.T) {
 	server := NewShimServer(context.Background(), ":12340", "0.0.1.dev2", NewDummyRunner(), nil, nil, nil, nil)
 	requestBody := `{

runner/internal/shim/api/schemas.go

@@ -9,6 +9,10 @@ import (
 type HealthcheckResponse struct {
 	Service string `json:"service"`
 	Version string `json:"version"`
+	// Optional host GPU driver info; empty on hosts without GPUs or if
+	// detection failed. Old servers ignore these fields.
+	GpuVendor        string `json:"gpu_vendor,omitempty"`
+	GpuDriverVersion string `json:"gpu_driver_version,omitempty"`
 }
 
 type ShutdownRequest struct {

runner/internal/shim/api/server.go

@@ -13,6 +13,7 @@ import (
 	"github.com/dstackai/dstack/runner/internal/shim"
 	"github.com/dstackai/dstack/runner/internal/shim/components"
 	"github.com/dstackai/dstack/runner/internal/shim/dcgm"
+	"github.com/dstackai/dstack/runner/internal/shim/host"
 )
 
 type TaskRunner interface {
@@ -22,6 +23,7 @@ type TaskRunner interface {
 	Remove(ctx context.Context, taskID string) error
 
 	Resources(context.Context) shim.Resources
+	Gpus(context.Context) []host.GpuInfo
 	TaskList() []*shim.TaskListItem
 	TaskInfo(taskID string) shim.TaskInfo
 }

runner/internal/shim/docker.go

@@ -324,6 +324,12 @@ func (d *DockerRunner) Resources(ctx context.Context) Resources {
 	}
 }
 
+// Gpus returns the GPUs detected at startup without collecting other host
+// resources, making it suitable for frequently called paths.
+func (d *DockerRunner) Gpus(ctx context.Context) []host.GpuInfo {
+	return d.gpus
+}
+
 func (d *DockerRunner) TaskList() []*TaskListItem {
 	tasks := d.tasks.List()
 	result := make([]*TaskListItem, 0, len(tasks))

runner/internal/shim/host/gpu.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
@@ -41,6 +42,10 @@ type GpuInfo struct {
 	// AMD: empty string
 	// Intel: accelerator index: ("0", "1", ...), as reported by `hl-smi -Q index`
 	Index string
+	// Version of the installed host driver, e.g., "570.86.15" (NVIDIA),
+	// "6.10.5" (AMD amdgpu), "2.0.0" (Tenstorrent TT-KMD).
+	// Empty string if detection failed. All GPUs on a host share the same driver.
+	DriverVersion string
 }
 
 func GetGpuInfo(ctx context.Context) []GpuInfo {
@@ -59,12 +64,23 @@ func GetGpuInfo(ctx context.Context) []GpuInfo {
 	return []GpuInfo{}
 }
 
+// normalizeDriverVersion filters out placeholder values SMI tools emit when a
+// query field is not available, e.g., "N/A" or "[Not Supported]".
+func normalizeDriverVersion(value string) string {
+	value = strings.TrimSpace(value)
+	switch strings.ToUpper(value) {
+	case "N/A", "[N/A]", "UNKNOWN", "[UNKNOWN]", "[NOT SUPPORTED]", "[NOT AVAILABLE]":
+		return ""
+	}
+	return value
+}
+
 func getNvidiaGpuInfo(ctx context.Context) []GpuInfo {
 	gpus := []GpuInfo{}
 
 	cmd := execute.ExecTask{
 		Command:     "nvidia-smi",
-		Args:        []string{"--query-gpu=name,memory.total,uuid", "--format=csv,noheader,nounits"},
+		Args:        []string{"--query-gpu=name,memory.total,uuid,driver_version", "--format=csv,noheader,nounits"},
 		StreamStdio: false,
 	}
 	res, err := cmd.Execute(ctx)
@@ -90,8 +106,8 @@ func getNvidiaGpuInfo(ctx context.Context) []GpuInfo {
 			log.Error(ctx, "cannot read csv", "err", err)
 			return gpus
 		}
-		if len(record) != 3 {
-			log.Error(ctx, "3 csv fields expected", "len", len(record))
+		if len(record) != 4 {
+			log.Error(ctx, "4 csv fields expected", "len", len(record))
 			return gpus
 		}
 		vram, err := strconv.Atoi(strings.TrimSpace(record[1]))
@@ -100,19 +116,43 @@ func getNvidiaGpuInfo(ctx context.Context) []GpuInfo {
 			vram = 0
 		}
 		gpus = append(gpus, GpuInfo{
-			Vendor: gpu.GpuVendorNvidia,
-			Name:   strings.TrimSpace(record[0]),
-			Vram:   vram,
-			ID:     strings.TrimSpace(record[2]),
+			Vendor:        gpu.GpuVendorNvidia,
+			Name:          strings.TrimSpace(record[0]),
+			Vram:          vram,
+			ID:            strings.TrimSpace(record[2]),
+			DriverVersion: normalizeDriverVersion(record[3]),
 		})
 	}
 	return gpus
 }
 
 type amdGpu struct {
-	Asic amdAsic `json:"asic"`
-	Vram amdVram `json:"vram"`
-	Bus  amdBus  `json:"bus"`
+	Asic   amdAsic   `json:"asic"`
+	Vram   amdVram   `json:"vram"`
+	Bus    amdBus    `json:"bus"`
+	Driver amdDriver `json:"driver"`
+}
+
+// amdDriver is the `driver` section of `amd-smi static --driver`.
+// Key names and value shapes differ between amd-smi versions, so it is parsed
+// defensively: an unexpected format leaves Version empty instead of failing
+// the whole GPU detection. Key matching is case-insensitive (encoding/json).
+type amdDriver struct {
+	Version string
+}
+
+func (d *amdDriver) UnmarshalJSON(data []byte) error {
+	var section struct {
+		Version       string `json:"version"`
+		DriverVersion string `json:"driver_version"`
+	}
+	// The error is ignored deliberately: an unexpected shape leaves Version empty.
+	_ = json.Unmarshal(data, &section)
+	d.Version = section.Version
+	if d.Version == "" {
+		d.Version = section.DriverVersion
+	}
+	return nil
 }
 
 // amd-smi >= 7.x wraps the array in {"gpu_data": [...]}
@@ -151,38 +191,53 @@ func parseAmdSmiOutput(data []byte) ([]amdGpu, error) {
 	return wrapped.GpuData, nil
 }
 
-func getAmdGpuInfo(ctx context.Context) []GpuInfo {
-	gpus := []GpuInfo{}
-
+func execAmdSmiStatic(ctx context.Context, withDriver bool) (string, error) {
 	ctx, cancel := context.WithTimeout(ctx, 2*time.Minute)
 	defer cancel()
 
+	args := []string{
+		"run",
+		"--rm",
+		"--device", "/dev/kfd",
+		"--device", "/dev/dri",
+		amdSmiImage,
+		"static", "--json", "--asic", "--vram", "--bus",
+	}
+	if withDriver {
+		args = append(args, "--driver")
+	}
 	cmd := execute.ExecTask{
-		Command: "docker",
-		Args: []string{
-			"run",
-			"--rm",
-			"--device", "/dev/kfd",
-			"--device", "/dev/dri",
-			amdSmiImage,
-			"static", "--json", "--asic", "--vram", "--bus",
-		},
+		Command:     "docker",
+		Args:        args,
 		StreamStdio: false,
 	}
 	res, err := cmd.Execute(ctx)
 	if err != nil {
-		log.Error(ctx, "failed to execute amd-smi", "err", err)
-		return gpus
+		return "", err
 	}
 	if res.ExitCode != 0 {
-		log.Error(
-			ctx, "failed to execute amd-smi",
-			"exitcode", res.ExitCode, "stdout", res.Stdout, "stderr", res.Stderr,
+		return "", fmt.Errorf(
+			"exitcode: %d, stdout: %s, stderr: %s", res.ExitCode, res.Stdout, res.Stderr,
 		)
+	}
+	return res.Stdout, nil
+}
+
+func getAmdGpuInfo(ctx context.Context) []GpuInfo {
+	gpus := []GpuInfo{}
+
+	stdout, err := execAmdSmiStatic(ctx, true)
+	if err != nil {
+		// Fall back for amd-smi versions without the --driver option.
+		log.Error(ctx, "failed to execute amd-smi with --driver, retrying without", "err", err)
+		stdout, err = execAmdSmiStatic(ctx, false)
+	}
+	if err != nil {
+		log.Error(ctx, "failed to execute amd-smi", "err", err)
 		return gpus
 	}
 
-	amdGpus, err := parseAmdSmiOutput([]byte(res.Stdout))
+	amdGpus, err := parseAmdSmiOutput([]byte(stdout))
 	if err != nil {
 		log.Error(ctx, "cannot read json", "err", err)
 		return gpus
@@ -198,6 +253,7 @@ func getAmdGpuInfo(ctx context.Context) []GpuInfo {
 			Name:           amdGpu.Asic.Name,
 			Vram:           amdGpu.Vram.Size.Value,
 			RenderNodePath: renderNodePath,
+			DriverVersion:  amdGpu.Driver.Version,
 		})
 	}
 	return gpus
@@ -413,6 +469,20 @@ func getGpusFromTtSmiSnapshot(snapshot *ttSmiSnapshot) []GpuInfo {
 	return gpus
 }
 
+// tenstorrentDriverVersionPath is the TT-KMD version file; it is what tt-smi
+// itself reads to report the driver version. It is a variable so tests can
+// override it.
+var tenstorrentDriverVersionPath = "/sys/module/tenstorrent/version"
+
+func getTenstorrentDriverVersion(ctx context.Context) string {
+	data, err := os.ReadFile(tenstorrentDriverVersionPath)
+	if err != nil {
+		log.Error(ctx, "failed to read tenstorrent driver version", "err", err)
+		return ""
+	}
+	return strings.TrimSpace(string(data))
+}
+
 func getTenstorrentGpuInfo(ctx context.Context) []GpuInfo {
 	gpus := []GpuInfo{}
 
@@ -447,7 +517,13 @@ func getTenstorrentGpuInfo(ctx context.Context) []GpuInfo {
 		return gpus
 	}
 
-	return getGpusFromTtSmiSnapshot(ttSmiSnapshot)
+	gpus = getGpusFromTtSmiSnapshot(ttSmiSnapshot)
+	if driverVersion := getTenstorrentDriverVersion(ctx); driverVersion != "" {
+		for i := range gpus {
+			gpus[i].DriverVersion = driverVersion
+		}
+	}
+	return gpus
 }
 
 func getAmdRenderNodePath(bdf string) (string, error) {