diff --git a/runner/cmd/shim/main.go b/runner/cmd/shim/main.go index 116e16b50c..50a163a0eb 100644 --- a/runner/cmd/shim/main.go +++ b/runner/cmd/shim/main.go @@ -269,7 +269,7 @@ func start(ctx context.Context, args shim.CLIArgs, serviceMode bool) (err error) var dcgmExporter *dcgm.DCGMExporter var dcgmWrapper dcgm.DCGMWrapperInterface - if gpu.GetGpuVendor() == gpu.GpuVendorNvidia { + if gpu.GetGpuVendor(ctx) == gpu.GpuVendorNvidia { dcgmExporterPath, err := dcgm.GetDCGMExporterExecPath(ctx) if err == nil { interval := time.Duration(args.DCGMExporter.Interval * int(time.Millisecond)) diff --git a/runner/internal/common/gpu/gpu.go b/runner/internal/common/gpu/gpu.go index 72ae83bb56..0a2c5afac8 100644 --- a/runner/internal/common/gpu/gpu.go +++ b/runner/internal/common/gpu/gpu.go @@ -1,8 +1,11 @@ package gpu import ( + "context" "errors" "os" + + "github.com/dstackai/dstack/runner/internal/common/log" ) type GpuVendor string @@ -15,26 +18,65 @@ const ( GpuVendorTenstorrent GpuVendor = "tenstorrent" ) -func GetGpuVendor() GpuVendor { - // FIXME: There might be errors other than os.ErrNotExist that are ignored silently. - // Propagate and log. - if _, err := os.Stat("/dev/kfd"); !errors.Is(err, os.ErrNotExist) { - return GpuVendorAmd - } - if _, err := os.Stat("/dev/nvidiactl"); !errors.Is(err, os.ErrNotExist) { +func GetGpuVendor(ctx context.Context) GpuVendor { + // Some devices can be detected unambiguously -- they have unique device paths. + // + // The order within this group does not matter -- false positives are unlikely, + // but more likely options should be checked first (e.g., NVIDIA before Tenstorrent). + // + // This group **must** stay before the ambiguous one below, otherwise a host with + // an AMD iGPU and an NVIDIA dGPU (a common combination) is detected as AMD, + // see: https://github.com/dstackai/dstack/issues/4085 + + // NVIDIA + if checkPath(ctx, "/dev/nvidiactl") { return GpuVendorNvidia } - if _, err := os.Stat("/dev/accel"); !errors.Is(err, os.ErrNotExist) { - return GpuVendorIntel + // NVIDIA on WSL2 + if checkPath(ctx, "/dev/dxg") && checkPath(ctx, "/usr/lib/wsl/lib/nvidia-smi") { + return GpuVendorNvidia } - if _, err := os.Stat("/dev/tenstorrent"); !errors.Is(err, os.ErrNotExist) { + // Tenstorrent + if checkPath(ctx, "/dev/tenstorrent") { return GpuVendorTenstorrent } - if _, err := os.Stat("/dev/dxg"); !errors.Is(err, os.ErrNotExist) { - // WSL2 - if _, err := os.Stat("/usr/lib/wsl/lib/nvidia-smi"); !errors.Is(err, os.ErrNotExist) { - return GpuVendorNvidia - } + + // The following devices are tricky -- the same paths are used for devices that we + // support and expect and devices that we don't support and don't want to support, such + // as AMD iGPUs or AMD/Intel NPUs integrated into CPUs. + // + // The order **does** matter, since both paths can be present on the same host, e.g., an + // AMD dGPU (/dev/kfd) on a host with an AMD NPU (/dev/accel). We decided to check for AMD + // first and Intel Gaudi last, because /dev/kfd is the more reliable signal of the two: + // * /dev/kfd is at least vendor-specific -- it identifies the vendor but not the device class + // * /dev/accel is standardized[1] vendor-agnostic path used by (including but not limited to): + // Intel (Habana Labs devices and NPU), AMD, Qualcomm, Rockchip -- it identifies neither, + // so mapping it to Intel Gaudi is a last-resort guess + // + // Apparently, there are also far more legit AMD accelerators (e.g., Instinct) deployed in + // the wild than Intel Gaudi accelerators, so the guess we make more often is the better one. + // + // [1]: https://github.com/torvalds/linux/blob/master/Documentation/accel/introduction.rst + + // AMD, **including** iGPU + if checkPath(ctx, "/dev/kfd") { + return GpuVendorAmd + } + // Intel/Habana Labs Gaudi OR some other compute accelerator + if checkPath(ctx, "/dev/accel") { + return GpuVendorIntel } + return GpuVendorNone } + +func checkPath(ctx context.Context, path string) bool { + _, err := os.Stat(path) + if err == nil { + return true + } + if !errors.Is(err, os.ErrNotExist) { + log.Error(ctx, "Failed to check path while detecting accelerator", "path", path, "err", err) + } + return false +} diff --git a/runner/internal/runner/metrics/metrics.go b/runner/internal/runner/metrics/metrics.go index 56c27a2bb1..34f22c731a 100644 --- a/runner/internal/runner/metrics/metrics.go +++ b/runner/internal/runner/metrics/metrics.go @@ -29,7 +29,7 @@ func NewMetricsCollector(ctx context.Context) (*MetricsCollector, error) { if err != nil { return nil, fmt.Errorf("get cgroup mount point: %w", err) } - gpuVendor := gpu.GetGpuVendor() + gpuVendor := gpu.GetGpuVendor(ctx) return &MetricsCollector{ cgroupMountPoint: cgroupMountPoint, gpuVendor: gpuVendor, diff --git a/runner/internal/shim/host/gpu.go b/runner/internal/shim/host/gpu.go index a9e7f97058..0526123da3 100644 --- a/runner/internal/shim/host/gpu.go +++ b/runner/internal/shim/host/gpu.go @@ -49,7 +49,7 @@ type GpuInfo struct { } func GetGpuInfo(ctx context.Context) []GpuInfo { - switch gpuVendor := gpu.GetGpuVendor(); gpuVendor { + switch gpuVendor := gpu.GetGpuVendor(ctx); gpuVendor { case gpu.GpuVendorNvidia: return getNvidiaGpuInfo(ctx) case gpu.GpuVendorAmd: