[runner] Improve GPU metrics collector (#2612)

* Use only vendor-specific tool, don't try them all
* AMD: handle arbitrary field number and order
* AMD: handle 'N/A' values (return an error)

Closes: #2551

Author: un-def

runner/cmd/shim/main.go

@@ -12,11 +12,11 @@ import (
 	"time"
 
 	"github.com/dstackai/dstack/runner/consts"
+	"github.com/dstackai/dstack/runner/internal/common"
 	"github.com/dstackai/dstack/runner/internal/log"
 	"github.com/dstackai/dstack/runner/internal/shim"
 	"github.com/dstackai/dstack/runner/internal/shim/api"
 	"github.com/dstackai/dstack/runner/internal/shim/dcgm"
-	"github.com/dstackai/dstack/runner/internal/shim/host"
 	"github.com/sirupsen/logrus"
 	"github.com/urfave/cli/v2"
 )
@@ -197,7 +197,7 @@ func start(ctx context.Context, args shim.CLIArgs, serviceMode bool) (err error)
 
 	var dcgmExporter *dcgm.DCGMExporter
 
-	if host.GetGpuVendor() == host.GpuVendorNvidia {
+	if common.GetGpuVendor() == common.GpuVendorNvidia {
 		dcgmExporterPath, err := dcgm.GetDCGMExporterExecPath(ctx)
 		if err == nil {
 			interval := time.Duration(args.DCGMExporter.Interval * int(time.Millisecond))

runner/internal/common/gpu.go

@@ -0,0 +1,32 @@
+package common
+
+import (
+	"errors"
+	"os"
+)
+
+type GpuVendor string
+
+const (
+	GpuVendorNone        GpuVendor = "none"
+	GpuVendorNvidia      GpuVendor = "nvidia"
+	GpuVendorAmd         GpuVendor = "amd"
+	GpuVendorIntel       GpuVendor = "intel"
+	GpuVendorTenstorrent GpuVendor = "tenstorrent"
+)
+
+func GetGpuVendor() GpuVendor {
+	if _, err := os.Stat("/dev/kfd"); !errors.Is(err, os.ErrNotExist) {
+		return GpuVendorAmd
+	}
+	if _, err := os.Stat("/dev/nvidiactl"); !errors.Is(err, os.ErrNotExist) {
+		return GpuVendorNvidia
+	}
+	if _, err := os.Stat("/dev/accel"); !errors.Is(err, os.ErrNotExist) {
+		return GpuVendorIntel
+	}
+	if _, err := os.Stat("/dev/tenstorrent"); !errors.Is(err, os.ErrNotExist) {
+		return GpuVendorTenstorrent
+	}
+	return GpuVendorNone
+}

runner/internal/metrics/metrics.go

@@ -3,28 +3,33 @@ package metrics
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"os/exec"
 	"strconv"
 	"strings"
 	"time"
 
+	"github.com/dstackai/dstack/runner/internal/common"
 	"github.com/dstackai/dstack/runner/internal/log"
 	"github.com/dstackai/dstack/runner/internal/schemas"
 )
 
 type MetricsCollector struct {
 	cgroupVersion int
+	gpuVendor     common.GpuVendor
 }
 
 func NewMetricsCollector() (*MetricsCollector, error) {
 	cgroupVersion, err := getCgroupVersion()
 	if err != nil {
 		return nil, err
 	}
+	gpuVendor := common.GetGpuVendor()
 	return &MetricsCollector{
 		cgroupVersion: cgroupVersion,
+		gpuVendor:     gpuVendor,
 	}, nil
 }
 
@@ -144,19 +149,24 @@ func (s *MetricsCollector) GetMemoryCacheBytes() (uint64, error) {
 }
 
 func (s *MetricsCollector) GetGPUMetrics() ([]schemas.GPUMetrics, error) {
-	metrics, err := s.GetNVIDIAGPUMetrics()
-	if err == nil {
-		return metrics, nil
+	var metrics []schemas.GPUMetrics
+	var err error
+	switch s.gpuVendor {
+	case common.GpuVendorNvidia:
+		metrics, err = s.GetNVIDIAGPUMetrics()
+	case common.GpuVendorAmd:
+		metrics, err = s.GetAMDGPUMetrics()
+	case common.GpuVendorIntel:
+		metrics, err = s.GetIntelAcceleratorMetrics()
+	case common.GpuVendorTenstorrent:
+		err = errors.New("tenstorrent metrics not suppored")
+	case common.GpuVendorNone:
+		// pass
 	}
-	metrics, err = s.GetAMDGPUMetrics()
-	if err == nil {
-		return metrics, nil
+	if metrics == nil {
+		metrics = []schemas.GPUMetrics{}
 	}
-	metrics, err = s.GetIntelAcceleratorMetrics()
-	if err == nil {
-		return metrics, nil
-	}
-	return []schemas.GPUMetrics{}, err
+	return metrics, err
 }
 
 func (s *MetricsCollector) GetNVIDIAGPUMetrics() ([]schemas.GPUMetrics, error) {
@@ -170,32 +180,65 @@ func (s *MetricsCollector) GetNVIDIAGPUMetrics() ([]schemas.GPUMetrics, error) {
 }
 
 func (s *MetricsCollector) GetAMDGPUMetrics() ([]schemas.GPUMetrics, error) {
-	metrics := []schemas.GPUMetrics{}
-
 	cmd := exec.Command("amd-smi", "monitor", "-vu", "--csv")
 	var out bytes.Buffer
 	cmd.Stdout = &out
 	if err := cmd.Run(); err != nil {
 		return nil, fmt.Errorf("failed to execute amd-smi: %w", err)
 	}
+	return s.getAMDGPUMetrics(out.String())
+}
+
+func (s *MetricsCollector) getAMDGPUMetrics(csv string) ([]schemas.GPUMetrics, error) {
+	lines := strings.Split(strings.TrimSpace(csv), "\n")
+	if len(lines) < 2 {
+		return nil, errors.New("too few lines in amd-smi output")
+	}
 
-	lines := strings.Split(strings.TrimSpace(out.String()), "\n")
+	gpuUtilIndex := -1
+	memUsedIndex := -1
+	for index, header := range strings.Split(lines[0], ",") {
+		switch header {
+		case "gfx":
+			gpuUtilIndex = index
+		case "vram_used":
+			memUsedIndex = index
+		}
+	}
+	if gpuUtilIndex == -1 {
+		return nil, errors.New("GPU utilization column not found")
+	}
+	if memUsedIndex == -1 {
+		return nil, errors.New("used VRAM column not found")
+	}
+
+	metrics := []schemas.GPUMetrics{}
 	for _, line := range lines[1:] {
 		fields := strings.Split(line, ",")
-		if len(fields) != 5 {
-			return nil, fmt.Errorf("unexpected number of columns in amd-smi output")
+		if len(fields) <= gpuUtilIndex || len(fields) <= memUsedIndex {
+			return nil, errors.New("too few columns in amd-smi output")
 		}
-		memUsed, err := strconv.ParseUint(strings.TrimSpace(fields[3]), 10, 64)
-		if err != nil {
-			return nil, fmt.Errorf("failed to parse VRAM used: %w", err)
+
+		gpuUtilRaw := strings.TrimSpace(fields[gpuUtilIndex])
+		if strings.ToUpper(gpuUtilRaw) == "N/A" {
+			return nil, errors.New("GPU utilization is N/A")
 		}
-		utilization, err := strconv.ParseUint(strings.TrimSpace(fields[1]), 10, 64)
+		gpuUtil, err := strconv.ParseUint(gpuUtilRaw, 10, 64)
 		if err != nil {
 			return nil, fmt.Errorf("failed to parse GPU utilization: %w", err)
 		}
+
+		memUsedRaw := strings.TrimSpace(fields[memUsedIndex])
+		if strings.ToUpper(memUsedRaw) == "N/A" {
+			return nil, errors.New("used VRAM is N/A")
+		}
+		memUsed, err := strconv.ParseUint(memUsedRaw, 10, 64)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse used VRAM: %w", err)
+		}
 		metrics = append(metrics, schemas.GPUMetrics{
 			GPUMemoryUsage: memUsed * 1024 * 1024,
-			GPUUtil:        utilization,
+			GPUUtil:        gpuUtil,
 		})
 	}
 

runner/internal/metrics/metrics_test.go

@@ -0,0 +1,47 @@
+package metrics
+
+import (
+	"testing"
+
+	"github.com/dstackai/dstack/runner/internal/schemas"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetAMDGPUMetrics_OK(t *testing.T) {
+	collector, err := NewMetricsCollector()
+	assert.NoError(t, err)
+
+	cases := []struct {
+		csv      string
+		expected []schemas.GPUMetrics
+	}{
+		// AMDSMI Tool: 24.7.1+0012a68 | AMDSMI Library version: 24.7.1.0 | ROCm version: 6.3.1
+		{
+			csv: "gpu,gfx,gfx_clock,vram_used,vram_total\n0,10,132,283,196300\n",
+			expected: []schemas.GPUMetrics{
+				{GPUUtil: 10, GPUMemoryUsage: 296747008},
+			},
+		},
+		// AMDSMI Tool: 25.3.0+ede62f2 | AMDSMI Library version: 25.3.0 | ROCm version: 6.4.0
+		{
+			csv: "gpu,gfx_clk,gfx,vram_used,vram_free,vram_total,vram_percent\n0,132,10,283,196309,196592,0.0\n",
+			expected: []schemas.GPUMetrics{
+				{GPUUtil: 10, GPUMemoryUsage: 296747008},
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		metrics, err := collector.getAMDGPUMetrics(tc.csv)
+		assert.NoError(t, err)
+		assert.Equal(t, tc.expected, metrics)
+	}
+}
+
+func TestGetAMDGPUMetrics_ErrorGPUUtilNA(t *testing.T) {
+	collector, err := NewMetricsCollector()
+	assert.NoError(t, err)
+	metrics, err := collector.getAMDGPUMetrics("gpu,gfx,gfx_clock,vram_used,vram_total\n0,N/A,N/A,283,196300\n")
+	assert.ErrorContains(t, err, "GPU utilization is N/A")
+	assert.Nil(t, metrics)
+}

runner/internal/shim/docker.go

@@ -30,6 +30,7 @@ import (
 	"github.com/docker/go-connections/nat"
 	"github.com/docker/go-units"
 	"github.com/dstackai/dstack/runner/consts"
+	"github.com/dstackai/dstack/runner/internal/common"
 	"github.com/dstackai/dstack/runner/internal/log"
 	"github.com/dstackai/dstack/runner/internal/shim/backends"
 	"github.com/dstackai/dstack/runner/internal/shim/host"
@@ -54,7 +55,7 @@ type DockerRunner struct {
 	dockerParams DockerParameters
 	dockerInfo   dockersystem.Info
 	gpus         []host.GpuInfo
-	gpuVendor    host.GpuVendor
+	gpuVendor    common.GpuVendor
 	gpuLock      *GpuLock
 	tasks        TaskStorage
 }
@@ -69,12 +70,12 @@ func NewDockerRunner(ctx context.Context, dockerParams DockerParameters) (*Docke
 		return nil, tracerr.Wrap(err)
 	}
 
-	var gpuVendor host.GpuVendor
+	var gpuVendor common.GpuVendor
 	gpus := host.GetGpuInfo(ctx)
 	if len(gpus) > 0 {
 		gpuVendor = gpus[0].Vendor
 	} else {
-		gpuVendor = host.GpuVendorNone
+		gpuVendor = common.GpuVendorNone
 	}
 	gpuLock, err := NewGpuLock(gpus)
 	if err != nil {
@@ -134,7 +135,7 @@ func (d *DockerRunner) restoreStateFromContainers(ctx context.Context) error {
 			log.Error(ctx, "failed to inspect container", "id", containerID, "task", taskID)
 		} else {
 			switch d.gpuVendor {
-			case host.GpuVendorNvidia:
+			case common.GpuVendorNvidia:
 				deviceRequests := containerFull.HostConfig.Resources.DeviceRequests
 				if len(deviceRequests) == 1 {
 					gpuIDs = deviceRequests[0].DeviceIDs
@@ -145,28 +146,28 @@ func (d *DockerRunner) restoreStateFromContainers(ctx context.Context) error {
 						"id", containerID, "task", taskID,
 					)
 				}
-			case host.GpuVendorAmd:
+			case common.GpuVendorAmd:
 				for _, device := range containerFull.HostConfig.Resources.Devices {
 					if host.IsRenderNodePath(device.PathOnHost) {
 						gpuIDs = append(gpuIDs, device.PathOnHost)
 					}
 				}
-			case host.GpuVendorTenstorrent:
+			case common.GpuVendorTenstorrent:
 				for _, device := range containerFull.HostConfig.Resources.Devices {
 					if strings.HasPrefix(device.PathOnHost, "/dev/tenstorrent/") {
 						// Extract the device ID from the path
 						deviceID := strings.TrimPrefix(device.PathOnHost, "/dev/tenstorrent/")
 						gpuIDs = append(gpuIDs, deviceID)
 					}
 				}
-			case host.GpuVendorIntel:
+			case common.GpuVendorIntel:
 				for _, envVar := range containerFull.Config.Env {
 					if indices, found := strings.CutPrefix(envVar, "HABANA_VISIBLE_DEVICES="); found {
 						gpuIDs = strings.Split(indices, ",")
 						break
 					}
 				}
-			case host.GpuVendorNone:
+			case common.GpuVendorNone:
 				gpuIDs = []string{}
 			}
 			ports = extractPorts(ctx, containerFull.NetworkSettings.Ports)
@@ -1014,12 +1015,12 @@ func configureGpuDevices(hostConfig *container.HostConfig, gpuDevices []GPUDevic
 	}
 }
 
-func configureGpus(config *container.Config, hostConfig *container.HostConfig, vendor host.GpuVendor, ids []string) {
+func configureGpus(config *container.Config, hostConfig *container.HostConfig, vendor common.GpuVendor, ids []string) {
 	// NVIDIA: ids are identifiers reported by nvidia-smi, GPU-<UUID> strings
 	// AMD: ids are DRI render node paths, e.g., /dev/dri/renderD128
 	// Tenstorrent: ids are device indices to be used with /dev/tenstorrent/<id>
 	switch vendor {
-	case host.GpuVendorNvidia:
+	case common.GpuVendorNvidia:
 		hostConfig.Resources.DeviceRequests = append(
 			hostConfig.Resources.DeviceRequests,
 			container.DeviceRequest{
@@ -1030,7 +1031,7 @@ func configureGpus(config *container.Config, hostConfig *container.HostConfig, v
 				DeviceIDs:    ids,
 			},
 		)
-	case host.GpuVendorAmd:
+	case common.GpuVendorAmd:
 		// All options are listed here: https://hub.docker.com/r/rocm/pytorch
 		// Only --device are mandatory, other seem to be performance-related.
 		// --device=/dev/kfd
@@ -1060,7 +1061,7 @@ func configureGpus(config *container.Config, hostConfig *container.HostConfig, v
 		// --security-opt=seccomp=unconfined
 		hostConfig.SecurityOpt = append(hostConfig.SecurityOpt, "seccomp=unconfined")
 		// TODO: in addition, for non-root user, --group-add=video, and possibly --group-add=render, are required.
-	case host.GpuVendorTenstorrent:
+	case common.GpuVendorTenstorrent:
 		// For Tenstorrent, simply add each device
 		for _, id := range ids {
 			devicePath := fmt.Sprintf("/dev/tenstorrent/%s", id)
@@ -1081,7 +1082,7 @@ func configureGpus(config *container.Config, hostConfig *container.HostConfig, v
 				Target: "/dev/hugepages-1G",
 			})
 		}
-	case host.GpuVendorIntel:
+	case common.GpuVendorIntel:
 		// All options are listed here:
 		// https://docs.habana.ai/en/latest/Installation_Guide/Additional_Installation/Docker_Installation.html
 		// --runtime=habana
@@ -1092,7 +1093,7 @@ func configureGpus(config *container.Config, hostConfig *container.HostConfig, v
 		hostConfig.CapAdd = append(hostConfig.CapAdd, "SYS_NICE")
 		// -e HABANA_VISIBLE_DEVICES=0,1,...
 		config.Env = append(config.Env, fmt.Sprintf("HABANA_VISIBLE_DEVICES=%s", strings.Join(ids, ",")))
-	case host.GpuVendorNone:
+	case common.GpuVendorNone:
 		// nothing to do
 	}
 }