metal: harden for ggml initialization failures (#15755)

* metal: harden for ggml initialization failures

ggml_metal_device_init performs a probe to verify the tensor API compiles.  On
some systems this passes, even though kernel coverage isn't complete, which
results in a later crash when compiling the real kernels.  This change adds a
single retry if any of the error strings match this failure mode to disable the
tensor API.  It also hardens an error case in the Go initDevices to detect
device initialization failures and panic instead of crashing later on a nil
array entry.

Fixes #15734

* review comments

* review comments

ml/backend/ggml/ggml.go

@@ -50,8 +50,18 @@ var initDevices = sync.OnceFunc(func() {
 	backends = make(map[C.ggml_backend_dev_t]C.ggml_backend_t)
 	for i := range C.ggml_backend_dev_count() {
 		d := C.ggml_backend_dev_get(i)
+		t := C.ggml_backend_dev_type(d)
+		name := C.GoString(C.ggml_backend_dev_name(d))
 
-		switch C.ggml_backend_dev_type(d) {
+		b := C.ggml_backend_dev_init(d, nil)
+		if b == nil {
+			slog.Error("failed to initialize ggml backend device", "device", name, "type", t)
+			panic(fmt.Sprintf("failed to initialize ggml backend device: %s", name))
+		}
+
+		backends[d] = b
+
+		switch t {
 		case C.GGML_BACKEND_DEVICE_TYPE_CPU:
 			if len(cpus) == 0 {
 				// only the first cpu device should be used
@@ -63,8 +73,6 @@ var initDevices = sync.OnceFunc(func() {
 			C.GGML_BACKEND_DEVICE_TYPE_IGPU:
 			gpus = append(gpus, d)
 		}
-
-		backends[d] = C.ggml_backend_dev_init(d, nil)
 	}
 })
 

ml/device.go

@@ -314,6 +314,11 @@ type DeviceInfo struct {
 
 	// Where backends were loaded from
 	LibraryPath []string
+
+	// RunnerEnvOverrides stores exceptional per-device runner environment
+	// overrides discovered during bootstrap. This is internal server state and
+	// is not serialized.
+	RunnerEnvOverrides map[string]string `json:"-"`
 }
 
 type SystemInfo struct {
@@ -519,16 +524,24 @@ func (f FlashAttentionType) String() string {
 }
 
 // Given the list of GPUs this instantiation is targeted for,
-// figure out the visible devices environment variables
-// Set mustFilter true to enable filtering of CUDA devices
-func GetVisibleDevicesEnv(l []DeviceInfo, mustFilter bool) map[string]string {
+// figure out the device environment variables and any recorded
+// per-device runner environment overrides. Set mustFilter true to enable
+// filtering of CUDA devices.
+func GetDevicesEnv(l []DeviceInfo, mustFilter bool) map[string]string {
 	if len(l) == 0 {
 		return nil
 	}
 	env := map[string]string{}
 	for _, d := range l {
 		d.updateVisibleDevicesEnv(env, mustFilter)
+		for k, v := range d.RunnerEnvOverrides {
+			if existing, ok := env[k]; ok && existing != v {
+				slog.Warn("conflicting device environment override", "key", k, "existing", existing, "new", v, "library", d.Library, "id", d.ID)
+			}
+			env[k] = v
+		}
 	}
+
 	return env
 }
 

ml/device_test.go

@@ -0,0 +1,60 @@
+package ml
+
+import (
+	"bytes"
+	"log/slog"
+	"strings"
+	"testing"
+)
+
+func TestMergeEnvWithRunnerEnvOverrides(t *testing.T) {
+	devices := []DeviceInfo{
+		{
+			DeviceID:           DeviceID{Library: "Metal", ID: "0"},
+			RunnerEnvOverrides: map[string]string{"GGML_METAL_TENSOR_DISABLE": "1"},
+		},
+		{
+			DeviceID: DeviceID{Library: "CUDA", ID: "3"},
+		},
+	}
+
+	env := GetDevicesEnv(devices, true)
+
+	if got, want := env["GGML_METAL_TENSOR_DISABLE"], "1"; got != want {
+		t.Fatalf("GGML_METAL_TENSOR_DISABLE = %q, want %q", got, want)
+	}
+
+	if got, want := env["CUDA_VISIBLE_DEVICES"], "3"; got != want {
+		t.Fatalf("CUDA_VISIBLE_DEVICES = %q, want %q", got, want)
+	}
+}
+
+func TestGetDevicesEnvWarnsOnConflictingOverrides(t *testing.T) {
+	var logs bytes.Buffer
+	oldLogger := slog.Default()
+	slog.SetDefault(slog.New(slog.NewTextHandler(&logs, &slog.HandlerOptions{Level: slog.LevelDebug})))
+	t.Cleanup(func() {
+		slog.SetDefault(oldLogger)
+	})
+
+	devices := []DeviceInfo{
+		{
+			DeviceID:           DeviceID{Library: "Metal", ID: "0"},
+			RunnerEnvOverrides: map[string]string{"TEST_OVERRIDE": "one"},
+		},
+		{
+			DeviceID:           DeviceID{Library: "Metal", ID: "1"},
+			RunnerEnvOverrides: map[string]string{"TEST_OVERRIDE": "two"},
+		},
+	}
+
+	env := GetDevicesEnv(devices, false)
+
+	if got, want := env["TEST_OVERRIDE"], "two"; got != want {
+		t.Fatalf("TEST_OVERRIDE = %q, want %q", got, want)
+	}
+
+	if !strings.Contains(logs.String(), "conflicting device environment override") {
+		t.Fatalf("expected warning log, got %q", logs.String())
+	}
+}