Enable MPS by default on macOS and add device override

Author: danias

.github/workflows/release.yml

@@ -29,7 +29,7 @@ jobs:
           - os: macos-15-intel
             target: x86_64-apple-darwin
             allow_failure: false
-            run_real_smoke: false
+            run_real_smoke: true
           - os: macos-14
             target: aarch64-apple-darwin
             allow_failure: false

README.md

@@ -17,28 +17,33 @@ The initial public model identifier is `bge-m3`.
 - Upstream model id: `BAAI/bge-m3`
 - Backend: `sentence-transformers`
 - Device:
-  - Apple Silicon macOS: `mps` when available, otherwise CPU
+  - macOS: `auto` by default, which selects `mps` when available, otherwise CPU
   - all other current targets: CPU
 - Provisioning: first-run download into a local cache directory
 
 The command and HTTP layers are written against an internal backend registry so additional models or inference backends can be added later without changing the user-facing contracts.
 
 ## Acceleration support
 
-Current hardware acceleration support is intentionally limited in `v0.1.2`:
+Current hardware acceleration support in `v0.1.3`:
 
 - `aarch64-apple-darwin`:
+  - defaults to `auto`
   - uses Apple Metal Performance Shaders (`mps`) automatically when available
   - falls back to CPU if MPS is unavailable
 - `x86_64-apple-darwin`:
-  - CPU only
+  - defaults to `auto`
+  - uses Apple Metal Performance Shaders (`mps`) automatically when available
+  - falls back to CPU if MPS is unavailable
 - `x86_64-unknown-linux-gnu`:
   - CPU only
 - `aarch64-unknown-linux-gnu`:
   - CPU only
 - `x86_64-pc-windows-msvc`:
   - CPU only
 
+Explicit device override is available on embedding, server, and daemon commands via `--device auto|cpu|mps`.
+
 The current release does not expose CUDA, ROCm, DirectML, or Intel GPU acceleration paths yet.
 
 ## Requirements
@@ -86,7 +91,7 @@ Example response:
   "embeddings": [[0.123, -0.456, 0.789]],
   "runtime": {
     "name": "bitloops-embeddings",
-    "version": "0.1.2"
+    "version": "0.1.3"
   }
 }
 ```
@@ -100,6 +105,13 @@ bitloops-embeddings embed \
   --output ./embedding.json
 ```
 
+Force CPU or request MPS explicitly:
+
+```bash
+bitloops-embeddings embed --model bge-m3 --input "Hello World" --device cpu
+bitloops-embeddings serve --model bge-m3 --device mps
+```
+
 Inspect model metadata without loading the model:
 
 ```bash
@@ -160,7 +172,7 @@ Response shape:
   "embeddings": [[0.123, -0.456, 0.789]],
   "runtime": {
     "name": "bitloops-embeddings",
-    "version": "0.1.2"
+    "version": "0.1.3"
   }
 }
 ```
@@ -258,6 +270,7 @@ Run the real-model smoke test against an installed console script or packaged ex
 
 ```bash
 python scripts/real_backend_smoke.py --binary bitloops-embeddings
+python scripts/real_backend_smoke.py --binary bitloops-embeddings --device mps
 ```
 
 ## GitHub Actions

scripts/real_backend_smoke.py

@@ -21,12 +21,19 @@ def main() -> None:
         required=True,
         help="Executable to invoke. This may be a console script name or an absolute path.",
     )
+    parser.add_argument(
+        "--device",
+        default="auto",
+        choices=("auto", "cpu", "mps"),
+        help="Inference device override to pass through to the runtime.",
+    )
     args = parser.parse_args()
 
     binary = args.binary
-    run_with_retries("embed smoke", lambda: run_embed_smoke(binary))
-    run_with_retries("server smoke", lambda: run_server_smoke(binary, reserve_free_port()))
-    run_with_retries("daemon smoke", lambda: run_daemon_smoke(binary))
+    device = args.device
+    run_with_retries("embed smoke", lambda: run_embed_smoke(binary, device))
+    run_with_retries("server smoke", lambda: run_server_smoke(binary, reserve_free_port(), device))
+    run_with_retries("daemon smoke", lambda: run_daemon_smoke(binary, device))
 
 
 def run_with_retries(name: str, operation) -> None:
@@ -48,9 +55,9 @@ def run_with_retries(name: str, operation) -> None:
             time.sleep(delay_seconds)
 
 
-def run_embed_smoke(binary: str) -> None:
+def run_embed_smoke(binary: str, device: str) -> None:
     completed = subprocess.run(
-        [binary, "embed", "--model", "bge-m3", "--input", "Hello World"],
+        [binary, "embed", "--model", "bge-m3", "--input", "Hello World", "--device", device],
         check=False,
         capture_output=True,
         text=True,
@@ -69,7 +76,7 @@ def run_embed_smoke(binary: str) -> None:
         raise RuntimeError("Embed smoke returned an empty embedding vector.")
 
 
-def run_server_smoke(binary: str, port: int) -> None:
+def run_server_smoke(binary: str, port: int, device: str) -> None:
     with tempfile.TemporaryDirectory(prefix="bitloops-embeddings-serve-logs-") as temp_dir:
         log_file = Path(temp_dir) / "serve.log"
         process = subprocess.Popen(
@@ -82,6 +89,8 @@ def run_server_smoke(binary: str, port: int) -> None:
                 "127.0.0.1",
                 "--port",
                 str(port),
+                "--device",
+                device,
                 "--log-file",
                 str(log_file),
             ],
@@ -109,7 +118,7 @@ def run_server_smoke(binary: str, port: int) -> None:
                 process.wait(timeout=5)
 
 
-def run_daemon_smoke(binary: str) -> None:
+def run_daemon_smoke(binary: str, device: str) -> None:
     with tempfile.TemporaryDirectory(prefix="bitloops-embeddings-daemon-logs-") as temp_dir:
         log_file = Path(temp_dir) / "daemon.log"
         process = subprocess.Popen(
@@ -118,6 +127,8 @@ def run_daemon_smoke(binary: str) -> None:
                 "daemon",
                 "--model",
                 "bge-m3",
+                "--device",
+                device,
                 "--log-file",
                 str(log_file),
             ],

src/bitloops_embeddings/backend/sentence_transformers_backend.py

@@ -7,7 +7,11 @@
 from threading import RLock
 from typing import Any
 
-from bitloops_embeddings.errors import BackendLoadError, InferenceError
+from bitloops_embeddings.errors import (
+    BackendLoadError,
+    InferenceError,
+    UnsupportedDeviceError,
+)
 from bitloops_embeddings.logging_utils import LOGGER_NAME, log_event
 
 
@@ -23,13 +27,15 @@ def __init__(
         upstream_model_id: str,
         cache_dir: Path,
         dimensions: int,
+        requested_device: str = "auto",
     ) -> None:
         self._model_id = model_id
         self._upstream_model_id = upstream_model_id
         self._cache_dir = cache_dir
         self._dimensions = dimensions
         self._model: Any = None
-        self._device = resolve_inference_device()
+        self._requested_device = requested_device
+        self._device = resolve_inference_device(requested_device=requested_device)
 
     @property
     def model_id(self) -> str:
@@ -188,28 +194,53 @@ def _configure_tqdm_lock_for_single_process() -> None:
     _TQDM_THREAD_LOCK_CONFIGURED = True
 
 
-def resolve_inference_device() -> str:
-    if platform.system() != "Darwin":
-        return "cpu"
+def resolve_inference_device_for_request(requested_device: str) -> str:
+    if requested_device == "auto":
+        return "mps" if _is_mps_available() else "cpu"
 
-    if platform.machine().lower() not in ("arm64", "aarch64"):
+    if requested_device == "cpu":
         return "cpu"
 
+    if requested_device == "mps":
+        unavailable_reason = _resolve_mps_unavailable_reason()
+        if unavailable_reason is None:
+            return "mps"
+        raise UnsupportedDeviceError(
+            f"MPS was requested but is unavailable: {unavailable_reason}"
+        )
+
+    raise UnsupportedDeviceError(
+        f"Unsupported device '{requested_device}'. Supported devices: auto, cpu, mps."
+    )
+
+
+def resolve_inference_device(requested_device: str = "auto") -> str:
+    return resolve_inference_device_for_request(requested_device)
+
+
+def _is_mps_available() -> bool:
+    return _resolve_mps_unavailable_reason() is None
+
+
+def _resolve_mps_unavailable_reason() -> str | None:
+    if platform.system() != "Darwin":
+        return "MPS is only available on macOS."
+
     try:
         import torch
     except ImportError:
-        return "cpu"
+        return "PyTorch is not installed."
 
     mps_backend = getattr(getattr(torch, "backends", None), "mps", None)
     if mps_backend is None:
-        return "cpu"
+        return "the installed PyTorch build does not expose torch.backends.mps."
 
     is_built = getattr(mps_backend, "is_built", None)
     if callable(is_built) and not is_built():
-        return "cpu"
+        return "the installed PyTorch build was not built with MPS support."
 
     is_available = getattr(mps_backend, "is_available", None)
     if callable(is_available) and is_available():
-        return "mps"
+        return None
 
-    return "cpu"
+    return "macOS 12.3 or later and an MPS-enabled GPU are required."

src/bitloops_embeddings/cli.py

@@ -40,6 +40,12 @@ class LogLevel(str, Enum):
     ERROR = "error"
 
 
+class Device(str, Enum):
+    AUTO = "auto"
+    CPU = "cpu"
+    MPS = "mps"
+
+
 class Transport(str, Enum):
     STDIO = "stdio"
 
@@ -66,6 +72,14 @@ def embed(
             writable=True,
         ),
     ] = None,
+    device: Annotated[
+        Device,
+        typer.Option(
+            "--device",
+            help="Inference device. auto prefers MPS when available, otherwise CPU.",
+            case_sensitive=False,
+        ),
+    ] = Device.AUTO,
     output: Annotated[
         Optional[Path],
         typer.Option(
@@ -81,7 +95,7 @@ def embed(
         raise typer.BadParameter("Only JSON output is supported in v1.")
 
     try:
-        backend = _build_backend(model=model, cache_dir=cache_dir)
+        backend = _build_backend(model=model, cache_dir=cache_dir, device=device)
         response = EmbeddingResponse(
             model_id=backend.model_id,
             dimensions=backend.dimensions,
@@ -110,6 +124,14 @@ def serve(
             writable=True,
         ),
     ] = None,
+    device: Annotated[
+        Device,
+        typer.Option(
+            "--device",
+            help="Inference device. auto prefers MPS when available, otherwise CPU.",
+            case_sensitive=False,
+        ),
+    ] = Device.AUTO,
     log_level: Annotated[
         LogLevel,
         typer.Option("--log-level", help="Server log verbosity.", case_sensitive=False),
@@ -130,7 +152,7 @@ def serve(
 ) -> None:
     try:
         configure_logging(log_level.value, log_file=log_file, prefer_os_log=True)
-        backend = _build_backend(model=model, cache_dir=cache_dir)
+        backend = _build_backend(model=model, cache_dir=cache_dir, device=device)
         backend.load()
         app_instance = create_app(backend, max_batch_size=max_batch_size)
         log_event(
@@ -164,6 +186,14 @@ def daemon(
             writable=True,
         ),
     ] = None,
+    device: Annotated[
+        Device,
+        typer.Option(
+            "--device",
+            help="Inference device. auto prefers MPS when available, otherwise CPU.",
+            case_sensitive=False,
+        ),
+    ] = Device.AUTO,
     log_level: Annotated[
         LogLevel,
         typer.Option("--log-level", help="Daemon log verbosity.", case_sensitive=False),
@@ -182,7 +212,7 @@ def daemon(
         configure_logging(log_level.value, log_file=log_file, prefer_os_log=True)
         if transport is not Transport.STDIO:
             raise typer.BadParameter("Only stdio transport is supported in v1.")
-        backend = _build_backend(model=model, cache_dir=cache_dir)
+        backend = _build_backend(model=model, cache_dir=cache_dir, device=device)
         raise typer.Exit(code=run_daemon(backend))
     except typer.BadParameter:
         raise
@@ -226,10 +256,15 @@ def describe(
         _exit_with_error(BitloopsEmbeddingsError(f"Unexpected runtime error: {exc}"))
 
 
-def _build_backend(*, model: str, cache_dir: Optional[Path]) -> EmbeddingBackend:
+def _build_backend(
+    *,
+    model: str,
+    cache_dir: Optional[Path],
+    device: Device = Device.AUTO,
+) -> EmbeddingBackend:
     resolved_cache_dir = ensure_cache_dir(resolve_cache_dir(cache_dir))
     spec = get_model_spec(model)
-    return spec.create_backend(resolved_cache_dir)
+    return spec.create_backend(resolved_cache_dir, requested_device=device.value)
 
 
 def _emit_json(payload: str, *, output: Optional[Path]) -> None:

src/bitloops_embeddings/errors.py

@@ -24,6 +24,11 @@ class UnsupportedModelError(BitloopsEmbeddingsError):
     default_status_code = 400
 
 
+class UnsupportedDeviceError(BitloopsEmbeddingsError):
+    default_code = "unsupported_device"
+    default_status_code = 400
+
+
 class BackendLoadError(BitloopsEmbeddingsError):
     default_code = "backend_load_error"
     default_status_code = 500