Avoid probing MPS twice after automatic device selection fails

kev/device.py

@@ -18,6 +18,8 @@
 # The child prints one line with this prefix. ROCm and other wheels write startup text to stdout around it;
 # a bare "cuda" line in that text is not a result.
 PROBE_PREFIX = "kev-device "
+# Announce the selected device before allocation, so a failed auto probe does not retry MPS.
+PROBE_START_PREFIX = "kev-probing "
 DEVICE_HELP = ("auto probes in a child process and falls back to cpu (CUDA/ROCm first, then MPS; "
                "each probe waits KEV_DEVICE_PROBE_TIMEOUT seconds, default 30); "
                "an explicit accelerator fails closed; cpu skips discovery")
@@ -42,7 +44,7 @@ def probe_timeout():
 def select(device="auto", *, dtype=torch.float32, timeout=None):
     """Resolve a device before loading weights. A fresh interpreter contains native driver crashes.
 
-    Explicit CPU never inspects accelerators. Auto tries CUDA/ROCm and, if that probe fails, MPS, then falls back to CPU.
+    Explicit CPU never inspects accelerators. Auto tries CUDA/ROCm, then MPS, then CPU. A failed MPS probe is not repeated.
     An explicit accelerator fails closed. This tests one allocation/matmul/sync, not every model kernel or whether the checkpoint fits.
     """
     if device != "auto" and device not in DEVICES:
@@ -57,6 +59,11 @@ def select(device="auto", *, dtype=torch.float32, timeout=None):
             return _run_probe(attempt, dtype, timeout)
         except (OSError, subprocess.SubprocessError, ValueError) as error:
             failures.append((attempt, error))
+            output = getattr(error, "output", None) or ""
+            if isinstance(output, bytes):  # TimeoutExpired keeps captured output as bytes, even with text=True.
+                output = output.decode("utf-8", errors="replace")
+            if attempt == "auto" and PROBE_START_PREFIX + "mps" in output.splitlines():
+                break
     message = "Device probe failed (" + "; ".join(f"{attempt}: {_probe_detail(error)}" for attempt, error in failures) + ")."
     if device != "auto":
         raise RuntimeError(f"{message} Fix the accelerator runtime or pass --device cpu.") from failures[-1][1]
@@ -87,6 +94,7 @@ def _probe(device, dtype):
     """Child-process entry point: no accelerator operations from here may run in the selecting process."""
     if device == "auto":
         device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+    print(PROBE_START_PREFIX + device, flush=True)
     if device != "cpu":
         with torch.no_grad():
             matrix = torch.ones(8, 8, device=device, dtype=dtype)

tests/test_device.py

@@ -355,13 +355,19 @@ def test_base_mmlu_probe_help_does_not_probe(monkeypatch):
     assert error.value.code == 0
 
 
-def test_auto_tries_mps_after_the_cuda_probe_fails(monkeypatch):
+@pytest.mark.parametrize("failure", [
+    subprocess.CalledProcessError(-11, "probe"),
+    subprocess.CalledProcessError(-11, "probe", output=device.PROBE_START_PREFIX + "cuda\n",
+                                  stderr="HIP error: invalid device function"),
+    subprocess.TimeoutExpired("probe", 30, output=(device.PROBE_START_PREFIX + "cuda\n").encode()),
+])
+def test_auto_tries_mps_after_the_cuda_probe_fails(monkeypatch, failure):
     calls = []
     def run(command, **kwargs):
         calls.append(command[-2])
         assert kwargs["timeout"] == 30.0
         if command[-2] == "auto":
-            raise subprocess.CalledProcessError(-11, "probe", stderr="HIP error: invalid device function")
+            raise failure
         return SimpleNamespace(stdout=reported("mps"))
     monkeypatch.setattr(subprocess, "run", run)
     with warnings.catch_warnings():
@@ -525,3 +531,42 @@ def run(command, **kwargs):
     monkeypatch.setattr(subprocess, "run", run)
     serve.main()
     assert seen == ["cuda"] and startup[0][0] == "cuda"
+
+
+@pytest.mark.parametrize("failure", ["error", "timeout", "crash"])
+def test_auto_does_not_repeat_a_failed_mps_probe(tmp_path, monkeypatch, failure):
+    """Run the real child with an MPS stand-in; each allocation records one probe attempt."""
+    fake_torch = '''
+import os, signal, time
+from contextlib import nullcontext
+from types import SimpleNamespace
+float32 = "float32"
+cuda = SimpleNamespace(is_available=lambda: False)
+backends = SimpleNamespace(mps=SimpleNamespace(is_available=lambda: True))
+no_grad = nullcontext
+def ones(*args, **kwargs):
+    with open(os.environ["PROBE_ATTEMPTS"], "a", encoding="utf-8") as out:
+        out.write(kwargs["device"] + "\\n")
+    failure = os.environ["PROBE_FAILURE"]
+    if failure == "timeout":
+        time.sleep(30)
+    elif failure == "crash":
+        os.kill(os.getpid(), signal.SIGTERM)
+    raise TypeError("unsupported MPS dtype")
+'''
+    (tmp_path / "torch.py").write_text(fake_torch, encoding="utf-8")
+    attempts = tmp_path / "attempts.txt"
+    monkeypatch.setenv("PYTHONPATH", str(tmp_path))
+    monkeypatch.setenv("PROBE_ATTEMPTS", str(attempts))
+    monkeypatch.setenv("PROBE_FAILURE", failure)
+    with pytest.warns(RuntimeWarning, match="Falling back to CPU") as caught:
+        assert device.select(timeout=2 if failure == "timeout" else 10) == "cpu"
+    assert attempts.read_text(encoding="utf-8").splitlines() == ["mps"]
+    assert "; mps:" not in str(caught[0].message)
+
+
+def test_probe_start_marker_does_not_report_success(monkeypatch):
+    monkeypatch.setattr(subprocess, "run", lambda *a, **k:
+                        SimpleNamespace(stdout=device.PROBE_START_PREFIX + "mps\n"))
+    with pytest.warns(RuntimeWarning, match="unexpected probe result"):
+        assert device.select() == "cpu"