diff --git a/kev/device.py b/kev/device.py index df0c2b2f..a99d6e62 100644 --- a/kev/device.py +++ b/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) diff --git a/tests/test_device.py b/tests/test_device.py index a146546a..8a74f5b2 100644 --- a/tests/test_device.py +++ b/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"