From 2779740382fe7e2538e0236590d3c6050d8a29f8 Mon Sep 17 00:00:00 2001 From: Cody Date: Mon, 22 Jun 2026 17:57:23 -0700 Subject: [PATCH] fix(cli): accept auto/mps devices in pixelrag index build --- index/src/pixelrag_index/pipelines.py | 5 ++++- tests/test_cli.py | 8 ++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/index/src/pixelrag_index/pipelines.py b/index/src/pixelrag_index/pipelines.py index f943ec8..0f4800f 100644 --- a/index/src/pixelrag_index/pipelines.py +++ b/index/src/pixelrag_index/pipelines.py @@ -217,7 +217,10 @@ def main(): ) parser.add_argument("--output", "-o", default=None, help="Output directory") parser.add_argument( - "--device", default=None, choices=["cpu", "cuda"], help="Embedding device" + "--device", + default=None, + choices=["auto", "cpu", "mps", "cuda"], + help="Embedding device (default: from config; auto-detects cuda/mps/cpu)", ) parser.add_argument( "--limit", "-n", type=int, default=None, help="Max documents to process" diff --git a/tests/test_cli.py b/tests/test_cli.py index 1c43df2..2d32016 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -32,6 +32,14 @@ def test_pixelrag_umbrella_help(): assert stage in out +def test_index_build_device_choices(): + # CLI must offer every device the embedder supports (auto/mps were missing). + # Assert the argparse-rendered choices token, not loose words in the help prose. + r = _run("pixelrag", "index", "build", "--help") + assert r.returncode == 0 + assert "{auto,cpu,mps,cuda}" in r.stdout + + def test_pixelrag_unknown_stage_errors(): r = _run("pixelrag", "definitely-not-a-stage") assert r.returncode != 0