Skip to content

Commit 0861720

Browse files
committed
fix(cli): accept auto/mps devices in pixelrag index build
1 parent 096adf5 commit 0861720

2 files changed

Lines changed: 12 additions & 1 deletion

File tree

index/src/pixelrag_index/pipelines.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,10 @@ def main():
212212
)
213213
parser.add_argument("--output", "-o", default=None, help="Output directory")
214214
parser.add_argument(
215-
"--device", default=None, choices=["cpu", "cuda"], help="Embedding device"
215+
"--device",
216+
default=None,
217+
choices=["auto", "cpu", "mps", "cuda"],
218+
help="Embedding device (default: from config; auto-detects cuda/mps/cpu)",
216219
)
217220
parser.add_argument(
218221
"--limit", "-n", type=int, default=None, help="Max documents to process"

tests/test_cli.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,14 @@ def test_pixelrag_umbrella_help():
3232
assert stage in out
3333

3434

35+
def test_index_build_device_choices():
36+
# CLI must offer every device the embedder supports (auto/mps were missing).
37+
# Assert the argparse-rendered choices token, not loose words in the help prose.
38+
r = _run("pixelrag", "index", "build", "--help")
39+
assert r.returncode == 0
40+
assert "{auto,cpu,mps,cuda}" in r.stdout
41+
42+
3543
def test_pixelrag_unknown_stage_errors():
3644
r = _run("pixelrag", "definitely-not-a-stage")
3745
assert r.returncode != 0

0 commit comments

Comments
 (0)